from pathlib import Path

import numpy as np
import pandas as pd
import yaml

from sklearn.preprocessing import LabelEncoder

from src.logger import setup_logger
from src.utils import yaml_config_loader

logger = setup_logger(__name__)


class AlarmDataProcessor:
    def __init__(
        self,
        config_path: str = "configs/config.yaml",
    ) -> None:
        self.config = self.load_config(config_path)
        self.data_paths = self.config["data_paths"]
        self.raw_data_path = self.data_paths["raw_data_path"]
        self.interim_data_path = Path(self.data_paths["interim_data_path"])
        self.external_data_path = Path(self.data_paths["external_data_path"])
        self.threshold = self.config["preprocessing"]["threshold"]
        self.label_encoders: dict = {}

    @staticmethod
    def load_config(config_path: str) -> dict:
        return yaml_config_loader(config_path)

    @staticmethod
    def load_data(file_path: str) -> pd.DataFrame:
        try:
            df = pd.read_csv(file_path)
            logger.info(f"> Data loaded successfully from {file_path}")
            return df
        except Exception as e:
            logger.error(f"> Error loading data: {e}")
            raise

    @staticmethod
    def convert_to_datetime(df: pd.DataFrame) -> pd.DataFrame:
        date_columns = [
            "First Occurence",
            "Last Occurence",
            "Event Time",
            "Clearance Time",
            "Acknoledgement Time",
            "FM Receive Time",
        ]

        date_format = "%m/%d/%Y %I:%M:%S %p"

        for col in date_columns:
            try:
                df[col] = pd.to_datetime(df[col].replace("", pd.NaT), format=date_format, errors="coerce")
            except KeyError:
                logger.warning(f"Column {col} not found in the data...")
                df[col] = np.nan
            except Exception as e:
                logger.error(f"Error converting column {col} to datetime: {e}")
                raise

        # Adding additional time-related features
        df["first_occurrence_minute"] = df["First Occurence"].dt.minute
        df["first_occurrence_hour"] = df["First Occurence"].dt.hour
        df["first_occurrence_day_of_week"] = df["First Occurence"].dt.dayofweek
        df["first_occurrence_month"] = df["First Occurence"].dt.month

        df["last_occurrence_minute"] = df["Last Occurence"].dt.minute
        df["last_occurrence_hour"] = df["Last Occurence"].dt.hour
        df["last_occurrence_day_of_week"] = df["Last Occurence"].dt.dayofweek
        df["last_occurrence_month"] = df["Last Occurence"].dt.month

        logger.info(" >> Datetime conversion and date features engineering done")

        return df

    @staticmethod
    def drop_unnecessary_columns(df: pd.DataFrame) -> pd.DataFrame:
        columns_to_drop = [
            "Status",
            "Cell Name",
            "Ticket Id",
            "Vendor",
            "Event Time",
            "Clearance User",
            "Acknoledgement User",
            "Region",
            "Zone",
        ]
        df.drop(columns=columns_to_drop, inplace=True, errors="ignore")
        logger.info(" >> Unnecessary columns dropped")
        return df

    @staticmethod
    def handle_null_values(df: pd.DataFrame) -> pd.DataFrame:
        mask = df["Acknoledgement Time"].isnull() & df["Clearance Time"].isnull()
        df = df[~mask]
        df = df.dropna(subset=["Clearance Time"])
        logger.info(" >> Null values handled")
        return df

    @staticmethod
    def remove_outliers(df: pd.DataFrame) -> pd.DataFrame:
        df.drop(df[df["First Occurence"].dt.year > 2024].index, inplace=True)
        df.loc[:, "Clearance Duration"] = (df["Clearance Time"] - df["First Occurence"]).dt.total_seconds()
        df.drop(df[df["Clearance Duration"] < 0].index, axis=0, inplace=True)
        df_outliers = df[df["Clearance Duration"] >= 4.204676e06]
        df.drop(df_outliers.index, axis=0, inplace=True)
        logger.info(" >> Outliers removed")
        return df

    def create_target_column(self, df: pd.DataFrame) -> pd.DataFrame:
        df.loc[:, "Target"] = df.apply(
            lambda row: (
                "Irrelevant"
                if (row["Clearance Time"] - row["First Occurence"] <= pd.Timedelta(minutes=self.threshold))
                else "Relevant"
            ),
            axis=1,
        )
        logger.info(" >> Target column created")
        return df

    @staticmethod
    def convert_to_categorical(df: pd.DataFrame) -> pd.DataFrame:
        categorical_columns = [
            "Severity",
            "Technical ID",
            "Alarm Name",
            "Site Name",
            "Site ID",
            "NE Name",
            "Source",
            "NE Type",
            "Technology",
            "Alarm Category",
            "Parent Node",
            "Target",
        ]

        # convert derived datetime columns to categorical if needed
        derived_datetime_columns = [
            "first_occurrence_minute",
            "first_occurrence_hour",
            "first_occurrence_day_of_week",
            "first_occurrence_month",
            "last_occurrence_minute",
            "last_occurrence_hour",
            "last_occurrence_day_of_week",
            "last_occurrence_month",
        ]

        categorical_columns.extend(derived_datetime_columns)

        for col in categorical_columns:
            if col in df.columns:
                df[col] = df[col].astype("category")

        # if "Target" in df.columns:
        #     categorical_columns.append("Target")

        # df.loc[:, categorical_columns] = df[categorical_columns].astype("category")
        logger.info(" >> Categorical conversion done")
        return df

    def cyclical_encoding(self, df: pd.DataFrame) -> pd.DataFrame:
        df["first_occurrence_minute_sin"] = np.sin(2 * np.pi * df["first_occurrence_minute"] / 60)
        df["first_occurrence_minute_cos"] = np.cos(2 * np.pi * df["first_occurrence_minute"] / 60)
        df["first_occurrence_hour_sin"] = np.sin(2 * np.pi * df["first_occurrence_hour"] / 24)
        df["first_occurrence_hour_cos"] = np.cos(2 * np.pi * df["first_occurrence_hour"] / 24)
        df["first_occurrence_day_of_week_sin"] = np.sin(2 * np.pi * df["first_occurrence_day_of_week"] / 7)
        df["first_occurrence_day_of_week_cos"] = np.cos(2 * np.pi * df["first_occurrence_day_of_week"] / 7)
        df["first_occurrence_month_sin"] = np.sin(2 * np.pi * df["first_occurrence_month"] / 12)
        df["first_occurrence_month_cos"] = np.cos(2 * np.pi * df["first_occurrence_month"] / 12)

        df["last_occurrence_minute_sin"] = np.sin(2 * np.pi * df["last_occurrence_minute"] / 60)
        df["last_occurrence_minute_cos"] = np.cos(2 * np.pi * df["last_occurrence_minute"] / 60)
        df["last_occurrence_hour_sin"] = np.sin(2 * np.pi * df["last_occurrence_hour"] / 24)
        df["last_occurrence_hour_cos"] = np.cos(2 * np.pi * df["last_occurrence_hour"] / 24)
        df["last_occurrence_day_of_week_sin"] = np.sin(2 * np.pi * df["last_occurrence_day_of_week"] / 7)
        df["last_occurrence_day_of_week_cos"] = np.cos(2 * np.pi * df["last_occurrence_day_of_week"] / 7)
        df["last_occurrence_month_sin"] = np.sin(2 * np.pi * df["last_occurrence_month"] / 12)
        df["last_occurrence_month_cos"] = np.cos(2 * np.pi * df["last_occurrence_month"] / 12)

        logger.info(" >> Cyclical encoding done")
        return df

    def encode_labels(self, df: pd.DataFrame) -> pd.DataFrame:
        categorical_columns = df.select_dtypes(include=["category"]).columns

        encoders_info = {}
        for col in categorical_columns:
            le = LabelEncoder()
            df[col] = le.fit_transform(df[col])
            self.label_encoders[col] = le

            labels = le.classes_.tolist()
            encoded_labels = le.transform(le.classes_).tolist()

            zip_data = dict(zip(labels, encoded_labels))
            encoders_info[col] = zip_data

        self.save_encoded_labels(encoders_info, self.interim_data_path.parent / "label_encoders_info.yaml")

        logger.info(" >> Label encoding done")
        return df

    def save_encoded_labels(self, encoders_info: dict, file_path: Path) -> None:
        try:
            file_path.parent.mkdir(exist_ok=True, parents=True)
            with open(file_path, "w") as file:
                yaml.dump(encoders_info, file)
            logger.info(f"> Encoded labels saved successfully to {file_path}")
        except Exception as e:
            logger.error(f"Error saving encoded labels: {e}")
            raise

    def preprocess_data(self, df: pd.DataFrame, is_training: bool = True) -> pd.DataFrame:
        df = self.convert_to_datetime(df)
        df = self.drop_unnecessary_columns(df)
        if is_training:
            df = self.handle_null_values(df)
            df = self.remove_outliers(df)
            df = self.create_target_column(df)
        df = self.cyclical_encoding(df)
        df = self.convert_to_categorical(df)
        if is_training:
            df = self.encode_labels(df)
        else:
            df = self.encode_labels_from_yaml(df)
        logger.info("> Data preprocessed successfully")
        return df

    def encode_labels_from_yaml(self, df: pd.DataFrame) -> pd.DataFrame:
        encoders_info = yaml_config_loader(self.interim_data_path.parent / "label_encoders_info.yaml")

        for feature in df.columns:
            if feature in encoders_info:
                df[feature] = df[feature].apply(
                    lambda x: (
                        encoders_info[feature].get(x)
                        if x in encoders_info[feature] or x is None or x == np.nan
                        else logger.warning(f"Value {x} not found in categories for {feature}") or None
                    )
                )
            # df[feature] = df[feature].map(encoders_info[feature])

        logger.info(" >> Label encoding done")
        return df

    def save_data(self, df: pd.DataFrame, file_path: Path) -> None:
        try:
            file_path.parent.mkdir(exist_ok=True, parents=True)
            df.to_csv(file_path, index=False)
            logger.info(f"> Data saved successfully to {file_path}")
        except Exception as e:
            logger.error(f"Error saving data: {e}")
            raise

    def run_training(self) -> None:
        df = self.load_data(self.raw_data_path)
        df = self.preprocess_data(df, is_training=True)
        self.save_data(df, self.interim_data_path)

    def run_inference(self, df: pd.DataFrame) -> pd.DataFrame:
        df = self.preprocess_data(df, is_training=False)
        return df


if __name__ == "__main__":
    processor = AlarmDataProcessor()
    processor.run_training()
