package com.birdsounds.identify


import android.content.Context
import android.graphics.Bitmap
import android.location.Location
import android.media.Image
import android.util.Log
import android.widget.ImageView
import org.apache.commons.math3.stat.descriptive.summary.Sum
import org.tensorflow.lite.Interpreter
import org.tensorflow.lite.support.image.ImageProcessor
import org.tensorflow.lite.support.image.TensorImage
import org.tensorflow.lite.support.image.ops.ResizeOp
import org.tensorflow.lite.support.label.Category
import org.tensorflow.lite.support.tensorbuffer.TensorBuffer
import org.tensorflow.lite.task.core.BaseOptions
import org.tensorflow.lite.task.vision.classifier.Classifications
import org.tensorflow.lite.task.vision.classifier.ImageClassifier
import org.tensorflow.lite.task.vision.classifier.ImageClassifier.ImageClassifierOptions
import java.io.File
import java.nio.ByteBuffer
import java.nio.FloatBuffer
import java.nio.channels.FileChannel
import java.nio.file.StandardOpenOption
import java.time.LocalDate
import java.time.ZoneId
import java.util.Calendar
import java.util.Calendar.WEEK_OF_YEAR
import java.util.Date


var local_species_spectrogram = "";


fun ByteBuffer.toUByteArray(): UByteArray {
    // Get remaining bytes in the buffer
    val bytes = ByteArray(this.remaining())

    // Copy bytes from the buffer
    this.get(bytes)

    // Convert to UByteArray
    return bytes.toUByteArray()
}

class SoundClassifierSpectrogramBased(context_in: MainActivity) {
    private var c_threshold: Float = 0.0f;
    private var model_ready: Boolean = false;
    private var context: MainActivity = context_in;
    var norm_factor: Float = (1 shl 15).toFloat();
    var score_string = "";
    var image_resize_op =
        ImageProcessor.Builder().add(ResizeOp(128, 512, ResizeOp.ResizeMethod.BILINEAR)).build();
    private lateinit var interpreter_spectrogram: Interpreter;
    private var spec_inp_dim: IntArray = intArrayOf(); // [512]
    private var spec_out_dim: IntArray = intArrayOf(); // [128,1]
    private lateinit var specInputBuffer: TensorBuffer;
    private lateinit var specOutputBuffer: TensorBuffer;


    private lateinit var interpreter_geo: Interpreter;
    private var geo_inp_dim: IntArray = intArrayOf();
    private var geo_out_dim: IntArray = intArrayOf();

    private lateinit var interpreter_classify: Interpreter;
    private lateinit var interpreter_imc: ImageClassifier;
    private var classify_inp_dim: IntArray = intArrayOf();
    private var classify_out_dim: IntArray = intArrayOf();
    private lateinit var classify_input_buffer: FloatBuffer;
    private lateinit var classify_input_array: FloatArray;

    private lateinit var classify_output_buffer: FloatBuffer;
    private lateinit var classify_output_array: FloatArray;

    private lateinit var labels_geo: Array<String>
    private lateinit var labels_classify: Array<String>

    private lateinit var geo_input: FloatBuffer
    private lateinit var geo_output: FloatArray

    private var geo_prob_map = mutableMapOf<String, Float>();
    var short_long_map = mutableMapOf<String, String>()


    init {;
        loadMetaLabels()
        loadLabels()
        loadShortLongMap()
        setupSpectrogramGenerator()
        setupInterpreter()
        setupMetaInterpreter()
        this.model_ready = true;
    }


    private fun setupSpectrogramGenerator() {
        interpreter_spectrogram =
            load_interpreter(Settings.spectrogram_generate_model_file["file"].toString())

        var spec_inp = this.interpreter_spectrogram.getInputTensor(/* inputIndex = */ 0)
        specInputBuffer = TensorBuffer.createFixedSize(spec_inp.shape(), spec_inp.dataType())


        var spec_out = this.interpreter_spectrogram.getOutputTensor(/* inputIndex = */ 0)
        specOutputBuffer = TensorBuffer.createFixedSize(spec_out.shape(), spec_out.dataType())
    }

    fun generateSpectrogram(input_array: FloatArray): ShortArray {
        specInputBuffer.buffer.clear()
        specOutputBuffer.buffer.clear()


        val array_to_load = input_array.copyOf(512).map { it / norm_factor }.toFloatArray()
        specInputBuffer.loadArray(array_to_load)
//        Log.d(TAG, "Loaded Array: ${array_to_load.min()} ${array_to_load.max()}")

        interpreter_spectrogram.run(specInputBuffer.buffer, specOutputBuffer.buffer);

        val output = specOutputBuffer.intArray.map { it.toShort() }.toShortArray()

        return output
    }

    private fun setupMetaInterpreter() {
        interpreter_geo = load_interpreter(Settings.spectrogram_geo_model_file["file"].toString())

        geo_inp_dim = this.interpreter_geo.getInputTensor(0).shape()
        Log.i(TAG, "TFLite model input shape: ${geo_inp_dim.contentToString()}")
        geo_input = FloatBuffer.allocate(3);


        geo_out_dim = this.interpreter_geo.getOutputTensor(0).shape()
        Log.i(TAG, "TFLite model input shape: ${geo_out_dim.contentToString()}")
        geo_output = FloatArray(geo_out_dim[0]) { 1f };
    }

    private fun setupInterpreter() {


        interpreter_classify =
            load_interpreter(Settings.spectrogram_classify_model_file["file"].toString())

        interpreter_imc = ImageClassifier.createFromFileAndOptions(
            File(get_file_path(Settings.spectrogram_classify_model_file["file"].toString())),
            ImageClassifierOptions.builder().setBaseOptions(
                BaseOptions.builder().setNumThreads(2).build()
            ).setMaxResults(10).build()
        )

        classify_inp_dim = this.interpreter_classify.getInputTensor(0).shape();
        Log.i(TAG, "TFLite model input shape: ${classify_inp_dim.contentToString()}")

        classify_out_dim = this.interpreter_classify.getOutputTensor(0).shape();
        Log.i(TAG, "TFLite model output shape: ${classify_out_dim.contentToString()}")


        val n_floats = classify_inp_dim.reduce(Int::times)
        classify_input_buffer = FloatBuffer.allocate(n_floats)
        classify_input_array = FloatArray(n_floats);

        classify_output_buffer = FloatBuffer.allocate(classify_out_dim[1]);
        classify_output_array = FloatArray(classify_out_dim[1]) { 0f };

    }


    private fun loadMetaLabels() {
        labels_geo = load_labels(Settings.spectrogram_geo_label_file["file"].toString())
    }

    private fun loadLabels() {
        labels_classify = load_labels(Settings.spectrogram_classify_label_file["file"].toString())
    }

    private fun loadShortLongMap() {
        var labx = load_labels(Settings.short_to_long_map["file"].toString());
        for (x in labx) {
            var ff = x.split(",")
            short_long_map[ff[0]] = ff[1]
        }
    }

    private fun geo_prob_clamp(i: Float): Float {
        return when {
            i > 0.01 -> 1.0f
            i > 0.008 -> 0.8f
            i > 0.001 -> 0.5f
            else -> 0.0f
        }
    }

    fun runMetaInterpreter(location: Location) {

        var localDate: LocalDate = LocalDate.now()
        val date: Date = Date.from(localDate.atStartOfDay(ZoneId.systemDefault()).toInstant())
        val calendar = Calendar.getInstance()
        calendar.time = date


        val lon = this.asFloatArray(location.longitude.toFloat());
        val lat = this.asFloatArray(location.latitude.toFloat());
        val week_of_year = this.asFloatArray(calendar.get(WEEK_OF_YEAR).coerceAtMost(52).toFloat())

        val input_array = arrayOf(lon, week_of_year, lat)


        var outputTensor = interpreter_geo.getOutputTensor(0);
        var fixedTensor =
            TensorBuffer.createFixedSize(outputTensor.shape(), outputTensor.dataType())

        val outputs: MutableMap<Int, Any> = mutableMapOf()
        outputs[0] = fixedTensor.buffer;
        interpreter_geo.runForMultipleInputsOutputs(input_array, outputs)

        fixedTensor.floatArray.zip(labels_geo).forEach { (score, label) ->
            geo_prob_map[label] = geo_prob_clamp(score)
        }

        val cloned = fixedTensor.floatArray.clone()
        val sorted_indices = cloned.withIndex().sortedByDescending { it.value }.map { it.index }

        local_species_spectrogram = "Most likely:\n"
        for (i in 0..10) {
            var c_species = labels_geo[sorted_indices[i]]
            local_species_spectrogram += "          " + short_long_map.get(c_species).toString();
            local_species_spectrogram += "\n";
        }

        Log.d(TAG, "Likely: " + local_species_spectrogram)
    }

    fun is_model_ready(): Boolean {
        return this.model_ready;
    }


    private fun get_file_path(file: String): String {
        return this.context.getDir(
            "", Context.MODE_PRIVATE
        ).absolutePath + "/" + file.toString();
    }

    private fun load_labels(file: String): Array<String> {
        val label_file_path = this.get_file_path(file);
        val labelFile = File(label_file_path)
        return labelFile.readLines().toTypedArray()
    }


    private fun load_interpreter(file: String): Interpreter {

        val modelFilePath = get_file_path(file);

        val i = Log.i(TAG, "Trying to create TFLite buffer from $modelFilePath")
        val modelFile = File(modelFilePath)
        val tfliteBuffer: ByteBuffer =
            FileChannel.open(modelFile.toPath(), StandardOpenOption.READ).use { channel ->
                channel.map(FileChannel.MapMode.READ_ONLY, 0, channel.size())
            }
        Log.i(TAG, "Done creating TFLite buffer from $modelFilePath")

        return Interpreter(tfliteBuffer, Interpreter.Options())

    }


    private fun asFloatArray(d: Float): FloatArray {
        return floatArrayOf(d)
    }

    fun run_interpreter(createBitmap: Bitmap) {
        val fromBitmap = TensorImage.fromBitmap(createBitmap);
        val out: TensorImage = image_resize_op.process(fromBitmap);


        this.context.runOnUiThread {
            this.context.set_image_view(out.bitmap)
        }


        var classify: MutableList<Classifications> = interpreter_imc.classify(out)

        var redone_scores = FloatArray(classify_output_array.size) { 0f };
        var cats: MutableList<Category>? = classify.get(0).categories;


        score_string = "";
        cats?.forEach {
            val c_name = it.label
            val index = it.index
            val score = it.score
            var cur_score = 0f;
            if (geo_prob_map.containsKey(c_name)) {
                cur_score = score * geo_prob_map[c_name]!!
            }
            if (c_name == "bird1") {
                cur_score = score
                c_threshold = Settings.threshold_bird_singing / 100f;
            }
            else
            {
                c_threshold =  Settings.threshold_spectrogram / 100f
            }
            redone_scores[index] = cur_score




            if (cur_score > c_threshold) {
                var use_name = "";
                if (short_long_map.containsKey(c_name)) {
                    use_name = short_long_map[c_name].toString()
                } else {
                    use_name = c_name
                };

                score_string += use_name + "_" + use_name;
                score_string += ","
                score_string += cur_score.toString();
                score_string += ';'
            }


        }



//
//
//        val outList = redone_scores.withIndex().sortedByDescending { it -> it.value };
//        for (i in 0 until 6) {
//            var cs: IndexedValue<Float> = outList[i]
//            Log.d(TAG, i.toString() + " " + labels_classify[cs.index] + " " + cs.value.toString())
//        }
//
//        Log.d(TAG, classify_output_buffer.toString())
//        Log.d(TAG, classify_output_array.sum().toString())
    }


}