import argparse
from pathlib import Path
from datetime import datetime
from typing import List
from core.Logger import logger
from core.DataLoader import DataLoader
from core.FeatureProcessor import FeatureProcessor
from MLModel.Models import AutogluonModel
import pandas as pd
from core.Settings import (
    MODEL_DIRECTORY,
    REGISTRY_PATH
)

TIMESTAMP = datetime.now().strftime(format="%Y-%m-%d_%H-%M-%S")

# Can move to a hyperparmeters file
# If anything we might want to have a file that can be loaded and sent to this script
HYPERPARAMETERS = {
    'problem_type': 'regression',
    'eval_metric': 'mean_absolute_error', 
    'time_limit': 60, 
    'presets': 'medium_quality', 
    'excluded_model_types': None
}

# FOR TESTING
train_filepath = "./model_build_data/change_data/rdsap_full/train_validation_data.parquet"
test_filepath = "./model_build_data/change_data/rdsap_full/test_data.parquet"
target_column = "RDSAP_CHANGE"
model_type = "autogluon"
hyperparameter = HYPERPARAMETERS
subsample_factor = 200


def ingest_arguments() -> argparse.Namespace:
    """
    Helper function to take in arguments from script start
    """

    parser = argparse.ArgumentParser(description='Inputs for training script')

    parser.add_argument('--train-filepath', type=str, help='Location of Parquet dataset to load for training', required=True)
    parser.add_argument('--test-filepath', type=str, help='Location of Parquet dataset to load for testing', required=True)
    parser.add_argument('--model-type', type=str, help='The type of model to train', choices=["autogluon"], default="autogluon")
    parser.add_argument('--target-column', type=str, help='The response variable', choices=["RDSAP_CHANGE"], default='RDSAP_CHANGE')

    args = parser.parse_args()

    return args
            

def training(
        train_filepath: str, 
        test_filepath: str, 
        target_column: str = "RDSAP_CHANGE", 
        model_type: str = "autogluon", 
        hyperparameter: dict = HYPERPARAMETERS
        ) -> None:
    """
    Pipeline to run training on the dataset
    """

    logger.info('--- Loading data ---')
    dataloader = DataLoader()
    train_df = dataloader.load(filepath=train_filepath)
    test_df = dataloader.load(filepath=test_filepath)
 
    logger.info('--- Feature processing ---')

    feature_processor = FeatureProcessor()

    subsample_amount = round(len(train_df)/subsample_factor)

    train_df = feature_processor.process(train_df, target_column=target_column, subsample_amount=subsample_amount)
    test_df = feature_processor.process(test_df, target_column=target_column)

    logger.info('--- Build Model ---')
    if model_type == "autogluon":
        model_root = f"{target_column}-{HYPERPARAMETERS['presets']}-{HYPERPARAMETERS['time_limit']}-{TIMESTAMP}".lower()
        output_base = Path(MODEL_DIRECTORY) / model_type / model_root 

        model_folder = "model"
        metrics_folder = "metrics"

        model = AutogluonModel(
            output_filepath = output_base / model_folder
            )
    else:
        logger.error("No alternative model implemented yet")
        exit(1)
    
    model.train_model(
        data=train_df, 
        target_column=target_column, 
        hyperparameters=hyperparameter
        )
    
    logger.info("--- Save Model ---")
    model.save_model(output_filepath=model.output_filepath)

    logger.info('--- Generate evaluation metrics ---')
    metrics_df = model.model_evaluation(
        validation_data=test_df, 
        target_column=target_column,
        metrics_location = output_base / metrics_folder
        )
    
    # TODO: introduce a seperate script for model optimisation, and from there, optimise for deployment
    # Imagining for now that the model trained here is the best model amongst all models built

    logger.info("--- Optimising model for deployment ---")
    optimised_folder = "deployment"        
    deployment_model_path = model.optimise_model_for_deployment(deployment_path= output_base / optimised_folder)
    logger.info("Optimised version of best model can be found at: {deployment_model_path}")

    # TODO: Need a model registry - for now have this as a CSV
    # Save this in the model directory
    logger.info("--- Append registry with new model ---")
    
    if REGISTRY_PATH.exists():
        logger.info("Registry file found - Loading into Dataframe")
        registry_df = pd.read_csv(REGISTRY_PATH, index_col=None)
    else:
        registry_df = pd.DataFrame(columns=['model_type', 'model_name', 'model_location', 'mean_absolute_error', 'root_mean_squared_error', 'mean_squared_error', 'r2', 'pearsonr', 'median_absolute_error', 'mape', 'best_model'])

    model_details_df = pd.DataFrame(
        [{
            'model_type': model_type, 
            'model_name': model_root, 
            'model_location': deployment_model_path
        }]
        )
    
    registry_row = pd.concat([model_details_df, metrics_df], axis=1)
    registry_df = pd.concat([registry_df, registry_row], axis=0).reset_index(drop=True)

    # TODO: will need a rebuild script metric script -i.e. if we add new metrics, we will want to load models and regenerate new metrics
    # TODO: decide metric to optimise to
    registry_df = registry_df.sort_values("mean_absolute_error", ascending=False).reset_index(drop=True)
    registry_df['best_model'] = [False]*len(registry_df)
    registry_df.loc[0, 'best_model'] = True

    logger.info("--- Saving new model to registry ---")
    registry_df.to_csv(REGISTRY_PATH, index=False)

    logger.info("--- Training Pipeline Complete --- ")


if __name__ == "__main__":

    logger.info('---Begin Pipeline---')

    logger.info('---Ingest Arguments---')
    args = ingest_arguments()

    # TODO: Ingest hyper parameters from somewhere - currently change at the top of script
    training(
        train_filepath=args.train_filepath, 
        test_filepath=args.test_filepath, 
        target_column=args.target_column, 
        model_type=args.model_type
        )