import Foundation

public enum PPGPeakDetector {
  public static func detect(
    signal: [Double],
    timestampsSeconds: [Double],
    validityMask: [Bool],
    configuration: PPGProcessingConfiguration
  ) throws -> PeakDetectionResult {
    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 positive = detectPolarity(
      signal: signal,
      timestampsSeconds: timestampsSeconds,
      validityMask: validityMask,
      polarity: 1,
      configuration: configuration
    )
    let negative = detectPolarity(
      signal: signal,
      timestampsSeconds: timestampsSeconds,
      validityMask: validityMask,
      polarity: -1,
      configuration: configuration
    )

    if negative.confidence > positive.confidence + 1e-12 {
      return negative
    }
    if abs(negative.confidence - positive.confidence) <= 1e-12,
       negative.acceptedIntervals.count > positive.acceptedIntervals.count {
      return negative
    }
    return positive
  }

  private struct Candidate {
    var index: Int
    var amplitude: Double
    var prominence: Double
  }

  private static func detectPolarity(
    signal: [Double],
    timestampsSeconds: [Double],
    validityMask: [Bool],
    polarity: Int,
    configuration: PPGProcessingConfiguration
  ) -> PeakDetectionResult {
    let values = signal.map { ($0.isFinite ? $0 : 0) * Double(polarity) }
    let timestampIntervals = zip(timestampsSeconds, timestampsSeconds.dropFirst())
      .map { $1 - $0 }
      .filter { $0.isFinite && $0 > 0 }
    let sampleInterval = SignalStatistics.median(timestampIntervals)
    guard sampleInterval > 0 else {
      return PeakDetectionResult(peaks: [], intervals: [], selectedPolarity: polarity, confidence: 0)
    }
    let sampleRate = 1 / sampleInterval

    let validValues = values.indices.compactMap { validityMask[$0] ? values[$0] : nil }
    let center = SignalStatistics.median(validValues)
    let robustScale = max(1.4826 * SignalStatistics.mad(validValues, around: center), 1e-6)
    let amplitudeThreshold = center + max(configuration.minimumPeakAmplitude, robustScale * 0.18)
    let prominenceRadius = max(Int((0.45 * sampleRate).rounded()), 2)

    var candidates: [Candidate] = []
    candidates.reserveCapacity(max(signal.count / 15, 8))
    for index in 1..<(values.count - 1) {
      guard validityMask[index],
            validityMask[index - 1],
            validityMask[index + 1],
            values[index] >= amplitudeThreshold,
            values[index] > values[index - 1],
            values[index] >= values[index + 1]
      else { continue }

      let leftStart = max(0, index - prominenceRadius)
      let rightEnd = min(values.count - 1, index + prominenceRadius)
      let leftMinimum = values[leftStart..<index].min() ?? values[index]
      let rightMinimum = values[(index + 1)...rightEnd].min() ?? values[index]
      let prominence = min(values[index] - leftMinimum, values[index] - rightMinimum)
      guard prominence >= configuration.minimumPeakProminence else { continue }
      candidates.append(Candidate(index: index, amplitude: values[index], prominence: prominence))
    }

    let minimumDistance = max(
      Int((configuration.refractoryPeriodSeconds * sampleRate).rounded()),
      2
    )
    var selected: [Candidate] = []
    for candidate in candidates {
      if let last = selected.last, candidate.index - last.index < minimumDistance {
        let candidateStrength = candidate.prominence + 0.25 * candidate.amplitude
        let lastStrength = last.prominence + 0.25 * last.amplitude
        if candidateStrength > lastStrength {
          selected[selected.count - 1] = candidate
        }
      } else {
        selected.append(candidate)
      }
    }

    let peaks = selected.map { candidate -> DetectedPeak in
      let index = candidate.index
      let left = values[index - 1]
      let middle = values[index]
      let right = values[index + 1]
      let denominator = left - 2 * middle + right
      let fractionalOffset: Double
      if abs(denominator) > 1e-12 {
        fractionalOffset = min(max(0.5 * (left - right) / denominator, -0.5), 0.5)
      } else {
        fractionalOffset = 0
      }
      let refinedTimestamp = timestampsSeconds[index] + fractionalOffset * sampleInterval
      return DetectedPeak(
        sampleIndex: index,
        timestampSeconds: refinedTimestamp,
        amplitude: candidate.amplitude,
        prominence: candidate.prominence
      )
    }

    var intervals = zip(peaks, peaks.dropFirst()).enumerated().map { sourceIndex, pair in
      let durationMilliseconds = (pair.1.timestampSeconds - pair.0.timestampSeconds) * 1_000
      let lowerIndex = min(pair.0.sampleIndex, pair.1.sampleIndex)
      let upperIndex = max(pair.0.sampleIndex, pair.1.sampleIndex)
      let crossesInvalidRegion = validityMask[lowerIndex...upperIndex].contains(false)
      return PulseInterval(
        sourceIndex: sourceIndex,
        startSeconds: pair.0.timestampSeconds,
        endSeconds: pair.1.timestampSeconds,
        durationMilliseconds: durationMilliseconds,
        rejectionReason: crossesInvalidRegion ? .invalidSignalRegion : nil
      )
    }

    let minimumInterval = 1_000 / configuration.maximumFrequencyHz
    let maximumInterval = 1_000 / configuration.minimumFrequencyHz
    for index in intervals.indices where intervals[index].rejectionReason == nil {
      let duration = intervals[index].durationMilliseconds
      if !duration.isFinite || duration < minimumInterval || duration > maximumInterval {
        intervals[index].rejectionReason = .outsidePhysiologicalRange
      }
    }

    // Score every pulse independently. A short movement or pressure change can
    // therefore reject only the affected beats instead of invalidating the
    // complete recording. Cycles are normalized to 50 points before their
    // shape is compared with neighbouring cycles, following the published
    // morphology-based quality approach used for camera PPG.
    let cycleShapes = intervals.indices.map { index in
      normalizedCycle(
        values: values,
        lowerIndex: peaks[index].sampleIndex,
        upperIndex: peaks[index + 1].sampleIndex,
        outputCount: 50
      )
    }
    let medianProminence = max(SignalStatistics.median(peaks.map(\.prominence)), 1e-9)
    for index in intervals.indices {
      let lowerPeakQuality = SignalStatistics.clamp01(peaks[index].prominence / medianProminence)
      let upperPeakQuality = SignalStatistics.clamp01(peaks[index + 1].prominence / medianProminence)
      let prominenceQuality = min(lowerPeakQuality, upperPeakQuality)
      let neighbourhood = max(0, index - 3)...min(intervals.count - 1, index + 3)
      let correlations = neighbourhood.compactMap { neighbour -> Double? in
        guard neighbour != index,
              let shape = cycleShapes[index],
              let other = cycleShapes[neighbour]
        else { return nil }
        return shapeCorrelation(shape, other)
      }
      let morphologyQuality: Double
      if correlations.isEmpty {
        morphologyQuality = intervals.count < 3 ? 0.75 : 0
      } else {
        // Use the best two local matches so one adjacent corrupted cycle does
        // not drag a clean beat below the acceptance threshold.
        let strongest = correlations.sorted(by: >).prefix(2)
        morphologyQuality = SignalStatistics.clamp01(
          (SignalStatistics.mean(Array(strongest)) + 1) / 2
        )
      }
      let technicalQuality = SignalStatistics.clamp01(
        0.65 * morphologyQuality + 0.35 * prominenceQuality
      )
      intervals[index].qualityScore = technicalQuality
      if intervals[index].rejectionReason == nil,
         intervals.count >= 3,
         technicalQuality < 0.42 {
        intervals[index].rejectionReason = .lowMorphologyQuality
      }
    }

    let physiologicallyPlausible = intervals
      .filter(\.isAccepted)
      .map(\.durationMilliseconds)
    if physiologicallyPlausible.count >= 3 {
      let median = SignalStatistics.median(physiologicallyPlausible)
      let robustDeviation = 1.4826 * SignalStatistics.mad(physiologicallyPlausible, around: median)
      let allowedDeviation = max(0.28 * median, 3.5 * robustDeviation, 90)
      for index in intervals.indices where intervals[index].rejectionReason == nil {
        if abs(intervals[index].durationMilliseconds - median) > allowedDeviation {
          intervals[index].rejectionReason = .localOutlier
        }
      }
    }

    let accepted = intervals.filter(\.isAccepted)
    let acceptedDurations = accepted.map(\.durationMilliseconds)
    let intervalAcceptance = intervals.isEmpty
      ? 0
      : Double(accepted.count) / Double(intervals.count)
    let countScore = min(Double(accepted.count) / 8, 1)
    let regularity: Double
    if acceptedDurations.count >= 2 {
      let median = max(SignalStatistics.median(acceptedDurations), 1)
      let robustCV = 1.4826 * SignalStatistics.mad(acceptedDurations, around: median) / median
      regularity = SignalStatistics.clamp01(1 - robustCV / 0.20)
    } else {
      regularity = 0
    }
    let prominenceScore = SignalStatistics.clamp01(medianProminence / 1.5)
    let validCoverage = Double(validityMask.filter { $0 }.count) / Double(max(validityMask.count, 1))
    let confidence = 0.27 * countScore
      + 0.23 * intervalAcceptance
      + 0.20 * regularity
      + 0.20 * prominenceScore
      + 0.10 * validCoverage

    return PeakDetectionResult(
      peaks: peaks,
      intervals: intervals,
      selectedPolarity: polarity,
      confidence: confidence
    )
  }

  private static func normalizedCycle(
    values: [Double],
    lowerIndex: Int,
    upperIndex: Int,
    outputCount: Int
  ) -> [Double]? {
    guard lowerIndex >= 0,
          upperIndex < values.count,
          upperIndex - lowerIndex >= 3,
          outputCount >= 3
    else { return nil }
    let span = Double(upperIndex - lowerIndex)
    let sampled = (0..<outputCount).map { outputIndex -> Double in
      let position = Double(lowerIndex)
        + Double(outputIndex) * span / Double(outputCount - 1)
      let left = min(max(Int(position.rounded(.down)), lowerIndex), upperIndex)
      let right = min(left + 1, upperIndex)
      let fraction = position - Double(left)
      return values[left] + fraction * (values[right] - values[left])
    }
    let center = SignalStatistics.mean(sampled)
    let scale = SignalStatistics.standardDeviation(sampled)
    guard scale.isFinite, scale > 1e-8 else { return nil }
    return sampled.map { ($0 - center) / scale }
  }

  private static func shapeCorrelation(_ lhs: [Double], _ rhs: [Double]) -> Double? {
    guard lhs.count == rhs.count, lhs.count >= 3 else { return nil }
    let value = zip(lhs, rhs).reduce(0.0) { $0 + $1.0 * $1.1 } / Double(lhs.count)
    return value.isFinite ? min(max(value, -1), 1) : nil
  }
}