from __future__ import annotations

import argparse
import json
from pathlib import Path

import torch
import torch.nn.functional as F

from .common import read_json, seed_everything, write_json
from .energy import (
    log_sinkhorn,
    relation_field_energy,
    retrieval_metrics,
    standardized_relation,
)
from .io import load_feature_pair, select_rows


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--manifest", default="artifacts/manifest.json")
    parser.add_argument("--vision", default="artifacts/vision.pt")
    parser.add_argument("--text", default="artifacts/text.pt")
    parser.add_argument("--text-orbits")
    parser.add_argument("--split", choices=["val", "test"], default="test")
    parser.add_argument("--samples", type=int, default=256)
    parser.add_argument(
        "--reservoir",
        choices=["same_set_shuffled", "unpaired"],
        default="same_set_shuffled",
    )
    parser.add_argument("--steps", type=int, default=1_200)
    parser.add_argument("--lr", type=float, default=0.12)
    parser.add_argument("--temperature-start", type=float, default=0.8)
    parser.add_argument("--temperature-end", type=float, default=0.12)
    parser.add_argument("--conditional-weight", type=float, default=0.08)
    parser.add_argument("--entropy-weight-start", type=float, default=0.005)
    parser.add_argument("--entropy-weight-end", type=float, default=0.055)
    parser.add_argument("--device", default="cuda:1")
    parser.add_argument("--seed", type=int, default=20260729)
    parser.add_argument(
        "--output", default="artifacts/energy_coupling_test.pt"
    )
    parser.add_argument(
        "--metrics-output", default="artifacts/energy_coupling_test.json"
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    seed_everything(args.seed)
    manifest = read_json(args.manifest)
    vision, text, vision_lookup, text_lookup = load_feature_pair(
        args.vision, args.text
    )
    rows = manifest[args.split][: args.samples]
    visual = select_rows(
        vision["features"], vision_lookup, rows
    ).to(args.device)
    paired_text = select_rows(
        text["features"], text_lookup, rows
    ).to(args.device)
    text_population = select_rows(
        text["features"], text_lookup, manifest["text_only_train"]
    ).to(args.device)
    if args.text_orbits:
        state = torch.load(
            args.text_orbits, map_location="cpu", weights_only=False
        )
        lookup = {
            int(row): index for index, row in enumerate(state["rows"])
        }
        orbit_mean = F.normalize(
            state["features"].float().mean(1), dim=-1
        )
        paired_text = select_rows(
            orbit_mean, lookup, rows
        ).to(args.device)
        text_population = select_rows(
            orbit_mean, lookup, manifest["text_only_train"]
        ).to(args.device)

    generator = torch.Generator(device=args.device).manual_seed(args.seed)
    target: torch.Tensor | None = None
    if args.reservoir == "same_set_shuffled":
        permutation = torch.randperm(
            len(paired_text), generator=generator, device=args.device
        )
        anchors = paired_text[permutation]
        target = torch.argsort(permutation)
    else:
        anchor_indices = torch.randperm(
            len(text_population),
            generator=generator,
            device=args.device,
        )[: len(visual)]
        anchors = text_population[anchor_indices]

    initial_permutation = torch.randperm(
        len(visual), generator=generator, device=args.device
    )
    logits = torch.nn.Parameter(
        0.02
        * torch.randn(
            len(visual),
            len(visual),
            generator=generator,
            device=args.device,
        )
    )
    with torch.no_grad():
        logits[torch.arange(len(visual), device=args.device), initial_permutation] += 3
    visual_relation, visual_standardized = standardized_relation(visual)
    optimizer = torch.optim.Adam([logits], lr=args.lr)
    history: list[dict] = []

    for step in range(args.steps + 1):
        progress = min(step / max(args.steps, 1), 1.0)
        temperature = max(
            args.temperature_end,
            args.temperature_start
            + progress
            * (args.temperature_end - args.temperature_start),
        )
        coupling = log_sinkhorn(logits, temperature, iterations=15)
        particles = F.normalize(coupling @ anchors, dim=-1)
        relation, conditional = relation_field_energy(
            visual_relation, visual_standardized, particles
        )
        entropy = -(
            coupling * coupling.clamp_min(1e-12).log()
        ).sum(-1).mean()
        entropy_weight = (
            args.entropy_weight_start
            + progress
            * (args.entropy_weight_end - args.entropy_weight_start)
        )
        loss = (
            relation
            + args.conditional_weight * conditional
            + entropy_weight * entropy
        )
        if step % 100 == 0 or step == args.steps:
            record = {
                "step": step,
                "total": float(loss.detach()),
                "relation": float(relation.detach()),
                "conditional": float(conditional.detach()),
                "entropy": float(entropy.detach()),
                "temperature": temperature,
                "paired_evaluation_only": retrieval_metrics(
                    particles.detach(), paired_text
                ),
            }
            if target is not None:
                prediction = coupling.argmax(-1)
                record["exact_coupling_evaluation_only"] = {
                    "accuracy": float((prediction == target).float().mean()),
                    "true_mass": float(
                        coupling[
                            torch.arange(
                                len(coupling), device=args.device
                            ),
                            target,
                        ].mean()
                    ),
                }
            history.append(record)
            print(json.dumps(record))
        if step == args.steps:
            break
        optimizer.zero_grad(set_to_none=True)
        loss.backward()
        torch.nn.utils.clip_grad_norm_([logits], 5.0)
        optimizer.step()

    result = {
        "protocol": (
            "The optimizer sees only frozen within-modality states and "
            "energy. Pair identity and paired text are used only for "
            "evaluation diagnostics."
        ),
        "mode": "mass_conserving_language_state_coupling",
        "reservoir": args.reservoir,
        "split": args.split,
        "rows": rows,
        "args": vars(args),
        "history": history,
    }
    state = {
        **result,
        "anchors": anchors.detach().cpu(),
        "final_coupling": coupling.detach().cpu(),
        "final_particles": particles.detach().cpu(),
    }
    Path(args.output).parent.mkdir(parents=True, exist_ok=True)
    torch.save(state, args.output)
    write_json(args.metrics_output, result)
    print(f"Wrote {args.output} and {args.metrics_output}")


if __name__ == "__main__":
    main()