VesiscanClinicalAndroid/app/src/main/java/com/medithings/vesiscan/managers/NirsCalcEngine.kt

package com.medithings.vesiscan.managers

import com.medithings.vesiscan.models.NirsConstants
import com.medithings.vesiscan.models.NirsMagResponse
import com.medithings.vesiscan.models.NirsOxyResult
import kotlin.math.abs
import kotlin.math.ln
import kotlin.math.max
import kotlin.math.pow

class NirsCalcEngine {

    // Gain Calibration
    private var gainValues: IntArray = intArrayOf()
    val hasGain: Boolean get() = gainValues.size >= 24

    fun setGainFromMag(response: NirsMagResponse) {
        gainValues = response.gainValues
    }

    fun clearGain() {
        gainValues = intArrayOf()
    }

    // AGC DAC -> Software Amplification
    private fun agcGain(dacValue: Int): Double {
        return 10.0.pow(((dacValue.toDouble() * 3.3 / 4096.0 * -80.0) + 88.0) / 20.0)
    }

    // DeGain Calculation
    fun degain(rawLed: IntArray): DoubleArray {
        if (rawLed.size < 24 || !hasGain) return DoubleArray(24) { 0.0 }

        val result = DoubleArray(24)
        for (i in 0 until 24) {
            val raw = rawLed[i].toDouble()
            val sag = gainValues[i]
            val pkg = i / 6
            val wl = i % 6

            val softAmp = agcGain(sag)
            val hAmp = NirsConstants.hardwareAmp[pkg]
            val spResponse = NirsConstants.spectralResponse[wl]
            val dpf = NirsConstants.dpf[wl]

            val denominator = softAmp * hAmp
            if (denominator > 0) {
                var dg = raw / denominator
                dg /= spResponse
                dg *= dpf
                result[i] = dg
            }
        }
        return result
    }

    // Butterworth IIR Filter
    private var filterState: Array<DoubleArray> = emptyArray()
    private var filterInitialized = false
    var isFilterEnabled = true

    fun resetFilter() {
        filterState = Array(24) { doubleArrayOf(0.0, 0.0, 0.0, 0.0) }
        filterInitialized = false
    }

    private fun applyFilter(x: Double, state: DoubleArray): Double {
        val b = NirsConstants.filterB
        val a = NirsConstants.filterA

        val y = b[0] * x + b[1] * state[0] + b[2] * state[1] - a[1] * state[2] - a[2] * state[3]

        state[1] = state[0]
        state[0] = x
        state[3] = state[2]
        state[2] = y

        return y
    }

    fun filterDeGain(degainValues: DoubleArray): DoubleArray {
        if (degainValues.size < 24) return degainValues

        if (filterState.isEmpty()) resetFilter()

        if (!filterInitialized) {
            for (i in 0 until 24) {
                val v = degainValues[i]
                filterState[i] = doubleArrayOf(v, v, v, v)
            }
            filterInitialized = true
        }

        val filtered = DoubleArray(24)
        for (i in 0 until 24) {
            filtered[i] = applyFilter(degainValues[i], filterState[i])
        }
        return filtered
    }

    // Oxygenation Calculation
    fun calcOxygenation(degainValues: DoubleArray): NirsOxyResult {
        if (degainValues.size < 24) return NirsOxyResult.zero

        val c = NirsConstants
        val muA = DoubleArray(c.numWavelengths)

        for (i in 0 until c.numWavelengths) {
            val i2cm = degainValues[i]
            val i5cm = degainValues[18 + i]

            val i2norm = i2cm / (c.ledPower[0] / c.referencePower)
            val i5norm = i5cm / (c.ledPower[3] / c.referencePower)

            val mix = c.srsMixNear * i2norm + c.srsMixFar * i5norm

            if (mix > 0) {
                val slope = -ln(mix) / c.srsPathLength
                muA[i] = -slope
            }
        }

        val solution = solveNormalEq6x3(muA)

        var hbO2 = max(0.0, solution.first) * c.oxyScale
        var hb = max(0.0, solution.second) * c.oxyScale

        val tHb = hbO2 + hb
        val stO2 = if (tHb > 0) (hbO2 / tHb) * 100.0 else 0.0

        return NirsOxyResult(hbO2 = hbO2, hb = hb, tHb = tHb, stO2 = stO2)
    }

    // Least Squares Solver (6x3)
    private fun solveNormalEq6x3(muA: DoubleArray): Triple<Double, Double, Double> {
        val eps = arrayOf(
            NirsConstants.extinctionHbO2,
            NirsConstants.extinctionHb,
            NirsConstants.extinctionH2O
        )

        val n = Array(3) { DoubleArray(3) }
        for (i in 0 until 3) {
            for (j in 0 until 3) {
                var sum = 0.0
                for (k in 0 until 6) sum += eps[i][k] * eps[j][k]
                n[i][j] = sum
            }
        }

        val b = DoubleArray(3)
        for (i in 0 until 3) {
            var sum = 0.0
            for (k in 0 until 6) sum += eps[i][k] * muA[k]
            b[i] = sum
        }

        val det = n[0][0] * (n[1][1] * n[2][2] - n[1][2] * n[2][1]) -
                n[0][1] * (n[1][0] * n[2][2] - n[1][2] * n[2][0]) +
                n[0][2] * (n[1][0] * n[2][1] - n[1][1] * n[2][0])

        if (abs(det) < 1e-12) return Triple(0.0, 0.0, 0.0)

        val inv = Array(3) { DoubleArray(3) }
        inv[0][0] = (n[1][1] * n[2][2] - n[1][2] * n[2][1]) / det
        inv[0][1] = -(n[0][1] * n[2][2] - n[0][2] * n[2][1]) / det
        inv[0][2] = (n[0][1] * n[1][2] - n[0][2] * n[1][1]) / det
        inv[1][0] = -(n[1][0] * n[2][2] - n[1][2] * n[2][0]) / det
        inv[1][1] = (n[0][0] * n[2][2] - n[0][2] * n[2][0]) / det
        inv[1][2] = -(n[0][0] * n[1][2] - n[0][2] * n[1][0]) / det
        inv[2][0] = (n[1][0] * n[2][1] - n[1][1] * n[2][0]) / det
        inv[2][1] = -(n[0][0] * n[2][1] - n[0][1] * n[2][0]) / det
        inv[2][2] = (n[0][0] * n[1][1] - n[0][1] * n[1][0]) / det

        val c = DoubleArray(3) { i ->
            inv[i][0] * b[0] + inv[i][1] * b[1] + inv[i][2] * b[2]
        }

        return Triple(c[0], c[1], c[2])
    }

    // Full Processing Pipeline
    data class ProcessResult(
        val degain: DoubleArray,
        val degainFiltered: DoubleArray,
        val oxy: NirsOxyResult,
        val oxyFiltered: NirsOxyResult
    )

    fun process(rawLed: IntArray): ProcessResult {
        val dg = degain(rawLed)
        val oxyRaw = calcOxygenation(dg)

        val dgf: DoubleArray
        val oxyF: NirsOxyResult
        if (isFilterEnabled) {
            dgf = filterDeGain(dg)
            oxyF = calcOxygenation(dgf)
        } else {
            dgf = dg
            oxyF = oxyRaw
        }

        return ProcessResult(dg, dgf, oxyRaw, oxyF)
    }
}