import Collections
import Foundation

/// A dictionary of timestamps to normalized amplitude values.
// public typealias WaveformDataMap = [TimeInterval: CGFloat]
public typealias WaveformDataMap = OrderedDictionary<TimeInterval, CGFloat>

public extension WaveformDataMap {
    static let empty: Self = [:]
    var sortedPairs: [(key: TimeInterval, value: CGFloat)] {
        sorted(by: <)
    }
}

public struct WaveformData: Equatable {
    public enum EmptyWaveformStandInStyle {
        case empty
        case random(_ scale: Float)
        case uniformLevel(_ level: Float)
        case noiseySinwave
        case custom((TimeInterval) -> Float)

        func valueAtTime(_ t: TimeInterval) -> Float {
            let value: Float
            switch self {
            case .empty:
                value = 0.1

            case .random(let scale):
                value = Float.random(in: 0.0 ... 1.0) * scale

            case .uniformLevel(let level):
                value = level

            case .noiseySinwave:
                let randomNoise = Float.random(in: 0.0 ... 1.0) * 0.1
                let sinwaveOverDuration = ((Float(sin(t)) * 0.25) + 0.75)
                value = sinwaveOverDuration * randomNoise

            case .custom(let process):
                value = process(t)
            }

            return min(max(0.0, value), 1.0)
        }
    }

    public static let empty: Self = WaveformData(normalizedArray: [], totalDuration: .zero, emptyWaveformStandInStyle: .empty)

    public struct TimedAmplitedValue: Equatable {
        let timestamp: TimeInterval
        let normalizedAmplitude: CGFloat
    }

    public let totalDuration: TimeInterval
    public let waveformDataMap: WaveformDataMap
    public let timeSortedAmplitudeValues: [TimedAmplitedValue]

    public init(normalizedCGFloatArray: [CGFloat], totalDuration: TimeInterval, emptyWaveformStandInStyle: EmptyWaveformStandInStyle) {
        let mappedArray = normalizedCGFloatArray.map { Float($0) }
        self.init(normalizedArray: mappedArray, totalDuration: totalDuration, emptyWaveformStandInStyle: emptyWaveformStandInStyle)
    }

    public init(normalizedArray: [Float], totalDuration: TimeInterval, emptyWaveformStandInStyle: EmptyWaveformStandInStyle) {
        var newMap: WaveformDataMap = normalizedArray
            .enumerated()
            .reduce(.empty) { result, pair in
                var result = result
                let (index, value) = pair
                if normalizedArray.count > 1 {
                    let ratio = Double(index) / Double(normalizedArray.count - 1)
                    let timestamp = totalDuration * ratio
                    result[timestamp] = CGFloat(value)
                } else {
                    result[0] = CGFloat(value)
                }
                return result
            }

        if totalDuration > 0.0, newMap.isEmpty {
            /*
                 No waveform data for song
                 Put in symbolic wave for UI purposes only
             */
            var result: WaveformDataMap = [:]
            let samplesPerSecond: TimeInterval = 5.0
            for i in 0 ... Int(totalDuration * samplesPerSecond) {
                let timeKey = TimeInterval(i) / samplesPerSecond
                let amplitude = emptyWaveformStandInStyle.valueAtTime(timeKey)
                result[timeKey] = CGFloat(amplitude)
            }
            newMap = result
        }

        self.totalDuration = totalDuration
        self.waveformDataMap = newMap
        self.timeSortedAmplitudeValues = newMap
            .map { (key: TimeInterval, value: CGFloat) in
                TimedAmplitedValue(timestamp: key, normalizedAmplitude: value)
            }
            .sorted { $0.timestamp < $1.timestamp }
    }

    public func centerRange(_ windowDuration: TimeInterval) -> ClosedRange<TimeInterval> {
        let start = 0.0
        let end = totalDuration
        let middle = totalDuration / 2
        let halfWindowDuration = windowDuration / 2.0
        let centerStart = middle - halfWindowDuration
        let centerEnd = middle + halfWindowDuration
        return max(start, centerStart) ... min(end, centerEnd)
    }

    public func resample(targetCount: UInt) -> WaveformDataMap {
        let values = timeSortedAmplitudeValues

        // Special case handling: zero
        guard let firstValue = values.first, targetCount > 0 else {
            return .empty
        }

        // Special case handling: one
        guard values.count > 1 && targetCount > 1 else {
            return [firstValue.timestamp: firstValue.normalizedAmplitude]
        }

        // Get the time range
        let startTime = firstValue.timestamp
        let endTime = (values.last ?? firstValue).timestamp
        let duration = endTime - startTime

        var result: WaveformDataMap = .empty

        // Calculate new evenly spaced timestamps
        for i in 0 ..< targetCount {
            let ratio = i == targetCount - 1 ? 1.0 : Double(i) / Double(targetCount - 1)
            let newTime = startTime + ratio * duration

            // Find the surrounding original points for interpolation
            var lowerIndex = 0
            while lowerIndex < values.count - 1 && values[lowerIndex + 1].timestamp <= newTime {
                lowerIndex += 1
            }

            // Direct hit on an existing point
            if values[lowerIndex].timestamp == newTime {
                result[newTime] = values[lowerIndex].normalizedAmplitude
                continue
            }

            // We need to interpolate
            if lowerIndex < values.count - 1 {
                let lowerTime = values[lowerIndex].timestamp
                let upperTime = values[lowerIndex + 1].timestamp
                let lowerValue = values[lowerIndex].normalizedAmplitude
                let upperValue = values[lowerIndex + 1].normalizedAmplitude

                // Linear interpolation
                let t = (newTime - lowerTime) / (upperTime - lowerTime)
                let interpolatedValue = CGFloat(lowerValue + CGFloat(t) * (upperValue - lowerValue))
                result[newTime] = interpolatedValue
            } else if let lastPair = values.last {
                // Edge case - use the last value
                result[newTime] = lastPair.normalizedAmplitude
            }
        }

        return result
    }

    public func fitIn(width: CGFloat, barWidth: CGFloat, spacing: CGFloat) -> WaveformDataMap {
        let totalBarWidth = barWidth + spacing
        let barsThatFit: CGFloat = width / totalBarWidth
        let barCount = UInt((barsThatFit.isNaN || barsThatFit.isInfinite) ? 0 : barsThatFit)

        return resample(targetCount: barCount)
    }
}