import logging
import pickle
import time
from pathlib import Path

import numpy as np
import torch
from sklearn.metrics import (
    average_precision_score,
    precision_recall_curve,
    roc_auc_score,
    roc_curve,
)
from sklearn.svm import OneClassSVM

from base.base_dataset import BaseADDataset
from networks.main import build_autoencoder


class OCSVM(object):
    """A class for One-Class SVM models."""

    def __init__(self, kernel="rbf", nu=0.1, hybrid=False, latent_space_dim=128):
        """Init OCSVM instance."""
        self.kernel = kernel
        self.nu = nu
        self.rho = None
        self.gamma = None

        self.model = OneClassSVM(kernel=kernel, nu=nu)

        self.hybrid = hybrid
        self.latent_space_dim = latent_space_dim
        self.ae_net = None  # autoencoder network for the case of a hybrid model
        self.linear_model = (
            None  # also init a model with linear kernel if hybrid approach
        )

        self.results = {
            "train_time": None,
            "test_time": None,
            "test_auc_exp_based": None,
            "test_roc_exp_based": None,
            "test_prc_exp_based": None,
            "test_ap_exp_based": None,
            "test_scores_exp_based": None,
            "test_auc_manual_based": None,
            "test_roc_manual_based": None,
            "test_prc_manual_based": None,
            "test_ap_manual_based": None,
            "test_scores_manual_based": None,
            "train_time_linear": None,
            "test_time_linear": None,
            "test_auc_linear": None,
        }

    def train(
        self,
        dataset: BaseADDataset,
        device: str = "cpu",
        n_jobs_dataloader: int = 0,
        k_fold_idx: int = None,
        batch_size: int = 32,
    ):
        """Trains the OC-SVM model on the training data."""
        logger = logging.getLogger()

        if k_fold_idx is not None:
            train_loader, _ = dataset.loaders_k_fold(
                fold_idx=k_fold_idx,
                batch_size=batch_size,
                num_workers=n_jobs_dataloader,
            )
        else:
            train_loader, _, _ = dataset.loaders(
                batch_size=batch_size, num_workers=n_jobs_dataloader
            )

        # Get data from loader
        X = ()
        for data in train_loader:
            inputs, _, _, _, _, _ = data  # Updated unpacking
            inputs = inputs.to(device)
            if self.hybrid:
                inputs = self.ae_net.encoder(inputs)
            X_batch = inputs.view(inputs.size(0), -1)
            X += (X_batch.cpu().data.numpy(),)
        X = np.concatenate(X)

        # Training
        logger.info("Starting training...")

        # Select model via hold-out test set of 1000 samples
        gammas = np.logspace(-7, 2, num=10, base=2)
        best_auc = 0.0

        # Sample hold-out set from test set
        if k_fold_idx is not None:
            _, test_loader = dataset.loaders_k_fold(
                fold_idx=k_fold_idx,
                batch_size=batch_size,
                num_workers=n_jobs_dataloader,
            )
        else:
            _, test_loader, _ = dataset.loaders(
                batch_size=batch_size, num_workers=n_jobs_dataloader
            )

        # Sample hold-out set from test set
        X_test = ()
        labels_exp = []
        labels_manual = []

        for data in test_loader:
            inputs, label_exp, label_manual, _, _, _ = data  # Updated unpacking
            inputs = inputs.to(device)
            label_exp = label_exp.to(device)
            label_manual = label_manual.to(device)

            if self.hybrid:
                inputs = self.ae_net.encoder(inputs)
            X_batch = inputs.view(inputs.size(0), -1)
            X_test += (X_batch.cpu().data.numpy(),)
            labels_exp += label_exp.cpu().data.numpy().astype(np.int64).tolist()
            labels_manual += label_manual.cpu().data.numpy().astype(np.int64).tolist()

        X_test = np.concatenate(X_test)
        labels_exp = np.array(labels_exp)
        labels_manual = np.array(labels_manual)

        # Use experiment-based labels for model selection (could also use manual-based)
        labels = labels_exp

        # Count samples for validation split (-1: anomaly, 1: normal, 0: unknown)
        n_test = len(X_test)
        n_normal = np.sum(labels == 1)
        n_outlier = np.sum(labels == -1)
        n_val = int(0.1 * n_test)  # Use 10% of test data for validation

        # Only consider labeled samples for validation
        valid_mask = labels != 0
        X_test_valid = X_test[valid_mask]
        labels_valid = labels[valid_mask]

        # Calculate validation split sizes
        n_val_normal = int(n_val * (n_normal / (n_normal + n_outlier)))
        n_val_outlier = int(n_val * (n_outlier / (n_normal + n_outlier)))

        # Create validation set with balanced normal/anomaly ratio
        perm = np.random.permutation(len(X_test_valid))
        X_val = np.concatenate(
            (
                X_test_valid[perm][labels_valid[perm] == 1][:n_val_normal],
                X_test_valid[perm][labels_valid[perm] == -1][:n_val_outlier],
            )
        )
        val_labels = np.array(
            [0] * n_val_normal + [1] * n_val_outlier
        )  # Convert to binary (0: normal, 1: anomaly)

        # Model selection loop
        i = 1
        for gamma in gammas:
            # Model candidate
            model = OneClassSVM(
                kernel=self.kernel,
                nu=self.nu,
                gamma=gamma,
            )

            # Train
            start_time = time.time()
            model.fit(X)
            train_time = time.time() - start_time

            # Test on small hold-out set from test set
            scores = (-1.0) * model.decision_function(X_val)
            scores = scores.flatten()

            # Compute AUC with binary labels
            auc = roc_auc_score(val_labels, scores)

            logger.info(
                f"  | Model {i:02}/{len(gammas):02} | Gamma: {gamma:.8f} | Train Time: {train_time:.3f}s "
                f"| Val AUC: {100.0 * auc:.2f} |"
            )

            if auc > best_auc:
                best_auc = auc
                self.model = model
                self.gamma = gamma
                self.results["train_time"] = train_time

            i += 1

        # If hybrid, also train a model with linear kernel
        if self.hybrid:
            self.linear_model = OneClassSVM(kernel="linear", nu=self.nu)
            start_time = time.time()
            self.linear_model.fit(X)
            train_time = time.time() - start_time
            self.results["train_time_linear"] = train_time

        logger.info(
            f"Best Model: | Gamma: {self.gamma:.8f} | AUC: {100.0 * best_auc:.2f}"
        )
        logger.info("Training Time: {:.3f}s".format(self.results["train_time"]))
        logger.info("Finished training.")

    def test(
        self,
        dataset: BaseADDataset,
        device: str = "cpu",
        n_jobs_dataloader: int = 0,
        k_fold_idx: int = None,
        batch_size: int = 32,
    ):
        """Tests the OC-SVM model on the test data."""
        logger = logging.getLogger()

        if k_fold_idx is not None:
            _, test_loader = dataset.loaders_k_fold(
                fold_idx=k_fold_idx,
                batch_size=batch_size,
                num_workers=n_jobs_dataloader,
            )
        else:
            _, test_loader, _ = dataset.loaders(
                batch_size=batch_size, num_workers=n_jobs_dataloader
            )

        # Get data from loader
        idx_label_score_exp = []
        idx_label_score_manual = []
        X = ()
        idxs = []
        labels_exp = []
        labels_manual = []

        for data in test_loader:
            inputs, label_exp, label_manual, _, idx, _ = data  # Updated unpacking
            inputs, label_exp, label_manual, idx = (
                inputs.to(device),
                label_exp.to(device),
                label_manual.to(device),
                idx.to(device),
            )

            if self.hybrid:
                inputs = self.ae_net.encoder(inputs)
            X_batch = inputs.view(inputs.size(0), -1)
            X += (X_batch.cpu().data.numpy(),)
            idxs += idx.cpu().data.numpy().astype(np.int64).tolist()
            labels_exp += label_exp.cpu().data.numpy().astype(np.int64).tolist()
            labels_manual += label_manual.cpu().data.numpy().astype(np.int64).tolist()

        X = np.concatenate(X)
        labels_exp = np.array(labels_exp)
        labels_manual = np.array(labels_manual)

        # Count and log label stats for exp_based
        n_exp_normal = np.sum(labels_exp == 1)
        n_exp_anomaly = np.sum(labels_exp == -1)
        n_exp_unknown = np.sum(labels_exp == 0)
        logger.info(
            f"Exp-based labels: normal(1)={n_exp_normal}, "
            f"anomaly(-1)={n_exp_anomaly}, unknown(0)={n_exp_unknown}"
        )

        # Count and log label stats for manual_based
        n_manual_normal = np.sum(labels_manual == 1)
        n_manual_anomaly = np.sum(labels_manual == -1)
        n_manual_unknown = np.sum(labels_manual == 0)
        logger.info(
            f"Manual-based labels: normal(1)={n_manual_normal}, "
            f"anomaly(-1)={n_manual_anomaly}, unknown(0)={n_manual_unknown}"
        )

        # Testing
        logger.info("Starting testing...")
        start_time = time.time()
        scores = (-1.0) * self.model.decision_function(X)
        self.results["test_time"] = time.time() - start_time
        scores = scores.flatten()

        # Save triples of (idx, label, score) for both label types
        idx_label_score_exp += list(zip(idxs, labels_exp.tolist(), scores.tolist()))
        idx_label_score_manual += list(
            zip(idxs, labels_manual.tolist(), scores.tolist())
        )

        self.results["test_scores_exp_based"] = idx_label_score_exp
        self.results["test_scores_manual_based"] = idx_label_score_manual

        # --- Evaluation for exp_based (only labeled samples) ---
        valid_mask_exp = labels_exp != 0
        if np.any(valid_mask_exp):
            labels_exp_binary = (labels_exp[valid_mask_exp] == -1).astype(int)
            scores_exp_valid = scores[valid_mask_exp]

            self.results["test_auc_exp_based"] = roc_auc_score(
                labels_exp_binary, scores_exp_valid
            )
            self.results["test_roc_exp_based"] = roc_curve(
                labels_exp_binary, scores_exp_valid
            )
            self.results["test_prc_exp_based"] = precision_recall_curve(
                labels_exp_binary, scores_exp_valid
            )
            self.results["test_ap_exp_based"] = average_precision_score(
                labels_exp_binary, scores_exp_valid
            )

            logger.info(
                "Test AUC (exp_based): {:.2f}%".format(
                    100.0 * self.results["test_auc_exp_based"]
                )
            )
        else:
            logger.info("Test AUC (exp_based): N/A (no labeled samples)")

        # --- Evaluation for manual_based (only labeled samples) ---
        valid_mask_manual = labels_manual != 0
        if np.any(valid_mask_manual):
            labels_manual_binary = (labels_manual[valid_mask_manual] == -1).astype(int)
            scores_manual_valid = scores[valid_mask_manual]

            self.results["test_auc_manual_based"] = roc_auc_score(
                labels_manual_binary, scores_manual_valid
            )
            self.results["test_roc_manual_based"] = roc_curve(
                labels_manual_binary, scores_manual_valid
            )
            self.results["test_prc_manual_based"] = precision_recall_curve(
                labels_manual_binary, scores_manual_valid
            )
            self.results["test_ap_manual_based"] = average_precision_score(
                labels_manual_binary, scores_manual_valid
            )

            logger.info(
                "Test AUC (manual_based): {:.2f}%".format(
                    100.0 * self.results["test_auc_manual_based"]
                )
            )
        else:
            logger.info("Test AUC (manual_based): N/A (no labeled samples)")

        # If hybrid, also test model with linear kernel
        if self.hybrid:
            start_time = time.time()
            scores_linear = (-1.0) * self.linear_model.decision_function(X)
            self.results["test_time_linear"] = time.time() - start_time
            scores_linear = scores_linear.flatten()
            # Save linear model results for both label types
            # --- exp_based ---
            valid_mask_exp_linear = labels_exp != 0
            if np.any(valid_mask_exp_linear):
                labels_exp_binary_linear = (
                    labels_exp[valid_mask_exp_linear] == -1
                ).astype(int)
                scores_exp_linear_valid = scores_linear[valid_mask_exp_linear]
                self.results["test_auc_linear_exp_based"] = roc_auc_score(
                    labels_exp_binary_linear, scores_exp_linear_valid
                )
                self.results["test_roc_linear_exp_based"] = roc_curve(
                    labels_exp_binary_linear, scores_exp_linear_valid
                )
                self.results["test_prc_linear_exp_based"] = precision_recall_curve(
                    labels_exp_binary_linear, scores_exp_linear_valid
                )
                self.results["test_ap_linear_exp_based"] = average_precision_score(
                    labels_exp_binary_linear, scores_exp_linear_valid
                )
            else:
                self.results["test_auc_linear_exp_based"] = None
                self.results["test_roc_linear_exp_based"] = None
                self.results["test_prc_linear_exp_based"] = None
                self.results["test_ap_linear_exp_based"] = None

            # --- manual_based ---
            valid_mask_manual_linear = labels_manual != 0
            if np.any(valid_mask_manual_linear):
                labels_manual_binary_linear = (
                    labels_manual[valid_mask_manual_linear] == -1
                ).astype(int)
                scores_manual_linear_valid = scores_linear[valid_mask_manual_linear]
                self.results["test_auc_linear_manual_based"] = roc_auc_score(
                    labels_manual_binary_linear, scores_manual_linear_valid
                )
                self.results["test_roc_linear_manual_based"] = roc_curve(
                    labels_manual_binary_linear, scores_manual_linear_valid
                )
                self.results["test_prc_linear_manual_based"] = precision_recall_curve(
                    labels_manual_binary_linear, scores_manual_linear_valid
                )
                self.results["test_ap_linear_manual_based"] = average_precision_score(
                    labels_manual_binary_linear, scores_manual_linear_valid
                )
            else:
                self.results["test_auc_linear_manual_based"] = None
                self.results["test_roc_linear_manual_based"] = None
                self.results["test_prc_linear_manual_based"] = None
                self.results["test_ap_linear_manual_based"] = None
            # Log exp_based results for linear model
            if self.results["test_auc_linear_exp_based"] is not None:
                logger.info(
                    "Test AUC linear model (exp_based): {:.2f}%".format(
                        100.0 * self.results["test_auc_linear_exp_based"]
                    )
                )
            else:
                logger.info(
                    "Test AUC linear model (exp_based): N/A (no labeled samples)"
                )

            # Log manual_based results for linear model
            if self.results["test_auc_linear_manual_based"] is not None:
                logger.info(
                    "Test AUC linear model (manual_based): {:.2f}%".format(
                        100.0 * self.results["test_auc_linear_manual_based"]
                    )
                )
            else:
                logger.info(
                    "Test AUC linear model (manual_based): N/A (no labeled samples)"
                )

            logger.info(
                "Test Time linear model: {:.3f}s".format(
                    self.results["test_time_linear"]
                )
            )

        # Log results for both label types
        if self.results.get("test_auc_exp_based") is not None:
            logger.info(
                "Test AUC (exp_based): {:.2f}%".format(
                    100.0 * self.results["test_auc_exp_based"]
                )
            )
        else:
            logger.info("Test AUC (exp_based): N/A (no labeled samples)")

        if self.results.get("test_auc_manual_based") is not None:
            logger.info(
                "Test AUC (manual_based): {:.2f}%".format(
                    100.0 * self.results["test_auc_manual_based"]
                )
            )
        else:
            logger.info("Test AUC (manual_based): N/A (no labeled samples)")
        logger.info("Test Time: {:.3f}s".format(self.results["test_time"]))
        logger.info("Finished testing.")

    def load_ae(self, model_path, net_name, device="cpu"):
        """Load pretrained autoencoder from model_path for feature extraction in a hybrid OC-SVM model."""

        model_dict = torch.load(model_path, map_location="cpu")
        ae_net_dict = model_dict["ae_net_dict"]

        if self.ae_net is None:
            self.ae_net = build_autoencoder(net_name, rep_dim=self.latent_space_dim)

        # update keys (since there was a change in network definition)
        ae_keys = list(self.ae_net.state_dict().keys())
        for i in range(len(ae_net_dict)):
            k, v = ae_net_dict.popitem(False)
            new_key = ae_keys[i]
            ae_net_dict[new_key] = v
            i += 1

        self.ae_net.load_state_dict(ae_net_dict)
        if device != "cpu":
            self.ae_net.to(torch.device(device))
        self.ae_net.eval()

    def save_model(self, export_path):
        """Save OC-SVM model to export_path."""
        with open(export_path, "wb") as fp:
            pickle.dump({"model": self.model, "linear_model": self.linear_model}, fp)

    def load_model(self, import_path):
        """Load OC-SVM model from import_path."""
        with open(import_path, "rb") as fp:
            models = pickle.load(fp)
            self.model = models["model"]
            self.linear_model = models["linear_model"]

    def save_results(self, export_pkl):
        with open(export_pkl, "wb") as fp:
            # json.dump(self.results, fp)
            pickle.dump(self.results, fp)
