#!/usr/bin/env python3
"""Render TerrainAtlas305GameTests native exports. No synthesized terrain or image smoothing.

python scripts/terrain_atlas.py --output build/atlas305 <native-json> [<native-json> ...]
See docs/terrain-atlas.md for acquisition, provenance and sampling limits.
"""
import argparse
import hashlib
import json
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.colors import TwoSlopeNorm
import numpy as np

LABELS = {
    "plateaus27": "Plateaux 24 / 27",
    "global301": "Bruit 3D global · 30.1",
    "stretch304": "Étirement vertical · 30.4",
    "current305": "Plateaux + déformation 3D · 30.5",
    "current306": "Relief légèrement renforcé · 30.6",
}


def save(fig, directory, name):
    fig.savefig(directory / f"{name}.png", dpi=150, facecolor="white")
    fig.savefig(directory / f"{name}.pdf", facecolor="white")
    plt.close(fig)


def render(source, output):
    raw = source.read_bytes()
    data = json.loads(raw)
    if data["schema"] != 1 or len(data["variants"]) != 4 or any(k not in LABELS for k in data["variants"]):
        raise ValueError(f"Unsupported atlas schema/variants: {source}")
    labels = {key: LABELS[key] for key in data["variants"]}
    target = output / source.stem
    target.mkdir(parents=True, exist_ok=True)
    title = f"Sanctuary · diamètre {data['diameter']} · graine {data['seed']}"
    provenance = (f"Champ natif Minecraft {data['minecraft']} · {data['version']} · {data['revision']}\n"
                  "Avant eau, matériaux de surface, végétation et structures. Le plafond de construction reste Y640.")
    variants = list(data["variants"].values())
    maps = []
    for variant in variants:
        columns = np.asarray(variant["columns"], dtype=float)
        xs, zs = np.unique(columns[:, 0]), np.unique(columns[:, 1])
        if columns.shape[0] != len(xs) * len(zs):
            raise ValueError("Incomplete native map grid")
        maps.append(columns.reshape(len(zs), len(xs), 5))
    step = data["map_step"]
    extent = (xs[0] - step / 2, xs[-1] + step / 2, zs[-1] + step / 2, zs[0] - step / 2)
    fig, axes = plt.subplots(3, 4, figsize=(18, 14), layout="constrained", sharex=True, sharey=True)
    fig.suptitle(title + "\nRelief et épaisseur de roche", fontsize=18)
    metrics = [(2, "Altitude du sommet (Y)", "terrain", 0, data["domain_y"][1]),
               (3, "Épaisseur solide cumulée (blocs)", "viridis", 0, 300),
               (4, "Roche continue sous le sommet (blocs)", "magma", 0, 100)]
    for row, (index, label, cmap, vmin, vmax) in enumerate(metrics):
        for col, (key, grid) in enumerate(zip(labels, maps)):
            ax = axes[row, col]
            values = np.ma.masked_where(grid[:, :, 2] == 0, grid[:, :, index])
            image = ax.imshow(values, extent=extent, origin="upper", interpolation="nearest",
                              cmap=cmap, vmin=vmin, vmax=vmax, aspect="equal")
            ax.set_facecolor("#e8edf0")
            ax.set_xlabel("X (blocs)")
            if col == 0:
                ax.set_ylabel("Z (blocs) · nord en haut")
            if row == 0:
                ax.set_title(labels[key], fontsize=11)
            ax.plot(0, 0, "+", color="black", ms=8, label="Origine")
        fig.colorbar(image, ax=axes[row, :], shrink=.85, label=label, extend="max")
    fig.supxlabel(f"Pas horizontal : {step} blocs · vertical : {data['vertical_step']} blocs · gris = aucun solide échantillonné.\n" + provenance, fontsize=9)
    save(fig, target, "maps")

    section_count = len(variants[0]["sections"])
    fig, axes = plt.subplots(4, section_count, figsize=(20, 12), layout="constrained", sharex=True, sharey=True)
    fig.suptitle(title + "\nCoupes de densité · proportions spatiales conservées", fontsize=18)
    for row, (key, variant) in enumerate(zip(labels, variants)):
        for col, section in enumerate(variant["sections"]):
            ax = axes[row, col]
            density = np.asarray(section["density"], dtype=float) / data["density_scale"]
            dy = section["step"]
            x = section["min_x"] + np.arange(density.shape[1]) * dy
            y = section["max_y"] - np.arange(density.shape[0]) * dy
            image = ax.imshow(density, extent=(x[0]-dy/2, x[-1]+dy/2, y[-1]-dy/2, y[0]+dy/2),
                              interpolation="nearest", cmap="RdBu_r", norm=TwoSlopeNorm(0, -.3, .3), aspect="equal")
            ax.contour(x, y, density, levels=[0], colors="#262626", linewidths=.35)
            if row == 0:
                ax.set_title(f"Plan vertical Z = {section['z']}")
            if col == 0:
                ax.set_ylabel(labels[key] + "\nY (blocs)", fontsize=10)
            ax.set_xlabel("X (blocs)")
            ax.set_ylim(0, data["domain_y"][1] + 1)
    fig.colorbar(image, ax=axes, shrink=.6, label="Densité sans unité · bleu : vide (≤ 0) · rouge : roche (> 0)", extend="both")
    fig.supxlabel("Coupes échantillonnées tous les 2 blocs, sans interpolation d’affichage ; isosurface zéro en noir.\n" + provenance, fontsize=9)
    save(fig, target, "sections")

    upper = data["variants"].get("current306", {}).get("upper_sections", [])
    if upper:
        fig, axes = plt.subplots(1, len(upper), figsize=(5*len(upper), 6), layout="constrained", squeeze=False)
        fig.suptitle(title + "\nMorceaux de terrain supérieur · densité native", fontsize=15)
        for ax, cut in zip(axes[0], upper):
            values = np.array(cut["solid"])
            dy=cut["step"]
            ax.imshow(values, cmap="Greys", vmin=0, vmax=1, interpolation="nearest", aspect="equal",
                      extent=(cut["min_x"]-1,cut["min_x"]+values.shape[1]*dy-1,cut["max_y"]-values.shape[0]*dy+1,cut["max_y"]+1))
            ax.set_title(f"Z = {cut['z']}"); ax.set_xlabel("X");ax.set_ylabel("Y")
        fig.supxlabel("Noir : roche · blanc : air. Aucune interpolation d’affichage. Avant végétation et structures.",fontsize=9)
        save(fig,target,"upper-terrain")

    summary = {"seed": data["seed"], "diameter": data["diameter"], "variants": {}}
    base = maps[0][:, :, 2]
    common = np.logical_and.reduce([m[:, :, 2] >= 200 for m in maps])
    for key, grid in zip(labels, maps):
        land = grid[:, :, 2] > 0
        uplands = base >= 200
        summary["variants"][key] = {
            "sampled_land_columns": int(land.sum()),
            "max_sampled_y": float(grid[:, :, 2].max()),
            "original_uplands_within_2_blocks_fraction": float((np.abs(grid[:, :, 2]-base)[uplands] <= 2).mean()),
            "common_upper_columns": int(common.sum()),
            "median_roof_thickness_common_upper_columns": float(np.median(grid[:, :, 4][common])),
            "roof_at_most_8_blocks_common_upper_fraction": float((grid[:, :, 4][common] <= 8).mean()),
        }
    (target / "summary.json").write_text(json.dumps(summary, indent=2) + "\n")
    (target / "provenance.json").write_text(json.dumps({
        "source": source.name, "source_sha256": hashlib.sha256(raw).hexdigest(),
        "script_sha256": hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
        "matplotlib": matplotlib.__version__, "numpy": np.__version__,
        "version": data["version"], "revision": data["revision"], "process_id": data["process_id"],
        "scope": data["scope"], "map_step": step, "vertical_step": data["vertical_step"],
        "density_quantization": 1 / data["density_scale"], "density_plot_saturation": [-.3, .3],
    }, indent=2) + "\n")
    print(json.dumps(summary))
    return summary


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("sources", nargs="+", type=Path)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    summaries = [render(source, args.output) for source in args.sources]
    (args.output / "summary.json").write_text(json.dumps(summaries, indent=2) + "\n")


if __name__ == "__main__":
    main()