import Algorithms
import Foundation
import PFFFT

struct Parameters {
    /// The number of beats to simulate.
    let numBeats: Int = 256

    /// 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 = 60.0

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

    /// 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: Int = 0

    /// 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
}

class PfftSetupCache: @unchecked Sendable {
    private var cache: [Int: OpaquePointer?] = [:]

    private let queue = DispatchQueue(label: String(describing: PfftSetupCache.self), attributes: .concurrent)

    func get(for nrr: Int) -> OpaquePointer? {
        var setup: OpaquePointer??
        queue.sync {
            setup = cache[nrr]
        }
        if setup == nil {
            queue.sync(flags: .barrier) {
                setup = cache[nrr]
                if setup == nil {
                    setup = pffftd_new_setup(Int32(nrr), PFFFT_REAL)
                    cache[nrr] = setup
                }
            }
        }
        return setup!
    }

    deinit {
        for (_, setup) in cache {
            if setup != nil {
                pffftd_destroy_setup(setup)
            }
        }
    }

    static let shared = PfftSetupCache()
}

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

func rrprocess(params: Parameters, nrr: Int) -> [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 rrmean = 60.0 / params.hrMean
    let rrstd = 60.0 * params.hrStd / (params.hrMean * params.hrMean)

    let sf = Double(params.sfInternal)
    let df = sf / Double(nrr)

    let fft = PfftSetupCache.shared.get(for: nrr)

    var rr = withUnsafeTemporaryAllocation(byteCount: (nrr + 2) * MemoryLayout<Double>.stride, alignment: 32) {
        let swc = $0.bindMemory(to: Double.self)

        for i in 0 ..< nrr / 2 + 1 {
            let w = df * Double(i) * 2.0 * .pi
            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 * Double.random(in: 0.0 ..< 1.0)

            swc[i * 2] = sw * cos(ph)
            swc[i * 2 + 1] = sw * sin(ph)
        }

        // pack Nyquist frequency real to imaginary of DC
        swc[1] = swc[nrr]

        return withUnsafeTemporaryAllocation(byteCount: nrr * MemoryLayout<Double>.stride, alignment: 32) {
            let outptr = $0.bindMemory(to: Double.self)

            pffftd_transform_ordered(fft, swc.baseAddress, outptr.baseAddress, nil, PFFFT_BACKWARD)
            return outptr.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
}


// func compute(params: Parameters) {
//     // adjust extrema parameters for mean heart rate
//     let hrFact = sqrt(params.hrMean / 60)
//     let hrFact2 = sqrt(hrFact)

//     let bi = params.b.map { $0 * hrFact }

//     /// XXX: Discrepancy here between Java/C and Matlab, the former uses 1.0 for ti[4] adjustment.
//     let ti = zip([hrFact2, hrFact, 1, hrFact, hrFact2], params.theta).map(*)

//     let ai = params.a

//     let x0 = SIMD3<Double>(1.0, 0.0, 0.04) // XXX: Convert to init from vector3d
//     let rseed = params.seed

//     // calculate time scales
//     let h = 1.0 / Double(params.sfInternal)
//     let tstep = 1.0 / Double(params.sfEcg)

//     // calculate length of the RR time series
//     let rrmean = (60.0 / params.hrMean)

//     let numRr = Int(pow(2.0, ceil(log2(Double(params.numBeats * params.sfInternal) * rrmean))))

//     var rr = [Double](repeating: 0.0, count: numRr)

//     // TODO: check sfInternal is integer multple of sfEcg

//     // define frequency parameters for rr process
// }
