using System.Collections.Generic;
using Unity.Sentis;
using UnityEngine;
using System.IO;
using FF = Unity.Sentis.Functional;
/*
 *  MovileNetV2 Inference Script
 *  ============================
 *  
 *  Place this script on the Main Camera
 *  
 *  Drag an image to the inputImage field
 *  
 *  When run the prediction of what the image is will output to the console window.
 *  You can modify the script to make it do something more interesting.
 * 
 */


public class RunMobileNet : MonoBehaviour
{
    //draw the sentis file here:
    public ModelAsset modelAsset;

    const string modelName = "mobilenet_v2.sentis";

    //The image to classify here:
    public Texture2D inputImage;

    //Link class_desc.txt here:
    public TextAsset labelsAsset;

    //All images are resized to these values to go into the model
    const int imageHeight = 224;
    const int imageWidth = 224;

    const BackendType backend = BackendType.GPUCompute;

    private IWorker engine;
    private string[] labels;

    //Used to normalise the input RGB values
    TensorFloat mulRGB = new TensorFloat(new TensorShape(1, 3, 1, 1), new float[] { 1 / 0.229f, 1 / 0.224f, 1 / 0.225f });
    TensorFloat shiftRGB = new TensorFloat(new TensorShape(1, 3, 1, 1), new float[] { 0.485f, 0.456f, 0.406f });

    void Start()
    {

        //Parse neural net labels
        labels = labelsAsset.text.Split('\n');

        //Load model from file or asset
        //var model = ModelLoader.Load(Path.Join(Application.streamingAssetsPath, modelName));
        var model = ModelLoader.Load(modelAsset);

        //We modify the model to normalise the input RGB values and select the highest prediction
        //probability and item number
        var model2 = FF.Compile(
            input =>
            {
                var probability = model.Forward(NormaliseRGB(input))[0];
                return (FF.ReduceMax(probability, 1), FF.ArgMax(probability, 1));
            },
            model.inputs[0]
        );

        //Setup the engine to run the model
        engine = WorkerFactory.CreateWorker(backend, model2);

        //Execute inference
        ExecuteML();
    }

    public void ExecuteML()
    {
        //Preprocess image for input
        using var input = TextureConverter.ToTensor(inputImage, imageWidth, imageHeight, 3);
        
        //Execute neural net
        engine.Execute(input);

        //Read output tensor
        var probability = engine.PeekOutput("output_0") as TensorFloat;
        var item = engine.PeekOutput("output_1") as TensorInt;
        item.CompleteOperationsAndDownload();
        probability.CompleteOperationsAndDownload();

        //Select the best output class and print the results
        var ID = item[0];
        var accuracy = probability[0];

        //The result is output to the console window
        int percent = Mathf.FloorToInt(accuracy * 100f + 0.5f);
        Debug.Log($"Prediction: {labels[ID]} {percent}﹪");

        //Clean memory
        Resources.UnloadUnusedAssets();
    }

    //This scales and shifts the RGB values for input into the model
    FunctionalTensor NormaliseRGB(FunctionalTensor image)
    {
        return (image - FunctionalTensor.FromTensor(shiftRGB)) * FunctionalTensor.FromTensor(mulRGB);
    }
    
    private void OnDestroy()
    {
        mulRGB?.Dispose();
        shiftRGB?.Dispose();
        engine?.Dispose();      
    }
}