///|
/// A station-tagged phase pick used by the associator.
pub struct PickObservation {
  station : String
  position : StationPosition
  pick : ArrivalPick
  phase : String
  weight : Double
}

///|
/// A predicted arrival and its residual for one station.
pub struct TravelTimeResidual {
  station : String
  phase : String
  observed : Double
  predicted : Double
  residual : Double
  weight : Double
}

///|
/// An associated event hypothesis with residual diagnostics.
pub struct AssociationHypothesis {
  origin_time : Double
  latitude : Double
  longitude : Double
  depth_km : Double
  p_picks : Int
  s_picks : Int
  residual_rms : Double
  score : Double
  observations : Array[TravelTimeResidual]
}

///|
/// Controls phase association and candidate evaluation.
pub struct AssociationConfig {
  p_velocity : Double
  s_velocity : Double
  time_tolerance : Double
  minimum_picks : Int
  minimum_score : Double
  max_depth_km : Double
}

///|
/// Return conservative local-network association defaults.
pub fn default_association_config() -> AssociationConfig {
  {
    p_velocity: 6.0,
    s_velocity: 3.5,
    time_tolerance: 2.0,
    minimum_picks: 3,
    minimum_score: 0.5,
    max_depth_km: 80.0,
  }
}

///|
/// Build a station-tagged observation.
pub fn PickObservation::new(
  station : String,
  position : StationPosition,
  pick : ArrivalPick,
  phase : String,
  weight : Double,
) -> PickObservation {
  { station, position, pick, phase, weight: weight.max(0.0) }
}

///|
/// Return the wave speed for a phase label.
pub fn association_phase_velocity(
  config : AssociationConfig,
  phase : String,
) -> Double {
  if phase == "S" || phase == "s" {
    config.s_velocity.max(0.01)
  } else {
    config.p_velocity.max(0.01)
  }
}

///|
/// Estimate an origin time from observed picks and a source guess.
pub fn estimate_origin_time(
  observations : Array[PickObservation],
  source : StationPosition,
  depth_km : Double,
  config : AssociationConfig,
) -> Double {
  let estimates = Array::new()
  for observation in observations {
    let distance = surface_distance_km(source, observation.position)
    let velocity = association_phase_velocity(config, observation.phase)
    let path = (distance * distance + depth_km * depth_km).sqrt()
    estimates.push(observation.pick.time - path / velocity)
  }
  median_samples(estimates)
}

///|
/// Predict an arrival time from a source and a station observation.
pub fn predict_pick_time(
  origin_time : Double,
  source : StationPosition,
  depth_km : Double,
  observation : PickObservation,
  config : AssociationConfig,
) -> Double {
  let distance = surface_distance_km(source, observation.position)
  let path = (distance * distance + depth_km * depth_km).sqrt()
  origin_time + path / association_phase_velocity(config, observation.phase)
}

///|
/// Evaluate residuals for a source hypothesis.
pub fn evaluate_association(
  observations : Array[PickObservation],
  source : StationPosition,
  depth_km : Double,
  origin_time : Double,
  config : AssociationConfig,
) -> AssociationHypothesis {
  let residuals = Array::new(capacity=observations.length())
  let mut weighted_square = 0.0
  let mut total_weight = 0.0
  let mut p_count = 0
  let mut s_count = 0
  for observation in observations {
    let predicted = predict_pick_time(
      origin_time, source, depth_km, observation, config,
    )
    let residual = observation.pick.time - predicted
    let weight = observation.weight.max(1.0e-6)
    weighted_square += residual * residual * weight
    total_weight += weight
    if observation.phase == "S" || observation.phase == "s" {
      s_count += 1
    } else {
      p_count += 1
    }
    residuals.push({
      station: observation.station,
      phase: observation.phase,
      observed: observation.pick.time,
      predicted,
      residual,
      weight,
    })
  }
  let rms = if total_weight <= 0.0 {
    0.0
  } else {
    (weighted_square / total_weight).sqrt()
  }
  {
    origin_time,
    latitude: source.latitude,
    longitude: source.longitude,
    depth_km,
    p_picks: p_count,
    s_picks: s_count,
    residual_rms: rms,
    score: association_score(observations.length(), rms, config),
    observations: residuals,
  }
}

///|
/// Score an association from pick count and RMS residual.
pub fn association_score(
  pick_count : Int,
  residual_rms : Double,
  config : AssociationConfig,
) -> Double {
  let count_score = if config.minimum_picks <= 0 {
    1.0
  } else {
    (pick_count.to_double() / config.minimum_picks.to_double()).min(1.0)
  }
  let time_score = @math.exp(
    -residual_rms.abs() / config.time_tolerance.max(1.0e-6),
  )
  count_score * time_score
}

///|
/// Reject observations with residuals outside a robust tolerance.
pub fn inlier_observations(
  hypothesis : AssociationHypothesis,
  tolerance : Double,
) -> Array[TravelTimeResidual] {
  let result = Array::new()
  let limit = tolerance.abs()
  for observation in hypothesis.observations {
    if observation.residual.abs() <= limit {
      result.push(observation)
    }
  }
  result
}

///|
/// Refine origin time by a weighted mean of the current residuals.
pub fn refine_origin_time(
  hypothesis : AssociationHypothesis,
  damping : Double,
) -> Double {
  let mut correction = 0.0
  let mut weight = 0.0
  for observation in hypothesis.observations {
    let w = observation.weight.max(1.0e-6)
    correction += observation.residual * w
    weight += w
  }
  let factor = damping.max(0.0).min(1.0)
  if weight <= 0.0 {
    hypothesis.origin_time
  } else {
    hypothesis.origin_time + factor * correction / weight
  }
}

///|
/// Iteratively refine origin time and retain the best residual result.
pub fn refine_association(
  observations : Array[PickObservation],
  source : StationPosition,
  depth_km : Double,
  initial_origin : Double,
  config : AssociationConfig,
  iterations : Int,
) -> AssociationHypothesis {
  let count = if iterations <= 0 { 1 } else { iterations }
  let mut hypothesis = evaluate_association(
    observations, source, depth_km, initial_origin, config,
  )
  for _ in 0.. Array[Double] {
  observations.map(observation => {
    let distance = surface_distance_km(source, observation.position)
    let path = (distance * distance + depth_km * depth_km).sqrt()
    observation.pick.time -
    path / association_phase_velocity(config, observation.phase)
  })
}

///|
/// Associate picks for a fixed source using the best origin-time seed.
pub fn associate_fixed_source(
  observations : Array[PickObservation],
  source : StationPosition,
  depth_km : Double,
  config : AssociationConfig,
) -> AssociationHypothesis? {
  if observations.length() < config.minimum_picks {
    return None
  }
  let seeds = origin_time_seeds(observations, source, depth_km, config)
  if seeds.is_empty() {
    return None
  }
  let mut best = refine_association(
    observations,
    source,
    depth_km,
    seeds[0],
    config,
    4,
  )
  for index = 1; index < seeds.length(); index = index + 1 {
    let candidate = refine_association(
      observations,
      source,
      depth_km,
      seeds[index],
      config,
      4,
    )
    if candidate.score > best.score {
      best = candidate
    }
  }
  if best.score >= config.minimum_score {
    Some(best)
  } else {
    None
  }
}

///|
/// Deduplicate observations from the same station and phase.
pub fn deduplicate_observations(
  observations : Array[PickObservation],
  separation_seconds : Double,
) -> Array[PickObservation] {
  let result : Array[PickObservation] = Array::new()
  for observation in observations {
    let mut duplicate = false
    for existing in result {
      if existing.station == observation.station &&
        existing.phase == observation.phase &&
        (existing.pick.time - observation.pick.time).abs() <=
        separation_seconds.abs() {
        duplicate = true
        break
      }
    }
    if !duplicate {
      result.push(observation)
    }
  }
  result
}

///|
/// Return the weighted median pick time for one phase.
pub fn phase_time_median(
  observations : Array[PickObservation],
  phase : String,
) -> Double {
  let values = Array::new()
  for observation in observations {
    if observation.phase == phase {
      values.push(observation.pick.time)
    }
  }
  median_samples(values)
}

///|
/// Return the median absolute residual of an association.
pub fn association_median_absolute_residual(
  hypothesis : AssociationHypothesis,
) -> Double {
  median_samples(hypothesis.observations.map(item => item.residual.abs()))
}

///|
/// Return a concise association quality label.
pub fn association_quality_label(hypothesis : AssociationHypothesis) -> String {
  if hypothesis.score >= 0.85 && hypothesis.residual_rms <= 0.5 {
    "excellent"
  } else if hypothesis.score >= 0.6 {
    "usable"
  } else if hypothesis.score >= 0.35 {
    "review"
  } else {
    "reject"
  }
}

///|
/// Serialize association residuals for analyst review.
pub fn association_to_csv(hypothesis : AssociationHypothesis) -> String {
  let mut output = "station,phase,observed,predicted,residual,weight\n"
  for item in hypothesis.observations {
    output = output +
      item.station +
      "," +
      item.phase +
      "," +
      item.observed.to_string() +
      "," +
      item.predicted.to_string() +
      "," +
      item.residual.to_string() +
      "," +
      item.weight.to_string() +
      "\n"
  }
  output
}

///|
/// Format the headline of an association hypothesis.
pub fn association_summary_text(hypothesis : AssociationHypothesis) -> String {
  "origin=" +
  hypothesis.origin_time.to_string() +
  " lat=" +
  hypothesis.latitude.to_string() +
  " lon=" +
  hypothesis.longitude.to_string() +
  " depth=" +
  hypothesis.depth_km.to_string() +
  " p_picks=" +
  hypothesis.p_picks.to_string() +
  " s_picks=" +
  hypothesis.s_picks.to_string() +
  " rms=" +
  hypothesis.residual_rms.to_string() +
  " score=" +
  hypothesis.score.to_string() +
  " quality=" +
  association_quality_label(hypothesis)
}