package com.evetensorflowyolo.tflite;

import android.content.res.AssetManager;
import android.util.Log;

import java.io.IOException;

public class DetectorFactory {
    public static com.evetensorflowyolo.tflite.YoloV5Classifier getDetector(
            final AssetManager assetManager,
            final String modelFilename)
            throws IOException {
        String labelFilename = null;
        boolean isQuantized = false;
        int inputSize = 0;
        int[] output_width = new int[] { 0 };
        int[][] masks = new int[][] { { 0 } };
        int[] anchors = new int[] { 0 };

        if (modelFilename.equals("ammoTypes.tflite")) {
            Log.v("Love", "set file name for model");
            labelFilename = "file:///android_asset/ammo_types.txt";
            Log.v("Love", "set file name for classes");
            isQuantized = false;
            inputSize = 512;
            output_width = new int[] { 80, 40, 20 };
            masks = new int[][] { { 0, 1, 2 }, { 3, 4, 5 }, { 6, 7, 8 } };
            anchors = new int[] {
                    10, 13, 16, 30, 33, 23, 30, 61, 62, 45, 59, 119, 116, 90, 156, 198, 373, 326
            };
        } else if (modelFilename.equals("ObjectDetection.tflite")) {
            Log.v("Love", "set file name for model");
            labelFilename = "file:///android_asset/coco.txt";
            Log.v("Love", "set file name for classes");
            isQuantized = false;
            inputSize = 640;
            output_width = new int[] { 80, 40, 20 };
            masks = new int[][] { { 0, 1, 2 }, { 3, 4, 5 }, { 6, 7, 8 } };
            anchors = new int[] {
                    10, 13, 16, 30, 33, 23, 30, 61, 62, 45, 59, 119, 116, 90, 156, 198, 373, 326
            };

        } else if (modelFilename.equals("efficientdet3.tflite")) {
            Log.v("Love", "set file name for model");
            labelFilename = "file:///android_asset/coco.txt";
            Log.v("Love", "set file name for classes");
            isQuantized = false;
            inputSize = 224;
            output_width = new int[] { 80, 40, 20 };
            masks = new int[][] { { 0, 1, 2 }, { 3, 4, 5 }, { 6, 7, 8 } };
            anchors = new int[] {
                    10, 13, 16, 30, 33, 23, 30, 61, 62, 45, 59, 119, 116, 90, 156, 198, 373, 326
            };
        } else if (modelFilename.equals("EquipmentCheck.tflite")) {
            Log.v("Love", "set file name for model");
            labelFilename = "file:///android_asset/equipment_check.txt";
            Log.v("Love", "set file name for classes");
            isQuantized = false;
            inputSize = 512;
            output_width = new int[] { 80, 40, 20 };
            masks = new int[][] { { 0, 1, 2 }, { 3, 4, 5 }, { 6, 7, 8 } };
            anchors = new int[] {
                    10, 13, 16, 30, 33, 23, 30, 61, 62, 45, 59, 119, 116, 90, 156, 198, 373, 326
            };
        } else if (modelFilename.equals("ProjAmmoDetection.tflite")) {
            Log.v("Love", "set file name for model");
            labelFilename = "file:///android_asset/proj_ammo.txt";
            Log.v("Love", "set file name for classes");
            isQuantized = false;
            inputSize = 224;
            output_width = new int[] { 80, 40, 20 };
            masks = new int[][] { { 0, 1, 2 }, { 3, 4, 5 }, { 6, 7, 8 } };
            anchors = new int[] {
                    10, 13, 16, 30, 33, 23, 30, 61, 62, 45, 59, 119, 116, 90, 156, 198, 373, 326
            };
        } else if (modelFilename.equals("ObjectDetection2.tflite")) {
            Log.v("Love", "set file name for model");
            labelFilename = "file:///android_asset/coco.txt";
            Log.v("Love", "set file name for classes");
            isQuantized = false;
            inputSize = 224;
            output_width = new int[] { 80, 40, 20 };
            masks = new int[][] { { 0, 1, 2 }, { 3, 4, 5 }, { 6, 7, 8 } };
            anchors = new int[] {
                    10, 13, 16, 30, 33, 23, 30, 61, 62, 45, 59, 119, 116, 90, 156, 198, 373, 326
            };

        } else if (modelFilename.equals("WordsDegradation.tflite")) {
            Log.v("Love", "set file name for model");
            labelFilename = "file:///android_asset/words_degradation.txt";
            Log.v("Love", "set file name for classes");
            isQuantized = false;
            inputSize = 224;
            output_width = new int[] { 80, 40, 20 };
            masks = new int[][] { { 0, 1, 2 }, { 3, 4, 5 }, { 6, 7, 8 } };
            anchors = new int[] {
                    10, 13, 16, 30, 33, 23, 30, 61, 62, 45, 59, 119, 116, 90, 156, 198, 373, 326
            };
        } else if (modelFilename.equals("yolov5s1.tflite")) {
            labelFilename = "file:///android_asset/coco.txt";
            isQuantized = false;
            inputSize = 320;
            output_width = new int[] { 40, 20, 10 };
            masks = new int[][] { { 0, 1, 2 }, { 3, 4, 5 }, { 6, 7, 8 } };
            anchors = new int[] {
                    10, 13, 16, 30, 33, 23, 30, 61, 62, 45, 59, 119, 116, 90, 156, 198, 373, 326
            };
        } else if (modelFilename.equals("yolov5s-int8.tflite")) {
            labelFilename = "file:///android_asset/customclasses.txt";
            isQuantized = true;
            inputSize = 416;
            output_width = new int[] { 40, 20, 10 };
            masks = new int[][] { { 0, 1, 2 }, { 3, 4, 5 }, { 6, 7, 8 } };
            anchors = new int[] {
                    10, 13, 16, 30, 33, 23, 30, 61, 62, 45, 59, 119, 116, 90, 156, 198, 373, 326
            };
        }
        Log.v("Love", "about to call Yolov5 classifier");
        return com.evetensorflowyolo.tflite.YoloV5Classifier.create(assetManager, modelFilename, labelFilename,
                isQuantized,
                inputSize);
    }

}