"""towerbase — unified TowerBase ABC, backbone factory, and modular training utilities.

This module is the structural backbone of the HyperTower v3 architecture.
It provides the abstract tower interface plus the training/eval helpers that
operate on any list of TowerBase instances.

Design principles
-----------------
* No concrete tower classes are defined here (ImageEncoder, ClinicalEncoder, etc.
  live in their respective tower files).
* No imports from any tower file — this module is self-contained with respect to
  the tower layer.  Tower files import from here; this file does not import from them.
* Removing any tower file leaves this module fully intact.
* Duck-typing via optional TowerBase methods (e.g. cd_warmup_embedding) replaces
  isinstance checks so new tower types never require changes here.
"""

from __future__ import annotations

import math
from abc import ABC, abstractmethod
from random import random
from dataclasses import dataclass, field
from typing import Optional

import numpy as np
import torch
import torch.nn.functional as F
from torch import nn
from torch.utils.data import DataLoader
from torchvision import transforms

from v3.classes.backbones import BACKBONES, list_names, load_backbone_weights
from v3.classes.bridges import Bridge
from random import random


# ---------------------------------------------------------------------------
# Cross-tower communication context
# ---------------------------------------------------------------------------


@dataclass
class EarlyPassContext:
    eye_train: list[dict]
    bilat_train: list[dict]
    bilat_val: list[dict]
    bilat_test: list[dict]
    image_preprocessor: object
    image_cache: object
    device: torch.device
    store: dict = field(default_factory=dict)  # cross-tower key-value store


# ---------------------------------------------------------------------------
# Backbone factory
# ---------------------------------------------------------------------------


def build_backbone(name: str, freeze_ratio: float = 0.0, augment: bool = True):
    """
    Operational builder:
      - instantiate with DEFAULT weights
      - strip classifier → features
      - apply ratio-based freezing over coarse blocks
      - return (model, out_dim, transform)
    """
    key = (name or "").lower()
    if key not in BACKBONES:
        raise ValueError(
            f"Unsupported backbone '{name}'. Valid options: {list_names()}"
        )

    spec = BACKBONES[key]
    m = spec.ctor(weights=spec.weights_default)
    out_dim, m = spec.strip(m)
    load_backbone_weights(key, m)

    mean = getattr(spec.weights_default, "meta", {}).get("mean", (0.485, 0.456, 0.406))
    std = getattr(spec.weights_default, "meta", {}).get("std", (0.229, 0.224, 0.225))
    crop = 299 if key == "inception_v3" else 224

    if augment:
        transform = transforms.Compose(
            [
                transforms.Resize(256),
                transforms.CenterCrop(crop),
                transforms.RandomHorizontalFlip(),
                transforms.RandomVerticalFlip(),
                transforms.RandomRotation(15),
                transforms.ColorJitter(0.1, 0.1, 0.1, 0.05),
                transforms.ToTensor(),
                transforms.Normalize(mean=mean, std=std),
            ]
        )
    else:
        transform = transforms.Compose(
            [
                transforms.Resize(256),
                transforms.CenterCrop(crop),
                transforms.ToTensor(),
                transforms.Normalize(mean=mean, std=std),
            ]
        )

    fr = max(0.0, min(1.0, float(freeze_ratio)))
    blocks = spec.blocks(m)
    n = len(blocks)
    freeze_n = int(math.floor(n * fr))
    for b in blocks[:freeze_n]:
        for p in b.parameters():
            p.requires_grad = False

    return m, out_dim, transform


# ---------------------------------------------------------------------------
# TowerBase ABC
# ---------------------------------------------------------------------------


class TowerBase(ABC):
    """Abstract base class for a HyperTower tower.

    Concrete sub-classes must implement ``embed_dims``, ``embed_batch``, and
    ``prepare_fold``.  Everything else has a sensible default no-op.
    """

    # ------------------------------------------------------------------
    # Required interface
    # ------------------------------------------------------------------

    @property
    @abstractmethod
    def embed_dims(self) -> list[int]:
        """Ordered list of embedding dimensionalities contributed to the bridge.

        Most towers contribute one embedding (e.g. GeometryTower → [geom_hidden]).
        ImageClinicalTower contributes two (image + clinical → [img_dim, cd_dim]).
        """

    @abstractmethod
    def embed_batch(
        self,
        batch: dict,
        *,
        device: torch.device,
        slot: int = 1,
    ) -> list[torch.Tensor]:
        """Return a list of embeddings for one eye slot in *batch*.

        Parameters
        ----------
        batch  : dict  — batch produced by a SlotDataset loader
        device : torch.device
        slot   : 1 (OD / image_1 / matrix_1) or 2 (OS / image_2 / matrix_2)

        Returns
        -------
        list of Tensor — same length and order as ``embed_dims``
        """

    @abstractmethod
    def prepare_fold(
        self,
        *,
        eye_train: list,
        bilat_train: list,
        bilat_val: list,
        bilat_test: list,
        image_preprocessor,
        image_cache,
        device: torch.device,
        args,
    ) -> None:
        """Called once per fold before the main epoch loop."""

    def early_pass(self, context: EarlyPassContext) -> None:
        """Optional: called once per fold before loaders are built."""
        pass

    def get_sample(self, entry) -> "torch.Tensor | dict":  # noqa: ARG002
        """Optional: called by HTDataset to retrieve one sample for this tower.

        entry : ShellEntry — carries entity_id, label, side, paired.

        Single mode (entry.paired=False): return one Tensor.
        Paired mode (entry.paired=True):  return {"a": Tensor, "b": Tensor}.

        Default raises NotImplementedError.  Towers that participate in the
        v4 HTDataset pipeline must implement this.
        """
        raise NotImplementedError(
            f"{type(self).__name__}.get_sample() is not implemented. "
            "Implement it to use this tower with HTDataset."
        )

    # ------------------------------------------------------------------
    # Optional interface
    # ------------------------------------------------------------------

    def augment_samples(self, samples: list) -> list:
        """Optional: add modality-specific keys to sample dicts before loaders are built.

        Called by the orchestrator on each sample list (eye_train, bilat_train,
        bilat_val, bilat_test) *after* ``prepare_fold`` completes.

        The default implementation is a no-op.  GeometryTower overrides this to
        inject ``seg_map_1`` / ``seg_map_2`` numpy arrays so the loader can deliver
        them as tensors alongside the image and clinical slots.
        """
        return samples

    def cd_warmup_embedding(
        self,
        batch: dict,
        *,
        device: torch.device,
    ) -> Optional[torch.Tensor]:
        """Return the clinical embedding for cd_warmup phase, or None if not applicable.

        ClinicalDataTower overrides this to return its encoder output.
        All other towers return None (the default).

        This replaces isinstance(t, ClinicalDataTower) checks in train_towers_epoch,
        so new tower types never require changes to towerbase.py.
        """
        return None

    def set_phase(self, phase: str) -> None:
        """Control requires_grad on this tower's parameters for *phase*.

        Phases: ``cd_warmup``, ``tower_warmup``, ``fused_warmup``, ``main``.
        Default: no-op (tower parameters always trainable unless overridden).
        """

    @property
    def total_epochs(self) -> int:
        """How many epochs this tower participates in the main loop."""
        return 0

    @property
    def train_loader(self) -> Optional[DataLoader]:
        """Single-eye training loader, or None if not applicable."""
        return None

    @property
    def cd_only_loader(self) -> Optional[DataLoader]:
        """Clinical-data-only loader for cd_warmup phase, or None."""
        return None

    def finalize_fold(
        self,
        *,
        bridge,
        bilat_train_loader: Optional[DataLoader] = None,
        val_loader: Optional[DataLoader] = None,
        device: torch.device,
        args,
    ) -> None:
        """Optional post-epoch-loop operations (e.g. fused head training)."""


# ---------------------------------------------------------------------------
# Shared training utilities
# ---------------------------------------------------------------------------


def _set_requires_grad(module: nn.Module, enabled: bool) -> None:
    for p in module.parameters():
        p.requires_grad = enabled


def _to_label_tensor(labels, device: torch.device) -> torch.Tensor:
    if torch.is_tensor(labels):
        return labels.to(device=device, dtype=torch.long)
    return torch.as_tensor(labels, dtype=torch.long, device=device)


# ---------------------------------------------------------------------------
# Modular tower training / evaluation (multi-tower interface)
# ---------------------------------------------------------------------------


def train_towers_epoch(
    towers: "list[TowerBase]",
    bridge: Bridge,
    loader: DataLoader,
    optimizer,
    device: torch.device,
    *,
    phase: str,
    bcd_prob: float = 0.5,
    tower_loss_mode: str = "bcd",
) -> tuple[float, float]:
    """Train one epoch using the modular tower interface.

    All towers and the bridge are set to the given phase via their
    ``set_phase`` methods.  BCD / fused loss semantics mirror the
    existing ``train_single_epoch`` logic:

      cd_warmup    — clinical encoder aux head only (slot index 1)
      tower_warmup — img + cd aux heads equally
      fused_warmup — fused bridge output only
      main         — BCD (randomly img or cd aux) vs fused, per tower_loss_mode

    Duck typing: towers that implement ``cd_warmup_embedding`` participate in
    cd_warmup; all others are skipped for that phase.  No isinstance checks.
    """
    for t in towers:
        t.set_phase(phase)
    bridge.set_phase(phase)

    total_loss = total_correct = total_n = 0

    for batch in loader:
        y = batch.get("label_1")
        if y is None:
            continue

        # ---- cd_warmup: train clinical encoder via its aux head ----
        if phase == "cd_warmup":
            z_cd = None
            for t in towers:
                z = t.cd_warmup_embedding(batch, device=device)
                if z is not None:
                    z_cd = z
                    break
            if z_cd is None:
                continue
            logits = bridge.aux_heads[1](z_cd)
            y_t = _to_label_tensor(y, device)
            loss = F.cross_entropy(logits, y_t)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            bs = y_t.shape[0]
            total_loss += float(loss.item()) * bs
            total_correct += int((logits.argmax(1) == y_t).sum())
            total_n += bs
            continue

        # ---- collect all slot-1 embeddings from every tower ----
        all_embs = []
        for t in towers:
            all_embs.extend(t.embed_batch(batch, device=device, slot=1))

        if any(e is None for e in all_embs):
            continue

        y_t = _to_label_tensor(y, device)

        if phase == "tower_warmup":
            # Average loss across all available aux heads
            aux_logits = [bridge.aux_heads[i](emb) for i, emb in enumerate(all_embs)]
            loss = sum(F.cross_entropy(l, y_t) for l in aux_logits) / len(aux_logits)
            # Softmax average for metrics
            logits = sum(F.softmax(l, dim=1) for l in aux_logits) / len(aux_logits)

        elif phase == "fused_warmup":
            logits_fused, _ = bridge.fuse(all_embs)
            loss = F.cross_entropy(logits_fused, y_t)
            logits = logits_fused

        else:  # main
            if tower_loss_mode == "all":
                logits_fused, aux = bridge.fuse(all_embs)
                loss = F.cross_entropy(logits_fused, y_t)
                for aux_l in aux:
                    loss = loss + F.cross_entropy(aux_l, y_t)
                logits = logits_fused
            elif random() < bcd_prob:
                # Randomly pick ONE tower to train (Generalized BCD)
                idx = int(random() * len(all_embs))
                logits = bridge.aux_heads[idx](all_embs[idx])
                loss = F.cross_entropy(logits, y_t)
            else:
                logits_fused, _ = bridge.fuse(all_embs)
                loss = F.cross_entropy(logits_fused, y_t)
                logits = logits_fused

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        bs = y_t.shape[0]
        total_loss += float(loss.item()) * bs
        total_correct += int((logits.argmax(1) == y_t).sum())
        total_n += bs

    return (
        total_loss / total_n if total_n else float("nan"),
        total_correct / total_n if total_n else float("nan"),
    )


def collect_probs_towers(
    towers: "list[TowerBase]",
    bridge: Bridge,
    loader: DataLoader,
    device: torch.device,
    *,
    tower_mode: str = "ensemble",
) -> tuple[np.ndarray, np.ndarray]:
    """Evaluate using the modular tower interface.

    Returns ``(y_true, probs)`` with patient-level probabilities
    (OD + OS averaged for ensemble mode, OD-only for single mode).
    """
    for t in towers:
        if isinstance(t, nn.Module):
            t.eval()
    bridge.eval()

    y_chunks, p_chunks = [], []
    with torch.no_grad():
        for batch in loader:
            y = batch.get("label_1")
            if y is None:
                continue
            if not (
                torch.is_tensor(batch.get("image_1"))
                and torch.is_tensor(batch.get("image_2"))
            ):
                continue

            all_embs_od = []
            all_embs_os = []
            for t in towers:
                all_embs_od.extend(t.embed_batch(batch, device=device, slot=1))
                all_embs_os.extend(t.embed_batch(batch, device=device, slot=2))

            if any(e is None for e in all_embs_od + all_embs_os):
                continue

            logits_od, _ = bridge.fuse(all_embs_od)
            logits_os, _ = bridge.fuse(all_embs_os)

            if tower_mode == "ensemble":
                probs = 0.5 * (
                    F.softmax(logits_od, dim=1) + F.softmax(logits_os, dim=1)
                )
            else:
                probs = F.softmax(logits_od, dim=1)

            y_chunks.append(_to_label_tensor(y, device).cpu().numpy())
            p_chunks.append(probs.cpu().numpy())

    for t in towers:
        if isinstance(t, nn.Module):
            t.train()
    bridge.train()

    if not y_chunks:
        return np.array([], dtype=np.int64), np.zeros((0, 0), dtype=np.float32)
    return np.concatenate(y_chunks), np.concatenate(p_chunks, axis=0)