"""
training.py — Training and evaluation routines.
"""

import math
import time
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import random

from experiment_config import ExperimentConfig
from data import get_dataloaders, selfsup_mask_collate

# ── Reproducibility ───────────────────────────────────────────────────────────

def set_seed(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    random.seed(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

# ── Optimizer ─────────────────────────────────────────────────────────────────

def build_optimizer(model, config):
    """
    Adam optimizer + CosineAnnealingLR.
    For RamanSelfSupModel after encoder frozen: only trainable params.
    """
    # Collect trainable parameters (works for both frozen and unfrozen models)
    trainable = [p for p in model.parameters() if p.requires_grad]

    opt   = torch.optim.Adam(trainable, lr=config.lr, weight_decay=config.weight_decay)
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=config.epochs)
    return opt, sched

# ── Domain-flip helper ────────────────────────────────────────────────────────

def flip_channel(x):
    """
    Simulate instrument variation between source and target domains.
    Applies random per-channel gain + bias (Raman intensity scaling proxy).
    Operates on a single image [3, H, W].
    """
    gain = torch.empty(3, 1, 1).uniform_(0.8, 1.2)
    bias = torch.empty(3, 1, 1).normal_(0, 0.05)
    return (x * gain + bias).clamp(0.0, 1.0)

# ── Training loop ─────────────────────────────────────────────────────────────

def train_one_epoch(model, loaders, opt, config, device, epoch):
    """
    Single training epoch with model-type dispatch.

    loaders: dict with keys 'source_train', 'target_train'
    Returns dict: {loss, class_loss, strategy_loss, accuracy}
    """
    model.train()
    total_loss, total_cls, total_strat = 0.0, 0.0, 0.0
    correct, total = 0, 0

    # Zip source and target loaders to get paired batches
    src_iter = iter(loaders['source_train'])
    tgt_iter = iter(loaders['target_train'])

    num_batches = min(
        len(loaders['source_train']), len(loaders['target_train'])
    )

    for _ in range(num_batches):
        src_batch = next(src_iter)
        tgt_batch = next(tgt_iter)

        images_s, labels_s, _ = src_batch
        images_t, labels_t, _  = tgt_batch

        images_s = images_s.to(device)
        labels_s = labels_s.to(device)
        images_t = images_t.to(device)
        labels_t = labels_t.to(device)

        # Combine for mixed-batch processing
        images_all = torch.cat([images_s, images_t], dim=0)
        labels_all = torch.cat([labels_s, labels_t], dim=0)
        # Domain labels: first half = source (0), second half = target (1)
        domain_labels = torch.cat([
            torch.zeros(images_s.size(0), dtype=torch.long, device=device),
            torch.ones(images_t.size(0), dtype=torch.long, device=device),
        ])

        model_name = model.__class__.__name__

        if model_name == 'DANNModel':
            # DANN: gradient reversal discriminator
            class_logits, domain_logits = model(images_all, return_domain=True)
            loss, cls_l, strat_l = model.add_strategy_loss(
                class_logits, domain_logits,
                labels_all, domain_labels, device
            )

        elif model_name == 'CORALModel':
            # CORAL: covariance alignment on bottleneck
            z = model.encode(images_all)
            class_logits = model.classifier(z)
            loss, cls_l, strat_l = model.add_strategy_loss(
                class_logits, z, labels_all, domain_labels, device
            )

        elif model_name == 'RamanSelfSupModel':
            # Raman self-supervised: masked spectrum reconstruction
            batch_for_ss = [
                (img, l) for img, l in zip(images_all, labels_all)
            ]
            masked_spec, targets_ss, labels_ss = selfsup_mask_collate(
                batch_for_ss, mask_ratio=config.mask_ratio
            )
            masked_spec = masked_spec.to(device)
            targets_ss  = targets_ss.to(device)
            labels_ss   = labels_ss.to(device)

            # forward() always returns a 2-tuple (cls_logits, recon)
            cls_logits, recon = model(masked_spectrum=masked_spec)
            loss, cls_l, strat_l = model.add_strategy_loss(
                cls_logits, recon, targets_ss, labels_ss, device
            )

        else:  # BaseModel — standard cross-entropy only
            logits = model(images_all)
            loss   = F.cross_entropy(logits, labels_all)
            cls_l  = loss.detach()
            strat_l = torch.tensor(0.0, device=device)

        # ── Gradient clipping + NaN guard ──────────────────────
        if torch.isnan(loss) or loss.item() > 100:
            print(
                f'  WARN: NaN/divergence at epoch={epoch}, loss={loss.item():.4f}'
            )
            return {
                'loss': float('nan'), 'class_loss': float('nan'),
                'strategy_loss': float('nan'), 'accuracy': 0.0
            }

        opt.zero_grad()
        loss.backward()
        nn.utils.clip_grad_norm_(model.parameters(), config.gradient_clip_norm)
        opt.step()

        # Metrics (only on source portion for accuracy)
        with torch.no_grad():
            if model_name == 'RamanSelfSupModel':
                src_logits, _ = model(x=images_s)
            else:
                src_logits = model(images_s)
            correct += (src_logits.argmax(1) == labels_s).sum().item()
            total   += labels_s.size(0)

        total_loss  += loss.item()
        total_cls   += cls_l.item() if isinstance(cls_l, torch.Tensor) else cls_l
        total_strat += strat_l.item() if isinstance(strat_l, torch.Tensor) else strat_l

    n = num_batches
    return {
        'loss':          total_loss / n,
        'class_loss':    total_cls  / n,
        'strategy_loss': total_strat / n,
        'accuracy':      correct / total,
    }

# ── Evaluation ────────────────────────────────────────────────────────────────

@torch.no_grad()
def evaluate(model, loader, device):
    """
    Standard eval loop. Computes classification accuracy as primary_metric.
    loader: any single DataLoader (val or test).
    Returns dict with keys: {loss, accuracy, primary_metric}
    """
    model.eval()
    total_loss, correct, total = 0.0, 0, 0
    n_batches = 0

    for images, labels, *_ in loader:
        images = images.to(device)
        labels = labels.to(device)

        # Handle RamanSelfSupModel which returns 2-tuple
        model_name = model.__class__.__name__
        if model_name == 'RamanSelfSupModel':
            logits, _ = model(x=images)
        else:
            logits = model(images)

        loss   = F.cross_entropy(logits, labels)
        total_loss += loss.item()
        correct    += (logits.argmax(1) == labels).sum().item()
        total      += labels.size(0)
        n_batches  += 1

    acc = correct / total if total > 0 else 0.0
    return {
        'loss':          total_loss / max(n_batches, 1),
        'accuracy':      acc,
        'primary_metric': acc,   # == accuracy; higher is better
    }

# ── Time guard ────────────────────────────────────────────────────────────────

def time_guard(start_time, time_budget_minutes):
    elapsed_minutes = (time.time() - start_time) / 60.0
    if elapsed_minutes >= time_budget_minutes:
        print(
            f'  [TIME] Budget exceeded: {elapsed_minutes:.1f}m >= '
            f'{time_budget_minutes}m'
        )
        return False
    return True

# ── DANN alpha schedule ───────────────────────────────────────────────────────

def dann_alpha(epoch, num_epochs):
    """
    DANN annealing schedule from Ganin & Lempitsky 2016:
      α_p = 2 / (1 + exp(-10·p)) − 1   where p = epoch / num_epochs
    """
    p = epoch / max(num_epochs - 1, 1)
    return float(2.0 / (1.0 + math.exp(-10.0 * p)) - 1.0)