import Foundation

public struct PRVCalculationResult: Equatable, Sendable {
  public var metrics: PRVMetrics
  public var intervals: [PulseInterval]

  public init(metrics: PRVMetrics, intervals: [PulseInterval]) {
    self.metrics = metrics
    self.intervals = intervals
  }
}

public enum PPGPRVCalculator {
  public static func calculate(
    intervals inputIntervals: [PulseInterval],
    configuration: PPGProcessingConfiguration
  ) throws -> PRVCalculationResult {
    let configuration = try configuration.validated()
    var intervals = inputIntervals
    let minimumInterval = 1_000 / configuration.maximumFrequencyHz
    let maximumInterval = 1_000 / configuration.minimumFrequencyHz

    for index in intervals.indices where intervals[index].rejectionReason == nil {
      let value = intervals[index].durationMilliseconds
      if !value.isFinite || value < minimumInterval || value > maximumInterval {
        intervals[index].rejectionReason = .outsidePhysiologicalRange
      }
    }

    // Filter against a moving neighbourhood (±13 intervals), not one global
    // median. This follows gradual physiological changes while isolating a
    // missed/double peak or a short movement artifact. Decisions use a frozen
    // snapshot so one rejected beat cannot cascade into later rejections.
    let initiallyAccepted = Set(intervals.indices.filter { intervals[$0].isAccepted })
    if initiallyAccepted.count >= 3 {
      for index in intervals.indices where initiallyAccepted.contains(index) {
        let lower = max(0, index - 13)
        let upper = min(intervals.count - 1, index + 13)
        let localValues = (lower...upper).compactMap { neighbour -> Double? in
          guard neighbour != index, initiallyAccepted.contains(neighbour) else { return nil }
          return intervals[neighbour].durationMilliseconds
        }
        guard localValues.count >= 2 else { continue }
        let center = SignalStatistics.median(localValues)
        let robustDeviation = 1.4826 * SignalStatistics.mad(localValues, around: center)
        let allowedDeviation = max(center * 0.20, 3.5 * robustDeviation, 70)
        let value = intervals[index].durationMilliseconds
        let resemblesMissedBeat = abs(value - 2 * center) <= max(0.18 * center, 70)
        let resemblesDoublePeak = abs(2 * value - center) <= max(0.18 * center, 70)
        if abs(value - center) > allowedDeviation
            || resemblesMissedBeat
            || resemblesDoublePeak {
          intervals[index].rejectionReason = .localOutlier
        }
      }
    }

    let accepted = intervals.filter(\.isAccepted)
    let values = accepted.map(\.durationMilliseconds)
    let rejectedCount = intervals.count - accepted.count
    let rejectedRatio = intervals.isEmpty
      ? 1
      : Double(rejectedCount) / Double(intervals.count)

    let median = values.isEmpty ? nil : SignalStatistics.median(values)
    let mean = values.isEmpty ? nil : SignalStatistics.mean(values)
    let sdnn = values.count >= 2 ? SignalStatistics.standardDeviation(values) : nil
    let coefficientOfVariation: Double?
    if let median, let sdnn, median > 0 {
      coefficientOfVariation = sdnn / median
    } else {
      coefficientOfVariation = nil
    }

    var successiveDifferences: [Double] = []
    for pair in zip(accepted, accepted.dropFirst()) {
      // Do not bridge over a rejected or missing beat: it would create a fictitious PRV difference.
      guard pair.1.sourceIndex == pair.0.sourceIndex + 1 else { continue }
      let difference = pair.1.durationMilliseconds - pair.0.durationMilliseconds
      successiveDifferences.append(difference)
    }
    let squaredSuccessiveDifferences = successiveDifferences.map { $0 * $0 }
    let rmssd = squaredSuccessiveDifferences.isEmpty
      ? nil
      : SignalStatistics.mean(squaredSuccessiveDifferences).squareRoot()
    let pnn50Percent = successiveDifferences.isEmpty
      ? nil
      : 100 * Double(successiveDifferences.filter { abs($0) > 50 }.count)
        / Double(successiveDifferences.count)
    let rangeIpiMs: Double? = values.count >= 2
      ? (values.max() ?? 0) - (values.min() ?? 0)
      : nil
    let histogram = Dictionary(grouping: values) { Int(floor($0 / 50)) }
    let modalBin = histogram.max { lhs, rhs in
      if lhs.value.count == rhs.value.count { return lhs.key > rhs.key }
      return lhs.value.count < rhs.value.count
    }
    let modeIpiMs = modalBin.map { SignalStatistics.mean($0.value) }
    let amo50Percent = modalBin.map {
      100 * Double($0.value.count) / Double(values.count)
    }

    let reliable = values.count >= configuration.minimumIntervalsForPRV
      && rejectedRatio <= configuration.maximumRejectedIntervalRatioForPRV
      && rmssd != nil
      && sdnn != nil

    let metrics = PRVMetrics(
      estimatedRmssdMs: rmssd,
      estimatedSdnnMs: sdnn,
      meanIpiMs: mean,
      medianIpiMs: median,
      pnn50Percent: pnn50Percent,
      rangeIpiMs: rangeIpiMs,
      modeIpiMs: modeIpiMs,
      amo50Percent: amo50Percent,
      coefficientOfVariation: coefficientOfVariation,
      validIntervalCount: values.count,
      rejectedIntervalCount: rejectedCount,
      rejectedIntervalRatio: rejectedRatio,
      reliable: reliable
    )
    return PRVCalculationResult(metrics: metrics, intervals: intervals)
  }

  public static func calculate(
    intervalsMilliseconds: [Double],
    configuration: PPGProcessingConfiguration
  ) throws -> PRVMetrics {
    var elapsed = 0.0
    let intervals = intervalsMilliseconds.enumerated().map { index, duration in
      let start = elapsed
      elapsed += duration / 1_000
      return PulseInterval(
        sourceIndex: index,
        startSeconds: start,
        endSeconds: elapsed,
        durationMilliseconds: duration
      )
    }
    return try calculate(intervals: intervals, configuration: configuration).metrics
  }
}