use crate::utils::qdf::check_quantamental_dataframe;
use crate::utils::qdf::pivots::*;
use crate::utils::qdf::reduce_df::*;
use polars::prelude::*;
use std::collections::HashMap;
const TOLERANCE: f64 = 1e-8;

fn _norm_apply_nan_mask(
    weights_dfw: &DataFrame,
    nan_mask_dfw: &DataFrame,
) -> Result<DataFrame, PolarsError> {
    let mut norm_dfw = DataFrame::new(vec![])?;
    for col in weights_dfw.get_column_names().clone() {
        let weight_col = weights_dfw.column(col)?;
        let nan_mask_col = nan_mask_dfw.column(col)?;
        let nan_mask_series = (!(nan_mask_col.bool()?)).into_column();
        let float_col = nan_mask_series.cast(&DataType::Float64)?;
        let float_series = (weight_col * &float_col)?;
        let sum = float_series.f64()?.sum().unwrap();
        if sum < TOLERANCE {
            // get the sum of the nan_mask_series
            let nan_mask_sum = nan_mask_series
                .cast(&DataType::Float64)?
                .f64()?
                .sum()
                .unwrap();
            // if the sum of the nan_mask_series is close to the len of the series, then ok
            let weight_sum = weight_col.f64()?.sum().unwrap();
            let ratio = (nan_mask_sum - weight_col.len() as f64).abs() / weight_col.len() as f64;
            assert!(
                ratio < TOLERANCE,
                "The sum of the updated series is zero: {:?} w ratio: {:?} and the weights summing to: {:?}",
                sum, nan_mask_sum, weight_sum
            );
        }
        norm_dfw.with_column(float_series)?;
    }
    Ok(norm_dfw)
}

fn normalize_weights_with_nan_mask(
    weights_dfw: DataFrame,
    nan_mask_dfw: DataFrame,
) -> Result<DataFrame, PolarsError> {
    // Ensure dimensions and columns match
    assert_eq!(weights_dfw.shape(), nan_mask_dfw.shape());
    assert!(
        weights_dfw
            .get_column_names()
            .iter()
            .all(|col| nan_mask_dfw.get_column_names().contains(col)),
        "The columns in weights_dfw and nan_mask_dfw do not match."
    );

    let mut norm_dfw = _norm_apply_nan_mask(&weights_dfw, &nan_mask_dfw)?;

    let sums_of_rows = norm_dfw
        .sum_horizontal(polars::frame::NullStrategy::Ignore)?
        .unwrap()
        .into_column();

    for col in weights_dfw.get_column_names() {
        let weight_series = norm_dfw.column(col)?;
        let divided_series = weight_series.divide(&sums_of_rows)?;
        let normalized_series = divided_series.as_series().unwrap().clone();
        norm_dfw.replace(col, normalized_series)?;
    }

    Ok(norm_dfw)
}

fn _form_agg_data_dfw(dfw: &DataFrame, agg_targs: &Vec<String>) -> Result<DataFrame, PolarsError> {
    let mut data_dfw = DataFrame::new(vec![])?;
    for agg_targ in agg_targs {
        match dfw.column(agg_targ) {
            Ok(agg_targ_series) => {
                // If the column exists, clone it and add it to data_dfw
                data_dfw.with_column(agg_targ_series.clone())?;
            }
            Err(_) => {
                // If the column does not exist, create a series full of NaNs
                let nan_series =
                    Series::full_null(agg_targ.into(), dfw.height(), &DataType::Float64);
                data_dfw.with_column(nan_series)?;
            }
        }
    }
    Ok(data_dfw)
}

fn _form_agg_nan_mask_dfw(data_dfw: &DataFrame) -> Result<DataFrame, PolarsError> {
    // Create a NaN mask DataFrame
    let mut nan_mask_dfw = DataFrame::new(vec![])?;
    for column in data_dfw.get_columns() {
        let nan_mask_series = column.is_null().into_series();
        nan_mask_dfw.with_column(nan_mask_series)?;
    }
    Ok(nan_mask_dfw)
}

fn _form_agg_nan_mask_series(nan_mask_dfw: &DataFrame) -> Result<Series, PolarsError> {
    // perform a row-wise OR
    let columns = nan_mask_dfw.get_columns();
    let mut combined = columns[0].bool()?.clone();
    for column in columns.iter().skip(1) {
        combined = &combined | column.bool()?;
    }
    Ok(combined.into_series())
}

fn _form_agg_weights_dfw(
    agg_weights_map: &HashMap<String, Vec<f64>>,
    data_dfw: DataFrame,
) -> Result<DataFrame, PolarsError> {
    let mut weights_dfw = DataFrame::new(vec![])?;
    for (agg_targ, weight_signs) in agg_weights_map.iter() {
        let wgt = weight_signs[0] * weight_signs[1];
        let wgt_series = Series::new(agg_targ.into(), vec![wgt; data_dfw.height()]);
        weights_dfw.with_column(wgt_series)?;
    }
    Ok(weights_dfw)
}

fn _assert_weights_and_data_match(
    data_dfw: &DataFrame,
    weights_dfw: &DataFrame,
) -> Result<(), PolarsError> {
    for col in weights_dfw.get_column_names() {
        if !data_dfw.get_column_names().contains(&col) {
            return Err(PolarsError::ComputeError(
                format!(
                    "weights_dfw and data_dfw do not have the same set of columns: {}",
                    col
                )
                .into(),
            ));
        }
    }
    Ok(())
}

fn perform_single_group_agg(
    dfw: &DataFrame,
    agg_on: &String,
    agg_targs: &Vec<String>,
    agg_weights_map: &HashMap<String, Vec<f64>>,
    normalize_weights: bool,
    complete: bool,
) -> Result<Column, PolarsError> {
    let data_dfw = _form_agg_data_dfw(dfw, agg_targs)?;
    let nan_mask_dfw = _form_agg_nan_mask_dfw(&data_dfw)?;
    let nan_mask_series = _form_agg_nan_mask_series(&nan_mask_dfw)?;
    let weights_dfw = _form_agg_weights_dfw(agg_weights_map, data_dfw.clone())?;
    let weights_dfw = match normalize_weights {
        true => normalize_weights_with_nan_mask(weights_dfw, nan_mask_dfw)?,
        false => weights_dfw,
    };

    println!("weights_dfw: {:?}", weights_dfw);
    // assert weights and data have the same set of columns
    _assert_weights_and_data_match(&data_dfw, &weights_dfw)?;

    // let output_dfw = weights_dfw * data_dfw;
    let mut output_columns = Vec::with_capacity(data_dfw.get_column_names().len());
    for col_name in data_dfw.get_column_names() {
        let data_series = data_dfw.column(col_name)?;
        let weight_series = weights_dfw.column(col_name)?;
        let mut multiplied = (data_series * weight_series)?;
        multiplied.rename(col_name.to_string().into());
        output_columns.push(multiplied);
    }

    let mut sum_series = output_columns[0].clone();
    for i in 1..output_columns.len() {
        let filled_column = output_columns[i].fill_null(FillNullStrategy::Zero)?;
        sum_series = (&sum_series + &filled_column)?;
    }

    if complete {
        sum_series = sum_series.zip_with(
            nan_mask_series.bool()?,
            &Series::full_null(agg_on.clone().into(), sum_series.len(), sum_series.dtype())
                .into_column(),
        )?;
    }
    sum_series.rename(agg_on.clone().into());
    // new_dfw.with_column(sum_series)?;

    Ok(sum_series)
}

fn perform_multiplication(
    dfw: &DataFrame,
    mult_targets: &HashMap<String, Vec<String>>,
    weights_map: &HashMap<String, HashMap<String, Vec<f64>>>,
    complete: bool,
    normalize_weights: bool,
) -> Result<DataFrame, PolarsError> {
    // let real_date = dfw.column("real_date".into())?.clone();
    // let mut new_dfw = DataFrame::new(vec![real_date])?;
    let mut new_dfw = DataFrame::new(vec![])?;
    assert!(!mult_targets.is_empty(), "agg_targs is empty");
    for (agg_on, agg_targs) in mult_targets.iter() {
        // perform_single_group_agg
        let cols_len = new_dfw.get_column_names().len();
        let new_col = perform_single_group_agg(
            dfw,
            agg_on,
            agg_targs,
            &weights_map[agg_on],
            normalize_weights,
            complete,
        )?;
        // assert that the number of columns has grown
        assert!(
            new_col.len() != 0,
            "The new DataFrame is empty after aggregation."
        );
        new_dfw.with_column(new_col)?;
        // if the height of new_dfw is 0,
        let new_cols_len = new_dfw.get_column_names().len();

        assert!(
            new_cols_len > cols_len,
            "The number of columns did not grow after aggregation."
        );
    }
    // get the real_date column from dfw
    let real_date = dfw.column("real_date".into())?.clone();
    new_dfw.with_column(real_date.clone())?;

    // Placeholder logic to return a valid DataFrame for now
    // Replace this with the actual implementation
    Ok(new_dfw)
}

fn get_agg_on_agg_targs(cids: Vec<String>, xcats: Vec<String>) -> (Vec<String>, Vec<String>) {
    let _agg_xcats_for_cid = agg_xcats_for_cid(cids.clone(), xcats.clone());
    let (agg_on, agg_targ) = if _agg_xcats_for_cid {
        (cids.clone(), xcats.clone())
    } else {
        (xcats.clone(), cids.clone())
    };
    // assert that if agg_xcats_for_cid is true, agg_on = cids
    match _agg_xcats_for_cid {
        true => {
            assert_eq!(agg_on, cids);
            assert_eq!(agg_targ, xcats);
        }
        false => {
            assert_eq!(agg_on, xcats);
            assert_eq!(agg_targ, cids);
        }
    }
    (agg_on, agg_targ)
}

/// Get the mapping of aggregation targets for the implied mode of aggregation.
/// # Returns
/// * `HashMap<String, Vec<String>>` - A mapping of cid/xcat to the list of tickers to be aggregated.
fn get_mul_targets(
    cids: Vec<String>,
    xcats: Vec<String>,
) -> Result<HashMap<String, Vec<String>>, Box<dyn std::error::Error>> {
    let _agg_xcats_for_cid = agg_xcats_for_cid(cids.clone(), xcats.clone());
    let mut mul_targets = HashMap::new();

    let (agg_on, agg_targ) = get_agg_on_agg_targs(cids.clone(), xcats.clone());

    for agg_o in agg_on {
        let mut targets = Vec::new();
        for agg_t in &agg_targ {
            let ticker = match _agg_xcats_for_cid {
                true => format!("{}_{}", agg_o, agg_t),
                false => format!("{}_{}", agg_t, agg_o),
            };
            targets.push(ticker);
        }
        if !targets.is_empty() {
            mul_targets.insert(agg_o.clone(), targets);
        }
    }
    // check if mul_targets is empty
    assert!(
        !mul_targets.is_empty(),
        "The mul_targets is empty. Please check the input DataFrame."
    );
    Ok(mul_targets)
}

fn form_weights_and_signs_map(
    cids: Vec<String>,
    xcats: Vec<String>,
    weights: Option<Vec<f64>>,
    signs: Option<Vec<f64>>,
) -> Result<HashMap<String, HashMap<String, Vec<f64>>>, Box<dyn std::error::Error>> {
    let _agg_xcats_for_cid = agg_xcats_for_cid(cids.clone(), xcats.clone());

    let (agg_on, agg_targ) = get_agg_on_agg_targs(cids.clone(), xcats.clone());

    // if weights are None, create a vector of 1s of the same length as agg_targ
    let weights = weights.unwrap_or(vec![1.0 / agg_targ.len() as f64; agg_targ.len()]);
    let signs = signs.unwrap_or(vec![1.0; agg_targ.len()]);

    // check that the lengths of weights and signs match the length of agg_targ
    check_weights_signs_lengths(
        weights.clone(),
        signs.clone(),
        _agg_xcats_for_cid,
        agg_targ.len(),
    )?;

    let mut weights_map = HashMap::new();

    for agg_o in agg_on {
        let mut agg_t_map = HashMap::new();
        for (i, agg_t) in agg_targ.iter().enumerate() {
            let ticker = match _agg_xcats_for_cid {
                true => format!("{}_{}", agg_o, agg_t),
                false => format!("{}_{}", agg_t, agg_o),
            };
            let weight_signs = vec![weights[i], signs[i]];
            agg_t_map.insert(ticker, weight_signs);
        }
        weights_map.insert(agg_o.clone(), agg_t_map);
    }
    Ok(weights_map)
}

fn check_weights_signs_lengths(
    weights_vec: Vec<f64>,
    signs_vec: Vec<f64>,
    _agg_xcats_for_cid: bool,
    agg_targ_len: usize,
) -> Result<(), Box<dyn std::error::Error>> {
    // for vx, vname in ...
    let agg_targ = match _agg_xcats_for_cid {
        true => "xcats",
        false => "cids",
    };
    for (vx, vname) in vec![
        (weights_vec.clone(), "weights"),
        (signs_vec.clone(), "signs"),
    ] {
        for (i, v) in vx.iter().enumerate() {
            if *v == 0.0 {
                return Err(format!("The {} at index {} is 0.0", vname, i).into());
            }
        }
        if vx.len() != agg_targ_len {
            return Err(format!(
                "The length of {} ({}) does not match the length of {} ({})",
                vname,
                vx.len(),
                agg_targ,
                agg_targ_len
            )
            .into());
        }
    }
    Ok(())
}
fn rename_result_dfw_cols(
    xcats: Vec<String>,
    cids: Vec<String>,
    new_xcat: String,
    new_cid: String,
    dfw: &mut DataFrame,
) -> Result<(), Box<dyn std::error::Error>> {
    let dfw_cols: Vec<String> = dfw
        .get_column_names()
        .iter()
        .filter(|&s| *s != "real_date") // Exclude 'real_date'
        .map(|s| s.to_string())
        .collect();
    let _agg_xcats_for_cid = agg_xcats_for_cid(cids.clone(), xcats.clone());
    Ok(for col in dfw_cols {
        let new_name = match _agg_xcats_for_cid {
            true => format!("{}_{}", col, new_xcat),
            false => format!("{}_{}", new_cid, col),
        };
        // rename the column
        dfw.rename(&col, new_name.into())?;
    })
}

/// Flags if the xcats are aggregated for a given cid.
/// If true, the xcats are aggregated for each cid, creating a new xcat.
/// If false, the cids are aggregated for each xcat, creating a new cid.
#[allow(unused_variables)]
fn agg_xcats_for_cid(cids: Vec<String>, xcats: Vec<String>) -> bool {
    // if there is more than 1 xcat, return xcats.len() > 1
    xcats.len() > 1
}

/// Weighted linear combinations of cross sections or categories
/// # Arguments
/// * `df` - QDF DataFrame
/// * `xcats` - List of category names or a single category name
/// * `cids` - List of cross section names or None
/// * `weights` - List of weights or a string indicating a weight `xcat`
/// * `normalize_weights` - Normalize weights to sum to 1 before applying
/// * `signs` - List of signs for each category (+1 or -1)
/// * `start` - Start date for the analysis
/// * `end` - End date for the analysis
/// * `blacklist` - Dictionary of blacklisted categories
/// * `complete_xcats` - If True, complete xcats with missing values
/// * `complete_cids` - If True, complete cids with missing values
/// * `new_xcat` - Name of the new xcat
/// * `new_cid` - Name of the new cid
///
/// # Returns
/// * `DataFrame` - DataFrame with the linear composite
pub fn linear_composite(
    df: &DataFrame,
    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<HashMap<String, Vec<String>>>,
    complete_xcats: bool,
    complete_cids: bool,
    new_xcat: Option<String>,
    new_cid: Option<String>,
) -> Result<DataFrame, Box<dyn std::error::Error>> {
    // Check if the DataFrame is a Quantamental DataFrame
    check_quantamental_dataframe(df)?;
    let rdf = reduce_dataframe(
        df.clone(),
        Some(cids.clone()),
        Some(xcats.clone()),
        Some(vec!["value".to_string()]),
        start.clone(),
        end.clone(),
        false,
    )
    .unwrap();

    let new_xcat = new_xcat.unwrap_or_else(|| "COMPOSITE".to_string());
    let new_cid = new_cid.unwrap_or_else(|| "GLB".to_string());

    let dfw = pivot_dataframe_by_ticker(rdf.clone(), Some("value".to_string())).unwrap();

    let mul_targets = get_mul_targets(cids.clone(), xcats.clone())?;
    let weights_map = form_weights_and_signs_map(cids.clone(), xcats.clone(), weights, signs)?;

    for (ticker, targets) in mul_targets.iter() {
        println!("ticker: {}, targets: {:?}", ticker, targets);
    }
    for (agg_on, agg_t_map) in weights_map.iter() {
        println!("agg_on: {}, agg_t_map: {:?}", agg_on, agg_t_map);
    }
    // assert!(0==1, "Debugging weights_map: {:?}", weights_map);
    // Perform the multiplication
    let complete = complete_xcats || complete_cids;
    let mut dfw = perform_multiplication(
        &dfw,
        &mul_targets,
        &weights_map,
        complete,
        normalize_weights,
    )?;

    rename_result_dfw_cols(xcats, cids, new_xcat, new_cid, &mut dfw)?;

    let dfw = pivot_wide_dataframe_to_qdf(dfw, Some("value".to_string()))?;
    Ok(dfw)
}
