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 (
    projection_quantile_target,
    prototype_manifold_energy,
    relation_field_energy,
    retrieval_metrics,
    sliced_distribution_energy,
    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",
        help=(
            "Optional multi-description feature cache. When supplied, each "
            "language particle is the normalized mean of its observation "
            "orbit."
        ),
    )
    parser.add_argument("--prototypes", default="artifacts/gw.pt")
    parser.add_argument("--split", choices=["val", "test"], default="test")
    parser.add_argument("--samples", type=int, default=256)
    parser.add_argument("--steps", type=int, default=100)
    parser.add_argument("--lr", type=float, default=0.03)
    parser.add_argument("--projections", type=int, default=128)
    parser.add_argument("--relation-weight-start", type=float, default=0.4)
    parser.add_argument("--relation-weight-end", type=float, default=2.0)
    parser.add_argument("--conditional-weight-start", type=float, default=0.05)
    parser.add_argument("--conditional-weight-end", type=float, default=0.2)
    parser.add_argument("--distribution-weight", type=float, default=80.0)
    parser.add_argument("--manifold-weight", type=float, default=1.0)
    parser.add_argument("--noise", type=float, default=0.01)
    parser.add_argument(
        "--shuffle-visual-energy",
        action="store_true",
        help=(
            "Evaluation control: permute image particles before constructing "
            "the energy while leaving evaluation rows unchanged."
        ),
    )
    parser.add_argument("--device", default="cuda:1")
    parser.add_argument("--seed", type=int, default=20260729)
    parser.add_argument(
        "--output", default="artifacts/energy_free_test.pt"
    )
    parser.add_argument(
        "--metrics-output", default="artifacts/energy_free_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:
        orbit_state = torch.load(
            args.text_orbits, map_location="cpu", weights_only=False
        )
        orbit_lookup = {
            int(row): index
            for index, row in enumerate(orbit_state["rows"])
        }
        orbit_mean = F.normalize(
            orbit_state["features"].float().mean(1), dim=-1
        )
        paired_text = select_rows(
            orbit_mean, orbit_lookup, rows
        ).to(args.device)
        text_population = select_rows(
            orbit_mean, orbit_lookup, manifest["text_only_train"]
        ).to(args.device)
    if args.shuffle_visual_energy:
        control_generator = torch.Generator(
            device=args.device
        ).manual_seed(args.seed + 10_000)
        visual = visual[
            torch.randperm(
                len(visual),
                generator=control_generator,
                device=args.device,
            )
        ]
    prototype_state = torch.load(
        args.prototypes, map_location="cpu", weights_only=False
    )
    prototypes = prototype_state["text_centers"].to(args.device)
    if prototypes.shape[-1] != text_population.shape[-1]:
        raise ValueError(
            "Prototype dimension does not match text features; use the GW "
            "cache built from the selected text backbone."
        )

    generator = torch.Generator(device=args.device).manual_seed(args.seed)
    initial_indices = torch.randperm(
        len(text_population),
        generator=generator,
        device=args.device,
    )[: len(visual)]
    initial = text_population[initial_indices].clone()
    initial = initial + args.noise * torch.randn(
        initial.shape, generator=generator, device=args.device
    )
    particles = torch.nn.Parameter(F.normalize(initial, dim=-1))
    directions, target_quantiles = projection_quantile_target(
        text_population,
        particles=len(particles),
        projections=args.projections,
        generator=generator,
    )
    visual_relation, visual_standardized = standardized_relation(visual)
    optimizer = torch.optim.Adam([particles], lr=args.lr)

    def energy_values() -> tuple[torch.Tensor, ...]:
        relation, conditional = relation_field_energy(
            visual_relation, visual_standardized, particles
        )
        distribution = sliced_distribution_energy(
            particles, directions, target_quantiles
        )
        manifold = prototype_manifold_energy(particles, prototypes)
        return relation, conditional, distribution, manifold

    history: list[dict] = []
    initial_cpu = F.normalize(particles.detach(), dim=-1).cpu()
    for step in range(args.steps + 1):
        relation, conditional, distribution, manifold = energy_values()
        progress = min(step / max(args.steps, 1), 1.0)
        relation_weight = (
            args.relation_weight_start
            + progress
            * (args.relation_weight_end - args.relation_weight_start)
        )
        conditional_weight = (
            args.conditional_weight_start
            + progress
            * (
                args.conditional_weight_end
                - args.conditional_weight_start
            )
        )
        loss = (
            relation_weight * relation
            + conditional_weight * conditional
            + args.distribution_weight * distribution
            + args.manifold_weight * manifold
        )
        if step % 20 == 0 or step == args.steps:
            history.append(
                {
                    "step": step,
                    "total": float(loss.detach()),
                    "relation": float(relation.detach()),
                    "conditional": float(conditional.detach()),
                    "distribution": float(distribution.detach()),
                    "manifold": float(manifold.detach()),
                    "paired_evaluation_only": retrieval_metrics(
                        particles.detach(), paired_text
                    ),
                }
            )
            print(json.dumps(history[-1]))
        if step == args.steps:
            break
        optimizer.zero_grad(set_to_none=True)
        loss.backward()
        torch.nn.utils.clip_grad_norm_([particles], 2.0)
        optimizer.step()
        with torch.no_grad():
            particles.copy_(F.normalize(particles, dim=-1))

    with torch.no_grad():
        oracle_relation, oracle_conditional = relation_field_energy(
            visual_relation, visual_standardized, paired_text
        )
        oracle_distribution = sliced_distribution_energy(
            paired_text, directions, target_quantiles
        )
        oracle_manifold = prototype_manifold_energy(
            paired_text, prototypes
        )
    result = {
        "protocol": (
            "No cross-modal map and no image-text pair is used by the energy "
            "or optimizer. Paired text is loaded only for trajectory and "
            "oracle diagnostics; the final step is fixed by CLI arguments."
        ),
        "mode": "free_language_latent_particles",
        "split": args.split,
        "rows": rows,
        "vision_model": vision["model"],
        "text_model": text["model"],
        "text_observation": (
            "multi-description orbit mean"
            if args.text_orbits
            else "single description"
        ),
        "args": vars(args),
        "history": history,
        "oracle_energy_components": {
            "relation": float(oracle_relation),
            "conditional": float(oracle_conditional),
            "distribution": float(oracle_distribution),
            "manifold": float(oracle_manifold),
        },
    }
    state = {
        **result,
        "initial_particles": initial_cpu,
        "final_particles": F.normalize(
            particles.detach(), dim=-1
        ).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()