import pickle

from pathlib import Path
from typing import List, Optional

import numpy as np
import pandas as pd

from catboost import CatBoostClassifier
from pydantic import BaseModel, Field

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

logger = setup_logger(__name__)


class AlarmDataModel(BaseModel):
    severity: Optional[str] = Field(None, alias="Severity")
    alarm_name: Optional[str] = Field(None, alias="Alarm Name")
    technical_ID: Optional[str] = Field(None, alias="Technical ID")
    site_name: Optional[str] = Field(None, alias="Site Name")
    site_ID: Optional[str] = Field(None, alias="Site ID")
    NE_name: Optional[str] = Field(None, alias="NE Name")
    source: Optional[str] = Field(None, alias="Source")
    NE_type: Optional[str] = Field(None, alias="NE Type")
    technology: Optional[str] = Field(None, alias="Technology")
    alarm_category: Optional[str] = Field(None, alias="Alarm Category")
    first_occurence: Optional[str] = Field(None, alias="First Occurence")
    last_occurence: Optional[str] = Field(None, alias="Last Occurence")
    parent_node: Optional[str] = Field(None, alias="Parent Node")

    class Config:
        populate_by_name = True


class InferenceModel:
    def __init__(self, config_path: str = "configs/config.yaml") -> None:
        self.config = yaml_config_loader(config_path)
        self.inference_config = self.config["inference"]
        self.model = self.load_model(self.inference_config["model_path"])
        self.inference_data_dir = Path(self.inference_config["inference_data_dir"])
        self.encoded_labels = yaml_config_loader(self.inference_config["encoded_labels"])
        self.features: list = self.inference_config["features"]

    @staticmethod
    def load_model(model_path: str) -> CatBoostClassifier:
        with open(model_path, "rb") as file:
            model = pickle.load(file)
        return model

    @staticmethod
    def get_data_as_df(
        data: List[AlarmDataModel] | AlarmDataModel | pd.DataFrame | dict | list,
    ) -> pd.DataFrame:
        if isinstance(data, AlarmDataModel):
            data = [data]
            df = pd.DataFrame([alarm.model_dump(by_alias=True) for alarm in data])
            logger.info(f"Converting data to DataFrame with {len(data)} rows")
        elif isinstance(data, pd.DataFrame):
            df = data
            logger.info(f"Data is already a DataFrame with {len(df)} rows")
        elif isinstance(data, list) and all(isinstance(item, AlarmDataModel) for item in data):
            df = pd.DataFrame([alarm.model_dump(by_alias=True) for alarm in data])
            logger.info(f"Converting list of AlarmDataModel to DataFrame with {len(data)} rows")
        elif isinstance(data, list) and all(isinstance(item, dict) for item in data):
            df = pd.DataFrame(data)
            logger.info(f"Converting list of dicts to DataFrame with {len(data)} rows")
        elif isinstance(data, dict):
            df = pd.DataFrame([data])
            logger.info("Converting single dict to DataFrame with 1 row")
        else:
            logger.error("Unsupported data type for conversion to DataFrame")
            raise ValueError("Unsupported data type for conversion to DataFrame")
        return df

    def preprocess(self, df: pd.DataFrame) -> pd.DataFrame:
        processor = AlarmDataProcessor()
        df = processor.run_inference(df=df)
        df.drop(columns=["First Occurence", "Last Occurence"], inplace=True, errors="ignore")
        for col in df.columns:
            if df[col].dtype == "category":
                df[col] = pd.to_numeric(df[col], errors="coerce")
        return df

    def load_model_from_mlflow(self) -> None:
        pass

    def save_new_data(
        self,
        data: List[AlarmDataModel] | AlarmDataModel | pd.DataFrame,
        save_dir: Path,
        csv_file_name: str,
    ) -> None:
        if isinstance(data, AlarmDataModel):
            data = [data]
            df = self.get_data_as_df(data)
        elif isinstance(data, pd.DataFrame):
            df = data
        else:
            df = self.get_data_as_df(data)

        try:
            save_dir.mkdir(parents=True, exist_ok=True)
            csv_file_path = save_dir / csv_file_name

            if csv_file_path.exists() and csv_file_path.stat().st_size > 0:
                df.to_csv(csv_file_path, mode="a", header=False, index=False)
                logger.info(f"New data appended to > {csv_file_path}")
            else:
                df.to_csv(csv_file_path, index=False)
                logger.info(f"New data saved to > {csv_file_path}")

        except Exception as e:
            logger.error(f"Error saving new data: {e}")

    def predict(self, data: List[AlarmDataModel] | AlarmDataModel | pd.DataFrame) -> np.ndarray:
        if isinstance(data, pd.DataFrame):
            df = data
        else:
            df = self.get_data_as_df(data)

        try:
            predictions = self.model.predict(df)
            logger.info(f"Predictions made on {len(predictions)} rows")
        except Exception as e:
            logger.error(f"Error making predictions: {e}")
            raise e
        return predictions

    def run_inference(self, data: List[AlarmDataModel] | AlarmDataModel) -> pd.DataFrame:
        self.save_new_data(data, self.inference_data_dir, "raw_inference_data.csv")
        raw_df = self.get_data_as_df(data)
        df = self.preprocess(raw_df[self.features])
        self.save_new_data(df, self.inference_data_dir, "processed_inference_data.csv")

        predictions = pd.DataFrame(self.predict(df), columns=["predictions"])
        df_with_predictions = pd.concat([raw_df, predictions], axis=1)
        df_with_predictions["predictions"] = df_with_predictions["predictions"].map({0: "Irrelevant", 1: "Relevant"})
        self.save_new_data(df_with_predictions, self.inference_data_dir, "predictions.csv")
        return predictions


if __name__ == "__main__":
    model = InferenceModel()
    data = [
        AlarmDataModel(
            severity="Cleared",
            alarm_name="NodeB Unavailable",
            technicam_ID=22214,
            site_name="Dummar",
            site_id="DA0005C1",
            NE_name="BSDMMDN1",
            source="BSDMMDN1",
            NE_type="BSC6900GU",
            technology="M2000",
            first_occurence="12/01/2015 02:41:22 PM",
            last_occurence="12/01/2015 02:41:22 PM",
            alarm_category="Clear",
        ),
        AlarmDataModel(
            severity="Critical",
            alarm_name="NodeB Unavailable",
            technicam_ID=22213,
            # site_id='DA0005C1',
            # NE_name='BSDMMDN1',
            # source='BSDMMDN1',
            # NE_type='BSC6900GU',
            # technology='M2000',
            first_occurence="12/01/2015 02:41:22 PM",
            # last_occurence='12/01/2015 02:41:22 PM',
            alarm_category="Clear",
        ),
        AlarmDataModel(
            severity="Minor",
            alarm_name="New type of unknown alarm",
            technicam_ID=22214,
            site_name="New site",
            site_id="DA0005C1",
            NE_name="A new network element name",
            source="BSDMMDN1",
            NE_type="BSC6900GU",
            technology="M2000",
            first_occurence="12/01/2015 02:41:22 PM",
            last_occurence="12/01/2015 02:41:22 PM",
            alarm_category="Maintenance",
        ),
    ]

    predictions = model.run_inference(data)
    print(predictions)
