import Foundation

public struct FrequencyBPMEstimate: Codable, Equatable, Sendable {
  public var bpm: Double?
  public var autocorrelation: Double
  public var confidence: Double

  public init(bpm: Double?, autocorrelation: Double, confidence: Double) {
    self.bpm = bpm
    self.autocorrelation = autocorrelation.isFinite ? autocorrelation : 0
    self.confidence = SignalStatistics.clamp01(confidence)
  }
}

public enum PPGBPMEstimator {
  public static func estimate(
    signal: [Double],
    timestampsSeconds: [Double],
    validityMask: [Bool],
    peakDetection: PeakDetectionResult,
    selectedChannel: PPGSignalChannel,
    configuration: PPGProcessingConfiguration
  ) throws -> BPMEstimate {
    let configuration = try configuration.validated()
    guard signal.count == timestampsSeconds.count,
          signal.count == validityMask.count,
          signal.count >= 5
    else {
      throw PPGSignalError.insufficientSamples(required: 5, actual: signal.count)
    }

    let acceptedIntervals = peakDetection.acceptedIntervals.map(\.durationMilliseconds)
    let temporalBPM: Double? = acceptedIntervals.count >= 4
      ? 60_000 / SignalStatistics.median(acceptedIntervals)
      : nil
    let frequency = frequencyEstimate(
      signal: signal,
      timestampsSeconds: timestampsSeconds,
      validityMask: validityMask,
      minimumFrequencyHz: configuration.minimumFrequencyHz,
      maximumFrequencyHz: configuration.maximumFrequencyHz
    )

    // Autocorrelation can lock onto the second harmonic (2×) or a sub-harmonic
    // (0.5×). Prefer the equivalent candidate closest to the clean beat-to-beat
    // estimate before judging the two methods as inconsistent.
    let adjustedFrequencyBPM = temporalBPM.flatMap { temporal -> Double? in
      guard let raw = frequency.bpm else { return nil }
      return [raw, raw / 2, raw * 2]
        .filter {
          $0 >= configuration.minimumFrequencyHz * 60
            && $0 <= configuration.maximumFrequencyHz * 60
        }
        .min(by: { abs($0 - temporal) < abs($1 - temporal) })
    } ?? frequency.bpm
    let agreement = temporalBPM.flatMap { temporal in
      adjustedFrequencyBPM.map { abs(temporal - $0) }
    }
    let agreementScore = agreement.map {
      SignalStatistics.clamp01(1 - $0 / max(configuration.maximumMethodDisagreementBPM, 1e-9))
    } ?? 0

    let finalBPM: Double?
    if let temporalBPM,
       let frequencyBPM = adjustedFrequencyBPM,
       let agreement,
       agreement <= configuration.maximumMethodDisagreementBPM {
      finalBPM = 0.65 * temporalBPM + 0.35 * frequencyBPM
    } else if let temporalBPM,
              acceptedIntervals.count >= 8,
              peakDetection.confidence >= 0.55 {
      // A long, regular sequence of detected beats is sufficient for mean HR.
      // Keep PRV subject to the stricter downstream interval-quality checks.
      finalBPM = temporalBPM
    } else if let frequencyBPM = adjustedFrequencyBPM,
              frequency.confidence >= 0.55 {
      finalBPM = frequencyBPM
    } else {
      finalBPM = nil
    }

    let countScore = min(Double(acceptedIntervals.count) / 10, 1)
    let confidence = 0.40 * peakDetection.confidence
      + 0.25 * frequency.confidence
      + 0.25 * agreementScore
      + 0.10 * countScore

    return BPMEstimate(
      bpm: finalBPM,
      temporalBpm: temporalBPM,
      frequencyBpm: adjustedFrequencyBPM,
      methodAgreementBpm: agreement,
      validPeakCount: peakDetection.peaks.count,
      rejectedPeakCount: peakDetection.rejectedIntervalCount,
      confidence: confidence,
      selectedChannel: selectedChannel,
      selectedPolarity: peakDetection.selectedPolarity
    )
  }

  public static func frequencyEstimate(
    signal: [Double],
    timestampsSeconds: [Double],
    validityMask: [Bool],
    minimumFrequencyHz: Double,
    maximumFrequencyHz: Double
  ) -> FrequencyBPMEstimate {
    guard signal.count == timestampsSeconds.count,
          signal.count == validityMask.count,
          signal.count >= 12,
          minimumFrequencyHz > 0,
          maximumFrequencyHz > minimumFrequencyHz
    else {
      return FrequencyBPMEstimate(bpm: nil, autocorrelation: 0, confidence: 0)
    }

    let intervals = zip(timestampsSeconds, timestampsSeconds.dropFirst())
      .map { $1 - $0 }
      .filter { $0.isFinite && $0 > 0 }
    let sampleInterval = SignalStatistics.median(intervals)
    guard sampleInterval > 0 else {
      return FrequencyBPMEstimate(bpm: nil, autocorrelation: 0, confidence: 0)
    }
    let sampleRate = 1 / sampleInterval
    let minimumLag = max(Int((sampleRate / maximumFrequencyHz).rounded(.down)), 2)
    let maximumLag = min(Int((sampleRate / minimumFrequencyHz).rounded(.up)), signal.count / 2)
    guard maximumLag > minimumLag else {
      return FrequencyBPMEstimate(bpm: nil, autocorrelation: 0, confidence: 0)
    }

    let validValues = signal.indices.compactMap {
      validityMask[$0] && signal[$0].isFinite ? signal[$0] : nil
    }
    let center = SignalStatistics.mean(validValues)
    var correlations: [Int: Double] = [:]
    for lag in minimumLag...maximumLag {
      var cross = 0.0
      var leftPower = 0.0
      var rightPower = 0.0
      var pairCount = 0
      for index in 0..<(signal.count - lag) {
        guard validityMask[index], validityMask[index + lag] else { continue }
        let left = signal[index] - center
        let right = signal[index + lag] - center
        guard left.isFinite, right.isFinite else { continue }
        cross += left * right
        leftPower += left * left
        rightPower += right * right
        pairCount += 1
      }
      guard pairCount >= 10, leftPower > 1e-12, rightPower > 1e-12 else { continue }
      correlations[lag] = cross / (leftPower * rightPower).squareRoot()
    }

    guard !correlations.isEmpty,
          let strongest = correlations.max(by: { $0.value < $1.value }),
          strongest.value > 0.10
    else {
      return FrequencyBPMEstimate(bpm: nil, autocorrelation: 0, confidence: 0)
    }

    let localPeaks = correlations.keys.sorted().filter { lag in
      let value = correlations[lag] ?? -.infinity
      let left = correlations[lag - 1] ?? -.infinity
      let right = correlations[lag + 1] ?? -.infinity
      return value >= left && value >= right
    }
    let strongThreshold = max(strongest.value * 0.92, 0.10)
    let chosenLag = localPeaks.first(where: { (correlations[$0] ?? 0) >= strongThreshold })
      ?? strongest.key
    let centerCorrelation = correlations[chosenLag] ?? strongest.value
    let leftCorrelation = correlations[chosenLag - 1] ?? centerCorrelation
    let rightCorrelation = correlations[chosenLag + 1] ?? centerCorrelation
    let denominator = leftCorrelation - 2 * centerCorrelation + rightCorrelation
    let fractionalOffset: Double
    if abs(denominator) > 1e-12 {
      fractionalOffset = min(
        max(0.5 * (leftCorrelation - rightCorrelation) / denominator, -0.5),
        0.5
      )
    } else {
      fractionalOffset = 0
    }
    let refinedLag = Double(chosenLag) + fractionalOffset
    guard refinedLag > 0 else {
      return FrequencyBPMEstimate(bpm: nil, autocorrelation: centerCorrelation, confidence: 0)
    }
    let bpm = 60 * sampleRate / refinedLag
    let confidence = SignalStatistics.clamp01((centerCorrelation - 0.10) / 0.75)
    return FrequencyBPMEstimate(
      bpm: bpm.isFinite ? bpm : nil,
      autocorrelation: centerCorrelation,
      confidence: confidence
    )
  }
}