import os

import mlflow
import numpy as np
from loguru import logger
import re
import pandas as pd

from scripts.constants.app_constants import Constants
from scripts.core.model_factory.helper import training_data_generator
from scripts.core.model_factory.pycaret_util import GetBestModel
import pytz
from dateutil import tz
from datetime import datetime
from scripts.constants.app_configuration import REQUIRED_TZ, MlflowMetaData, PycaretMetaData, PARAMETER_INPUT, \
    PARAMETER_OUTPUT

mlflow_tracking_uri = MlflowMetaData.MLFLOW_TRACKING_URI
AZURE_STORAGE_CONNECTION_STRING = MlflowMetaData.AZURE_STORAGE_CONNECTION_STRING
AZURE_STORAGE_ACCESS_KEY = MlflowMetaData.AZURE_STORAGE_ACCESS_KEY

os.environ["MLFLOW_TRACKING_USERNAME"] = MlflowMetaData.MLFLOW_TRACKING_USERNAME
os.environ["MLFLOW_TRACKING_PASSWORD"] = MlflowMetaData.MLFLOW_TRACKING_PASSWORD
os.environ["AZURE_STORAGE_CONNECTION_STRING"] = AZURE_STORAGE_CONNECTION_STRING
os.environ["AZURE_STORAGE_ACCESS_KEY"] = AZURE_STORAGE_ACCESS_KEY
mlflow.set_tracking_uri(mlflow_tracking_uri)
mlflow.set_registry_uri(mlflow_tracking_uri)
client = mlflow.tracking.MlflowClient()


class MlFlowUtil:
    @staticmethod
    def get_last_run_time_diff(run_info, date_param):
        try:
            logger.info(f"Checking the time difference in {date_param}")
            df_time = run_info.copy()
            df_time['end_time'] = pd.to_datetime(df_time['end_time']).dt.tz_convert(REQUIRED_TZ)
            to_zone = tz.gettz(REQUIRED_TZ)
            df_time["days"] = df_time['end_time'].dt.date
            df_time["hours"] = df_time['end_time'].dt.hour
            last_model_time = list(df_time['end_time'])[0].to_pydatetime()
            today = datetime.now(pytz.utc)
            central_current = today.astimezone(to_zone)
            time_diff = central_current - last_model_time
            if date_param.lower() == "days":
                return int(time_diff.days)
            elif date_param.lower() == "hours":
                return int(time_diff.total_seconds() // 3600)
            elif date_param.lower() == "minutes":
                return int(time_diff.total_seconds() // 60)
            else:
                logger.info("No Valid Date format was given")
                return 0

        except Exception as e:
            logger.warning(f"Exception while checking the last run time of the model - {e}")
            return 0

    @staticmethod
    def log_model(model, model_name):
        try:
            mlflow.sklearn.log_model(model, model_name)
            logger.info("logged the model")
            return True
        except Exception as e:
            logger.exception(str(e))

    @staticmethod
    def log_metrics(metrics):
        try:
            updated_metric = {}
            for key, value in metrics.items():
                key = re.sub(r"[([{})\]]", "", key)
                updated_metric[key] = value
            mlflow.log_metrics(updated_metric)
            return True
        except Exception as e:
            logger.exception(str(e))

    @staticmethod
    def log_hyper_param(hyper_params):
        try:
            mlflow.log_params(hyper_params)
            return True
        except Exception as e:
            logger.exception(str(e))

    @staticmethod
    def set_tag(child_run_id, key, value):
        try:
            client.set_tag(run_id=child_run_id, key=key, value=value)
        except Exception as e:
            logger.exception(f"Exception while setting the tag - {e}")

    @staticmethod
    def remove_file_if_exists(path):
        if os.path.exists(path):
            os.remove(path)

    def log_artifact(self, df):
        try:
            self.remove_file_if_exists(Constants.PROFILE_FILE_NAME)
            df.to_csv(Constants.PROFILE_FILE_NAME,index=False)
            logger.info(f"Logging the Artifact: {Constants.PROFILE_FILE_NAME}")
            mlflow.log_artifact(Constants.PROFILE_FILE_NAME)
            self.remove_file_if_exists(Constants.PROFILE_FILE_NAME)
            return True
        except Exception as e:
            logger.exception(str(e))


class ModelReTrainer:
    def __init__(self, experiment_name, parent_run_name, model_save_name):
        self.experiment_name = experiment_name
        self.parent_run_name = parent_run_name
        self.model_save_name = model_save_name
        self._mfu_ = MlFlowUtil()
        self.list_of_models = PycaretMetaData.LIST_OF_MODELS.split(',')
        self.fine_tune_tech = PycaretMetaData.SEARCH_LIBRARY
        self.comparison_metric = PycaretMetaData.COMPARISON_METRIC
        self.retrain_param_unit = MlflowMetaData.RETRAIN_PARAM_UNIT
        self.retrain_param_value = int(MlflowMetaData.RETRAIN_PARAM_VALUE)

    def check_create_experiment(self):
        """
        check if experiment exists, if not creates a new experiment
        :return: experiment_id of the experiment
        """
        experiment_info = mlflow.get_experiment_by_name(self.experiment_name)
        if experiment_info is None:
            logger.info(f"No experiment found with name {self.experiment_name}, So creating one")
            mlflow.create_experiment(self.experiment_name)
        else:
            logger.info(f"Proceeding with existing Experiment {self.experiment_name}")
        mlflow.set_experiment(experiment_name=self.experiment_name)
        experiment_info = mlflow.get_experiment_by_name(self.experiment_name)
        experiment_id = experiment_info.experiment_id
        return experiment_id

    def check_create_parent_run(self, experiment_id):
        """
        check if a parent run exists in the experiment, if not create it with the mentioned parent run name
        :param experiment_id: Experiment id
        :return: returns the parent run id
        """
        parent_runs_df = mlflow.search_runs(experiment_id)
        run_key = 'tags.mlflow.runName'
        if run_key in parent_runs_df.columns:
            parent_runs_df = parent_runs_df[parent_runs_df[run_key] == self.parent_run_name]
        else:
            parent_runs_df = parent_runs_df.iloc[:0]
        if not parent_runs_df.empty:
            logger.info(f"Proceeding with existing Parent Run {self.parent_run_name}")
            return list(parent_runs_df['run_id'])[0]
        # no parent run found
        logger.info(f"No Parent Run present {self.parent_run_name}")
        with mlflow.start_run(experiment_id=experiment_id, run_name=self.parent_run_name) as run:
            logger.info(f"Creating the parent Run {self.parent_run_name} with Parent Run Id {run.info.run_id}")
            return run.info.run_id

    def check_create_child_run(self, experiment_id, parent_run_id):
        """
        check if a child run exists in the experiment id under the parent run id
        if exists take the child run id which has the model saved and validate when was it lastly trained.
        Based on the lastly trained see if you have to retrain or not. if retrain create a new child run
        else if no child run exists under the parent run id of experiment id, create a new child run
        :param experiment_id: experiment id
        :param parent_run_id: parent run id
        :return: child run id, retrain flag
        """
        child_runs_df = mlflow.search_runs(experiment_id, filter_string=f"tags.mlflow.parentRunId='{parent_run_id}'")
        if not child_runs_df.empty:
            logger.info(f"Already Child runs are present for Parent Run Id {parent_run_id}")
            child_run_id, retrain = self.get_latest_child_run(experiment_id, parent_run_id, child_runs_df)
            return child_run_id, retrain
        else:
            logger.info(f"Child runs are not present for Parent Run Id {parent_run_id}")
            with mlflow.start_run(experiment_id=experiment_id, run_id=parent_run_id, nested=True):
                with mlflow.start_run(experiment_id=experiment_id, nested=True) as child_run:
                    return child_run.info.run_id, True

    def get_latest_child_run(self, experiment_id, parent_run_id, runs_df):
        """
        Check if child runs are present. if not create a new child run. Otherwise, validate the last run time and
        create a new child run if retraining needed or take the last child run id which has model saved
        :param experiment_id: experiment id
        :param parent_run_id: parent run id
        :param runs_df: the child runs of the parent id
        :return: last child run id, retrain flag
        """
        history_key = 'tags.mlflow.log-model.history'
        if history_key in runs_df.columns:
            runs_df = runs_df[runs_df[history_key].notna()]
        else:
            runs_df = runs_df.iloc[:0]
        if runs_df.empty:
            logger.info("Existing Child Runs doesn't contain any model to run. So creating new child run")
            with mlflow.start_run(experiment_id=experiment_id, run_id=parent_run_id, nested=True):
                with mlflow.start_run(experiment_id=experiment_id, nested=True) as child_run:
                    return child_run.info.run_id, True
        latest_child_run_id = list(runs_df['run_id'])[0]
        latest_run_info = runs_df.iloc[:1]
        retrain = False
        day_check_flag = self.check_existing_model_retrain(latest_child_run_id, latest_run_info, retrain)
        if day_check_flag:
            with mlflow.start_run(experiment_id=experiment_id, run_id=parent_run_id, nested=True):
                with mlflow.start_run(experiment_id=experiment_id, nested=True) as child_run:
                    return child_run.info.run_id, True
        return latest_child_run_id, retrain

    @staticmethod
    def load_model_pyfunc(model_path):
        """
        Function to load the model from mlflow artifact path
        :param model_path: model path on mlflow
        :return: loaded model
        """
        try:
            model = mlflow.pyfunc.load_model(model_path)
            logger.info("loading the model")
            return model
        except Exception as e:
            logger.exception(str(e))

    def check_existing_model_retrain(self, latest_child_run_id, child_run_info, retrain):
        """
        If retrain is True, it returns true as retraining is required.
        If retrain is False, it checks the time difference between the last child run and the current time and returns
        true or false depending on the time difference
        :param latest_child_run_id: last child run id
        :param child_run_info: last child run info
        :param retrain: retrain flag
        :return: final retrain flag
        """
        if retrain:
            logger.info("Retraining Needed...")
            return True
        else:
            logger.info(f"Already trained model is present, checking the age of the existing model of run id "
                        f"{latest_child_run_id}")
            time_diff = self._mfu_.get_last_run_time_diff(child_run_info, self.retrain_param_unit)
            logger.info(f"Time difference is {time_diff} {self.retrain_param_unit}")
            if time_diff >= self.retrain_param_value:
                logger.info(f"Retraining needed as the last trained model time exceeds the mentioned time difference "
                            f"{self.retrain_param_value} {self.retrain_param_unit}")
                return True
            else:
                logger.info(f"Retraining not needed as the last trained model time doesnt exceeds the mentioned time "
                            f"difference {self.retrain_param_value} {self.retrain_param_unit}")
                return False

    def forming_loading_path(self, latest_run_id):
        """
        Creates the path from the child run id
        :param latest_run_id: latest child run id
        :return: the path to the model
        """
        try:
            model_name = self.model_save_name
            return f"runs:/{latest_run_id}/{model_name}"
        except Exception as e:
            logger.exception(f"Exception while forming loading path  - {e}")

    @staticmethod
    def get_train_data():
        df = training_data_generator()
        # df = pd.read_csv(r"D:\Projects-2024\jan\jub-r5-new-beta-ratio\r5-reaction-golden-batch\train-data.csv")
        raw_df = df.copy(deep=True)
        
        df["timestamp"] = pd.to_datetime(df["timestamp"])  

        df["Alcohol flow"] = df["Totalizer Alcohol"].diff()
        df["Alcohol flow"].fillna(value=0, inplace=True)

        critical_params = ['Alcohol flow', 'Column C2 Bottom Temperature','Carburetor C4 temperature PV', 
            'Reactor 6A Vapour temperature','Reactor 6A Inlet temperature', 'Reactor 6A Outlet temperature',
            'Reactor 6B Vapour temperature', 'Reactor 6B Inlet temperature','Reactor 6B Outlet temperature',
            'Reactor 6C Vapour temperature',  'Reactor 6C Inlet temperature', 'Reactor 6C Outlet temperature',
            'Reactor 6D Vapour temperature', 'Reactor 6D Inlet temperature', 'Reactor 6D Outlet temperature',
            'Column_C9_Top temperature', 'Column_C9A_Top temperature','Product Output flow to Tank']

        df_temp= df[critical_params]
        df = df_temp[(df_temp["Alcohol flow"]>500)&(df_temp["Product Output flow to Tank"]>1500)]
       
        total_cols = list(PARAMETER_INPUT)
        total_cols.extend(iter(PARAMETER_OUTPUT))
        df = df[total_cols]
        raw_cols = list(raw_df.columns)
        raw_cols.remove('timestamp')
        raw_cols.remove('epochtime')
        raw_df = raw_df[raw_cols]
        return df, raw_df

    def model_trainer(self, experiment_id, parent_run_id, child_run_id):
        """
        Using the experiment id, parent run id and child run id, it will train the model
        :param experiment_id: experiment id
        :param parent_run_id: parent run id
        :param child_run_id: child run id
        :return: the final model
        """
        with mlflow.start_run(experiment_id=experiment_id, run_id=parent_run_id, nested=True):
            with mlflow.start_run(experiment_id=experiment_id, run_id=child_run_id, nested=True):
                train_df, raw_df = self.get_train_data()

                _gbm_ = GetBestModel(train_df, raw_df, PARAMETER_OUTPUT[0], self.list_of_models)
                model, model_name, metrics, hyperparams, profile_df = _gbm_. \
                    compare_get_best_model(self.fine_tune_tech, self.comparison_metric)
                self._mfu_.log_model(model=model, model_name=self.model_save_name)
                self._mfu_.log_metrics(metrics=metrics)
                self._mfu_.log_hyper_param(hyper_params=hyperparams)
                self._mfu_.set_tag(child_run_id=child_run_id, key="algorithm", value=model_name)
                self._mfu_.log_artifact(profile_df)
        return model

    def get_latest_model(self):
        """
        This is the Main function which will return the latest model
        :return:
        """
        experiment_id = self.check_create_experiment()
        parent_run_id = self.check_create_parent_run(experiment_id)
        child_run_id, retrain = self.check_create_child_run(experiment_id, parent_run_id)
        logger.info(f"Retrain flag is {retrain}")
        if retrain:
            logger.info("Retraining needed")
            self.model_trainer(experiment_id, parent_run_id, child_run_id)
            logger.info("New model trained successfully")
        else:
            logger.info(f"No retraining needed. proceeding to load the last child run model {child_run_id}")
        logger.info(f"Loading the model from the child run id {child_run_id}")
        final_model = self.load_model_pyfunc(model_path=self.forming_loading_path(latest_run_id=child_run_id))
        logger.info(f"Loading Profile from the run-id {child_run_id}")
        profile_df = pd.read_csv(client.download_artifacts(child_run_id, Constants.PROFILE_FILE_NAME))
        return final_model, profile_df
