"""
Training pipeline for active learning head-only fine-tuning.

Loads a pre-trained MobileNetV3 backbone (frozen), replaces the classification
head, and fine-tunes on versioned crop datasets using AdamW + early stopping.
All training is FP32 for Tesla P40 compatibility (CC 5.2, no Tensor Cores).
"""

import logging
import os
import random
from pathlib import Path
from typing import Dict, Optional, Tuple

import numpy as np

logger = logging.getLogger(__name__)


class Trainer:
    """Head-only fine-tuning pipeline for MobileNetV3 classifier."""

    def __init__(self, config):
        self._cfg = config
        self._models_path = Path(config.get("storage.models_path", "/models"))
        self._candidate_path = self._models_path / "candidate"
        self._candidate_path.mkdir(parents=True, exist_ok=True)

        al = config.get_section("active_learning")
        tr = al.get("training", {})
        self._epochs: int = int(tr.get("epochs", 20))
        self._batch_size: int = int(tr.get("batch_size", 32))
        self._lr: float = float(tr.get("learning_rate", 1e-3))
        self._weight_decay: float = float(tr.get("weight_decay", 1e-2))
        self._patience: int = int(tr.get("early_stopping_patience", 5))
        self._lr_factor: float = float(tr.get("lr_factor", 0.5))
        self._lr_patience: int = int(tr.get("lr_patience", 3))

        aug = al.get("augmentation", {})
        self._aug_flip: bool = bool(aug.get("horizontal_flip", True))
        self._aug_jitter: bool = bool(aug.get("color_jitter", True))
        self._aug_affine: bool = bool(aug.get("affine", True))
        self._aug_degrees: int = int(aug.get("affine_degrees", 10))
        self._aug_scale: float = float(aug.get("affine_scale", 0.1))

        seed = int(config.get("active_learning.seed", 42))
        self._seed = seed

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def train(self, dataset_path: str, version: str) -> Optional[str]:
        """
        Fine-tune the classification head on the versioned dataset.

        Returns the path to the best checkpoint or None on failure.
        """
        try:
            import torch
            import torch.nn as nn
            from torch.utils.data import DataLoader

            self._seed_everything(self._seed)
            self._assert_cuda_compatibility(torch)

            train_loader, val_loader = self._build_dataloaders(dataset_path, torch)
            model = self._build_model(nn)
            best_path = self._run_training(model, train_loader, val_loader, version, torch, nn)
            return best_path
        except ImportError as exc:
            logger.error("PyTorch not available: %s", exc)
            return None
        except Exception as exc:
            logger.error("Training failed: %s", exc, exc_info=True)
            return None

    # ------------------------------------------------------------------
    # Model construction
    # ------------------------------------------------------------------

    def _build_model(self, nn):
        """Load MobileNetV3-Small backbone, freeze it, attach trainable head."""
        import torchvision.models as models

        model = models.mobilenet_v3_small(weights=models.MobileNet_V3_Small_Weights.IMAGENET1K_V1)

        # Freeze backbone
        for param in model.features.parameters():
            param.requires_grad = False

        # Replace classification head
        model.classifier = nn.Sequential(
            nn.Linear(576, 256),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(256, 2),
        )
        return model

    # ------------------------------------------------------------------
    # Data loading
    # ------------------------------------------------------------------

    def _build_dataloaders(self, dataset_path: str, torch) -> Tuple:
        from torch.utils.data import DataLoader
        import torchvision.transforms as T

        train_transform = self._build_train_transform(T)
        val_transform = self._build_val_transform(T)

        train_dataset = _CropDataset(
            Path(dataset_path) / "crops", split="train",
            transform=train_transform, seed=self._seed,
        )
        val_dataset = _CropDataset(
            Path(dataset_path) / "crops", split="val",
            transform=val_transform, seed=self._seed,
        )

        g = torch.Generator()
        g.manual_seed(self._seed)

        train_loader = DataLoader(
            train_dataset,
            batch_size=self._batch_size,
            shuffle=True,
            num_workers=2,
            pin_memory=True,
            generator=g,
        )
        val_loader = DataLoader(
            val_dataset,
            batch_size=self._batch_size,
            shuffle=False,
            num_workers=2,
            pin_memory=True,
        )
        logger.info(
            "Dataset: %d train / %d val samples",
            len(train_dataset), len(val_dataset),
        )
        return train_loader, val_loader

    def _build_train_transform(self, T):
        ops = [T.Resize((224, 224))]
        if self._aug_flip:
            ops.append(T.RandomHorizontalFlip())
        if self._aug_jitter:
            ops.append(T.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1))
        if self._aug_affine:
            ops.append(T.RandomAffine(
                degrees=self._aug_degrees,
                scale=(1.0 - self._aug_scale, 1.0 + self._aug_scale),
            ))
        ops += [T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])]
        return T.Compose(ops)

    def _build_val_transform(self, T):
        return T.Compose([
            T.Resize((224, 224)),
            T.ToTensor(),
            T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ])

    # ------------------------------------------------------------------
    # Training loop
    # ------------------------------------------------------------------

    def _run_training(self, model, train_loader, val_loader, version: str, torch, nn) -> Optional[str]:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        model = model.to(device)
        logger.info("Training on device: %s", device)

        class_weights = self._compute_class_weights(train_loader, torch, device)
        criterion = nn.CrossEntropyLoss(weight=class_weights)

        optimizer = torch.optim.AdamW(
            filter(lambda p: p.requires_grad, model.parameters()),
            lr=self._lr,
            weight_decay=self._weight_decay,
            betas=(0.9, 0.999),
        )
        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
            optimizer, mode="max", factor=self._lr_factor, patience=self._lr_patience,
        )

        best_f1 = -1.0
        best_path: Optional[str] = None
        no_improve = 0

        for epoch in range(1, self._epochs + 1):
            train_loss = self._train_epoch(model, train_loader, criterion, optimizer, device, torch)
            val_f1 = self._val_epoch(model, val_loader, device, torch)
            scheduler.step(val_f1)

            ckpt_path = self._candidate_path / f"{version}_epoch_{epoch:02d}.pt"
            torch.save(model.state_dict(), ckpt_path)
            logger.info(
                "Epoch %d/%d  train_loss=%.4f  val_f1=%.4f",
                epoch, self._epochs, train_loss, val_f1,
            )

            if val_f1 > best_f1:
                best_f1 = val_f1
                best_path = str(self._candidate_path / f"{version}_best.pt")
                torch.save(model.state_dict(), best_path)
                no_improve = 0
            else:
                no_improve += 1
                if no_improve >= self._patience:
                    logger.info("Early stopping at epoch %d (patience=%d)", epoch, self._patience)
                    break

        logger.info("Training complete. Best val_f1=%.4f  checkpoint=%s", best_f1, best_path)
        return best_path

    @staticmethod
    def _train_epoch(model, loader, criterion, optimizer, device, torch) -> float:
        model.train()
        total_loss = 0.0
        for images, labels in loader:
            images, labels = images.to(device), labels.to(device)
            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            total_loss += loss.item() * len(images)
        return total_loss / max(len(loader.dataset), 1)

    @staticmethod
    def _val_epoch(model, loader, device, torch) -> float:
        from sklearn.metrics import f1_score

        model.eval()
        all_preds, all_labels = [], []
        with torch.no_grad():
            for images, labels in loader:
                images = images.to(device)
                outputs = model(images)
                preds = outputs.argmax(dim=1).cpu().numpy()
                all_preds.extend(preds)
                all_labels.extend(labels.numpy())

        return float(f1_score(all_labels, all_preds, average="macro", zero_division=0))

    @staticmethod
    def _compute_class_weights(train_loader, torch, device):
        """Compute inverse-frequency class weights to handle imbalance."""
        counts = np.zeros(2, dtype=np.int64)
        for _, labels in train_loader:
            for lbl in labels.numpy():
                counts[int(lbl)] += 1
        total = counts.sum()
        if total == 0 or counts.min() == 0:
            return None
        weights = total / (2.0 * counts.astype(np.float32))
        return torch.tensor(weights, dtype=torch.float32, device=device)

    # ------------------------------------------------------------------
    # Misc
    # ------------------------------------------------------------------

    @staticmethod
    def _assert_cuda_compatibility(torch):
        if torch.cuda.is_available():
            cuda_ver = torch.version.cuda or ""
            torch_ver = torch.__version__
            logger.info("CUDA version: %s  PyTorch: %s", cuda_ver, torch_ver)
            major = int(cuda_ver.split(".")[0]) if cuda_ver else 0
            if major > 11:
                logger.warning(
                    "CUDA %s detected; system requires CUDA ≤ 11.8 for Tesla P40", cuda_ver
                )

    @staticmethod
    def _seed_everything(seed: int):
        random.seed(seed)
        np.random.seed(seed)
        try:
            import torch
            torch.manual_seed(seed)
            if torch.cuda.is_available():
                torch.cuda.manual_seed_all(seed)
                torch.backends.cudnn.deterministic = True
                torch.backends.cudnn.benchmark = False
        except ImportError:
            pass


# ------------------------------------------------------------------
# Dataset helper
# ------------------------------------------------------------------

class _CropDataset:
    """Reads face crop images from class_0/ and class_1/ subdirectories."""

    _EXTENSIONS = {".jpg", ".jpeg", ".png"}

    def __init__(self, crops_root: Path, split: str, transform, seed: int):
        all_samples = self._discover(crops_root)
        train_samples, val_samples = _split_samples(all_samples, 0.80, seed)
        self._samples = train_samples if split == "train" else val_samples
        self._transform = transform

    def __len__(self) -> int:
        return len(self._samples)

    def __getitem__(self, idx: int):
        from PIL import Image
        path, label = self._samples[idx]
        image = Image.open(path).convert("RGB")
        if self._transform:
            image = self._transform(image)
        return image, label

    def _discover(self, crops_root: Path):
        samples = []
        for label, subdir in ((0, "class_0"), (1, "class_1")):
            d = crops_root / subdir
            if d.is_dir():
                for p in sorted(d.iterdir()):
                    if p.suffix.lower() in self._EXTENSIONS:
                        samples.append((p, label))
        return samples


def _split_samples(samples, train_frac: float, seed: int):
    """Stratified split mirroring LabelIngestor logic."""
    rng = random.Random(seed)
    by_label: Dict[int, list] = {}
    for path, label in samples:
        by_label.setdefault(label, []).append((path, label))

    train, val = [], []
    for group in by_label.values():
        shuffled = list(group)
        rng.shuffle(shuffled)
        split = max(1, int(len(shuffled) * train_frac))
        train.extend(shuffled[:split])
        val.extend(shuffled[split:])

    return train, val