#!/usr/bin/env python3
"""Render a Fig.3-style alignment plot from dLLM alignment.json.

The source paper's Fig. 3 visualizes pairwise gradient cosine similarities and
draws dashed interval boundaries. This renderer adds an adjacent-step cosine
curve below the heatmap because that is the most direct boundary diagnostic for
our mask-ratio grid.
"""

from __future__ import annotations

import argparse
import html
import json
from pathlib import Path


def color_for(value: float) -> str:
    value = max(-1.0, min(1.0, value))
    if value >= 0:
        t = value
        r = int(255 * (1 - 0.08 * t))
        g = int(255 * (1 - 0.62 * t))
        b = int(255 * (1 - 0.72 * t))
    else:
        t = -value
        r = int(255 * (1 - 0.70 * t))
        g = int(255 * (1 - 0.45 * t))
        b = int(255 * (1 - 0.05 * t))
    return f"#{r:02x}{g:02x}{b:02x}"


def adjacent_cosines(ratios: list[float], matrix: list[list[float]]) -> list[dict[str, float]]:
    out = []
    for i in range(len(ratios) - 1):
        out.append({
            "left": ratios[i],
            "right": ratios[i + 1],
            "mid": (ratios[i] + ratios[i + 1]) / 2,
            "cosine": matrix[i][i + 1],
        })
    return out


def best_adjacent_cut(adjacent: list[dict[str, float]]) -> int | None:
    if not adjacent:
        return None
    return min(range(len(adjacent)), key=lambda i: adjacent[i]["cosine"])


def write_json(path: Path, payload: dict) -> None:
    path.write_text(json.dumps(payload, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")


def write_markdown(path: Path, payload: dict) -> None:
    lines = [
        "# Fig.3-Style dLLM Alignment Plot",
        "",
        f"- Source alignment: `{payload['source_alignment']}`",
        f"- Figure: `{payload['figure']}`",
        f"- Candidate boundary: `{payload['candidate_boundary_label']}`",
        "",
        "## Adjacent Mask-Ratio Cosine Similarity",
        "",
        "| Step pair | Cosine similarity |",
        "| --- | ---: |",
    ]
    for row in payload["adjacent_cosines"]:
        lines.append(f"| {row['left']:g} -> {row['right']:g} | {row['cosine']:.3f} |")
    lines.extend([
        "",
        "Interpretation: lower adjacent cosine values indicate a sharper local",
        "change between neighboring denoising objectives and are natural candidate",
        "interval boundaries. The dashed line in the SVG marks the lowest adjacent",
        "cosine split.",
        "",
    ])
    path.write_text("\n".join(lines), encoding="utf-8")


def write_svg(path: Path, payload: dict) -> None:
    ratios = payload["ratios"]
    matrix = payload["cosine_similarity"]
    adjacent = payload["adjacent_cosines"]
    cut_index = payload["candidate_boundary_index"]

    n = len(ratios)
    requested_cell = int(payload.get("cell_size", 0) or 0)
    cell = requested_cell if requested_cell > 0 else (38 if n <= 12 else 14 if n <= 50 else 8)
    show_cell_text = bool(payload.get("show_cell_text", n <= 16))
    label_every = int(payload.get("label_every", 0) or (1 if n <= 16 else 5 if n <= 60 else 10))
    label_w = 92
    top = 54
    right_pad = 36
    heat_w = n * cell
    heat_h = n * cell
    gap = 54
    curve_h = 160
    legend_h = 46
    width = label_w + heat_w + right_pad
    height = top + heat_h + gap + curve_h + legend_h
    x0 = label_w
    y0 = top

    title = payload["title"]
    parts = [
        f'<svg xmlns="http://www.w3.org/2000/svg" width="{width}" height="{height}" viewBox="0 0 {width} {height}">',
        '<rect width="100%" height="100%" fill="#ffffff"/>',
        f'<text x="{x0}" y="24" font-family="Arial, sans-serif" font-size="17" font-weight="700">{html.escape(title)}</text>',
        f'<text x="{x0}" y="43" font-family="Arial, sans-serif" font-size="12" fill="#444">Pairwise gradient cosine by mask ratio; dashed line marks lowest adjacent-step cosine.</text>',
    ]

    for i, ratio in enumerate(ratios):
        if i % label_every != 0 and i != n - 1:
            continue
        x = x0 + i * cell + cell / 2
        parts.append(f'<text x="{x}" y="{y0 - 9}" text-anchor="middle" font-family="Arial, sans-serif" font-size="10">{ratio:g}</text>')
        y = y0 + i * cell + cell / 2 + 4
        parts.append(f'<text x="{x0 - 10}" y="{y}" text-anchor="end" font-family="Arial, sans-serif" font-size="10">{ratio:g}</text>')

    for row_idx, row in enumerate(matrix):
        for col_idx, value in enumerate(row):
            x = x0 + col_idx * cell
            y = y0 + row_idx * cell
            parts.append(f'<rect x="{x}" y="{y}" width="{cell}" height="{cell}" fill="{color_for(value)}" stroke="#ffffff" stroke-width="1"/>')
            if show_cell_text:
                parts.append(f'<text x="{x + cell / 2}" y="{y + cell / 2 + 4}" text-anchor="middle" font-family="Arial, sans-serif" font-size="9" fill="#111">{value:.2f}</text>')

    if cut_index is not None:
        cut_x = x0 + (cut_index + 1) * cell
        cut_y = y0 + (cut_index + 1) * cell
        parts.append(f'<line x1="{cut_x}" y1="{y0}" x2="{cut_x}" y2="{y0 + heat_h}" stroke="#111" stroke-width="2" stroke-dasharray="6 5"/>')
        parts.append(f'<line x1="{x0}" y1="{cut_y}" x2="{x0 + heat_w}" y2="{cut_y}" stroke="#111" stroke-width="2" stroke-dasharray="6 5"/>')

    # Adjacent-step cosine curve.
    curve_x0 = x0
    curve_y0 = y0 + heat_h + gap
    curve_w = heat_w
    min_v = min(-0.05, min(row["cosine"] for row in adjacent) - 0.04)
    max_v = max(1.0, max(row["cosine"] for row in adjacent) + 0.04)

    def sx(idx: int) -> float:
        if len(adjacent) == 1:
            return curve_x0 + curve_w / 2
        return curve_x0 + idx * curve_w / (len(adjacent) - 1)

    def sy(value: float) -> float:
        return curve_y0 + curve_h - (value - min_v) / (max_v - min_v) * curve_h

    parts.append(f'<text x="{curve_x0}" y="{curve_y0 - 18}" font-family="Arial, sans-serif" font-size="14" font-weight="700">Adjacent mask-ratio cosine</text>')
    parts.append(f'<line x1="{curve_x0}" y1="{curve_y0}" x2="{curve_x0}" y2="{curve_y0 + curve_h}" stroke="#222" stroke-width="1"/>')
    parts.append(f'<line x1="{curve_x0}" y1="{curve_y0 + curve_h}" x2="{curve_x0 + curve_w}" y2="{curve_y0 + curve_h}" stroke="#222" stroke-width="1"/>')
    for tick in [0.0, 0.25, 0.5, 0.75, 1.0]:
        if min_v <= tick <= max_v:
            y = sy(tick)
            parts.append(f'<line x1="{curve_x0 - 4}" y1="{y}" x2="{curve_x0 + curve_w}" y2="{y}" stroke="#dddddd" stroke-width="1"/>')
            parts.append(f'<text x="{curve_x0 - 10}" y="{y + 4}" text-anchor="end" font-family="Arial, sans-serif" font-size="10">{tick:.2f}</text>')

    points = []
    for i, row in enumerate(adjacent):
        points.append(f"{sx(i):.2f},{sy(row['cosine']):.2f}")
    parts.append(f'<polyline points="{" ".join(points)}" fill="none" stroke="#1f5fbf" stroke-width="2.5"/>')

    point_label_every = max(1, len(adjacent) // 12)
    for i, row in enumerate(adjacent):
        x = sx(i)
        y = sy(row["cosine"])
        fill = "#d62728" if i == cut_index else "#1f5fbf"
        parts.append(f'<circle cx="{x}" cy="{y}" r="4.2" fill="{fill}" stroke="#ffffff" stroke-width="1"/>')
        if i % point_label_every == 0 or i == len(adjacent) - 1 or i == cut_index:
            parts.append(f'<text x="{x}" y="{curve_y0 + curve_h + 16}" text-anchor="middle" font-family="Arial, sans-serif" font-size="9">{row["left"]:g}-{row["right"]:g}</text>')
            parts.append(f'<text x="{x}" y="{y - 8}" text-anchor="middle" font-family="Arial, sans-serif" font-size="9">{row["cosine"]:.2f}</text>')

    if cut_index is not None:
        x = sx(cut_index)
        parts.append(f'<line x1="{x}" y1="{curve_y0}" x2="{x}" y2="{curve_y0 + curve_h}" stroke="#d62728" stroke-width="1.5" stroke-dasharray="5 4"/>')
        parts.append(f'<text x="{x + 8}" y="{curve_y0 + 15}" font-family="Arial, sans-serif" font-size="11" fill="#d62728">candidate split</text>')

    legend_y = curve_y0 + curve_h + 38
    parts.append(f'<text x="{x0}" y="{legend_y}" font-family="Arial, sans-serif" font-size="11" fill="#444">Heatmap: red = positive alignment, white = near zero, yellow/blue = negative. Curve: cosine between adjacent ratios.</text>')
    parts.append("</svg>")
    path.write_text("\n".join(parts) + "\n", encoding="utf-8")


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--title", default="dLLM mask-ratio gradient alignment")
    parser.add_argument("--cell-size", type=int, default=0)
    parser.add_argument("--label-every", type=int, default=0)
    parser.add_argument("--show-cell-text", action="store_true")
    args = parser.parse_args()

    source = json.loads(args.input.read_text(encoding="utf-8"))
    ratios = [float(x) for x in source["ratios"]]
    matrix = source["cosine_similarity"]
    adjacent = adjacent_cosines(ratios, matrix)
    cut_index = best_adjacent_cut(adjacent)
    if cut_index is None:
        boundary = "n/a"
    else:
        boundary = f"{adjacent[cut_index]['left']:g} -> {adjacent[cut_index]['right']:g}"

    args.output_dir.mkdir(parents=True, exist_ok=True)
    figure = args.output_dir / "fig3_style_adjacent.svg"
    summary = args.output_dir / "fig3_style_adjacent.json"
    report = args.output_dir / "fig3_style_adjacent.md"
    payload = {
        "source_alignment": str(args.input),
        "figure": str(figure),
        "title": args.title,
        "cell_size": args.cell_size,
        "label_every": args.label_every,
        "show_cell_text": args.show_cell_text,
        "ratios": ratios,
        "cosine_similarity": matrix,
        "adjacent_cosines": adjacent,
        "candidate_boundary_index": cut_index,
        "candidate_boundary_label": boundary,
    }
    write_svg(figure, payload)
    write_json(summary, payload)
    write_markdown(report, payload)
    print(json.dumps({
        "figure": str(figure),
        "summary": str(summary),
        "report": str(report),
        "candidate_boundary": boundary,
    }, indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())