package org.jdetect;

import org.opencv.core.*;
import org.opencv.dnn.Dnn;
import org.opencv.dnn.Net;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.List;
import java.nio.file.Files;
import java.nio.file.Paths;


public class ModelLoader {
    private static Net net;
    private static List<String> classes;
    private static List<String> outputLayers;

    public static synchronized void loadModel() throws IOException {
        if (net != null) return;

        // Load network
        InputStream weightsStream = ModelLoader.class.getResourceAsStream("/yolov4-tiny.weights");
        InputStream configStream = ModelLoader.class.getResourceAsStream("/yolov4-tiny.cfg");

        assert weightsStream != null;
        byte[] weights = weightsStream.readAllBytes();
        assert configStream != null;
        byte[] config = configStream.readAllBytes();

        // Create temp files
        String weightsPath = "temp_weights.weights";
        String configPath = "temp_config.cfg";

        Files.write(Paths.get(weightsPath), weights);
        Files.write(Paths.get(configPath), config);

        net = Dnn.readNetFromDarknet(configPath, weightsPath);

        // Load classes
        classes = new ArrayList<>();
        try (InputStream classesStream = ModelLoader.class.getResourceAsStream("/coco.names")) {
            assert classesStream != null;
            try (BufferedReader br = new BufferedReader(new InputStreamReader(classesStream))) {
                String line;
                while ((line = br.readLine()) != null) {
                    classes.add(line.trim());
                }
            }
        }

        // Get output layers
        List<String> layerNames = net.getLayerNames();
        MatOfInt unconnectedOutLayers = net.getUnconnectedOutLayers();
        int[] indices = unconnectedOutLayers.toArray();

        outputLayers = new ArrayList<>();
        for (int idx : indices) {
            outputLayers.add(layerNames.get(idx - 1));
        }
    }

    public static Net getNet() {
        return net;
    }

    public static List<String> getClasses() {
        return classes;
    }

    public static List<String> getOutputLayers() {
        return outputLayers;
    }
}