import numpy as np
from pandas import DataFrame
from temporalio import activity, workflow


with workflow.unsafe.imports_passed_through():
    from laborious.activities.base import BaseActivity
    from laborious.utils.repository.model_repository import MLFlowRepository
    from typing import Any
    from logging import Logger
    from sientia_do.notifications.handlers import NotificationHandler


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.
        """
        self.logger.info('Transforming data...')
        data = DataFrame(input_data['data'])
        model_name = input_data['model_name']
        model_retention = input_data['model_retention']

        self.logger.debug(data)

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

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

        self.logger.debug(response_data)

        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.
        """
        self.logger.info('Predicting data...')
        data = DataFrame(input_data['data'])
        model_name = input_data['model_name']
        model_retention = input_data['model_retention']

        self.logger.debug(data)

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

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

        self.logger.debug(response_data)

        return response_data