package xyz.r0r5chach.dermy.models;

import android.content.res.AssetFileDescriptor;
import android.content.res.Resources;
import android.graphics.Bitmap;

import org.tensorflow.lite.Interpreter;


import java.io.FileInputStream;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.nio.MappedByteBuffer;
import java.nio.channels.FileChannel;
import java.nio.channels.FileChannel.MapMode;

import xyz.r0r5chach.dermy.R;

public class Model {
    protected final Interpreter modelInterpreter;

    public Model(Resources resources, int modelId) throws IOException {
        modelInterpreter = new Interpreter(loadModel(resources, modelId));
    }

    public float[] runInference(Bitmap image, int classes) {
        ByteBuffer inputBuffer = preprocessImage(image);
        float[][] output = new float[1][classes];
        modelInterpreter.run(inputBuffer, output);

        return output[0];
    }

    public MappedByteBuffer loadModel(Resources resources, int modelId) throws IOException {
        AssetFileDescriptor fileDescriptor = resources.openRawResourceFd(modelId);
        FileInputStream inputStream = new FileInputStream(fileDescriptor.getFileDescriptor()); //FIXME: add try-with-resources

        FileChannel fileChannel = inputStream.getChannel();
        long startOffset = fileDescriptor.getStartOffset();
        long declaredLength = fileDescriptor.getDeclaredLength();

        return fileChannel.map(MapMode.READ_ONLY, startOffset, declaredLength);
    }

    public static ByteBuffer preprocessImage(Bitmap image) {
        ByteBuffer buffer = ByteBuffer.allocateDirect(4 * 280 * 280 * 3);
        buffer.order(ByteOrder.nativeOrder());

        int[] intValues = new int[280 * 280];

        image.getPixels(intValues, 0, image.getWidth(), 0, 0, image.getWidth(), image.getHeight());

        int pixel = 0;
        for (int i = 0; i < 280; ++i) {
            for (int j = 0; j < 280; ++j) {
                int val = intValues[pixel++];
                buffer.putFloat((((val >> 16) & 0xFF) - 127) / 128.0f);
                buffer.putFloat((((val >> 8) & 0xFF) - 127) / 128.0f);
                buffer.putFloat(((val & 0xFF) - 127) / 128.0f);
            }
        }

        return buffer;
    }

}
