//! Cyclical encoding.
//!
//! [`CyclicalEncoder`] maps cyclical features (hour, day-of-week, month)
//! to `sin` and `cos` components so that `23:00` and `01:00` are close
//! in the encoded space.

use crate::traits::{Error, Fit, Result, Transform};
use polars::prelude::*;

/// Encode cyclical features with sin/cos transformation.
///
/// For each column, computes `sin(2π * x / period)` and
/// `cos(2π * x / period)`, preserving the cyclic relationship.
///
/// # Example
///
/// ```rust
/// use featrs::time_series::cyclical::CyclicalEncoder;
/// use featrs::traits::{Fit, Transform};
/// use polars::prelude::{Column, DataFrame, NamedFrom, Series};
///
/// let col = Column::from(Series::new("hour".into(), &[0.0_f64, 6.0, 12.0, 18.0]));
/// let df = DataFrame::new(4, vec![col])?;
///
/// let mut enc = CyclicalEncoder::new(&["hour"], 24);
/// enc.fit(df.clone())?;
/// let encoded = enc.transform(df)?;
/// assert_eq!(encoded.height(), 4);
/// # Ok::<(), Box<dyn std::error::Error>>(())
/// ```
pub struct CyclicalEncoder {
    fitted: bool,
    columns: Vec<(String, f64)>,
}

impl CyclicalEncoder {
    /// Create an encoder that maps each column to `(sin, cos)` using a single
    /// shared integer period (e.g. `24` for hours, `12` for months).
    pub fn new(columns: &[&str], period: usize) -> Self {
        let period_f = period as f64;
        Self {
            fitted: false,
            columns: columns.iter().map(|s| (s.to_string(), period_f)).collect(),
        }
    }

    /// Create an encoder where each column gets its own integer period, given as
    /// `(column_name, period)` pairs.
    pub fn with_periods(columns: &[(&str, usize)]) -> Self {
        Self {
            fitted: false,
            columns: columns
                .iter()
                .map(|(s, p)| (s.to_string(), *p as f64))
                .collect(),
        }
    }
}

impl Fit<DataFrame> for CyclicalEncoder {
    type Output = ();

    fn fit(&mut self, x: DataFrame) -> Result<()> {
        for (col, _) in &self.columns {
            if x.column(col.as_str()).is_err() {
                return Err(Error::InvalidInput(format!(
                    "CyclicalEncoder: column '{}' not found.",
                    col
                )));
            }
        }
        self.fitted = true;
        Ok(())
    }
}

impl Transform<DataFrame> for CyclicalEncoder {
    type Output = DataFrame;

    fn transform(&self, x: DataFrame) -> Result<DataFrame> {
        if !self.fitted {
            return Err(Error::NotFitted("CyclicalEncoder".into()));
        }
        let mut out = x.clone();
        let two_pi = 2.0 * std::f64::consts::PI;

        for (col, period) in &self.columns {
            let s = out.column(col.as_str()).map_err(|e| {
                Error::InvalidInput(format!(
                    "CyclicalEncoder.transform: column '{}' not found. {}",
                    col, e
                ))
            })?;
            let s = s.clone();
            let ca = s.f64().map_err(|e| {
                Error::InvalidInput(format!(
                    "CyclicalEncoder.transform: column '{}' has dtype {}; expected Float64. {}",
                    col,
                    s.dtype(),
                    e
                ))
            })?;

            let sin_vals: ChunkedArray<Float64Type> = ca
                .iter()
                .map(|opt| opt.map(|v| (two_pi * v / period).sin()))
                .collect();
            let cos_vals: ChunkedArray<Float64Type> = ca
                .iter()
                .map(|opt| opt.map(|v| (two_pi * v / period).cos()))
                .collect();

            let sin_name = format!("{}_sin", col);
            let cos_name = format!("{}_cos", col);
            out.with_column(
                sin_vals
                    .into_series()
                    .with_name(sin_name.as_str().into())
                    .into(),
            )
            .map_err(|e| Error::Computation(e.to_string()))?;
            out.with_column(
                cos_vals
                    .into_series()
                    .with_name(cos_name.as_str().into())
                    .into(),
            )
            .map_err(|e| Error::Computation(e.to_string()))?;
        }

        Ok(out)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use approx::assert_relative_eq;

    #[test]
    fn test_cyclical_encoding() {
        let vals = Column::from(Series::new("hour".into(), &[0.0f64, 6.0, 12.0, 18.0]));
        let df = DataFrame::new(4, vec![vals]).unwrap();
        let mut enc = CyclicalEncoder::new(&["hour"], 24);

        enc.fit(df.clone()).unwrap();
        let result = enc.transform(df).unwrap();

        assert_eq!(result.width(), 3); // hour, hour_sin, hour_cos
        let sin = result.column("hour_sin").unwrap().f64().unwrap();
        let cos = result.column("hour_cos").unwrap().f64().unwrap();

        // hour 0: sin=0, cos=1
        assert_relative_eq!(sin.get(0).unwrap(), 0.0, epsilon = 1e-6);
        assert_relative_eq!(cos.get(0).unwrap(), 1.0, epsilon = 1e-6);
        // hour 6 (90°): sin=1, cos=0
        assert_relative_eq!(sin.get(1).unwrap(), 1.0, epsilon = 1e-6);
        assert_relative_eq!(cos.get(1).unwrap(), 0.0, epsilon = 1e-6);
    }
}
