use crate::utils::misc::*;
use crate::utils::qdf::core::*;
use polars::prelude::*;
use std::collections::HashMap;
use std::error::Error;

#[allow(unused_imports)]
use std::f32::consts::E;

/// The required columns for a Quantamental DataFrame.
const QDF_INDEX_COLUMNS: [&str; 3] = ["real_date", "cid", "xcat"];

#[allow(dead_code)]
fn get_ticker_from_qdf(df: &DataFrame) -> Result<String, Box<dyn Error>> {
    check_quantamental_dataframe(df)?;
    let cid = df.column("cid")?.get(0).unwrap().to_string();
    let xcat = df.column("xcat")?.get(0).unwrap().to_string();
    let ticker = format!("{}_{}", cid, xcat);
    Ok(ticker)
}

/// Pivots a dataframe to a format where each ticker is a column.
/// # Arguments:
/// * `df` - The dataframe to pivot.
/// * `metric` - The metric to pivot on. If None, defaults to "value".
/// # Returns:
/// * `DataFrame` - The pivoted dataframe.
#[allow(dead_code)]
pub fn pivot_dataframe_by_ticker(
    df: DataFrame,
    metric: Option<String>,
) -> Result<DataFrame, Box<dyn Error>> {
    check_quantamental_dataframe(&df)?;

    // if no metric is provided, set it to 'value'
    let mut metric = metric.unwrap_or("value".into());
    if !df
        .get_column_names()
        .contains(&&PlSmallStr::from_string(metric.clone()))
    {
        // set metric to the first non-index column
        metric = df.get_column_names()[3].to_string();
    }
    let mut unique_dates: Column = get_unique_dates(&df)?;
    unique_dates.rename("real_date".into());

    // keep only the index columns and the metric column
    let mut keep_cols = QDF_INDEX_COLUMNS
        .to_vec()
        .iter()
        .map(|s| s.to_string())
        .collect::<Vec<String>>();
    keep_cols.push(metric.clone());

    let mut out_dfs = split_df_by_tickers(&df, Some(vec![metric.clone()]))?;
    for (tkr, odf) in out_dfs.iter_mut() {
        odf.rename(&metric, tkr.into())?;
        *odf = odf.select(vec!["real_date", tkr])?;
    }

    // create a new dataframe with the unique dates

    let mut new_df = DataFrame::new(vec![unique_dates])?;

    for (_, odf) in out_dfs.iter_mut() {
        new_df = new_df.left_join(odf, ["real_date"], ["real_date"])?;
    }

    // left join all dataframes on the date column

    Ok(new_df)
}

/// Pivot ticker dataframe to qdf.
/// # Arguments:
/// * `df` - The dataframe to pivot.
/// * `metric` - The metric to pivot on. If None, defaults to "value".
/// # Returns:
/// * `DataFrame` - The pivoted dataframe.
#[allow(dead_code)]
pub fn pivot_wide_dataframe_to_qdf(
    df: DataFrame,
    metric: Option<String>,
) -> Result<DataFrame, Box<dyn Error>> {
    // df list of dataframes
    let mut df_vec: Vec<DataFrame> = Vec::new();
    // if no metric is provided, set it to 'value'
    let metric = metric.unwrap_or("value".into());
    for tkr in df.get_column_names() {
        if tkr == "real_date" {
            continue;
        }
        let mut new_df = DataFrame::new(vec![df.column("real_date")?.clone()])?;
        // copy the column to the new dataframe
        new_df.with_column(df.column(tkr)?.clone())?;
        new_df.rename(tkr, metric.clone().into())?;
        // add the cid and xcat columns
        let (cid, xcat) = split_ticker(tkr.to_string())?;
        new_df.with_column(Series::new("cid".into(), vec![cid]))?;
        new_df.with_column(Series::new("xcat".into(), vec![xcat]))?;
        df_vec.push(new_df);
    }
    // create a new dataframe, concatting row-wise
    let mut new_df = DataFrame::new(vec![])?;
    for df in df_vec.iter() {
        new_df.vstack_mut(df)?;
    }

    // return the new dataframe
    check_quantamental_dataframe(&new_df)?;
    Ok(new_df)
}

/// Splits a dataframe by ticker.
/// # Arguments:
/// * `df` - The dataframe to split.
/// * `metrics` - The metrics to split on. If None, defaults to all metrics.
/// # Returns:
/// * `HashMap<String, DataFrame>` - A hashmap where the key is the ticker and the value is the dataframe.
#[allow(dead_code)]
fn split_df_by_tickers(
    df: &DataFrame,
    metrics: Option<Vec<String>>,
) -> Result<HashMap<String, DataFrame>, Box<dyn Error>> {
    check_quantamental_dataframe(df)?;

    let mut df_outs = HashMap::new();

    let unique_tickers = get_unique_tickers(df)?;
    let mut icids = Vec::new();
    let mut ixcats = Vec::new();
    for ticker in unique_tickers.iter() {
        let (cid, xcat) = split_ticker(ticker.to_string())?;
        icids.push(cid);
        ixcats.push(xcat);
    }
    let metrics = metrics.unwrap_or_else(|| get_unique_metrics(df).unwrap());
    // let keep_cols:Vec<String> = QDF_INDEX_COLUMNS + metrics;
    let mut keep_cols = QDF_INDEX_COLUMNS
        .to_vec()
        .iter()
        .map(|s| s.to_string())
        .collect::<Vec<String>>();

    keep_cols.extend(metrics);

    for (cid, xcat) in icids.iter().zip(ixcats.iter()) {
        // Apply filter while borrowing df and avoid moving ownership.
        let filter = col("cid")
            .eq(lit(cid.clone()))
            .and(col("xcat").eq(lit(xcat.clone())));

        let df_out = df.clone().lazy().filter(filter).collect()?;
        // select keep_cols
        let df_out = df_out.select(keep_cols.clone())?;
        df_outs.insert(format!("{}_{}", cid, xcat), df_out);
    }

    Ok(df_outs)
}

/// Splits a QDF container data for a single ticker into a Vec of DataFrames (one per metric).
/// The resulting DataFrames will have the "real_date" column and the metric column.
/// # Arguments:
/// * `df` - The dataframe to split.
/// # Returns:
/// * `Vec<DataFrame>` - A vector of dataframes, each containing the "real_date" and a metric column.
#[allow(dead_code)]
fn single_ticker_qdf_to_timeseries(mut df: DataFrame) -> Result<Vec<DataFrame>, Box<dyn Error>> {
    let mut df_vec = Vec::new();

    // Since we own `df`, we can remove the "real_date" column, consuming it.
    let date_col = df.drop_in_place("real_date")?;

    // Assuming `get_unique_metrics` only needs a reference to `df`
    let metrics = get_unique_metrics(&df)?;

    for metric in metrics.iter() {
        // Remove the metric column from `df`, consuming it.
        let metric_col = df.drop_in_place(metric)?;

        // Create a new DataFrame with the date column and the metric column.
        let new_df = DataFrame::new(vec![
            date_col.clone(), // Clone because `date_col` is used multiple times
            metric_col,       // Move the metric column (no clone needed)
        ])?;
        df_vec.push(new_df);
    }
    // At this point, `df` has had its metric columns removed and is effectively consumed.
    // The "date_col" has been moved out of `df` as well.

    // If you need to use `df` after this function, consider redesigning the function
    // to avoid consuming it.

    Ok(df_vec)
}
