from temporalio import activity, workflow


with workflow.unsafe.imports_passed_through():
    from sientia_do.temporal.activities.base import BaseActivity
    from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
    from sientia_do.notifications.models import NotificationLevel
    from sientia_do.temporal.utils.logger import Logger
    from laborious.utils.repository.model_repository import MLFlowRepository
    from typing import Any
    import numpy as np
    from pandas import DataFrame
    import traceback


class MLFlow(BaseActivity):
    def __init__(self, mlflow_host: str, mlflow_port: int, mlflow_username: str,
                 mlflow_password: str, logger: Logger, notification_handler: NotificationHandler):
        BaseActivity.__init__(self, logger, notification_handler)
        self.mlflow_host = mlflow_host
        self.mlflow_port = mlflow_port
        self.mlflow_username = mlflow_username
        self.mlflow_password = mlflow_password

        self.model_monitoring_repository = MLFlowRepository(
            f"{mlflow_host}:{mlflow_port}", mlflow_username, mlflow_password
        )

    @activity.defn(name="request_transform")
    async def request_transform(self, input_data: dict[str, Any]) -> dict[str, Any]:
        """
        Access MLFlow model to get the transformed data.
        Args:
            - input_data (dict): The input data. Contains:
                - data (dict[str, Any]): The data to transform.
                - model_name (str): The name of the model.
                - model_retention (int): The retention time of the model, in minutes.
        Returns:
            dict[str, Any]: The transformed data.
        """
        metadata = input_data['metadata']
        self.debug('Transforming data...', metadata)
        data = DataFrame(input_data['data'])
        model_name = input_data['model_name']
        model_retention = input_data['model_retention']

        self.debug("Raw input data:", metadata)
        self.debug(data, metadata)

        # Sort by created_at in descending order and keep first occurrence of each variable/timestamp pair
        data = data.sort_values('created_at', ascending=False).drop_duplicates(
            subset=['variable', 'timestamp'], keep='first'
        )

        data = data.pivot(
            index='timestamp', columns='variable',
            values='value')
        data.fillna(np.nan, inplace=True)
        data.reset_index(inplace=True)
        data.columns.name = None

        self.debug("Processed input data:", metadata)
        self.debug(data, metadata)

        response_data = self.model_monitoring_repository.transform(
            model_name, data, model_retention)

        self.debug("Response data:", metadata)
        self.debug(response_data, metadata)

        return response_data

    @activity.defn(name="request_predict")
    async def request_predict(self, input_data: dict[str, Any]) -> dict[str, Any]:
        """
        Access MLFlow model to get the predicted data.
        Args:
            - input_data (dict): The input data. Contains:
                - data (dict[str, Any]): The data to predict.
                - model_name (str): The name of the model.
                - model_retention (int): The retention time of the model, in minutes.
        Returns:
            dict[str, Any]: The predicted data.
        """
        metadata = input_data['metadata']
        self.debug('Predicting data...', metadata)
        data = DataFrame(input_data['data'])
        model_name = input_data['model_name']
        model_retention = input_data['model_retention']

        self.debug(data, metadata)

        data.replace(np.nan, None, inplace=True)

        response_data = self.model_monitoring_repository.predict(
            model_name, data, model_retention)

        self.debug(response_data, metadata)

        return response_data

    @activity.defn(name="retrain_model")
    async def retrain_model(self, input_data: dict[str, Any]) -> dict[str, Any]:
        """
        Retrain the model.
        Args:
            - input_data (dict): The input data. Contains:
                - model_name (str): The name of the model.
                - data (dict[str, Any]): The data to retrain the model.
        """
        metadata = input_data['metadata']
        data = DataFrame(input_data['data'])
        model_name = input_data['model_name']

        self.info(f'Retraining model {model_name}...', metadata)

        timestamp = data['timestamp'].max()
        self.debug(f'Timestamp: {timestamp}', metadata)

        data.drop(columns=['model_id'], inplace=True, errors='ignore')
        data.drop(columns=['created_at'], inplace=True, errors='ignore')

        data = data.pivot(index='timestamp', columns='variable',
                          values='value')
        data.sort_index(inplace=True)
        data.reset_index(inplace=True)

        data = data.dropna()
        data.columns.name = None

        try:
            retrain_output, experiment = self.model_monitoring_repository.retrain_model(
                data=data,
                model_name=model_name
            )

            return {
                'status': retrain_output,
                'timestamp': timestamp,
                'experiment': experiment
            }
        except Exception as e:
            trace = traceback.format_exc()
            self.send_notification(
                metadata=metadata,
                notification_id='RETRAIN_MODEL_ERROR',
                message=f'Error retraining model {model_name}: {e}',
                block='retrain_model',
                level=NotificationLevel.ERROR,
                attachment_content=trace
            )
            self.error(trace, metadata=metadata)
            raise e

    @activity.defn(name="update_production_model")
    async def update_production_model(self, input_data: dict[str, Any]) -> dict[Any, Any]:
        """
        Update the production model.
        Args:
            - input_data (dict): The input data. Contains:
                - model_name (str): The name of the model.
                - experiment (str): The name of the experiment.
                - model_id (str): The id of the model.
                - timestamp (str): The timestamp of the model.
                - status (str): The status of the model.
        Returns:
            dict[Any, Any]: The report of the model.
        """
        metadata = input_data['metadata']
        model_name = input_data['model_name']
        model_id = input_data['model_id']
        experiment = input_data['experiment']
        timestamp = input_data['timestamp']
        status = input_data['status']

        self.info(
            f'Updating production model {model_name} from experiment {experiment}...', metadata)

        try:
            response = self.model_monitoring_repository.update_production_model(
                experiment=experiment,
                model_name=model_name
            )

            report = DataFrame([response])
            report['model_id'] = model_id
            report['model_name'] = model_name
            report['timestamp'] = timestamp
            report['status'] = status

            self.info(
                f'Production model {model_name} updated successfully', metadata)
            return report.to_dict()

        except Exception as e:
            trace = traceback.format_exc()
            self.send_notification(
                metadata=metadata,
                notification_id='UPDATE_PRODUCTION_MODEL_ERROR',
                message=f'Error updating production model {model_name}: {e}',
                block='update_production_model',
                level=NotificationLevel.ERROR,
                attachment_content=trace
            )
            self.error(trace, metadata=metadata)
            raise e