from pathlib import Path

import pandas as pd

from imblearn.over_sampling import SMOTE
from sklearn.feature_selection import SelectKBest, f_classif
from sklearn.model_selection import StratifiedKFold

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

logger = setup_logger(__name__)


class TrainingProcessor:
    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.training_processing_config = self.config["training_processing"]

        self.interim_data_path = self.data_paths["interim_data_path"]
        self.processed_data_dir = Path(self.data_paths["processed_data_dir"])

        self.seed: int = self.training_processing_config["seed"]
        self.features = self.training_processing_config["features"]

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

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

    def stratified_split(self, df: pd.DataFrame) -> tuple[pd.DataFrame, pd.DataFrame, pd.Series, pd.Series]:
        X, y = df.drop("Target", axis=1), df["Target"]
        skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
        for train_index, test_index in skf.split(X, y):
            X_train, X_test = X.iloc[train_index], X.iloc[test_index]
            y_train, y_test = y.iloc[train_index], y.iloc[test_index]
        logger.info("Stratified splitting done")

        self.save_data(X_train, self.processed_data_dir / "X_train.csv")
        self.save_data(X_test, self.processed_data_dir / "X_test.csv")
        self.save_data(y_train, self.processed_data_dir / "y_train.csv")
        self.save_data(y_test, self.processed_data_dir / "y_test.csv")
        logger.info("X_train, X_test, y_train, y_test saved successfully")

        return X_train, X_test, y_train, y_test

    def apply_smote(self, X: pd.DataFrame, y: pd.Series) -> tuple[pd.DataFrame, pd.Series]:
        smote = SMOTE(random_state=self.seed)
        X_res, y_res = smote.fit_resample(X, y)
        logger.info("SMOTE applied")
        return X_res, y_res

    def feature_selection(self, X: pd.DataFrame, y: pd.Series) -> pd.DataFrame:
        selector = SelectKBest(f_classif, k=10)
        X_new = selector.fit_transform(X, y)
        logger.info("Feature selection done")
        print(X.columns[selector.get_support()])
        return pd.DataFrame(X_new, columns=X.columns[selector.get_support()])

    def process_training_data(self) -> None:
        df = self.load_data()
        df = df[self.features + ["Target"]]

        X_train, X_test, y_train, y_test = self.stratified_split(df)
        X_train, y_train = self.apply_smote(X_train, y_train)
        # X_train = self.feature_selection(X_train, y_train)
        X_test = X_test[X_train.columns]
        self.save_data(X_train, self.processed_data_dir / "SMOTE/X_train_processed.csv")
        self.save_data(X_test, self.processed_data_dir / "SMOTE/X_test_processed.csv")
        self.save_data(y_train, self.processed_data_dir / "SMOTE/y_train_processed.csv")
        self.save_data(y_test, self.processed_data_dir / "SMOTE/y_test_processed.csv")
        logger.info("X_train, X_test, y_train, y_test processed are saved successfully")

    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


if __name__ == "__main__":
    processor = TrainingProcessor()
    processor.process_training_data()

    # df = pd.read_csv("data/processed/X_train.csv")
    # df.info()
    # print(df[df["Is Active"] == True])
    # print(df['Is Active'].value_counts())
