"""Trainable vision projector for CORTEX AI (experimental, stage 1).

Pipeline exercised here::

    IMAGE -> real vision encoder -> VisionProjector -> PrefixFusion
          -> small real frozen causal LM -> logits -> causal LM loss

The vision encoder and the causal LM stay frozen; **only** ``VisionProjector`` is trained.
This module is deliberately independent of the 1.65 T DeepSeek-V4-Pro checkpoint: it neither
reads nor modifies the model-*-of-*.safetensors shards, and the projector checkpoint it writes
is a separate file.

Scope note: a projector trained against a small stand-in LM demonstrates that the projector is
genuinely trainable and that the frozen-weight invariants hold. It does **not** demonstrate that
the trained projector works with the real DeepSeek-V4-Pro core, nor native multimodal
understanding of that checkpoint.
"""

from __future__ import annotations

import json
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any, Callable, Iterable, Sequence

import torch
from torch import Tensor, nn

from .core import PrefixFusion, VisionConfig, VisionProjector
from .native import NativeVisionFusion, build_native_fusion

IGNORE_INDEX = -100


@dataclass
class ProjectorTrainingConfig:
    """Explicit dimensions, dtype and optimisation settings for projector training."""

    vision_model: str = "google/siglip-base-patch16-224"
    language_model: str = "EleutherAI/pythia-14m"
    projector_type: str = "mlp"
    dtype: str = "float32"
    device: str = "cpu"
    learning_rate: float = 1e-3
    weight_decay: float = 0.0
    epochs: int = 1
    batch_size: int = 2
    max_steps: int = 40
    max_text_tokens: int = 24
    dataset: str = "clip-benchmark/wds_flickr30k"
    dataset_split: str = "test"
    seed: int = 0
    output_checkpoint: str = "checkpoints/vision_projector.pt"
    log_every: int = 10
    limit: int | None = 128

    def resolved_dtype(self) -> torch.dtype:
        return {"float32": torch.float32, "float16": torch.float16,
                "bfloat16": torch.bfloat16}[self.dtype]

    def resolved_device(self) -> torch.device:
        return torch.device(self.device)


@dataclass
class TrainingReport:
    """Recorded evidence from a training run."""

    vision_hidden_size: int
    lm_hidden_size: int
    image_token_count: int
    text_token_count: int
    fused_sequence_length: int
    projector_parameters: int
    dtype: str
    device: str
    steps: int
    loss_first: float
    loss_last: float
    loss_history: list[float] = field(default_factory=list)
    projector_changed: bool = False
    projector_delta: float = 0.0
    encoder_unchanged: bool = False
    lm_unchanged: bool = False
    checkpoint_path: str = ""

    def to_json(self) -> str:
        return json.dumps(asdict(self), indent=2)


def freeze_module(module: nn.Module) -> int:
    """Set requires_grad=False everywhere and return the frozen parameter count."""
    n = 0
    for p in module.parameters():
        p.requires_grad_(False)
        n += p.numel()
    module.eval()
    return n


def count_parameters(module: nn.Module, trainable_only: bool = False) -> int:
    return sum(p.numel() for p in module.parameters() if (p.requires_grad or not trainable_only))


def parameter_fingerprint(module: nn.Module) -> dict[str, Tensor]:
    """Detached clone of every parameter, for before/after comparison."""
    return {k: v.detach().clone() for k, v in module.state_dict().items()}


def fingerprint_delta(before: dict[str, Tensor], after: dict[str, Tensor]) -> float:
    """Total absolute change across all tensors in two fingerprints."""
    total = 0.0
    for k, v in before.items():
        if k not in after:
            continue
        total += (after[k].to(torch.float32) - v.to(torch.float32)).abs().sum().item()
    return total


def build_fusion_pipeline(
    vision_encoder: nn.Module,
    core: nn.Module,
    config: ProjectorTrainingConfig,
    image_processor: Callable[..., Any] | None = None,
) -> tuple[NativeVisionFusion, int, int]:
    """Assemble IMAGE -> encoder -> projector -> PrefixFusion -> frozen core."""
    vision_hidden = int(vision_encoder.config.hidden_size)
    lm_hidden = int(core.config.hidden_size)
    fusion = build_native_fusion(
        vision_encoder,
        core,
        vision_hidden_size=vision_hidden,
        cortex_hidden_size=lm_hidden,
        image_processor=image_processor,
        projector_type=config.projector_type,
    )
    # The encoder and the core are frozen; the projector is the only trained part.
    freeze_module(fusion.vision_encoder)
    freeze_module(fusion.core)
    fusion.projector.train()
    return fusion, vision_hidden, lm_hidden


def fused_labels(image_token_count: int, input_ids: Tensor) -> Tensor:
    """Labels aligned to the fused sequence.

    Causal-LM loss is reported at every position, so the image-token positions would otherwise
    be trained to predict ordinary text tokens from a vision prefix. They are masked out, so the
    gradient on the projector comes from predicting the caption tokens conditioned on the image.
    """
    ignore = torch.full(
        (input_ids.shape[0], image_token_count), IGNORE_INDEX, dtype=torch.long, device=input_ids.device
    )
    return torch.cat((ignore, input_ids), dim=1)


def fused_attention_mask(image_token_count: int, attention_mask: Tensor | None, batch_size: int) -> Tensor:
    """Attention mask aligned to the fused sequence.

    Image tokens are always attended (1), so the text mask must be extended rather than passed
    through unchanged. Text padding keeps its own 0 entries.
    """
    if attention_mask is None:
        return torch.ones((batch_size, image_token_count), dtype=torch.long)
    prefix = torch.ones(
        (attention_mask.shape[0], image_token_count), dtype=attention_mask.dtype, device=attention_mask.device
    )
    return torch.cat((prefix, attention_mask), dim=1)


def train_step(
    fusion: NativeVisionFusion,
    optimizer: torch.optim.Optimizer,
    pixel_values: Tensor,
    input_ids: Tensor,
    attention_mask: Tensor | None = None,
) -> float:
    """One optimisation step. Returns the scalar loss."""
    image_embeds = fusion.encode_pixels(pixel_values)
    text_embeds = fusion.core.get_input_embeddings()(input_ids)
    fused = fusion.fusion.fuse(text_embeds, image_embeds)
    mask = fused_attention_mask(image_embeds.shape[1], attention_mask, fused.shape[0])
    labels = fused_labels(image_embeds.shape[1], input_ids)

    out = fusion.core(
        inputs_embeds=fused, attention_mask=mask, labels=labels, use_cache=False
    )
    loss = out.loss
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()
    return float(loss.detach())


def build_optimizer(projector: nn.Module, config: ProjectorTrainingConfig) -> torch.optim.Optimizer:
    params = [p for p in projector.parameters() if p.requires_grad]
    if not params:
        raise RuntimeError("projector has no trainable parameters")
    return torch.optim.AdamW(params, lr=config.learning_rate, weight_decay=config.weight_decay)


def iterate_caption_batches(
    images: Sequence[Any],
    captions: Sequence[str],
    tokenizer: Any,
    image_processor: Any,
    config: ProjectorTrainingConfig,
) -> Iterable[tuple[Tensor, Tensor, Tensor]]:
    """Yield pre-processed (pixel_values, input_ids, attention_mask) batches."""
    bs = config.batch_size
    for start in range(0, len(images), bs):
        batch_images = list(images[start:start + bs])
        batch_captions = list(captions[start:start + bs])
        if not batch_images:
            continue
        pixel_values = image_processor(batch_images, return_tensors="pt").pixel_values
        enc = tokenizer(
            batch_captions,
            return_tensors="pt",
            padding=True,
            truncation=True,
            max_length=config.max_text_tokens,
        )
        yield pixel_values, enc["input_ids"], enc["attention_mask"]


def run_training(
    fusion: NativeVisionFusion,
    batches: Iterable[tuple[Tensor, Tensor, Tensor]],
    config: ProjectorTrainingConfig,
    projector_parameters: int,
    vision_hidden_size: int,
    lm_hidden_size: int,
    on_step: Callable[[int, float], None] | None = None,
) -> TrainingReport:
    """Train the projector over the supplied batches and report the evidence."""
    dtype = config.resolved_dtype()
    device = config.resolved_device()

    before_proj = parameter_fingerprint(fusion.projector)
    before_enc = parameter_fingerprint(fusion.vision_encoder)
    before_lm = parameter_fingerprint(fusion.core)

    optimizer = build_optimizer(fusion.projector, config)
    losses: list[float] = []
    steps = 0
    first_shape: tuple[int, int, int] | None = None

    for _epoch in range(config.epochs):
        for pixel_values, input_ids, attention_mask in batches:
            if steps >= config.max_steps:
                break
            pixel_values = pixel_values.to(device=device, dtype=dtype)
            input_ids = input_ids.to(device)
            attention_mask = attention_mask.to(device)

            if first_shape is None:
                with torch.no_grad():
                    n_img = fusion.encode_pixels(pixel_values).shape[1]
                first_shape = (n_img, input_ids.shape[1], n_img + input_ids.shape[1])

            loss = train_step(fusion, optimizer, pixel_values, input_ids, attention_mask)
            losses.append(loss)
            steps += 1
            if on_step is not None:
                on_step(steps, loss)
        if steps >= config.max_steps:
            break

    after_proj = parameter_fingerprint(fusion.projector)
    after_enc = parameter_fingerprint(fusion.vision_encoder)
    after_lm = parameter_fingerprint(fusion.core)

    n_img, n_txt, n_fused = first_shape or (0, 0, 0)
    report = TrainingReport(
        vision_hidden_size=vision_hidden_size,
        lm_hidden_size=lm_hidden_size,
        image_token_count=n_img,
        text_token_count=n_txt,
        fused_sequence_length=n_fused,
        projector_parameters=projector_parameters,
        dtype=config.dtype,
        device=config.device,
        steps=steps,
        loss_first=losses[0] if losses else float("nan"),
        loss_last=losses[-1] if losses else float("nan"),
        loss_history=losses,
        projector_delta=fingerprint_delta(before_proj, after_proj),
        encoder_unchanged=fingerprint_delta(before_enc, after_enc) == 0.0,
        lm_unchanged=fingerprint_delta(before_lm, after_lm) == 0.0,
    )
    report.projector_changed = report.projector_delta > 0.0
    return report


def save_projector(projector: nn.Module, config: ProjectorTrainingConfig, report: TrainingReport) -> Path:
    """Write the projector checkpoint plus its config, separate from any base model."""
    path = Path(config.output_checkpoint)
    path.parent.mkdir(parents=True, exist_ok=True)
    report.checkpoint_path = str(path)
    torch.save(
        {
            "state_dict": projector.state_dict(),
            "config": asdict(config),
            "report": json.loads(report.to_json()),
            "scope": "experimental stage 1: projector trained against a small stand-in LM",
        },
        path,
    )
    return path