import logging
import pickle
import time

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

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


class IsoForest(object):
    """A class for Isolation Forest models."""

    def __init__(
        self,
        hybrid=False,
        n_estimators=100,
        max_samples="auto",
        contamination=0.1,
        n_jobs=-1,
        seed=None,
        **kwargs,
    ):
        """Init Isolation Forest instance."""
        self.n_estimators = n_estimators
        self.max_samples = max_samples
        self.contamination = contamination
        self.n_jobs = n_jobs
        self.seed = seed

        self.model = IsolationForest(
            n_estimators=n_estimators,
            max_samples=max_samples,
            contamination=contamination,
            n_jobs=n_jobs,
            random_state=seed,
            **kwargs,
        )

        self.hybrid = hybrid
        self.ae_net = None  # autoencoder network for the case of a hybrid model

        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,
        }

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

        # do not drop last batch for non-SGD optimization shallow_ssad
        # drop_last necessary?
        # train_loader = DataLoader(
        #     dataset=dataset.train_set,
        #     batch_size=128,
        #     shuffle=True,
        #     num_workers=n_jobs_dataloader,
        #     drop_last=False,
        # )

        if k_fold_idx is not None:
            train_loader, _ = dataset.loaders_k_fold(
                fold_idx=k_fold_idx,
                batch_size=128,
                num_workers=n_jobs_dataloader,
            )
        else:
            train_loader, _, _ = dataset.loaders(
                batch_size=128, num_workers=n_jobs_dataloader
            )
        # Get data from loader
        X = ()
        for data in train_loader:
            inputs, _, _, _, _, _ = data
            inputs = inputs.to(device)
            if self.hybrid:
                inputs = self.ae_net.encoder(
                    inputs
                )  # in hybrid approach, take code representation of AE as features
            X_batch = inputs.view(
                inputs.size(0), -1
            )  # X_batch.shape = (batch_size, n_channels * height * width)
            X += (X_batch.cpu().data.numpy(),)
        X = np.concatenate(X)

        # Training
        logger.info("Starting training...")
        start_time = time.time()
        self.model.fit(X)
        train_time = time.time() - start_time
        self.results["train_time"] = train_time

        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,
    ):
        """Tests the Isolation Forest 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=128,
                num_workers=n_jobs_dataloader,
            )
        else:
            _, test_loader, _ = dataset.loaders(
                batch_size=128, 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
            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) in a list 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) ---
        # Filter out unknown labels and convert to binary (1: anomaly, 0: normal) for ROC
        valid_mask_exp = labels_exp != 0
        if np.any(valid_mask_exp):
            # Convert to binary labels for ROC (-1 → 1, 1 → 0)
            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):
            # Convert to binary labels for ROC (-1 → 1, 1 → 0)
            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)")

        logger.info("Test Time: {:.3f}s".format(self.results["test_time"]))
        logger.info("Finished testing.")

    def load_ae(self, dataset_name, model_path):
        """Load pretrained autoencoder from model_path for feature extraction in a hybrid Isolation Forest model."""

        model_dict = torch.load(model_path, map_location="cpu")
        ae_net_dict = model_dict["ae_net_dict"]
        if dataset_name in ["mnist", "fmnist", "cifar10"]:
            net_name = dataset_name + "_LeNet"
        else:
            net_name = dataset_name + "_mlp"

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

        # 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)
        self.ae_net.eval()

    def save_model(self, export_path):
        """Save Isolation Forest model to export_path."""
        with open(export_path, "wb") as fp:
            pickle.dump(self.model, fp)

    def load_model(self, import_path, device: str = "cpu"):
        """Load Isolation Forest model from import_path."""
        with open(import_path, "rb") as fp:
            self.model = pickle.load(fp)

    def save_results(self, export_pkl):
        """Save results dict to a JSON-file."""
        with open(export_pkl, "wb") as fp:
            # json.dump(self.results, fp)
            pickle.dump(self.results, fp)
