from utils.load_data import load_data
import pandas as pd
from training.training import run_single_asset_trainig
from utils.launch_wandb import launch_wandb, seperate_configs
from models.model_map import map_model_name_to_function
from default_config import get_default_config

def setup_pipeline(project_name:str, with_wandb: bool, sweep: bool):
    model_config, training_config, data_config = get_default_config()
    
    wandb = None
    if with_wandb: 
        wandb = launch_wandb(project_name=project_name, default_config=dict(**model_config, **training_config, **data_config), sweep=sweep)
        model_config, training_config, data_config = seperate_configs(wandb, model_config, training_config, data_config)

    model_config = map_model_name_to_function(model_config, data_config['method'])
    pipeline(project_name, wandb, sweep, model_config, training_config, data_config)  
    

def pipeline(project_name:str, wandb, sweep:bool, model_config:dict, training_config:dict, data_config:dict ):
    results = pd.DataFrame()

    for asset in data_config['all_assets']:
        print('--------\nPredicting: ', asset)
        all_predictions = pd.DataFrame()

        # 1. Load data
        data_params = data_config.copy()
        data_params['target_asset'] = asset

        X, y, target_returns = load_data(**data_params)

        # 2. Train Level-1 models
        current_result, current_predictions = run_single_asset_trainig(
            ticker_to_predict = asset,
            X = X,
            y = y,
            target_returns = target_returns,
            models = model_config['level_1_models'],
            method = data_config['method'],
            sliding_window_size = training_config['sliding_window_size'],
            retrain_every =  training_config['retrain_every'],
            scaler =  training_config['scaler'],
            wandb = wandb,
            project_name=project_name,
            sweep=sweep
        )
        results = pd.concat([results, current_result], axis=1)
        all_predictions = pd.concat([all_predictions, current_predictions], axis=1)

        if len(model_config['level_2_models']) > 0: 
            # 3. Train Level-2 (Ensemble) model
            
            ensemble_X = all_predictions
            if training_config['include_original_data_in_ensemble']:
                ensemble_X = pd.concat([ensemble_X, X], axis=1)

            ensemble_result, ensemble_preds = run_single_asset_trainig(
                ticker_to_predict = asset,
                X = ensemble_X,
                y = y,
                target_returns = target_returns,
                models = model_config['level_2_models'],
                method = data_config['method'],
                sliding_window_size = training_config['sliding_window_size'],
                retrain_every = training_config['retrain_every'],
                scaler = training_config['scaler'],
                wandb = wandb,
                project_name=project_name,
                sweep=sweep
            )

            results = pd.concat([results, ensemble_result], axis=1)
            all_predictions = pd.concat([all_predictions, ensemble_preds], axis=1)

    results.to_csv('results.csv')

    level1_columns = results[[column for column in results.columns if 'Ensemble' not in column]]
    ensemble_columns = results[[column for column in results.columns if 'Ensemble' in column]]

    print("Mean Sharpe ratio for Level-1 models: ", level1_columns.loc['sharpe'].mean())
    print("Mean Sharpe ratio for Level-2 (Ensemble) models: ", ensemble_columns.loc['sharpe'].mean())

    if sweep:
        if wandb.run is not None:
            wandb.finish()
    
if __name__ == '__main__':
    setup_pipeline(project_name='price-prediction', with_wandb = False, sweep = False)