#!/usr/bin/env python3
from __future__ import annotations

import argparse
from dataclasses import asdict

import torch

from alpamayo_lite.config import load_config
from alpamayo_lite.data import load_manifest
from alpamayo_lite.expert import ExpertArchitecture, KVCachedDiffusionExpert
from alpamayo_lite.prompts import build_expert_prefix, build_vlm_prefix
from alpamayo_lite.runtime import (
    empty_device_cache,
    parameter_counts,
    resolve_dtype,
    seed_everything,
    select_device,
    write_json,
)
from alpamayo_lite.vlm import (
    configure_stage1_trainability,
    extract_prefix_kv,
    generate_reasoning,
    load_vlm,
    vlm_text_config,
)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", default="configs/local_overfit.yaml")
    args = parser.parse_args()
    cfg = load_config(args.config)
    seed_everything(cfg.seed)
    device = select_device()
    dtype = resolve_dtype(cfg.vlm.dtype, device)
    sample = load_manifest(cfg.resolve(cfg.paths.manifest))[0]
    history, _ = sample.load_trajectory()

    processor, vlm = load_vlm(cfg.resolve(cfg.paths.model_dir), cfg.vlm, device, dtype)
    vlm.eval()
    baseline_text = generate_reasoning(
        vlm,
        processor,
        sample.load_image(),
        build_vlm_prefix(history),
        device,
        dtype,
        cfg.vlm.max_reasoning_tokens,
    )
    prefix = extract_prefix_kv(
        vlm,
        processor,
        sample.load_image(),
        build_expert_prefix(history, sample.reasoning),
        device,
        dtype,
    )

    stage2 = cfg.stage2
    architecture = ExpertArchitecture(
        hidden_size=stage2.expert_hidden_size,
        intermediate_size=stage2.expert_intermediate_size,
        num_layers=stage2.expert_num_layers,
        num_attention_heads=stage2.expert_num_attention_heads,
        num_key_value_heads=stage2.expert_num_key_value_heads,
        head_dim=stage2.expert_head_dim,
        action_projection_hidden_size=stage2.action_projection_hidden_size,
        action_fourier_features=stage2.action_fourier_features,
        non_causal_attention=stage2.expert_non_causal_attention,
    )
    expert = KVCachedDiffusionExpert(vlm_text_config(vlm), architecture).to(
        device=device, dtype=dtype
    )
    with torch.no_grad():
        predicted_velocity = expert(
            noisy_action=torch.randn(1, cfg.data.expert_waypoints, 2, device=device, dtype=dtype),
            timestep=torch.full((1, 1, 1), 0.5, device=device, dtype=dtype),
            prefix=prefix,
        )

    trainability = configure_stage1_trainability(vlm, cfg.stage1)
    result = {
        "device": str(device),
        "dtype": str(dtype),
        "vlm_parameters": parameter_counts(vlm),
        "stage1_trainability": trainability,
        "expert_parameters": parameter_counts(expert),
        "expert_architecture": asdict(architecture),
        "prefix_layers": prefix.num_layers,
        "prefix_sequence_length": prefix.sequence_length,
        "prefix_key_shape": list(prefix.keys[0].shape),
        "expert_output_shape": list(predicted_velocity.shape),
        "expert_output_finite": bool(torch.isfinite(predicted_velocity).all().item()),
        "baseline_text": baseline_text,
        "target_text": sample.reasoning,
    }
    write_json(cfg.root / "outputs" / "connection_test.json", result)
    print(result)
    del expert, vlm
    empty_device_cache(device)


if __name__ == "__main__":
    main()