#!/usr/bin/env python3
"""Compute ELF timestep gradient alignment.

This probes the continuous-time flow-matching objective used by ELF.  It is
not a masked-token objective: for each fixed timestep t, the script noices a
batch of encoded text latents with the same Gaussian noise, computes the
denoising velocity L2 loss, and compares parameter gradients across timesteps.
"""

from __future__ import annotations

import argparse
import json
import os
import re
import sys
import tempfile
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

import numpy as np


DEFAULT_TEXTS = [
    "The quick brown fox jumps over the lazy dog.",
    "A language model can generate text by gradually denoising a latent state.",
    "Mathematical reasoning often benefits from preserving intermediate steps.",
    "Diffusion models and autoregressive models expose different computation paths.",
]


def parse_values(spec: str) -> list[float]:
    """Parse comma values or start:step:end, inclusive within floating tolerance."""
    spec = spec.strip()
    if ":" not in spec:
        return [float(x.strip()) for x in spec.split(",") if x.strip()]
    start_s, step_s, end_s = spec.split(":")
    start, step, end = float(start_s), float(step_s), float(end_s)
    vals = []
    cur = start
    while cur <= end + step * 0.5:
        vals.append(round(cur, 10))
        cur += step
    return vals


def tree_path_to_str(path: Any) -> str:
    parts = []
    for item in path:
        key = getattr(item, "key", None)
        if key is None:
            key = getattr(item, "name", None)
        if key is None:
            key = str(item)
        parts.append(str(key))
    return "/".join(parts)


def selected_tree_stats(jax, jnp, grads: Any, pattern: str | None):
    leaves = []
    names = []
    regex = re.compile(pattern) if pattern else None
    for path, leaf in jax.tree_util.tree_flatten_with_path(grads)[0]:
        name = tree_path_to_str(path)
        if regex is None or regex.search(name):
            leaves.append(leaf)
            names.append(name)
    if not leaves:
        raise ValueError(f"No gradient leaves matched pattern: {pattern!r}")
    sq_norm = sum(jnp.vdot(x, x).real for x in leaves)
    return leaves, names, sq_norm


def tree_dot(jnp, left: list[Any], right: list[Any]):
    return sum(jnp.vdot(a, b).real for a, b in zip(left, right))


def write_report(output_dir: Path, payload: dict[str, Any]) -> None:
    t_values = payload["t_values"]
    rows = payload["cosine_similarity"]
    lines = [
        "# ELF Timestep Gradient Alignment",
        "",
        f"- Updated: `{payload['updated_at']}`",
        f"- Model: `{payload['model']}`",
        f"- Checkpoint: `{payload['checkpoint_path']}`",
        f"- Samples: `{payload['num_samples']}`",
        f"- Max length: `{payload['max_length']}`",
        f"- Gradient leaf regex: `{payload['grad_leaf_regex']}`",
        f"- Selected leaves: `{payload['selected_leaf_count']}`",
        f"- Selected params: `{payload['selected_param_count']}`",
        "",
        "## Loss By Timestep",
        "",
        "| t | loss |",
        "| ---: | ---: |",
    ]
    for t, loss in zip(t_values, payload["loss_by_t"], strict=True):
        lines.append(f"| {t:g} | {loss:.6f} |")
    lines.extend(["", "## Adjacent Cosines", "", "| step pair | cosine |", "| --- | ---: |"])
    for item in payload["adjacent_cosines"]:
        lines.append(f"| {item['from']:g} -> {item['to']:g} | {item['cosine']:.3f} |")
    lines.extend(["", "## Cosine Similarity", ""])
    header = "| t | " + " | ".join(f"{t:g}" for t in t_values) + " |"
    lines.append(header)
    lines.append("| --- | " + " | ".join("---:" for _ in t_values) + " |")
    for t, row in zip(t_values, rows, strict=True):
        lines.append("| " + f"{t:g}" + " | " + " | ".join(f"{v:.3f}" for v in row) + " |")
    lines.append("")
    (output_dir / "report.md").write_text("\n".join(lines), encoding="utf-8")


def plot_heatmap(output_dir: Path, payload: dict[str, Any]) -> None:
    try:
        import matplotlib

        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
    except Exception as exc:  # pragma: no cover - plotting is optional on clusters
        (output_dir / "plot_error.txt").write_text(str(exc), encoding="utf-8")
        return

    t_values = payload["t_values"]
    matrix = np.asarray(payload["cosine_similarity"], dtype=np.float32)
    fig, ax = plt.subplots(figsize=(8, 6))
    im = ax.imshow(matrix, vmin=-1, vmax=1, cmap="coolwarm", origin="lower")
    ax.set_title("ELF timestep gradient cosine")
    ax.set_xlabel("t")
    ax.set_ylabel("t")
    ax.set_xticks(range(len(t_values)))
    ax.set_yticks(range(len(t_values)))
    ax.set_xticklabels([f"{x:g}" for x in t_values], rotation=45, ha="right")
    ax.set_yticklabels([f"{x:g}" for x in t_values])
    fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
    fig.tight_layout()
    fig.savefig(output_dir / "heatmap.svg")
    fig.savefig(output_dir / "heatmap.png", dpi=180)
    plt.close(fig)


def resolve_elf_config(config_path: str, output_dir: Path) -> str:
    """Write a copy of an ELF config with relative nested paths absolutized."""
    import yaml

    path = Path(config_path).resolve()
    cfg = yaml.safe_load(path.read_text(encoding="utf-8")) or {}
    sampling_path = cfg.get("sampling_configs_path")
    if isinstance(sampling_path, str) and sampling_path and not os.path.isabs(sampling_path):
        # Official configs live under <elf-src>/configs/training_configs/*.yml,
        # while sampling_configs_path is written relative to <elf-src>.
        elf_src = path.parent.parent.parent
        cfg["sampling_configs_path"] = str((elf_src / sampling_path).resolve())
    output_dir.mkdir(parents=True, exist_ok=True)
    fd, tmp_name = tempfile.mkstemp(prefix="resolved_elf_config_", suffix=".yml", dir=output_dir)
    os.close(fd)
    Path(tmp_name).write_text(yaml.safe_dump(cfg, sort_keys=False), encoding="utf-8")
    return tmp_name


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--elf-src", default="reference/code/ELF/src")
    parser.add_argument("--config", default="reference/code/ELF/src/configs/training_configs/train_owt_ELF-B.yml")
    parser.add_argument("--checkpoint-path", default="embedded-language-flows/ELF-B-owt")
    parser.add_argument("--output-dir", required=True)
    parser.add_argument("--texts-file", default=None, help="Optional JSONL/text file; JSONL may use text/generated/output.")
    parser.add_argument("--num-samples", type=int, default=4)
    parser.add_argument("--max-length", type=int, default=128)
    parser.add_argument("--t-values", default="0.05:0.05:0.95")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--grad-leaf-regex", default="final_layer|proj_kernel|proj_bias")
    parser.add_argument("--self-cond-cfg-scale", type=float, default=3.0)
    parser.add_argument("--use-cpu-init", action="store_true")
    args = parser.parse_args()

    elf_src = Path(args.elf_src).resolve()
    sys.path.insert(0, str(elf_src))

    import jax
    import jax.numpy as jnp
    import optax
    from flax import traverse_util
    from transformers import AutoTokenizer

    from configs.config import apply_config_overrides, load_config_from_yaml
    from modules.model import ELF_models
    from modules.t5_encoder import get_encoder
    from utils.checkpoint_utils import load_checkpoint, load_encoder_checkpoint
    from utils.encoder_utils import encode_text
    from utils.sampling_utils import add_noise, net_out_to_v_x
    from utils.train_utils import TrainState

    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    config_path = resolve_elf_config(args.config, output_dir)
    config = load_config_from_yaml(config_path)
    config = apply_config_overrides(
        config,
        [
            f"max_length={args.max_length}",
            "global_batch_size=1",
            "batch_size=1",
            "use_wandb=false",
            "online_eval=false",
        ],
    )

    if args.texts_file:
        texts = []
        for line in Path(args.texts_file).read_text(encoding="utf-8").splitlines():
            if not line.strip():
                continue
            if line.lstrip().startswith("{"):
                obj = json.loads(line)
                text = obj.get("text") or obj.get("generated") or obj.get("output") or obj.get("input")
            else:
                text = line
            if text:
                texts.append(str(text))
            if len(texts) >= args.num_samples:
                break
    else:
        texts = DEFAULT_TEXTS[: args.num_samples]
    if not texts:
        raise ValueError("No input texts found.")

    rng = jax.random.PRNGKey(args.seed)
    tokenizer = AutoTokenizer.from_pretrained(config.tokenizer_name or config.encoder_model_name)
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token = tokenizer.eos_token
    encoded = tokenizer(
        texts,
        add_special_tokens=False,
        max_length=args.max_length,
        truncation=True,
        padding="max_length",
        return_tensors="np",
    )
    input_ids = jnp.asarray(encoded["input_ids"], dtype=jnp.int32)
    attention_mask = jnp.asarray(encoded["attention_mask"], dtype=jnp.float32)

    encoder_config, encoder_model, _ = get_encoder(config.encoder_model_name, jnp.float32)
    encoder_params = load_encoder_checkpoint(config.encoder_checkpoint)
    x0 = encode_text(
        input_ids=input_ids,
        attention_mask=attention_mask,
        encoder_apply_fn=encoder_model.apply,
        encoder_params=encoder_params,
        latent_mean=config.latent_mean,
        latent_std=config.latent_std,
    )

    model = ELF_models[config.model](
        text_encoder_dim=encoder_config.d_model,
        max_length=config.max_length,
        attn_drop=config.attn_dropout,
        proj_drop=config.proj_dropout,
        num_time_tokens=config.num_time_tokens,
        num_self_cond_cfg_tokens=config.num_self_cond_cfg_tokens,
        vocab_size=tokenizer.vocab_size,
        num_model_mode_tokens=config.num_model_mode_tokens,
        bottleneck_dim=config.bottleneck_dim,
    )

    init_rng, dropout_rng, noise_rng = jax.random.split(rng, 3)
    input_dim = encoder_config.d_model * (2 if config.self_cond_prob > 0 else 1)
    dummy_x = jnp.ones((1, config.max_length, input_dim), dtype=jnp.float32)
    dummy_t = jnp.ones((1,), dtype=jnp.float32)
    dummy_sc = jnp.ones((1,), dtype=jnp.float32) if config.num_self_cond_cfg_tokens > 0 else None
    variables = model.init(
        init_rng,
        dummy_x,
        dummy_t,
        deterministic=True,
        self_cond_cfg_scale=dummy_sc,
        decoder_step_active=jnp.array(False),
    )
    state = TrainState.create(
        apply_fn=model.apply,
        params=variables["params"],
        tx=optax.adamw(learning_rate=1e-4),
        dropout_rng=dropout_rng,
        ema_params1=variables["params"],
    )
    state, _ = load_checkpoint(args.checkpoint_path, state)
    params = state.ema_params1 if state.ema_params1 is not None else state.params

    noise = jax.random.normal(noise_rng, x0.shape, dtype=x0.dtype)
    loss_mask = attention_mask
    self_cond_cfg = (
        jnp.full((input_ids.shape[0],), args.self_cond_cfg_scale, dtype=jnp.float32)
        if config.num_self_cond_cfg_tokens > 0
        else None
    )
    t_values = parse_values(args.t_values)

    def loss_for_t(p, t_value):
        t = jnp.full((x0.shape[0],), t_value, dtype=jnp.float32)
        z = add_noise(x0, noise, t, config, cond_seq_mask=None)
        t_expanded = t.reshape(-1, 1, 1)
        v_target = (x0 - z) / jnp.maximum(1.0 - t_expanded, config.t_eps)
        if config.self_cond_prob > 0:
            z0 = jnp.concatenate([z, jnp.zeros_like(z)], axis=-1)
            net_init, _ = state.apply_fn(
                {"params": p},
                z0,
                t,
                deterministic=True,
                self_cond_cfg_scale=self_cond_cfg,
                decoder_step_active=jnp.array(False),
            )
            _, x_pred_init = net_out_to_v_x(net_init, z, t, config.t_eps)
            model_input = jnp.concatenate([z, jax.lax.stop_gradient(x_pred_init)], axis=-1)
        else:
            model_input = z
        net_out, _ = state.apply_fn(
            {"params": p},
            model_input,
            t,
            deterministic=True,
            self_cond_cfg_scale=self_cond_cfg,
            decoder_step_active=jnp.array(False),
        )
        v_pred, _ = net_out_to_v_x(net_out, z, t, config.t_eps)
        per_token = jnp.mean((v_pred - v_target) ** 2, axis=-1)
        return (per_token * loss_mask).sum() / jnp.maximum(loss_mask.sum(), 1.0)

    grad_fn = jax.value_and_grad(loss_for_t)
    selected_grads = []
    selected_norms = []
    losses = []
    selected_names = None
    selected_param_count = 0
    for t_value in t_values:
        loss, grads = grad_fn(params, jnp.asarray(t_value, dtype=jnp.float32))
        leaves, names, sq_norm = selected_tree_stats(jax, jnp, grads, args.grad_leaf_regex)
        if selected_names is None:
            selected_names = names
            flat_params = traverse_util.flatten_dict(params)
            regex = re.compile(args.grad_leaf_regex) if args.grad_leaf_regex else None
            selected_param_count = int(
                sum(
                    np.prod(np.asarray(value).shape)
                    for key, value in flat_params.items()
                    if regex is None or regex.search("/".join(str(x) for x in key))
                )
            )
        selected_grads.append([jax.device_get(x) for x in leaves])
        selected_norms.append(float(jax.device_get(jnp.sqrt(sq_norm + 1e-30))))
        losses.append(float(jax.device_get(loss)))

    n = len(t_values)
    cosine = [[0.0 for _ in range(n)] for _ in range(n)]
    for i in range(n):
        for j in range(n):
            dot = float(jax.device_get(tree_dot(jnp, selected_grads[i], selected_grads[j])))
            cosine[i][j] = dot / max(selected_norms[i] * selected_norms[j], 1e-30)

    adjacent = [
        {"from": t_values[i], "to": t_values[i + 1], "cosine": cosine[i][i + 1]}
        for i in range(n - 1)
    ]
    payload = {
        "updated_at": datetime.now(timezone.utc).isoformat(),
        "mode": "elf_flow_matching",
        "model": config.model,
        "config": os.path.abspath(args.config),
        "checkpoint_path": args.checkpoint_path,
        "num_samples": len(texts),
        "max_length": args.max_length,
        "t_values": t_values,
        "loss_by_t": losses,
        "cosine_similarity": cosine,
        "adjacent_cosines": adjacent,
        "grad_leaf_regex": args.grad_leaf_regex,
        "selected_leaf_count": len(selected_names or []),
        "selected_param_count": selected_param_count,
        "selected_leaf_names": selected_names or [],
        "seed": args.seed,
        "texts": texts,
        "command": " ".join(sys.argv),
    }
    (output_dir / "alignment.json").write_text(json.dumps(payload, indent=2), encoding="utf-8")
    write_report(output_dir, payload)
    plot_heatmap(output_dir, payload)
    print(json.dumps({k: payload[k] for k in ["mode", "model", "t_values", "loss_by_t", "adjacent_cosines"]}, indent=2))


if __name__ == "__main__":
    main()