package com.birdsounds.identify


import android.content.Context
import android.location.Location
import android.os.SystemClock
import android.preference.PreferenceManager
import android.util.Log
import org.tensorflow.lite.Interpreter

import java.io.BufferedReader
import java.io.File
import java.io.IOException
import java.io.InputStreamReader
import java.nio.ByteBuffer
import java.nio.FloatBuffer
import java.nio.channels.FileChannel
import java.nio.file.StandardOpenOption
import java.time.LocalDate
import java.util.Locale
import java.util.TimerTask
import java.util.Timer
import kotlin.concurrent.scheduleAtFixedRate
import kotlin.math.ceil
import kotlin.math.cos
import uk.me.berndporr.iirj.Butterworth
import kotlin.math.sin



var local_species = ""
class SoundClassifier(
    context: Context,
    private val options: Options = Options()
) {

    internal var mContext: Context



    val TAG = "Sound Classifier"

    init {
        this.mContext = context.applicationContext
    }

    class Options(
        val labelsBase: String = "labels",
        /** Path of the converted .tflite file, relative to the assets/ directory.  */
        val assetFile: String = "assets.txt",
        /** Path of the converted .tflite file, relative to the assets/ directory.  */
        /** The required audio sample rate in Hz.  */
        val sampleRate: Int = 48000,
        /** Multiplier for audio samples  */
        val warmupRuns: Int = 3,
        /** Probability value above which a class in the meta model is labeled as active (i.e., detected) the display. (default 0.01) */
        var metaProbabilityThreshold1: Float = 0.01f,  //min must be > 0
        var metaProbabilityThreshold2: Float = 0.008f,  //min must be > 0
        var metaProbabilityThreshold3: Float = 0.001f,  //min must be > 0
        /** Probability value above which a class is shown as image. (default 0.5) */
        var displayImageThreshold: Float = 0.65f,  //min must be > 0
    )


    /** Names of the model's output classes.  */
    lateinit var labelList: List<String>

    /** Names of the model's output classes.  */
    lateinit var assetList: List<String>

    /** How many milliseconds between consecutive model inference calls.  */
    private var inferenceInterval = 800L

    /** The TFLite interpreter instance.  */
    private lateinit var interpreter: Interpreter
    private lateinit var meta_interpreter: Interpreter

    /** Audio length (in # of PCM samples) required by the TFLite model.  */
    private var modelInputLength = 0

    /** input Length of the meta model  */
    private var metaModelInputLength = 0

    /** Number of output classes of the TFLite model.  */
    private var modelNumClasses = 0
    private var metaModelNumClasses = 0
    private var settings = Settings;

    /** Used to hold the real-time probabilities predicted by the model for the output classes.  */
    private lateinit var predictionProbs: FloatArray
    private lateinit var metaPredictionProbs: FloatArray

    /** Latest prediction latency in milliseconds.  */
    private var latestPredictionLatencyMs = 0f

    private var recognitionTask: TimerTask? = null

    /** Buffer that holds audio PCM sample that are fed to the TFLite model for inference.  */
    private lateinit var inputBuffer: FloatBuffer
    private lateinit var metaInputBuffer: FloatBuffer
    private var model_ready = false;
    init {;
        loadLabels(context)
        setupInterpreter(context)
        setupMetaInterpreter(context)
        warmUpModel()
        this.model_ready = true;
    }

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


    /** Retrieve labels from "labels.txt" file */
    private fun loadLabels(context: Context) {
        var filename = "labels_en.txt"

        try {
            val reader =
                BufferedReader(InputStreamReader(context.assets.open(filename)))
            val wordList = mutableListOf<String>()
            reader.useLines { lines ->
                lines.forEach {
                    wordList.add(it)
                }
            }
            labelList = wordList.map { it.toTitleCase() }
            Log.i(TAG, "Label list entries: ${labelList.size}")

        } catch (e: IOException) {
            Log.e(TAG, "Failed to read labels ${filename}: ${e.message}")
        }
    }


    private fun setupInterpreter(context: Context) {
        try {
            val modelFilePath =
                context.getDir(
                "",
                Context.MODE_PRIVATE
            ).absolutePath + "/" + settings.waveform_classify_model_file["file"].toString();

            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")

            interpreter = Interpreter(tfliteBuffer, Interpreter.Options())



        } catch (e: IOException) {
            Log.e(TAG, "Failed to load TFLite model - ${e.message}")
            return
        }
        // Inspect input and output specs.
        val inputShape = interpreter.getInputTensor(0).shape()
        Log.i(TAG, "TFLite model input shape: ${inputShape.contentToString()}")
        modelInputLength = inputShape[1]

        val outputShape = interpreter.getOutputTensor(0).shape()
        Log.i(TAG, "TFLite output shape: ${outputShape.contentToString()}")
        modelNumClasses = outputShape[1]
        if (modelNumClasses != labelList.size) {
            Log.e(
                TAG,
                "Mismatch between metadata number of classes (${labelList.size})" +
                        " and model output length ($modelNumClasses)"
            )
        }
        // Fill the array with NaNs initially.
        predictionProbs = FloatArray(modelNumClasses) { Float.NaN }

        inputBuffer = FloatBuffer.allocate(modelInputLength)

    }

    private fun setupMetaInterpreter(context: Context) {

        try {
            val metaModelFilePath = context.getDir(
                "",
                Context.MODE_PRIVATE
            ).absolutePath + "/" + settings.waveform_meta_model_file["file"].toString();
            Log.i(TAG, "Trying to create TFLite buffer from $metaModelFilePath")
            val metaModelFile = File(metaModelFilePath)
            val tfliteBuffer: ByteBuffer =
                FileChannel.open(metaModelFile.toPath(), StandardOpenOption.READ).use { channel ->
                    channel.map(FileChannel.MapMode.READ_ONLY, 0, channel.size())
                }
            Log.i(TAG, "Done creating TFLite buffer from $metaModelFilePath")

            meta_interpreter = Interpreter(tfliteBuffer, Interpreter.Options())
        } catch (e: IOException) {
            Log.e(TAG, "Failed to load TFLite model - ${e.message}")
            return
        }
        // Inspect input and output specs.
        val metaInputShape = meta_interpreter.getInputTensor(0).shape()
        Log.i(TAG, "TFLite meta model input shape: ${metaInputShape.contentToString()}")
        metaModelInputLength = metaInputShape[1]

        val metaOutputShape = meta_interpreter.getOutputTensor(0).shape()
        Log.i(TAG, "TFLite meta model output shape: ${metaOutputShape.contentToString()}")
        metaModelNumClasses = metaOutputShape[1]
        if (metaModelNumClasses != labelList.size) {
            Log.e(
                TAG,
                "Mismatch between metadata number of classes (${labelList.size})" +
                        " and meta model output length ($metaModelNumClasses)"
            )
        }
        // Fill the array with 1 initially.
        metaPredictionProbs = FloatArray(metaModelNumClasses) { 1f }

        metaInputBuffer = FloatBuffer.allocate(metaModelInputLength)

    }

    fun runMetaInterpreter(location: Location) {
        val dayOfYear = LocalDate.now().dayOfYear
        val week = ceil(dayOfYear * 48.0 / 366.0) //model year has 48 weeks
        lat = location.latitude.toFloat()
        lon = location.longitude.toFloat()

        val weekMeta = cos(Math.toRadians(week * 7.5)) + 1.0

        metaInputBuffer.put(0, lat)
        metaInputBuffer.put(1, lon)
        metaInputBuffer.put(2, weekMeta.toFloat())
        metaInputBuffer.rewind() // Reset position to beginning of buffer

        val metaOutputBuffer = FloatBuffer.allocate(metaModelNumClasses)
        metaOutputBuffer.rewind()
        meta_interpreter.run(metaInputBuffer, metaOutputBuffer)
        metaOutputBuffer.rewind()
        metaOutputBuffer.get(metaPredictionProbs) // Copy data to metaPredictionProbs.

        var cloned = metaPredictionProbs.clone()

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

        local_species = "Most likely:\n"
        for (i in 0..5) {
            local_species += "          " +labelList[sorted_indices[i]].split("_")[1]
            local_species += "\n";
        }



        for (i in metaPredictionProbs.indices) {
            metaPredictionProbs[i] =
                if (metaPredictionProbs[i] >= options.metaProbabilityThreshold1) {
                    1f
                } else if (metaPredictionProbs[i] >= options.metaProbabilityThreshold2) {
                    0.8f
                } else if (metaPredictionProbs[i] >= options.metaProbabilityThreshold3) {
                    0.5f
                } else {
                    0f
                }
        }
    }

    private fun warmUpModel() {
        generateDummyAudioInput(inputBuffer)
        for (n in 0 until options.warmupRuns) {

            // Create input and output buffers.
            val outputBuffer = FloatBuffer.allocate(modelNumClasses)
            inputBuffer.rewind()
            outputBuffer.rewind()
            interpreter.run(inputBuffer, outputBuffer)

        }
    }

    private fun generateDummyAudioInput(inputBuffer: FloatBuffer) {
        val twoPiTimesFreq = 2 * Math.PI.toFloat() * 1000f
        for (i in 0 until modelInputLength) {
            val x = i.toFloat() / (modelInputLength - 1)
            inputBuffer.put(i, sin(twoPiTimesFreq * x.toDouble()).toFloat())
        }
    }

    private fun String.toTitleCase() =
        splitToSequence("_")
            .map { it.replaceFirstChar { if (it.isLowerCase()) it.titlecase(Locale.ROOT) else it.toString() } }
            .joinToString("_")
            .trim()

    companion object {
        private const val TAG = "SoundClassifier"
        var lat: Float = 0.0f
        var lon: Float = 0.0f

        /** Number of nanoseconds in a millisecond  */
        private const val NANOS_IN_MILLIS = 1_000_000.toDouble()
    }

    fun executeScoring(
        short_array: ShortArray
    ): List<IndexedValue<Float>> {
        val highPass = 0
        val butterworth = Butterworth()
        butterworth.highPass(6, 48000.0, highPass.toDouble())

        val outputBuffer = FloatBuffer.allocate(modelNumClasses)

        for (i in 0 until modelInputLength) {
            val s = short_array[i]

            if (highPass == 0) inputBuffer.put(i, s.toFloat())
            else inputBuffer.put(i, butterworth.filter(s.toDouble()).toFloat())
        }

        inputBuffer.rewind()
        outputBuffer.rewind()
        interpreter.run(inputBuffer, outputBuffer)
        outputBuffer.rewind()
        outputBuffer.get(predictionProbs) // Copy data to predictionProbs.

        val probList = mutableListOf<Float>()
        for (i in predictionProbs.indices) {
            probList.add(metaPredictionProbs[i] / (1 + kotlin.math.exp(-predictionProbs[i])))  //apply sigmoid
        }
        val outList = probList.withIndex().sortedByDescending { it -> it.value }
//        Log.w(TAG, outList.toString())
//        Log.i(TAG,labelList[outList[0].index].toString())
        return outList



    }

    private fun startRecognition() {
        if (modelInputLength <= 0 || modelNumClasses <= 0) {
            Log.e(TAG, "Switches: Cannot start recognition because model is unavailable.")
            return
        }
        val sharedPref = PreferenceManager.getDefaultSharedPreferences(mContext)
//        val highPass = sharedPref.getInt("high_pass", 0)
        val highPass = 300;
        val butterworth = Butterworth()
        butterworth.highPass(6, 48000.0, highPass.toDouble())

        val circularBuffer = ShortArray(modelInputLength)

        var j = 0 // Indices for the circular buffer next write

        Log.w(TAG, "recognitionPeriod:" + inferenceInterval)
        recognitionTask = Timer().scheduleAtFixedRate(inferenceInterval, inferenceInterval) task@{
            val outputBuffer = FloatBuffer.allocate(modelNumClasses)
            val recordingBuffer = ShortArray(modelInputLength)

            // Load new audio samples
//            val sampleCounts = loadAudio(recordingBuffer)
            val sampleCounts = 0
            if (sampleCounts == 0) {
                return@task
            }


            // Copy new data into the circular buffer
            for (i in 0 until sampleCounts) {
                circularBuffer[j] = recordingBuffer[i]
                j = (j + 1) % circularBuffer.size
            }

            // Feed data to the input buffer.
            var samplesAreAllZero = true
            for (i in 0 until modelInputLength) {
                val s = circularBuffer[(i + j) % modelInputLength]
                if (samplesAreAllZero && s.toInt() != 0) {
                    samplesAreAllZero = false
                }
                if (highPass == 0) inputBuffer.put(i, s.toFloat())
                else inputBuffer.put(i, butterworth.filter(s.toDouble()).toFloat())
            }

            if (samplesAreAllZero) {
                Log.w(TAG, mContext.resources.getString(R.string.samples_zero))


                return@task
            }

            val t0 = SystemClock.elapsedRealtimeNanos()
            inputBuffer.rewind()
            outputBuffer.rewind()
            interpreter.run(inputBuffer, outputBuffer)
            outputBuffer.rewind()
            outputBuffer.get(predictionProbs) // Copy data to predictionProbs.

            val probList = mutableListOf<Float>()
//            if (mBinding.checkIgnoreMeta.isChecked){
//                for (value in predictionProbs) {
//                    probList.add(1 / (1 + kotlin.math.exp(-value)))  //apply sigmoid
//                }
//            } else {
            for (i in predictionProbs.indices) {
                probList.add(metaPredictionProbs[i] / (1 + kotlin.math.exp(-predictionProbs[i])))  //apply sigmoid
            }
//            }

//            if (mBinding.progressHorizontal.isIndeterminate){  //if start/stop button set to "running"
//                probList.withIndex().also {
//                    val max = it.maxByOrNull { entry -> entry.value }
//                    updateTextView(max, mBinding.text1)
//                    updateImage(max)
            //after finding the maximum probability and its corresponding label (max), we filter out that entry from the list of entries before finding the second highest probability (secondMax)
//                    val secondMax = it.filterNot { entry -> entry == max }.maxByOrNull { entry -> entry.value }
//                    updateTextView(secondMax,mBinding.text2)
//                }
//            }

            latestPredictionLatencyMs =
                ((SystemClock.elapsedRealtimeNanos() - t0) / 1e6).toFloat()
        }
    }
}
