"""
active_learning.py
==================
Simplified Active Learning loop for the LSTM-Autoencoder.

Strategy: uncertainty sampling near the decision boundary. Sessions with reconstruction error close to the threshold are the most uncertain - these get queued for human review.
The human labels them as normal (0) or attack (1). Labels are used to refine the threshold without retraining. This implements the "Human-in-the-Loop" component described in the research proposal <need to add reference>. #TODO

Author : K.A.D.S.D. Kandanaarachchi (2020/ICT/19)
Project: Detecting Anomalous REST API Traffic — IT4216
"""

import argparse
import json
import logging
from pathlib import Path

import numpy as np
import torch
from torch.utils.data import DataLoader, TensorDataset

from model import (
    build_model_cicids2018,
    build_model_csic2010,
    build_model_unsw,
)

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s  %(levelname)s  %(message)s",
    datefmt="%H:%M:%S",
)
log = logging.getLogger(__name__)


class UncertaintySampler:
    """
    Identifies the most uncertain predictions for human review.

    A prediction is uncertain when its reconstruction error falls within a margin band around the threshold:
        [threshold - margin, threshold + margin]

    Sessions inside this band are neither clearly normal nor clearly anomalous — these benefit most from a human label.
    """

    def __init__(self, threshold: float, margin: float = 0.1):
        self.threshold = threshold
        self.margin = margin

    def get_uncertain_indices(self, errors: np.ndarray) -> np.ndarray:
        """
        Return indices of samples within the uncertainty band.
        """
        lower = self.threshold - (self.threshold * self.margin)
        upper = self.threshold + (self.threshold * self.margin)
        uncertain = np.where((errors >= lower) & (errors <= upper))[0]
        return uncertain

    def get_confidence(self, errors: np.ndarray) -> np.ndarray:
        """
        Compute a [0, 1] confidence score for each prediction.
        Score of 1.0 = far from threshold (very confident).
        Score of 0.0 = exactly on threshold (maximally uncertain).
        """
        distance = np.abs(errors - self.threshold)
        confidence = np.clip(distance / (self.threshold * self.margin + 1e-9), 0, 1)
        return confidence


class ThresholdRefiner:
    """
    Updates the anomaly threshold based on human-labeled samples.

    Logic:
    - If a human labels a high-error session as NORMAL → threshold should move UP (we were too aggressive)
    - If a human labels a low-error session as ATTACK → threshold should move DOWN (we were too lenient)

    Uses an exponential moving average to update smoothly rather than jumping to a new value instantly.
    """

    def __init__(self, initial_threshold: float, learning_rate: float = 0.1):
        self.threshold = initial_threshold
        self.learning_rate = learning_rate
        self.history = [initial_threshold]
        self.n_updates = 0

    def update(self, errors: np.ndarray, labels: np.ndarray) -> float:
        """
        Update threshold based on labeled samples.

        Parameters
        ----------
        errors : reconstruction errors of reviewed samples
        labels : human labels (0=normal, 1=attack)

        Returns updated threshold.
        """
        if len(errors) == 0:
            return self.threshold

        # Find errors that were mislabeled by current threshold
        pred = (errors > self.threshold).astype(int)
        mislabeled = errors[pred != labels]

        if len(mislabeled) == 0:
            log.info("  No mislabeled samples — threshold unchanged")
            return self.threshold

        # False positives: predicted attack, actually normal
        # → threshold too low, should move up
        fp_errors = errors[(pred == 1) & (labels == 0)]

        # False negatives: predicted normal, actually attack
        # → threshold too high, should move down
        fn_errors = errors[(pred == 0) & (labels == 1)]

        adjustment = 0.0
        if len(fp_errors) > 0:
            # Move threshold up toward the mean FP error
            adjustment += self.learning_rate * (fp_errors.mean() - self.threshold)
        if len(fn_errors) > 0:
            # Move threshold down toward the mean FN error
            adjustment -= self.learning_rate * (self.threshold - fn_errors.mean())

        old_threshold = self.threshold
        self.threshold = float(
            np.clip(
                self.threshold + adjustment,
                self.threshold * 0.5,  # don't drop below 50% of original
                self.threshold * 2.0,  # don't exceed 200% of original
            )
        )

        self.history.append(self.threshold)
        self.n_updates += 1

        log.info(
            "  Threshold: %.6f → %.6f  (Δ=%.6f)",
            old_threshold,
            self.threshold,
            self.threshold - old_threshold,
        )
        log.info("  FP samples=%d  FN samples=%d", len(fp_errors), len(fn_errors))

        return self.threshold


class SimulatedOracle:
    """
    Simulates a human security analyst reviewing flagged sessions.

    In a real deployment this would be replaced by an actual analyst clicking 'normal' or 'attack' in a GUI/TUI.

    For research purposes we use the ground-truth labels to simulate (not so 'perfect') human judgment, then measure how much the threshold improves as a result.
    """

    def __init__(self, y_true: np.ndarray):
        self.y_true = y_true
        self.n_labeled = 0

    def label(self, indices: np.ndarray) -> np.ndarray:
        """
        Return ground-truth labels for the given indices.
        Simulates a human reviewing and labeling each session.
        """
        self.n_labeled += len(indices)
        return self.y_true[indices]


class ActiveLearningLoop:
    """
    Ties everything together for one full AL cycle.

    Each iteration:
    1. Compute reconstruction errors on unlabeled pool
    2. Identify uncertain samples (near threshold)
    3. Oracle labels the uncertain samples
    4. Refiner updates the threshold
    5. Measure how metrics improved
    """

    def __init__(
        self,
        model,
        sampler: UncertaintySampler,
        refiner: ThresholdRefiner,
        oracle: SimulatedOracle,
        dataset: str,
        device: str = "cpu",
    ):
        self.model = model
        self.sampler = sampler
        self.refiner = refiner
        self.oracle = oracle
        self.dataset = dataset
        self.device = device

    def compute_errors(self, X: np.ndarray) -> np.ndarray:
        """Run model inference and return reconstruction errors."""
        self.model.eval()

        if self.dataset == "csic2010":
            tensor = torch.tensor(X, dtype=torch.long)
        else:
            X = np.nan_to_num(X, nan=0.0, posinf=0.0, neginf=0.0)
            tensor = torch.tensor(X, dtype=torch.float32)

        loader = DataLoader(
            TensorDataset(tensor),
            batch_size=512,
        )
        errors = []
        with torch.no_grad():
            for (batch,) in loader:
                batch = batch.to(self.device)
                errors.extend(self.model.reconstruction_error(batch).cpu().numpy())
        return np.array(errors)

    def evaluate(
        self, errors: np.ndarray, y_true: np.ndarray, threshold: float
    ) -> dict:
        """Compute metrics at a given threshold."""
        from sklearn.metrics import (
            confusion_matrix,
            f1_score,
            precision_score,
            recall_score,
        )

        y_pred = (errors > threshold).astype(int)
        prec = precision_score(y_true, y_pred, zero_division=0)
        rec = recall_score(y_true, y_pred, zero_division=0)
        f1 = f1_score(y_true, y_pred, zero_division=0)  # Type checker issue. Code runs

        cm = confusion_matrix(y_true, y_pred)
        tn, fp, fn, tp = cm.ravel()
        fpr = fp / (fp + tn) if (fp + tn) > 0 else 0.0

        return {
            "threshold": round(threshold, 6),
            "precision": round(prec, 4),
            "recall": round(rec, 4),
            "f1": round(f1, 4),
            "fpr": round(fpr, 4),
            "tp": int(tp),
            "fp": int(fp),
            "tn": int(tn),
            "fn": int(fn),
        }

    def run(
        self,
        X_pool: np.ndarray,
        y_pool: np.ndarray,
        n_iterations: int = 5,
        batch_size: int = 50,
    ) -> list[dict]:
        """
        Run the full Active Learning loop.

        Parameters
        ----------
        X_pool      : unlabeled session pool
        y_pool      : ground truth (used only by oracle)
        n_iterations: number of AL rounds
        batch_size  : sessions to review per round

        Returns list of metric dicts — one per iteration.
        """
        log.info("Computing initial reconstruction errors...")
        errors = self.compute_errors(X_pool)

        history = []

        # Baseline metrics before any AL
        baseline = self.evaluate(errors, y_pool, self.refiner.threshold)
        baseline["iteration"] = 0
        baseline["n_labeled"] = 0
        baseline["n_uncertain"] = 0
        history.append(baseline)

        log.info(
            "Baseline — F1=%.4f  Prec=%.4f  Rec=%.4f  FPR=%.4f",
            baseline["f1"],
            baseline["precision"],
            baseline["recall"],
            baseline["fpr"],
        )

        for i in range(1, n_iterations + 1):
            log.info("\n── AL Iteration %d/%d ──────────────────", i, n_iterations)

            # Find uncertain samples
            uncertain_idx = self.sampler.get_uncertain_indices(errors)

            if len(uncertain_idx) == 0:
                log.info("  No uncertain samples found — stopping early")
                break

            # Select a batch to review
            if len(uncertain_idx) > batch_size:
                selected = np.random.choice(uncertain_idx, batch_size, replace=False)
            else:
                selected = uncertain_idx

            log.info(
                "  Uncertain samples: %d  →  reviewing: %d",
                len(uncertain_idx),
                len(selected),
            )

            # Oracle labels them
            labels = self.oracle.label(selected)

            # Refine threshold
            self.refiner.update(errors[selected], labels)
            self.sampler.threshold = self.refiner.threshold

            # Evaluate with new threshold
            metrics = self.evaluate(errors, y_pool, self.refiner.threshold)
            metrics["iteration"] = i
            metrics["n_labeled"] = self.oracle.n_labeled
            metrics["n_uncertain"] = len(uncertain_idx)
            history.append(metrics)

            log.info(
                "  After AL — F1=%.4f  Prec=%.4f  Rec=%.4f  FPR=%.4f",
                metrics["f1"],
                metrics["precision"],
                metrics["recall"],
                metrics["fpr"],
            )

        return history


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--dataset", required=True, choices=["csic2010", "cicids2018", "unsw"]
    )
    parser.add_argument("--iterations", type=int, default=5)
    parser.add_argument("--batch_size", type=int, default=50)
    parser.add_argument("--margin", type=float, default=0.1)
    parser.add_argument("--lr", type=float, default=0.1)
    parser.add_argument("--window", type=int, default=5)
    args = parser.parse_args()

    device = "cuda" if torch.cuda.is_available() else "cpu"
    data_dir = Path("data/processed")
    mdl_dir = Path("models")
    res_dir = Path("results")

    run_id = f"{args.dataset}_w{args.window}"

    log.info("=" * 55)
    log.info("ACTIVE LEARNING — %s   (window=%d)", args.dataset.upper(), args.window)
    log.info("=" * 55)

    # Load data — filenames are window-suffixed by preprocessing.py
    X_test = np.load(data_dir / f"X_test_{run_id}.npy")
    y_test = np.load(data_dir / f"y_test_{run_id}.npy")

    # Load threshold — cross-check the window it was trained with
    thresh_path = mdl_dir / f"threshold_{run_id}.json"
    thresh_info = json.load(open(thresh_path))
    threshold = thresh_info["threshold"]
    saved_window = thresh_info.get("window")
    if saved_window is not None and saved_window != args.window:
        raise ValueError(
            f"Window mismatch: {thresh_path.name} was trained with "
            f"window={saved_window}, but --window={args.window} was passed. "
            f"Pass --window {saved_window} to match the trained model."
        )
    log.info("Initial threshold: %.6f", threshold)

    # Build and load model
    if args.dataset == "csic2010":
        checkpoint = torch.load(
            mdl_dir / f"best_{run_id}.pt", map_location=device
        )
        embed_weight = checkpoint["embedding.weight"]
        vocab_size = embed_weight.shape[0]
        model = build_model_csic2010(vocab_size=vocab_size, seq_len=args.window)
    elif args.dataset == "cicids2018":
        model = build_model_cicids2018(n_features=X_test.shape[2], seq_len=args.window)
    else:
        model = build_model_unsw(n_features=X_test.shape[2], seq_len=args.window)

    model.load_state_dict(
        torch.load(mdl_dir / f"best_{run_id}.pt", map_location=device)
    )

    # Subsample pool for speed — use 10k sessions
    pool_size = min(10_000, len(X_test))
    idx = np.random.choice(len(X_test), pool_size, replace=False)
    X_pool = X_test[idx]
    y_pool = y_test[idx]

    log.info(
        "Pool size: %d  (normal=%d  attack=%d)",
        pool_size,
        (y_pool == 0).sum(),
        (y_pool == 1).sum(),
    )

    # Build AL components
    sampler = UncertaintySampler(threshold, margin=args.margin)
    refiner = ThresholdRefiner(threshold, learning_rate=args.lr)
    oracle = SimulatedOracle(y_pool)

    al_loop = ActiveLearningLoop(model, sampler, refiner, oracle, args.dataset, device)

    # Run
    history = al_loop.run(
        X_pool,
        y_pool,
        n_iterations=args.iterations,
        batch_size=args.batch_size,
    )

    # Print summary table
    log.info("\n%s", "=" * 55)
    log.info("ACTIVE LEARNING RESULTS — %s", args.dataset.upper())
    log.info("=" * 55)
    log.info(
        "%-5s %-8s %-8s %-8s %-8s %-8s %-8s",
        "Iter",
        "Thresh",
        "Prec",
        "Rec",
        "F1",
        "FPR",
        "Labeled",
    )
    log.info("-" * 55)
    for r in history:
        log.info(
            "%-5d %-8.5f %-8.4f %-8.4f %-8.4f %-8.4f %-8d",
            r["iteration"],
            r["threshold"],
            r["precision"],
            r["recall"],
            r["f1"],
            r["fpr"],
            r["n_labeled"],
        )

    # Save
    out = {
        "dataset": args.dataset,
        "iterations": args.iterations,
        "margin": args.margin,
        "history": history,
        "window": args.window,
    }
    out_path = res_dir / f"active_learning_{run_id}.json"
    with open(out_path, "w") as f:
        json.dump(out, f, indent=2)
    log.info("Saved → %s", out_path)


if __name__ == "__main__":
    main()