//
//  PiecewiseIntegrator.swift
//  ToneShaperPreview
//
//  Created by Joseph Pagliaro on 12/7/23.
//

import Accelerate 

class PiecewiseIntegrator {
    
    var userIFCurve:[CGPoint]
    var sampleCount: Int
    var rangeCount: Int
    var delta:Double
    
    var useSimpsons: Bool
    
    var partitions:[ClosedRange<Int>]
    var scale:Double
    let userIFValues:[Double]
    
    var currentPartition:ClosedRange<Int>
    var partialV:[Double]
    var lastIntegral:[Double]
    
    var currentIndex = -1
    
    init(userIFCurve:[CGPoint], sampleCount: Int, rangeCount: Int, delta:Double, useSimpsons: Bool) {
        self.userIFCurve = userIFCurve
        self.sampleCount = sampleCount
        self.rangeCount = rangeCount
        
        self.delta = delta
        
        self.useSimpsons = useSimpsons
        
        self.partitions = PiecewiseIntegrator.createPartitionRanges(sampleCount: sampleCount, rangeCount: rangeCount)
        self.scale = PiecewiseIntegrator.scale(sampleCount: sampleCount, userIFCurve: userIFCurve)
        self.userIFValues = PiecewiseIntegrator.userIFValues_by_interpolateUserIFCurve(userIFCurve: userIFCurve)
        
        self.currentPartition = partitions[0]
        self.partialV = PiecewiseIntegrator.sampleIFValuesForRange(scale: scale, sampleRange: currentPartition, userIFValues: userIFValues) // overlap
        self.lastIntegral = PiecewiseIntegrator.integrate(samples: partialV, stepSize: delta, useSimpsons: useSimpsons)
        
        print(self.partitions)
    }
    
    class func integrate(samples:[Double], stepSize:Double, useSimpsons:Bool) -> [Double] {
        
        let sampleCount = samples.count
        var step = stepSize
        
        var result = [Double](repeating: 0.0, count: sampleCount)
        
        if useSimpsons {
            vDSP_vsimpsD(samples, 1, &step, &result, 1, vDSP_Length(sampleCount))
        }
        else {
            vDSP_vtrapzD(samples, 1, &step, &result, 1, vDSP_Length(sampleCount))
        }
        
        return result
    }
    
    class func scale(sampleCount:Int, userIFCurve:[CGPoint]) -> Double {
        Double(sampleCount-1) / userIFCurve.last!.x
    }
    
    class func createPartitionRanges(sampleCount: Int, rangeCount: Int) -> [ClosedRange<Int>] {
        
        guard rangeCount > 0 else {
            return []
        }
        
        var ranges: [ClosedRange<Int>] = []
        
        var start = 0
        while start < sampleCount {
            let end = min(start + rangeCount, sampleCount) - 1
            let range = start...end
            ranges.append(range)
            start = end + 1
        }
        
        return ranges
    }
    
    class func userIFValues_by_interpolateUserIFCurve(userIFCurve:[CGPoint]) -> [Double] {
        
        let indices = userIFCurve.map { Double($0.x) }
        let values = userIFCurve.map { Double($0.y) }
        
        return vDSP.linearInterpolate(values: values, atIndices: indices)
    }
    
    class func sampleIFValuesForRange(scale:Double, sampleRange:ClosedRange<Int>, userIFValues:[Double]) -> [Double] {
            // scale sample range into user range
        let ta = Double(sampleRange.lowerBound) / scale 
        let tb = Double(sampleRange.upperBound) / scale 
        
        var valuesToInterpolate:[Double] = []
        var sampleIndices:[Double] = []
        
        func appendInterploatedValue(_ t:Double) {
            let delta = (t - t.rounded(.down))
            let index = Int(t.rounded(.down))
            if delta == 0 || (index+1 > userIFValues.count-1) { // index+1 may be out of bounds when delta = 0, or very nearly 0
                valuesToInterpolate.append(userIFValues[index])
            }
            else {
                let interpolated = userIFValues[index] * (1 - delta) + userIFValues[index+1] * delta
                valuesToInterpolate.append(interpolated)
            }
        }
        
        if ta == tb {
            appendInterploatedValue(ta) 
            sampleIndices.append(Double(sampleRange.lowerBound))
        }
        else {
            
                // start
            appendInterploatedValue(ta) 
            sampleIndices.append(Double(sampleRange.lowerBound))
            
                // middle, if any
            var lowerBound = Int(ta.rounded(.up))
            if lowerBound == Int(ta.rounded(.down)) {
                lowerBound += 1
            }
            var upperBound = Int(tb.rounded(.down))
            if upperBound == Int(tb.rounded(.up)) {
                upperBound -= 1
            }
            
            if lowerBound <= upperBound {
                valuesToInterpolate.append(contentsOf: Array(userIFValues[lowerBound...upperBound]))
                sampleIndices.append(contentsOf: (lowerBound...upperBound).map { Double($0) * scale })
            }
            
                // end
            appendInterploatedValue(tb) 
            sampleIndices.append(Double(sampleRange.upperBound))
            
        }
        
        sampleIndices = sampleIndices.map { $0 - sampleIndices[0]}
        
        return vDSP.linearInterpolate(values: valuesToInterpolate, atIndices: sampleIndices)
    }
    
    func nextIntegral() -> [Double]? {
        
        if currentIndex == self.partitions.count-1 {
            return nil
        }
        else {
            
            currentIndex += 1
            
            if currentIndex == 0 {
                return self.lastIntegral
            }
            else {
                currentPartition = partitions[currentIndex]
                partialV = [partialV.last!] + PiecewiseIntegrator.sampleIFValuesForRange(scale: scale, sampleRange: currentPartition, userIFValues: userIFValues) // overlap
                let lastIntegralValue = lastIntegral.last!
                lastIntegral = PiecewiseIntegrator.integrate(samples: partialV, stepSize: delta, useSimpsons: useSimpsons)
                lastIntegral.removeFirst()
                lastIntegral = vDSP.add(lastIntegralValue, lastIntegral)
                
                return self.lastIntegral
            }
        }
    }
}
