use pyo3::prelude::*;
use pyo3_polars::{PyDataFrame, PySeries};

/// Python wrapper for [`crate::panel`] module.
#[allow(deprecated)]
#[pymodule]
pub fn panel(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(historic_vol, m)?)?;
    m.add_function(wrap_pyfunction!(get_bdates_from_col_hv, m)?)?;
    m.add_function(wrap_pyfunction!(get_period_indices_hv, m)?)?;
    m.add_function(wrap_pyfunction!(linear_composite, m)?)?;
    Ok(())
}

/// Python wrapper for calculating historic volatility.
/// See [`crate::panel::historic_vol`] for full documentation.
#[pyfunction]
pub fn historic_vol(
    df: PyDataFrame,
    xcat: String,
    cids: Option<Vec<String>>,
    lback_periods: Option<usize>,
    lback_method: Option<String>,
    half_life: Option<f64>,
    start: Option<String>,
    end: Option<String>,
    est_freq: Option<String>,
    remove_zeros: Option<bool>,
    postfix: Option<String>,
    nan_tolerance: Option<f64>,
) -> PyResult<PyDataFrame> {
    Ok(PyDataFrame(
        crate::panel::historic_vol::historic_vol(
            df.into(),
            xcat,
            cids,
            lback_periods,
            lback_method,
            half_life,
            start,
            end,
            est_freq,
            remove_zeros,
            postfix,
            nan_tolerance,
        )
        .unwrap(),
    ))
}

#[pyfunction]
pub fn get_bdates_from_col_hv(dfw: PyDataFrame, est_freq: &str) -> PyResult<PySeries> {
    Ok(PySeries(
        crate::panel::historic_vol::get_bdates_from_col_hv(&dfw.into(), est_freq)
            .unwrap()
            .into(),
    ))
}

#[pyfunction]
pub fn get_period_indices_hv(dfw: PyDataFrame, est_freq: &str) -> PyResult<Vec<usize>> {
    Ok(crate::panel::historic_vol::get_period_indices_hv(&dfw.into(), est_freq).unwrap())
}

#[pyfunction(signature = (
    df,
    xcats,
    cids,
    weights = None,
    signs = None,
    weight_xcats = None,
    normalize_weights = false,
    start = None,
    end = None,
    blacklist = None,
    complete_xcats = false,
    complete_cids = false,
    new_xcat = None,
    new_cid = None
))]
pub fn linear_composite(
    df: PyDataFrame,
    xcats: Vec<String>,
    cids: Vec<String>,
    weights: Option<Vec<f64>>,
    signs: Option<Vec<f64>>,
    weight_xcats: Option<Vec<String>>,
    normalize_weights: bool,
    start: Option<String>,
    end: Option<String>,
    blacklist: Option<std::collections::HashMap<String, Vec<String>>>,
    complete_xcats: bool,
    complete_cids: bool,
    new_xcat: Option<String>,
    new_cid: Option<String>,
) -> PyResult<PyDataFrame> {
    Ok(PyDataFrame(
        crate::panel::linear_composite::linear_composite(
            &df.into(),
            xcats,
            cids,
            weights,
            signs,
            weight_xcats,
            normalize_weights,
            start,
            end,
            blacklist,
            complete_xcats,
            complete_cids,
            new_xcat,
            new_cid,
        )
        .unwrap(),
    ))
}
