import Foundation

enum SignalStatistics {
  static func mean(_ values: [Double]) -> Double {
    guard !values.isEmpty else { return 0 }
    return values.reduce(0, +) / Double(values.count)
  }

  static func median(_ values: [Double]) -> Double {
    guard !values.isEmpty else { return 0 }
    let sorted = values.sorted()
    let middle = sorted.count / 2
    if sorted.count.isMultiple(of: 2) {
      return (sorted[middle - 1] + sorted[middle]) / 2
    }
    return sorted[middle]
  }

  static func mad(_ values: [Double], around center: Double? = nil) -> Double {
    guard !values.isEmpty else { return 0 }
    let center = center ?? median(values)
    return median(values.map { abs($0 - center) })
  }

  static func standardDeviation(_ values: [Double], sample: Bool = true) -> Double {
    guard values.count > (sample ? 1 : 0) else { return 0 }
    let average = mean(values)
    let sumSquares = values.reduce(0) { partial, value in
      let delta = value - average
      return partial + delta * delta
    }
    let divisor = Double(values.count - (sample ? 1 : 0))
    return (sumSquares / divisor).squareRoot()
  }

  static func rootMeanSquare(_ values: [Double]) -> Double {
    guard !values.isEmpty else { return 0 }
    let squareMean = values.reduce(0) { $0 + $1 * $1 } / Double(values.count)
    return squareMean.squareRoot()
  }

  static func percentile(_ values: [Double], _ probability: Double) -> Double {
    guard !values.isEmpty else { return 0 }
    let sorted = values.sorted()
    let p = min(max(probability, 0), 1)
    let position = p * Double(sorted.count - 1)
    let lower = Int(position.rounded(.down))
    let upper = Int(position.rounded(.up))
    guard lower != upper else { return sorted[lower] }
    let fraction = position - Double(lower)
    return sorted[lower] + fraction * (sorted[upper] - sorted[lower])
  }

  static func clamp01(_ value: Double) -> Double {
    guard value.isFinite else { return 0 }
    return min(max(value, 0), 1)
  }

  static func reflectedIndex(_ index: Int, count: Int) -> Int {
    guard count > 1 else { return 0 }
    var candidate = index
    while candidate < 0 || candidate >= count {
      if candidate < 0 {
        candidate = -candidate - 1
      }
      if candidate >= count {
        candidate = 2 * count - candidate - 1
      }
    }
    return candidate
  }

  /// Fits a sinusoid at `frequencyHz` and reports its power relative to residual power in dB.
  static func sinusoidalSNRDB(
    values: [Double],
    timestamps: [Double],
    validityMask: [Bool],
    frequencyHz: Double
  ) -> Double? {
    guard values.count == timestamps.count,
          values.count == validityMask.count,
          frequencyHz.isFinite,
          frequencyHz > 0
    else { return nil }

    let validIndices = values.indices.filter {
      validityMask[$0] && values[$0].isFinite && timestamps[$0].isFinite
    }
    guard validIndices.count >= 12 else { return nil }

    let average = validIndices.reduce(0) { $0 + values[$1] } / Double(validIndices.count)
    var sinSin = 0.0
    var cosCos = 0.0
    var sinCos = 0.0
    var valueSin = 0.0
    var valueCos = 0.0
    let origin = timestamps[validIndices[0]]

    for index in validIndices {
      let angle = 2 * Double.pi * frequencyHz * (timestamps[index] - origin)
      let sine = sin(angle)
      let cosine = cos(angle)
      let value = values[index] - average
      sinSin += sine * sine
      cosCos += cosine * cosine
      sinCos += sine * cosine
      valueSin += value * sine
      valueCos += value * cosine
    }

    let determinant = sinSin * cosCos - sinCos * sinCos
    guard abs(determinant) > 1e-12 else { return nil }
    let sineCoefficient = (valueSin * cosCos - valueCos * sinCos) / determinant
    let cosineCoefficient = (valueCos * sinSin - valueSin * sinCos) / determinant

    var modeledPower = 0.0
    var residualPower = 0.0
    for index in validIndices {
      let angle = 2 * Double.pi * frequencyHz * (timestamps[index] - origin)
      let modeled = sineCoefficient * sin(angle) + cosineCoefficient * cos(angle)
      let value = values[index] - average
      modeledPower += modeled * modeled
      let residual = value - modeled
      residualPower += residual * residual
    }

    guard modeledPower > 1e-12 else { return -60 }
    let ratio = modeledPower / max(residualPower, 1e-12)
    return 10 * log10(ratio)
  }

  /// Compares average spectral power around the pulse fundamental (and its second harmonic)
  /// with the remaining physiological band. The finite bandwidth tolerates slow BPM changes.
  static func periodicBandSNRDB(
    values: [Double],
    timestamps: [Double],
    validityMask: [Bool],
    fundamentalFrequencyHz: Double,
    minimumFrequencyHz: Double,
    maximumFrequencyHz: Double
  ) -> Double? {
    guard values.count == timestamps.count,
          values.count == validityMask.count,
          values.count >= 30,
          fundamentalFrequencyHz.isFinite,
          fundamentalFrequencyHz > 0,
          minimumFrequencyHz > 0,
          maximumFrequencyHz > minimumFrequencyHz
    else { return nil }

    let validIndices = values.indices.filter {
      validityMask[$0] && values[$0].isFinite && timestamps[$0].isFinite
    }
    guard validIndices.count >= 30 else { return nil }
    let origin = timestamps[validIndices[0]]
    let duration = timestamps[validIndices.last!] - origin
    guard duration > 2 else { return nil }
    let average = validIndices.reduce(0) { $0 + values[$1] } / Double(validIndices.count)
    let frequencyStep = max(1 / duration, 0.02)
    let fundamentalHalfWidth = max(0.12, min(0.28, fundamentalFrequencyHz * 0.18))
    let secondHarmonic = 2 * fundamentalFrequencyHz
    let harmonicHalfWidth = min(0.36, 1.5 * fundamentalHalfWidth)

    var signalPowers: [Double] = []
    var noisePowers: [Double] = []
    var frequency = minimumFrequencyHz
    while frequency <= maximumFrequencyHz + frequencyStep * 0.5 {
      var real = 0.0
      var imaginary = 0.0
      for index in validIndices {
        let relativeTime = timestamps[index] - origin
        let window = 0.5 - 0.5 * cos(2 * Double.pi * relativeTime / duration)
        let centeredValue = (values[index] - average) * window
        let angle = -2 * Double.pi * frequency * relativeTime
        real += centeredValue * cos(angle)
        imaginary += centeredValue * sin(angle)
      }
      let power = real * real + imaginary * imaginary
      let nearFundamental = abs(frequency - fundamentalFrequencyHz) <= fundamentalHalfWidth
      let nearSecondHarmonic = secondHarmonic <= maximumFrequencyHz
        && abs(frequency - secondHarmonic) <= harmonicHalfWidth
      if nearFundamental || nearSecondHarmonic {
        signalPowers.append(power)
      } else {
        noisePowers.append(power)
      }
      frequency += frequencyStep
    }

    guard !signalPowers.isEmpty, !noisePowers.isEmpty else { return nil }
    let signalPower = mean(signalPowers)
    let noisePower = mean(noisePowers)
    guard signalPower > 1e-12 else { return -60 }
    return 10 * log10(signalPower / max(noisePower, 1e-12))
  }
}