use msyrs::download::jpmaqsdownload::{JPMaQSDownload, JPMaQSDownloadGetIndicatorArgs};
use polars::prelude::*;

fn main() {
    println!("Authentication to DataQuery API");
    let mut jpamqs_download = JPMaQSDownload::default();

    println!("Checking connection to DataQuery API");
    jpamqs_download.check_connection().unwrap();

    println!("Retrieving catalogue of tickers");
    let mut start = std::time::Instant::now();
    let tickers = jpamqs_download.get_catalogue().unwrap();
    println!(
        "Retrieved catalogue with {} tickers in {:?}",
        tickers.len(),
        start.elapsed()
    );

    let num_ticks = 20;
    let sel_tickers: Vec<String> = tickers
        .iter()
        .take(num_ticks)
        .map(|s| s.to_string())
        .collect();
    // let mut df_deets = Vec::new();

    println!("Retrieving indicators for {} tickers", sel_tickers.len());
    start = std::time::Instant::now();
    let mut res_df: DataFrame = jpamqs_download
        .get_indicators_qdf(JPMaQSDownloadGetIndicatorArgs {
            tickers: sel_tickers.clone(),
            start_date: "2024-11-05".to_string(),
            ..Default::default()
        })
        .unwrap();

    println!(
        "Retrieved indicators for {} tickers in {:?}",
        sel_tickers.len(),
        start.elapsed()
    );

    // append _ to every cid
    let cid_vec: Vec<String> = res_df
        .column("cid")
        .unwrap()
        .str()
        .unwrap()
        .into_iter()
        .map(|s| s.unwrap_or("").to_string())
        .collect();
    let xcat_vec: Vec<String> = res_df
        .column("xcat")
        .unwrap()
        .str()
        .unwrap()
        .into_iter()
        .map(|s| s.unwrap_or("").to_string())
        .collect();

    let mut tickers_set = std::collections::HashSet::new();
    for (cid, xcat) in cid_vec.iter().zip(xcat_vec.iter()) {
        tickers_set.insert(cid.to_string() + "_" + xcat);
    }

    // save this df to disk
    let file_name = "data/jpmaqs_indicators_qdf.csv";
    let file = std::fs::File::create(file_name).unwrap();
    let mut csv_writer = CsvWriter::new(file);
    csv_writer.finish(&mut res_df).unwrap();

    // print len of tickers_set
    println!("Unique tickers: {}", tickers_set.len());
    println!("DataFrame shape: {:?}", res_df.shape());
}