//! Gap fill from the HII/ThaiWater mirror telemetry (`data.fill_from_hii`).
//!
//! For stations whose HII twin is the same sensor, hours missing from the RID
//! series (or carrying no level) are filled from HII, shifted by the median
//! offset between the two series over their overlap. ThaiWater relays RID's
//! own telemetry, so this repairs holes in RID's report but cannot bridge
//! RID's overnight pause (the mirror arrives with the same delay).

use crate::config::config;
use crate::grid::{Row, floor_hour};
use chrono::NaiveDateTime;
use std::collections::{BTreeMap, BTreeSet, HashSet};

/// One `hii_waterlevel` row joined to its RID station code; NaN = missing.
#[derive(Clone, Debug, PartialEq)]
pub struct HiiRow {
    pub timestamp: NaiveDateTime,
    pub station_code: String,
    pub wl_msl: f64,
    pub discharge: f64,
}

/// `data._prepare_hii_levels`: rows without a level dropped, one row per
/// station-hour, the latest reading of the hour winning (stable on ties).
pub fn prepare(rows: &[HiiRow]) -> Vec<HiiRow> {
    let mut kept: Vec<&HiiRow> = rows.iter().filter(|r| !r.wl_msl.is_nan()).collect();
    kept.sort_by_key(|r| r.timestamp); // stable
    let mut by_hour: BTreeMap<(String, NaiveDateTime), HiiRow> = BTreeMap::new();
    for r in kept {
        let t = floor_hour(r.timestamp);
        by_hour.insert(
            (r.station_code.clone(), t),
            HiiRow {
                timestamp: t,
                ..r.clone()
            },
        );
    }
    by_hour.into_values().collect()
}

/// The HII station codes worth fetching for these RID rows.
pub fn fill_candidates(rid: &[Row]) -> Vec<String> {
    let present: HashSet<&str> = rid.iter().map(|r| r.station_code.as_str()).collect();
    config()
        .hii_fill_stations
        .iter()
        .filter(|c| present.contains(c.as_str()))
        .cloned()
        .collect()
}

/// Per-station result of the fill, for logs and the notifier.
#[derive(Clone, Debug, PartialEq)]
pub struct FillInfo {
    pub station_code: String,
    pub offset: f64,
    pub filled_hours: usize,
    /// The newest filled hour and its level.
    pub latest: Option<(NaiveDateTime, f64)>,
}

/// pandas `Series.median` (numpy: mean of the two middle values when even).
fn median(mut v: Vec<f64>) -> f64 {
    v.sort_by(f64::total_cmp);
    let n = v.len();
    if n % 2 == 1 {
        v[n / 2]
    } else {
        (v[n / 2 - 1] + v[n / 2]) / 2.0
    }
}

/// `data.fill_from_hii`: `rid` as loaded (already one row per station-hour),
/// `hii` from [`prepare`]. Returns the merged rows (fill rows replace RID rows
/// that have no level) and what was filled.
pub fn fill(rid: &[Row], hii: &[HiiRow]) -> (Vec<Row>, Vec<FillInfo>) {
    let cfg = config();
    let candidates: BTreeSet<String> = fill_candidates(rid).into_iter().collect();
    let mut mirrors: BTreeMap<&str, Vec<&HiiRow>> = BTreeMap::new();
    for h in hii {
        if candidates.contains(&h.station_code) {
            mirrors.entry(h.station_code.as_str()).or_default().push(h);
        }
    }
    let mut fills: Vec<Row> = Vec::new();
    let mut info = Vec::new();
    for (code, mirror) in mirrors {
        let base: BTreeMap<NaiveDateTime, f64> = rid
            .iter()
            .filter(|r| r.station_code == code)
            .map(|r| (r.timestamp, r.water_level))
            .collect();
        let diffs: Vec<f64> = mirror
            .iter()
            .filter_map(|m| match base.get(&m.timestamp) {
                Some(level) if !level.is_nan() => Some(m.wl_msl - level),
                _ => None,
            })
            .collect();
        if diffs.len() < cfg.hii_min_overlap_hours {
            continue;
        }
        let offset = median(diffs);
        let exact = cfg.hii_exact_mirrors.iter().any(|c| c == code);
        let before = fills.len();
        for m in mirror {
            if base.get(&m.timestamp).is_some_and(|l| !l.is_nan()) {
                continue;
            }
            fills.push(Row {
                timestamp: m.timestamp,
                station_code: code.to_string(),
                water_level: m.wl_msl - offset,
                discharge: if exact { m.discharge } else { f64::NAN },
            });
        }
        if fills.len() > before {
            info.push(FillInfo {
                station_code: code.to_string(),
                offset,
                filled_hours: fills.len() - before,
                latest: fills
                    .iter()
                    .skip(before)
                    .max_by_key(|r| r.timestamp)
                    .map(|r| (r.timestamp, r.water_level)),
            });
        }
    }
    if fills.is_empty() {
        return (rid.to_vec(), info);
    }
    // concat + drop_duplicates(keep="last"): a fill replaces a level-less RID row
    let mut merged: BTreeMap<(String, NaiveDateTime), Row> = BTreeMap::new();
    for r in rid.iter().chain(fills.iter()) {
        merged.insert((r.station_code.clone(), r.timestamp), r.clone());
    }
    (merged.into_values().collect(), info)
}

#[cfg(test)]
mod tests {
    use super::*;
    use chrono::{Duration, NaiveDate};

    fn t0() -> NaiveDateTime {
        NaiveDate::from_ymd_opt(2026, 10, 1)
            .unwrap()
            .and_hms_opt(0, 0, 0)
            .unwrap()
    }

    fn rid(code: &str, hours: std::ops::Range<i64>) -> Vec<Row> {
        hours
            .map(|h| Row {
                timestamp: t0() + Duration::hours(h),
                station_code: code.into(),
                water_level: 1.0 + h as f64 * 0.01,
                discharge: 50.0,
            })
            .collect()
    }

    fn mirror(code: &str, hours: std::ops::Range<i64>, offset: f64) -> Vec<HiiRow> {
        hours
            .map(|h| HiiRow {
                timestamp: t0() + Duration::hours(h),
                station_code: code.into(),
                wl_msl: 1.0 + h as f64 * 0.01 + offset,
                discharge: 51.0,
            })
            .collect()
    }

    #[test]
    fn trailing_hours_filled_with_the_overlap_offset() {
        let r = rid("P.1", 0..200);
        let h = prepare(&mirror("P.1", 0..203, 300.0));
        let (merged, info) = fill(&r, &h);
        assert_eq!(merged.len(), 203);
        assert_eq!(info[0].filled_hours, 3);
        let last = merged.last().unwrap();
        assert_eq!(last.timestamp, t0() + Duration::hours(202));
        assert!((last.water_level - 3.02).abs() < 1e-9);
        assert_eq!(last.discharge, 51.0); // exact mirror keeps discharge
    }

    #[test]
    fn short_overlap_or_other_sensor_is_not_merged() {
        let r = rid("P.1", 0..100);
        assert_eq!(
            fill(&r, &prepare(&mirror("P.1", 0..103, 300.0))).0.len(),
            100
        );
        let r = rid("P.77", 0..200);
        assert_eq!(
            fill(&r, &prepare(&mirror("P.77", 0..203, 300.0))).0.len(),
            200
        );
    }

    #[test]
    fn p81_is_bias_corrected_without_discharge() {
        let (merged, _) = fill(&rid("P.81", 0..200), &prepare(&mirror("P.81", 0..201, 5.0)));
        assert!(merged.last().unwrap().discharge.is_nan());
    }

    #[test]
    fn latest_reading_of_an_hour_wins() {
        let mut rows = mirror("P.1", 0..1, 300.0);
        rows.push(HiiRow {
            timestamp: t0() + Duration::minutes(40),
            wl_msl: 302.0,
            ..rows[0].clone()
        });
        rows.push(HiiRow {
            timestamp: t0() + Duration::minutes(10),
            wl_msl: 301.5,
            ..rows[0].clone()
        });
        rows.push(HiiRow {
            wl_msl: f64::NAN,
            timestamp: t0() + Duration::minutes(50),
            ..rows[0].clone()
        });
        let p = prepare(&rows);
        assert_eq!(p.len(), 1);
        assert_eq!(p[0].wl_msl, 302.0);
        assert_eq!(p[0].timestamp, t0());
    }
}