//! Long readings -> complete hourly grid (`features.make_hourly_grid`).

use crate::rolling::{Series, nan_series};
use chrono::{Duration, NaiveDateTime, Timelike};
use std::collections::BTreeMap;

/// One reading in long format; NaN = missing.
#[derive(Clone, Debug, PartialEq)]
pub struct Row {
    pub timestamp: NaiveDateTime,
    pub station_code: String,
    pub water_level: f64,
    pub discharge: f64,
}

pub fn floor_hour(t: NaiveDateTime) -> NaiveDateTime {
    t.with_nanosecond(0)
        .and_then(|t| t.with_second(0))
        .and_then(|t| t.with_minute(0))
        .unwrap_or(t)
}

/// Every hour from the first to the last reading; per station the observed
/// (unfilled) level and discharge, NaN where nothing was recorded.
pub struct Grid {
    pub times: Vec<NaiveDateTime>,
    pub level: BTreeMap<String, Series>,
    pub discharge: BTreeMap<String, Series>,
}

impl Grid {
    /// Timestamps are floored to the hour; for a duplicate (station, hour)
    /// the later row wins (pandas `drop_duplicates(keep="last")`).
    pub fn from_rows(rows: &[Row]) -> Grid {
        let mut cells: BTreeMap<(String, NaiveDateTime), (f64, f64)> = BTreeMap::new();
        for r in rows {
            cells.insert(
                (r.station_code.clone(), floor_hour(r.timestamp)),
                (r.water_level, r.discharge),
            );
        }
        let (Some(first), Some(last)) = (
            cells.keys().map(|k| k.1).min(),
            cells.keys().map(|k| k.1).max(),
        ) else {
            return Grid {
                times: vec![],
                level: BTreeMap::new(),
                discharge: BTreeMap::new(),
            };
        };
        let n = ((last - first).num_hours() + 1) as usize;
        let times: Vec<NaiveDateTime> = (0..n).map(|i| first + Duration::hours(i as i64)).collect();
        let mut level: BTreeMap<String, Series> = BTreeMap::new();
        let mut discharge: BTreeMap<String, Series> = BTreeMap::new();
        for ((code, t), (l, q)) in cells {
            let i = (t - first).num_hours() as usize;
            level.entry(code.clone()).or_insert_with(|| nan_series(n))[i] = l;
            discharge.entry(code).or_insert_with(|| nan_series(n))[i] = q;
        }
        Grid {
            times,
            level,
            discharge,
        }
    }

    pub fn len(&self) -> usize {
        self.times.len()
    }

    pub fn is_empty(&self) -> bool {
        self.times.is_empty()
    }

    pub fn index_of(&self, t: NaiveDateTime) -> Option<usize> {
        let first = *self.times.first()?;
        let i = (t - first).num_hours();
        (i >= 0 && (i as usize) < self.len() && self.times[i as usize] == t).then_some(i as usize)
    }

    /// The station's last hour with an observed level.
    pub fn last_observed(&self, station: &str) -> Option<usize> {
        self.level.get(station)?.iter().rposition(|x| !x.is_nan())
    }
}