#!/usr/bin/env python3
"""Runtime architecture probe for RAVEL/CLARA models.

Requires project ML dependencies:

    pip install -r requirements.txt

This script instantiates a selected pipeline model, runs one dummy forward pass,
and writes:

- actual tensor shapes collected by hooks/manual probes;
- parameter breakdown from named_parameters();
- trainable/frozen component summary.

It is intentionally separate from revision_priority0_audit.py because this one
downloads/loads Hugging Face model weights through from_pretrained().
"""

from __future__ import annotations

import argparse
import csv
import sys
from collections import defaultdict
from pathlib import Path
from typing import Any, Dict, Iterable, List, Tuple

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Probe runtime architecture and parameters.")
    parser.add_argument(
        "--pipeline",
        choices=["mvsa_single", "mvsa_multiple", "hfm", "reference_model"],
        default="mvsa_multiple",
    )
    parser.add_argument("--output-dir", default="ravel_revision_results/architecture")
    parser.add_argument("--batch-size", type=int, default=2)
    parser.add_argument("--seq-len", type=int, default=16)
    parser.add_argument("--image-size", type=int, default=224)
    parser.add_argument("--device", default="cpu")
    return parser.parse_args()


def write_csv(path: Path, rows: Iterable[Dict[str, Any]], fieldnames: List[str]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames, extrasaction="ignore")
        writer.writeheader()
        for row in rows:
            writer.writerow(row)


def component_for_name(name: str) -> str:
    while name.startswith("impl."):
        name = name[len("impl.") :]
    if name.startswith(("vision_lora", "vision.", "vision_encoder")):
        if "lora_" in name:
            return "Vision LoRA"
        return "Vision backbone frozen/trainable weights"
    if name.startswith(("text.", "text_encoder")):
        if "lora_" in name:
            return "Text LoRA"
        if ".layer." in name and "lora_" not in name:
            return "Text backbone/unfrozen DeBERTa layers"
        return "Text backbone other weights"
    if name.startswith("fusion."):
        if ".cross." in name or ".blocks." in name or ".layers." in name:
            return "Co-attention"
        if ".v_proj." in name or ".t_proj." in name:
            return "Projection layers"
        return "Fusion other"
    if name.startswith(("visual_pool_proj.", "text_pool_proj.")):
        return "Projection layers"
    if name.startswith(("visual_head.", "text_head.")):
        return "Unimodal heads"
    if name.startswith(("pred.", "prediction_head")):
        return "Primary prediction head"
    if name.startswith(("veri.", "verification")):
        return "Verification head"
    if name.startswith(("feed.", "feedback", "refinement.")):
        return "Refinement/feed module"
    if name.startswith("extra_mlp_control."):
        return "Parameter-matched control head"
    if name.startswith("final."):
        return "Final classifier"
    if name.startswith("unc."):
        return "Uncertainty estimator"
    if "vision_projection" in name or "text_projection" in name:
        return "Projection layers"
    if "co_attention" in name:
        return "Co-attention"
    return "Other"


def parameter_breakdown(model: Any) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]:
    by_component: Dict[str, Dict[str, int]] = defaultdict(lambda: {"total": 0, "trainable": 0})
    detail_rows: List[Dict[str, Any]] = []
    for name, parameter in model.named_parameters():
        total = int(parameter.numel())
        trainable = int(parameter.numel()) if parameter.requires_grad else 0
        component = component_for_name(name)
        by_component[component]["total"] += total
        by_component[component]["trainable"] += trainable
        detail_rows.append(
            {
                "name": name,
                "component": component,
                "shape": list(parameter.shape),
                "total_parameters": total,
                "trainable_parameters": trainable,
                "requires_grad": bool(parameter.requires_grad),
            }
        )

    rows: List[Dict[str, Any]] = []
    total_all = sum(item["total"] for item in by_component.values())
    trainable_all = sum(item["trainable"] for item in by_component.values())
    for component, item in sorted(by_component.items()):
        pct = (100.0 * item["trainable"] / item["total"]) if item["total"] else 0.0
        rows.append(
            {
                "component": component,
                "total_parameters": item["total"],
                "trainable_parameters": item["trainable"],
                "trainable_percentage": f"{pct:.6f}",
            }
        )
    rows.append(
        {
            "component": "Total",
            "total_parameters": total_all,
            "trainable_parameters": trainable_all,
            "trainable_percentage": f"{100.0 * trainable_all / total_all:.6f}" if total_all else "0.000000",
        }
    )
    return rows, detail_rows


def load_model(pipeline: str) -> Any:
    if pipeline == "mvsa_single":
        from src.mvsa_single_pipeline import CLARAModel, DEFAULT_MVSA_SINGLE_CONFIG

        return CLARAModel(dict(DEFAULT_MVSA_SINGLE_CONFIG))
    if pipeline == "mvsa_multiple":
        from src.mvsa_multiple_pipeline import CLARAModel, DEFAULT_MVSA_MULTIPLE_CONFIG

        return CLARAModel(dict(DEFAULT_MVSA_MULTIPLE_CONFIG))
    if pipeline == "hfm":
        from src.hfm_pipeline import CLARAModel, DEFAULT_HFM_CONFIG

        return CLARAModel(dict(DEFAULT_HFM_CONFIG))
    if pipeline == "reference_model":
        from src.model import CLARAConfig, CLARAModel

        return CLARAModel(CLARAConfig())
    raise ValueError(pipeline)


def shape_of(value: Any) -> str:
    if hasattr(value, "shape"):
        return str(list(value.shape))
    if isinstance(value, (tuple, list)):
        return "[" + ", ".join(shape_of(item) for item in value) + "]"
    if isinstance(value, dict):
        return "{" + ", ".join(f"{key}: {shape_of(item)}" for key, item in value.items()) + "}"
    if hasattr(value, "last_hidden_state"):
        return f"last_hidden_state={shape_of(value.last_hidden_state)}"
    return type(value).__name__


def collect_shapes(model: Any, pipeline: str, batch_size: int, seq_len: int, image_size: int, device: str) -> List[Dict[str, str]]:
    import torch

    model.to(device)
    model.eval()
    rows: List[Dict[str, str]] = []
    hooks = []

    def add_row(name: str, tensor: Any, note: str = "") -> None:
        rows.append({"name": name, "shape": shape_of(tensor), "note": note})

    def hook(name: str):
        def _hook(_module: Any, inputs: Tuple[Any, ...], output: Any) -> None:
            rows.append(
                {
                    "name": name,
                    "shape": shape_of(output),
                    "note": "forward hook output",
                }
            )
            if inputs:
                rows.append(
                    {
                        "name": name + ".input",
                        "shape": shape_of(inputs),
                        "note": "forward hook input",
                    }
                )

        return _hook

    for name in [
        "vision",
        "text",
        "fusion",
        "visual_head",
        "text_head",
        "pred",
        "veri",
        "feed",
        "refinement",
        "extra_mlp_control",
        "final",
        "unc",
    ]:
        module = getattr(model, name, None)
        if module is not None:
            hooks.append(module.register_forward_hook(hook(name)))

    pixel_values = torch.zeros(batch_size, 3, image_size, image_size, device=device)
    input_ids = torch.ones(batch_size, seq_len, dtype=torch.long, device=device)
    attention_mask = torch.ones(batch_size, seq_len, dtype=torch.long, device=device)
    if seq_len > 4:
        attention_mask[:, -2:] = 0
    add_row("input_images", pixel_values)
    add_row("input_ids", input_ids)
    add_row("attention_mask", attention_mask)

    with torch.no_grad():
        try:
            out = model(
                pixel_values=pixel_values,
                input_ids=input_ids,
                attention_mask=attention_mask,
                return_attention=True,
            )
        except TypeError:
            out = model(pixel_values=pixel_values, input_ids=input_ids, attention_mask=attention_mask)
        add_row("model_output", out)
        if isinstance(out, dict):
            for key in [
                "projected_visual_tokens",
                "projected_text_tokens",
                "fused",
                "visual_logits",
                "text_logits",
                "pred_logits",
                "disagreement",
                "logits",
            ]:
                if key in out:
                    add_row(key, out[key], "manual output key")

    for handle in hooks:
        handle.remove()
    return rows


def collect_attention_stats(
    model: Any,
    batch_size: int,
    seq_len: int,
    image_size: int,
    device: str,
) -> List[Dict[str, Any]]:
    import torch

    model.to(device)
    model.eval()
    torch.manual_seed(1234)
    pixel_values = torch.randn(batch_size, 3, image_size, image_size, device=device)
    input_ids = torch.randint(1, 1000, (batch_size, seq_len), dtype=torch.long, device=device)
    attention_mask = torch.ones(batch_size, seq_len, dtype=torch.long, device=device)
    if seq_len > 4:
        attention_mask[:, -2:] = 0

    with torch.no_grad():
        try:
            out = model(
                pixel_values=pixel_values,
                input_ids=input_ids,
                attention_mask=attention_mask,
                return_attention=True,
            )
        except TypeError:
            return []

    if not isinstance(out, dict):
        return []

    rows: List[Dict[str, Any]] = []
    padding_mask = ~attention_mask.bool()
    for direction, key in [("v2t", "attention_v2t"), ("t2v", "attention_t2v")]:
        for layer_idx, attn in enumerate(out.get(key, []) or []):
            if attn is None:
                continue
            attn = attn.detach()
            entropy = -(attn.clamp_min(1e-12) * attn.clamp_min(1e-12).log()).sum(dim=-1)
            for head_idx in range(attn.shape[1]):
                head = attn[:, head_idx]
                padding_mass: Any = ""
                if direction == "v2t" and padding_mask.any():
                    padding_mass = float(head.masked_select(padding_mask[:, None, :]).sum().item())
                rows.append(
                    {
                        "dataset": "",
                        "seed": "",
                        "sample_id": "dummy_batch",
                        "layer": layer_idx,
                        "direction": direction,
                        "head": head_idx,
                        "attention_mean": float(head.mean().item()),
                        "attention_std": float(head.std().item()),
                        "attention_entropy": float(entropy[:, head_idx].mean().item()),
                        "max_attention": float(head.max().item()),
                        "padding_attention_mass": padding_mass,
                    }
                )
    return rows


def collect_gradient_probe(
    model: Any,
    batch_size: int,
    seq_len: int,
    image_size: int,
    device: str,
) -> List[Dict[str, Any]]:
    import torch
    import torch.nn as nn

    try:
        from src.revised_ravel_model import token_loss
    except Exception:
        token_loss = None

    model.to(device)
    model.train()
    model.zero_grad(set_to_none=True)
    torch.manual_seed(4321)
    pixel_values = torch.randn(batch_size, 3, image_size, image_size, device=device)
    input_ids = torch.randint(1, 1000, (batch_size, seq_len), dtype=torch.long, device=device)
    attention_mask = torch.ones(batch_size, seq_len, dtype=torch.long, device=device)
    if seq_len > 4:
        attention_mask[:, -2:] = 0
    labels = torch.arange(batch_size, device=device) % int(getattr(model, "num_classes", 2))

    out = model(pixel_values=pixel_values, input_ids=input_ids, attention_mask=attention_mask)
    criterion = nn.CrossEntropyLoss()
    if isinstance(out, dict) and {"visual_logits", "text_logits", "pred_logits", "logits"}.issubset(out):
        if token_loss is None:
            loss = criterion(out["logits"], labels)
        else:
            loss, _ = token_loss(out, labels, criterion)
    elif isinstance(out, dict):
        loss = criterion(out["logits"], labels)
        if "pred_logits" in out:
            loss = loss + 0.1 * criterion(out["pred_logits"], labels)
        if "verify_logits" in out:
            loss = loss + 0.1 * criterion(out["verify_logits"], labels)
    else:
        return []
    loss.backward()

    rows: List[Dict[str, Any]] = []
    for name, parameter in model.named_parameters():
        grad = parameter.grad
        rows.append(
            {
                "component": component_for_name(name),
                "parameter_name": name,
                "requires_grad": bool(parameter.requires_grad),
                "gradient_present": grad is not None,
                "gradient_norm": "" if grad is None else float(grad.detach().norm().item()),
            }
        )
    model.zero_grad(set_to_none=True)
    return rows


def collect_modality_dependency_tests(
    model: Any,
    batch_size: int,
    seq_len: int,
    image_size: int,
    device: str,
) -> List[Dict[str, Any]]:
    import torch

    model.to(device)
    model.eval()
    torch.manual_seed(5678)

    pixel_values = torch.randn(batch_size, 3, image_size, image_size, device=device)
    replacement_pixels = torch.randn(batch_size, 3, image_size, image_size, device=device)
    blank_pixels = torch.zeros_like(pixel_values)

    input_ids = torch.randint(1, 1000, (batch_size, seq_len), dtype=torch.long, device=device)
    replacement_input_ids = torch.randint(
        1,
        1000,
        (batch_size, seq_len),
        dtype=torch.long,
        device=device,
    )
    attention_mask = torch.ones(batch_size, seq_len, dtype=torch.long, device=device)
    if seq_len > 4:
        attention_mask[:, -2:] = 0

    minimal_text_ids = torch.ones_like(input_ids)
    minimal_text_mask = torch.zeros_like(attention_mask)
    minimal_text_mask[:, 0] = 1

    def forward_probs(
        pixels: Any,
        ids: Any,
        mask: Any,
    ) -> Tuple[Any, Any]:
        with torch.no_grad():
            out = model(pixel_values=pixels, input_ids=ids, attention_mask=mask)
        if not isinstance(out, dict) or "visual_probs" not in out or "text_probs" not in out:
            raise RuntimeError("Model output does not include visual_probs/text_probs.")
        return out["visual_probs"].detach(), out["text_probs"].detach()

    base_visual, base_text = forward_probs(pixel_values, input_ids, attention_mask)

    cases = [
        (
            "replace_image_only",
            replacement_pixels,
            input_ids,
            attention_mask,
            "visual posterior changes more than text posterior",
        ),
        (
            "replace_text_only",
            pixel_values,
            replacement_input_ids,
            attention_mask,
            "text posterior changes more than visual posterior",
        ),
        (
            "blank_image",
            blank_pixels,
            input_ids,
            attention_mask,
            "visual posterior changes more than text posterior",
        ),
        (
            "empty_text",
            pixel_values,
            minimal_text_ids,
            minimal_text_mask,
            "text posterior changes more than visual posterior",
        ),
    ]

    rows: List[Dict[str, Any]] = []
    eps = 1e-8
    for test_case, pixels, ids, mask, expected in cases:
        visual_probs, text_probs = forward_probs(pixels, ids, mask)
        visual_change = float((base_visual - visual_probs).abs().mean().item())
        text_change = float((base_text - text_probs).abs().mean().item())
        if test_case in {"replace_image_only", "blank_image"}:
            passed = visual_change > text_change + eps and text_change < 1e-6
        else:
            passed = text_change > visual_change + eps and visual_change < 1e-6
        rows.append(
            {
                "test_case": test_case,
                "visual_posterior_change": visual_change,
                "text_posterior_change": text_change,
                "expected_behavior": expected,
                "passed": bool(passed),
            }
        )
    return rows


def main() -> None:
    args = parse_args()
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    model = load_model(args.pipeline)
    breakdown, detail = parameter_breakdown(model)
    write_csv(
        output_dir / f"parameter_breakdown_{args.pipeline}.csv",
        breakdown,
        ["component", "total_parameters", "trainable_parameters", "trainable_percentage"],
    )
    write_csv(
        output_dir / f"parameter_detail_{args.pipeline}.csv",
        detail,
        ["name", "component", "shape", "total_parameters", "trainable_parameters", "requires_grad"],
    )
    shapes = collect_shapes(
        model=model,
        pipeline=args.pipeline,
        batch_size=args.batch_size,
        seq_len=args.seq_len,
        image_size=args.image_size,
        device=args.device,
    )
    write_csv(output_dir / f"tensor_shapes_{args.pipeline}.csv", shapes, ["name", "shape", "note"])
    with (output_dir / f"tensor_shapes_{args.pipeline}.txt").open("w", encoding="utf-8") as f:
        for row in shapes:
            f.write(f"{row['name']}: {row['shape']}  # {row['note']}\n")
    attention_rows = collect_attention_stats(
        model=model,
        batch_size=args.batch_size,
        seq_len=args.seq_len,
        image_size=args.image_size,
        device=args.device,
    )
    write_csv(
        output_dir / f"attention_statistics_{args.pipeline}.csv",
        attention_rows,
        [
            "dataset",
            "seed",
            "sample_id",
            "layer",
            "direction",
            "head",
            "attention_mean",
            "attention_std",
            "attention_entropy",
            "max_attention",
            "padding_attention_mass",
        ],
    )
    gradient_rows = collect_gradient_probe(
        model=model,
        batch_size=args.batch_size,
        seq_len=args.seq_len,
        image_size=args.image_size,
        device=args.device,
    )
    write_csv(
        output_dir / f"gradient_probe_{args.pipeline}.csv",
        gradient_rows,
        ["component", "parameter_name", "requires_grad", "gradient_present", "gradient_norm"],
    )
    try:
        modality_rows = collect_modality_dependency_tests(
            model=model,
            batch_size=args.batch_size,
            seq_len=args.seq_len,
            image_size=args.image_size,
            device=args.device,
        )
    except RuntimeError:
        modality_rows = []
    write_csv(
        output_dir / f"modality_dependency_tests_{args.pipeline}.csv",
        modality_rows,
        [
            "test_case",
            "visual_posterior_change",
            "text_posterior_change",
            "expected_behavior",
            "passed",
        ],
    )
    print(f"Wrote runtime probe outputs to {output_dir}")


if __name__ == "__main__":
    main()