using UnityEngine;
using Unity.Barracuda;
using UnityEngine.UI;
using TMPro;

public class ModelInference : MonoBehaviour
{
    // Referensi ke model ONNX (drag & drop file ONNX di Inspector)
    public NNModel onnxModel;
    private Model runtimeModel;
    private IWorker worker;

    // Ukuran input model (sesuai dengan training, misalnya 224x224)
    public int inputWidth = 224;
    public int inputHeight = 224;

    // Referensi UI untuk menampilkan input dan hasil
    public RawImage inputDisplay;
    public TextMeshProUGUI resultText;

    void Start()
    {
        // Muat model dari asset
        runtimeModel = ModelLoader.Load(onnxModel);
        // Buat worker untuk melakukan inferensi (auto memilih antara CPU/GPU)
        worker = WorkerFactory.CreateWorker(WorkerFactory.Type.Auto, runtimeModel);
    }

    // Fungsi untuk melakukan inferensi pada Texture2D
    public void RunModel(Texture2D inputTexture)
    {
        // Tampilkan input (opsional)
        if (inputDisplay != null)
        {
            inputDisplay.texture = inputTexture;
        }

        // Resize texture ke ukuran input model
        Texture2D resized = ResizeTexture(inputTexture, inputWidth, inputHeight);

        // Konversi texture ke tensor dengan normalisasi piksel ke rentang [-1, 1]
        Color32[] pixels = resized.GetPixels32();
        float[] floatValues = new float[inputWidth * inputHeight * 3];

        for (int i = 0; i < pixels.Length; i++)
        {
            // Normalisasi: (value / 127.5f) - 1.0f
            floatValues[i * 3 + 0] = (pixels[i].r / 127.5f) - 1.0f;
            floatValues[i * 3 + 1] = (pixels[i].g / 127.5f) - 1.0f;
            floatValues[i * 3 + 2] = (pixels[i].b / 127.5f) - 1.0f;
        }

        // Buat tensor input dengan dimensi [batch, height, width, channels]
        Tensor inputTensor = new Tensor(1, inputHeight, inputWidth, 3, floatValues);

        // Jalankan inferensi
        worker.Execute(inputTensor);
        Tensor outputTensor = worker.PeekOutput();

        // Misalnya, output memiliki 2 nilai (anorganik dan organik)
        int predictedIndex = ArgMax(outputTensor.ToReadOnlyArray());
        string predictedLabel = predictedIndex == 0 ? "anorganik" : "organik";

        // Tampilkan hasil prediksi pada UI
        resultText.text = "Predicted: " + predictedLabel;

        // Dispose tensor untuk menghindari memory leak
        inputTensor.Dispose();
        outputTensor.Dispose();
    }

    // Fungsi untuk mendapatkan indeks dengan nilai maksimum
    private int ArgMax(float[] array)
    {
        int index = 0;
        float max = array[0];
        for (int i = 1; i < array.Length; i++)
        {
            if (array[i] > max)
            {
                max = array[i];
                index = i;
            }
        }
        return index;
    }

    // Fungsi untuk meresize Texture2D
    private Texture2D ResizeTexture(Texture2D source, int newWidth, int newHeight)
    {
        RenderTexture rt = RenderTexture.GetTemporary(newWidth, newHeight);
        RenderTexture.active = rt;
        Graphics.Blit(source, rt);
        Texture2D newTexture = new Texture2D(newWidth, newHeight);
        newTexture.ReadPixels(new Rect(0, 0, newWidth, newHeight), 0, 0);
        newTexture.Apply();
        RenderTexture.active = null;
        RenderTexture.ReleaseTemporary(rt);
        return newTexture;
    }

    // Fungsi untuk mengambil gambar menggunakan NativeCamera dan mengoreksi orientasinya
    public void TakePicture()
    {
        // Periksa izin kamera terlebih dahulu
        NativeCamera.Permission permission = NativeCamera.CheckPermission(true);
        if (permission == NativeCamera.Permission.Denied || permission == NativeCamera.Permission.ShouldAsk)
        {
            NativeCamera.RequestPermission(true);
        }

        // Panggil kamera untuk mengambil gambar dengan ukuran maksimal inputWidth
        NativeCamera.TakePicture((string path) =>
        {
            if (!string.IsNullOrEmpty(path))
            {
                // Muat gambar dari path yang diambil
                Texture2D capturedTexture = NativeCamera.LoadImageAtPath(path, inputWidth);
                if (capturedTexture != null)
                {
                    // Pastikan texture dari kamera bisa diakses (readable)
                    Texture2D readableTexture = MakeTextureReadable(capturedTexture);
                    // Koreksi orientasi gambar yang diambil dari kamera
                    Texture2D correctedTexture = CorrectCameraImage(readableTexture);
                    // Lakukan inferensi pada gambar yang telah dikoreksi
                    RunModel(correctedTexture);
                }
                else
                {
                    Debug.LogError("Gagal memuat gambar dari kamera");
                }
            }
        }, inputWidth); // inputWidth sebagai ukuran maksimum gambar
    }

    // Fungsi untuk melakukan inferensi pada sample texture dari Resources
    public void TestInference()
    {
        Texture2D sampleTexture = Resources.Load<Texture2D>("n3");
        RunModel(sampleTexture);
    }

    // Fungsi untuk mengoreksi orientasi gambar yang diambil dari NativeCamera
    // Contoh ini memutar gambar 90° searah jarum jam
    private Texture2D CorrectCameraImage(Texture2D original)
    {
        int width = original.width;
        int height = original.height;
        Texture2D rotated = new Texture2D(height, width, original.format, false);

        for (int i = 0; i < width; i++)
        {
            for (int j = 0; j < height; j++)
            {
                rotated.SetPixel(j, width - i - 1, original.GetPixel(i, j));
            }
        }
        rotated.Apply();
        return rotated;
    }

    // Fungsi untuk membuat texture menjadi readable
    private Texture2D MakeTextureReadable(Texture2D tex)
    {
        RenderTexture tmp = RenderTexture.GetTemporary(tex.width, tex.height, 0, RenderTextureFormat.Default, RenderTextureReadWrite.Linear);
        Graphics.Blit(tex, tmp);
        RenderTexture previous = RenderTexture.active;
        RenderTexture.active = tmp;
        Texture2D readableTex = new Texture2D(tex.width, tex.height);
        readableTex.ReadPixels(new Rect(0, 0, tmp.width, tmp.height), 0, 0);
        readableTex.Apply();
        RenderTexture.active = previous;
        RenderTexture.ReleaseTemporary(tmp);
        return readableTex;
    }

    void OnDestroy()
    {
        // Pastikan untuk dispose worker saat objek dihancurkan
        worker.Dispose();
    }
}