import Algorithms
import ComplexModule
import Foundation
import OdeInt
import PFFFT
import PFFFTLib
import RealModule

public struct Parameters {
    /// The ratio of power between low and high frequencies.
    let lfhfRatio: Double = 0.5

    /// The ECT amplitude in mV.
    let amplitude: Double = 1.4

    /// RNG seed value.
    let seed: UInt64 = 111

    /// Amplitude of the noise.
    let aNoise: Double = 0.0

    /// The angle of each attractor (P, Q, R, S, T) around the limit cycle, in radians.
    let theta: [Double] = [-60, -15, 0, 15, 90].map { $0 * .pi / 180 }

    /// Widths of the attractors (P, Q, R, S, T).
    let a: [Double] = [1.2, -5, 30, -7.5, 0.75]

    /// The position of attractors (P, Q, R, S, T) above or below the z=0 plane.
    let b: [Double] = [0.25, 0.1, 0.1, 0.1, 0.4]

    /// Mayer wave frequency in Hz.
    let flo = 0.1

    /// flo standard deviation.
    let flostd = 0.01

    /// Respiratory rate frequency in Hz.
    let fhi = 0.25

    /// fhi standard deviation.
    let fhistd = 0.01
}

public struct TimeParameters {
    /// The number of beats to simulate.
    let numBeats: Int = 8

    /// The ECG sampling frequency in Hz.
    let sfEcg: Int = 256

    /// The internal sampling frequency in Hz.
    var sfInternal: Int = 512

    /// The mean heart rate in beats per minute.
    let hrMean: Double = 75.0

    /// The standard deviation of the heart rate.
    let hrStd: Double = 1.0
}

func stdev(_ data: [Double]) -> Double {
    let n = Double(data.count)
    let mean = data.reduce(0.0, +) / n
    return sqrt(data.lazy.map { ($0 - mean) * ($0 - mean) }.reduce(0.0, +) / (n - 1))
}

public struct RRProcess: ~Copyable {
    let nrr: Int
    let sfInternal: Int

    let fft: FFT<Double>
    let spectrum: Buffer<Complex<Double>>
    let signal: Buffer<Double>

    // mean and standard deviation of RR intervals
    let rrMean: Double
    let rrStd: Double

    public init(params: TimeParameters) {
        sfInternal = params.sfInternal
        rrMean = 60.0 / params.hrMean
        rrStd = 60.0 * params.hrStd / (params.hrMean * params.hrMean)

        nrr = Int(pow(2.0, ceil(log2(Double(params.numBeats * sfInternal) * rrMean))))

        fft = try! FFT<Double>(n: nrr)
        spectrum = fft.makeSpectrumBuffer(extra: 1)
        signal = fft.makeSignalBuffer()
    }

    public func generate(params: Parameters) -> [Double] {
        let w1 = 2.0 * .pi * params.flo
        let w2 = 2.0 * .pi * params.fhi
        let c1 = 2.0 * .pi * params.flostd
        let c2 = 2.0 * .pi * params.fhistd

        let sig2 = 1.0
        let sig1 = params.lfhfRatio

        let sf = Double(sfInternal)

        let dw = (sf / Double(nrr)) * 2.0 * .pi

        var rng = Xoshiro256Plus(seed: params.seed)

        spectrum.mapInPlaceSwapLast { i in
            let w = dw * Double(i)

            let dw1 = w - w1
            let dw2 = w - w2
            let hw = sig1 * exp(-dw1 * dw1 / (2.0 * c1 * c1)) / sqrt(2.0 * .pi * c1 * c1)
                + sig2 * exp(-dw2 * dw2 / (2.0 * c2 * c2)) / sqrt(2.0 * .pi * c2 * c2)

            let sw = (sf / 2.0) * sqrt(hw)
            let ph = 2.0 * .pi * rng.nextDouble()

            return Complex(length: sw, phase: ph)
        }

        fft.inverse(spectrum: spectrum, signal: signal)

        var rr = signal.map { $0 * 1.0 / Double(nrr) }

        let xstd = stdev(rr)
        let ratio = rrStd / xstd

        for i in 0 ..< nrr {
            rr[i] = rr[i] * ratio + rrMean
        }
        return rr
    }
}

public struct Generator: ~Copyable {
    let rrp: RRProcess

    let hrFact: Double
    let hrFactSqrt: Double

    let dt: Double


    let timeParams: TimeParameters

    public init(params: TimeParameters) {
        rrp = RRProcess(params: params)

        // stretching factors for intervals based on Bazett's formula
        hrFact = sqrt(params.hrMean / 60)
        hrFactSqrt = sqrt(hrFact)

        // init time scales
        dt = 1.0 / Double(params.sfInternal)

        timeParams = params
    }

    public func compute(params: Parameters) -> [Double] {
        let ai = params.a
        let bi = params.b.map { $0 * hrFact }
        // adjust extrema parameters for mean heart rate
        let ti = zip([hrFactSqrt, hrFact, 1, hrFact, hrFactSqrt], params.theta).map(*)

        let sfInternal = timeParams.sfInternal

        let rr = rrp.generate(params: params)

        let fhi = params.fhi

        let dt = self.dt

        // generate piecewise RR
        let rrpc = [Double](unsafeUninitializedCapacity: rr.count * 2) { rrpc, count in
            var tecg = 0.0
            var i = 0
            var j = 0
            while i < rr.count {
                tecg += rr[i]
                j = Int((tecg / dt).rounded(.toNearestOrEven))
                for k in i ... j {
                    rrpc[k] = rr[i]
                }
                i = j + 1
            }
            count = i
        }
        print("rrpc: \(rrpc.count)")

        let nt = rr.count

        let x0 = SIMD3<Double>(1.0, 0.0, 0.04)

        var mip = 0
        var mt2 = 0.0

        let ts = (0 ..< nt).map { Double($0) * dt }
        print("ts: \(ts.count) \(ts.last!)")

        let result = SIMD3<Double>.integrate(over: ts, y0: x0, tol: 1e-6) { x, t in
            let ta = atan2(x[1], x[0])

            let r0 = 1.0
            let a0 = 1.0 - sqrt(x[0] * x[0] + x[1] * x[1]) / r0

            let ip = Int(floor(t * Double(sfInternal)))
            mip = max(ip, mip)
            mt2 = max(t, mt2)
            //print("ip: \(ip) mip: \(mip) mt2: \(mt2)")

            let w0 = 2 * .pi / rrpc[min(ip, rrpc.count - 1)]

            let zbase = 0.005 * sin(2 * .pi * fhi * t)

            var dxdt = SIMD3<Double>()

            dxdt[0] = a0 * x[0] - w0 * x[1]
            dxdt[1] = a0 * x[1] + w0 * x[0]

            dxdt[2] = 0.0

            for i in 0 ..< ti.count {
                let dt = remainder(ta - ti[i], 2 * .pi)
                let dt² = dt * dt

                dxdt[2] += -ai[i] * dt * exp(-0.5 * dt² / (bi[i] * bi[i]))
            }
            dxdt[2] += -1.0 * (x[2] - zbase)

            return dxdt
        }

        print("mip: \(mip)")

        // downsample to ECG sampling frequency
        let qstep = sfInternal / timeParams.sfEcg
        var zresult = stride(from: 0, to: nt, by: qstep).map { result[$0][2] }

        let (zmin, zmax) = zresult.minAndMax()!
        let zrange = zmax - zmin

        var rng = Xoshiro256Plus(seed: params.seed + 1)

        // Scale signal between -0.4 and 1.2 mV
        // add uniformly distributed measurement noise
        for i in 0 ..< zresult.count {
            zresult[i] = 1.6 * (zresult[i] - zmin) / zrange - 0.4
            zresult[i] += params.aNoise * (2.0 * rng.nextDouble() - 1.0)
        }

        // write zresult to text/CSV file
        let filename = "ecg.csv"
        try! zresult[0 ..< 800].lazy.map { String($0) }.joined(separator: "\n").write(to: URL(fileURLWithPath: filename), atomically: true, encoding: .utf8)
        return zresult
    }
}
