msyrs/src/utils/qdf/pivots.rs

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"];

/// Pivots a dataframe to a format where each ticker a column.
#[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 new_df = df.clone();
    // 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());

    return Err("Not implemented".into());

    // new_df = new_df.select(keep_cols)?;
    // let ticker_col = get_ticker_column_for_quantamental_dataframe(&new_df)?;

    // new_df.with_column(ticker_col)?;
    // // drop the cid and xcat columns
    // new_df = new_df.drop_many(&["cid".to_string(), "xcat".to_string()]);

    // let dates_col = df.column("real_date")?;

    // Ok(df)
}

/// Splits a dataframe by ticker.
#[allow(dead_code)]
fn split_df_by_tickers(df: &DataFrame) -> 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);
    }

    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()?;
        df_outs.insert(format!("{}_{}", cid, xcat), df_out);
    }

    Ok(df_outs)
}