"""Prediction-transfer alignment energy for the synthetic world.

The third rung of the energy ladder: no hand-designed geometry and no
pair-trained EBM. Each modality's own predictor induces a substitution
kernel over scenes -- which other scenes the world model finds compatible
-- and the alignment energy is the conjugacy defect of the two kernels
under a candidate coupling.

- Vision kernel: the masked-reconstruction tower inpaints a half-masked
  render; pixel distance between the inpainting and other scenes' renders
  gives K_V[i, i'].
- Text kernel: the causal LM scores scene j's caption sentences as a
  continuation of scene i's caption prefix; referential consistency makes
  same-content continuations likely, giving K_T.

Both kernels are unimodal-predictor functionals. Hidden pairs enter only
through the gate that scores orderings of the energy.
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import torch
import torch.nn.functional as F
from tqdm import tqdm

from .common import batch_indices, read_json, seed_everything, write_json
from .manifold_gate import standardize_relation
from .ricci_control import run_gates
from .synth_towers import TextTower, VisionTower, load_image, tokenize


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--data-dir", default="artifacts/synth_v0")
    parser.add_argument(
        "--vision-tower", default="artifacts/synth_v0/vision_tower_simmim.pt"
    )
    parser.add_argument("--text-tower", default="artifacts/synth_v0/text_tower.pt")
    parser.add_argument("--split", choices=["val", "test"], default="test")
    parser.add_argument("--samples", type=int, default=512)
    parser.add_argument("--mask-ratio", type=float, default=0.5)
    parser.add_argument("--kernel-temperature", type=float, default=0.03)
    parser.add_argument("--batch-size", type=int, default=256)
    parser.add_argument("--random-perms", type=int, default=300)
    parser.add_argument("--descent-restarts", type=int, default=5)
    parser.add_argument("--descent-max-steps", type=int, default=200000)
    parser.add_argument(
        "--descent-objective", default="mse", choices=["mse", "m30_total"]
    )
    parser.add_argument("--descent-verify-top", type=int, default=64)
    parser.add_argument("--device", default="cuda:3")
    parser.add_argument("--seed", type=int, default=20260731)
    parser.add_argument(
        "--output", default="artifacts/synth_v0/transfer_energy_gate.json"
    )
    parser.add_argument(
        "--kernels-output", default="artifacts/synth_v0/transfer_kernels.pt"
    )
    return parser.parse_args()


@torch.inference_mode()
def vision_kernel(
    rows: list[int], manifest: dict, args: argparse.Namespace
) -> torch.Tensor:
    state = torch.load(args.vision_tower, map_location="cpu", weights_only=False)
    saved = state["args"]
    model = VisionTower(
        manifest["image_size"], saved["patch"], saved["dim"], saved["depth"], saved["heads"]
    ).to(args.device)
    model.load_state_dict(state["model"])
    model.eval()
    image_dir = Path(manifest["image_dir"])
    tokens = (manifest["image_size"] // saved["patch"]) ** 2
    generator = torch.Generator(device=args.device).manual_seed(args.seed)

    references = []
    for indices in batch_indices(len(rows), args.batch_size):
        pixels = torch.stack(
            [load_image(image_dir / f"scene{rows[i]:06d}_v0.png") for i in indices]
        )
        references.append(pixels)
    references = torch.cat(references)  # [N, 3, H, W] view-0 renders on cpu

    inpainted = []
    for indices in tqdm(
        list(batch_indices(len(rows), args.batch_size)), desc="vision kernel"
    ):
        # Inpaint view 1 under a random half mask; compare to view-0 renders.
        pixels = torch.stack(
            [load_image(image_dir / f"scene{rows[i]:06d}_v1.png") for i in indices]
        ).to(args.device)
        mask = (
            torch.rand(len(pixels), tokens, generator=generator, device=args.device)
            < args.mask_ratio
        )
        encoded = model.encode(pixels, mask=mask)
        patches = model.reconstruction(encoded[:, 1:])  # [B, T, p*p*3]
        grid = manifest["image_size"] // saved["patch"]
        patch = saved["patch"]
        image = patches.reshape(len(pixels), grid, grid, patch, patch, 3)
        image = image.permute(0, 5, 1, 3, 2, 4).reshape(
            len(pixels), 3, manifest["image_size"], manifest["image_size"]
        )
        blended = torch.where(
            mask.reshape(len(pixels), 1, grid, 1, grid, 1)
            .expand(-1, 3, -1, patch, -1, patch)
            .reshape_as(image),
            image,
            pixels,
        )
        inpainted.append(blended.float().cpu())
    inpainted = torch.cat(inpainted)

    # Layout-invariant content distance: Chamfer between foreground patch
    # sets in pixel space. Renders of the same scene share patch CONTENT
    # (colors, local shape fragments) while positions are resampled, so
    # per-pixel comparison would measure layout instead.
    patch = saved["patch"]
    grid = manifest["image_size"] // patch

    def patch_sets(images: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        pieces = images.reshape(len(images), 3, grid, patch, grid, patch)
        pieces = pieces.permute(0, 2, 4, 1, 3, 5).reshape(len(images), grid * grid, -1)
        background = pieces.mean(-1)
        foreground = background > background.min(1, keepdim=True).values + 0.02
        return pieces, foreground

    inpaint_pieces, inpaint_fg = patch_sets(inpainted)
    reference_pieces, reference_fg = patch_sets(references)
    distances = torch.zeros(len(rows), len(rows))
    device = args.device
    reference_pieces = reference_pieces.to(device)
    reference_fg = reference_fg.to(device)
    for start in range(0, len(rows), 32):
        stop = min(start + 32, len(rows))
        block = inpaint_pieces[start:stop].to(device)
        block_fg = inpaint_fg[start:stop].to(device)
        cross = torch.cdist(block, reference_pieces.reshape(-1, block.shape[-1]))
        cross = cross.reshape(stop - start, block.shape[1], len(rows), -1)
        masked_ab = cross.masked_fill(
            ~reference_fg[None, None, :, :], float("inf")
        ).amin(-1)
        forward = (
            (masked_ab * block_fg[:, :, None]).sum(1)
            / block_fg.sum(1).clamp_min(1)[:, None]
        )
        masked_ba = cross.masked_fill(
            ~block_fg[:, :, None, None].expand_as(cross), float("inf")
        ).amin(1)
        backward = (
            (masked_ba * reference_fg[None, :, :]).sum(-1)
            / reference_fg.sum(-1).clamp_min(1)[None, :]
        )
        distances[start:stop] = (0.5 * (forward + backward)).float().cpu()
    distances[torch.isinf(distances) | torch.isnan(distances)] = distances[
        torch.isfinite(distances)
    ].max()
    return distances


@torch.inference_mode()
def text_kernel(
    rows: list[int], captions: list[list[str]], args: argparse.Namespace
) -> torch.Tensor:
    state = torch.load(args.text_tower, map_location="cpu", weights_only=False)
    saved = state["args"]
    vocab = state["vocab"]
    model = TextTower(
        len(vocab), saved["text_dim"], saved["depth"], saved.get("text_heads", 4), saved["context"]
    ).to(args.device)
    model.load_state_dict(state["model"])
    model.eval()

    prefixes = [tokenize(captions[row][0], vocab) for row in rows]
    # Continuations are the RELATION sentences only: they refer back to
    # groups the prefix must have introduced, so their likelihood grades
    # shared content instead of rewarding verbatim repetition.
    continuations = []
    for row in rows:
        sentences = captions[row][1].split(". ")
        relational = ". ".join(sentences[1:]) if len(sentences) > 1 else sentences[0]
        continuations.append(tokenize(relational, vocab)[1:])  # drop bos
    scores = torch.zeros(len(rows), len(rows))
    jobs = [(i, j) for i in range(len(rows)) for j in range(len(rows))]
    for indices in tqdm(
        list(batch_indices(len(jobs), args.batch_size)), desc="text kernel"
    ):
        batch = [jobs[k] for k in indices]
        sequences = [
            (prefixes[i] + continuations[j])[: saved["context"]] for i, j in batch
        ]
        longest = max(len(s) for s in sequences)
        tokens = torch.zeros(len(batch), longest, dtype=torch.long)
        for row, sequence in enumerate(sequences):
            tokens[row, : len(sequence)] = torch.tensor(sequence)
        tokens = tokens.to(args.device)
        hidden = model(tokens)
        logits = model.logits(hidden[:, :-1])
        log_probs = F.log_softmax(logits, dim=-1)
        for row, (i, j) in enumerate(batch):
            begin = len(prefixes[i]) - 1
            end = min(len(prefixes[i]) + len(continuations[j]), longest) - 1
            if end <= begin:
                scores[i, j] = 0.0
                continue
            targets = tokens[row, begin + 1 : end + 1]
            picked = log_probs[row, begin:end].gather(1, targets[:, None])
            scores[i, j] = float(picked.mean())
    return -scores  # negative mean log-likelihood as a distance


def main() -> None:
    args = parse_args()
    seed_everything(args.seed)
    manifest = read_json(Path(args.data_dir, "manifest.json"))
    captions = read_json(Path(args.data_dir, "captions.json"))["captions"]
    rows = manifest[args.split][: args.samples]

    visual_distance = vision_kernel(rows, manifest, args)
    text_distance = text_kernel(rows, captions, args)
    torch.save(
        {
            "rows": rows,
            "visual_distance": visual_distance,
            "text_distance": text_distance,
        },
        args.kernels_output,
    )

    # Substitution kernels as standardized relation channels: negative
    # distances, symmetrized, standardized off-diagonal. The gate machinery
    # then scores orderings exactly as for any relation field.
    visual_similarity = -visual_distance
    text_similarity = -0.5 * (text_distance + text_distance.T)
    visual_channels = standardize_relation(
        0.5 * (visual_similarity + visual_similarity.T).double()
    )[0][None]
    text_channels = standardize_relation(text_similarity.double())[0][None]
    generator = torch.Generator().manual_seed(args.seed)
    report = {
        "protocol": (
            "Substitution kernels are unimodal-predictor functionals "
            "(masked inpainting distance; caption continuation "
            "likelihood). Hidden pairs score orderings only."
        ),
        "split": args.split,
        "samples": len(rows),
        **run_gates(text_channels, visual_channels, args, generator),
    }
    verdict = {
        "true_z_mse": report["gate_a"]["random"]["mse"]["true_z"],
        "improving_fraction": report["gate_b"]["improving_fraction"],
        "descent_keeps": report["descent_from_true"]["final_accuracy"],
        "true_mse": report["gate_a"]["true"]["mse"],
        "best_random_descent": min(
            (r["final_objective"] for r in report["descent_from_random"]),
            default=None,
        ),
    }
    verdict["counterfeit_found"] = bool(
        verdict["best_random_descent"] is not None
        and verdict["best_random_descent"] < verdict["true_mse"]
    )
    report["verdict"] = verdict
    print(json.dumps({"verdict": verdict}))
    write_json(args.output, report)
    print(f"Wrote {args.output}")


if __name__ == "__main__":
    main()