"""Reproduce the analytic V-valley tutorial with public LESS APIs.

Keep the script in scripts/ and the downloaded case files in ../data/v-valley/.
Run with --output to select a directory for the current results and figures.
"""
from pathlib import Path
import argparse
import json

import less
import numpy as np


DATA = Path(__file__).resolve().parents[1] / "data" / "v-valley"


def scene_from_input(config, backend):
    scene = less.Scene(backend=backend)
    scene.sampler = less.Sampler(config.get("sampling_method", "independent"))
    scene.size = tuple(config["size_m"])
    scene.repetitive = config["periodic_wrap_limit"]
    scene.terrain = less.Terrain(property=less.Lambertian(reflectance=0))
    valley = less.Object("valley", mesh=str(DATA / "scene.obj"))
    valley.set_property("grey", less.Lambertian(
        wavelengths=config["wavelengths_nm"],
        reflectance=config["grey_reflectance"]))
    valley.set_property("black", less.Lambertian(reflectance=config["black_reflectance"]))
    scene.add(valley, positions=[[0, 0, 0]])
    scene.illumination = less.Illumination(
        source=less.Sun(zenith=config["sun_zenith_degrees"], wavelengths=config["wavelengths_nm"],
                        irradiance=config["sun_normal_irradiance_W_m2_nm"]),
        atmosphere=less.NoAtmosphere())
    return scene


def analytic(config):
    beta = np.deg2rad(config["slope_degrees"])
    theta = np.deg2rad(config["signed_view_degrees"])
    rho = np.asarray(config["grey_reflectance"])
    visible = np.clip((1 + np.tan(beta) * np.tan(theta)) / 2, 0, 1)
    brf = visible[:, None] * rho[None, :] * np.cos(beta)
    albedo = rho * np.cos(beta) / 2
    return brf, albedo


def native_identity(product, backend):
    identity = product.metadata["execution"]
    assert identity["backend"] == backend, identity
    assert identity["native"] is True, identity
    assert identity["executor"].startswith(backend + ".slang."), identity
    return {key: identity[key] for key in
            ("backend", "workflow", "executor", "native", "random_seed")}


def run(config):
    reference_brf, reference_albedo = analytic(config)
    thresholds = config["thresholds"]
    records = []
    assert config["size_m"] == [2, 2] and config["valley_floor_m"] == 0.1
    assert config["slope_degrees"] == 45 and config["sun_zenith_degrees"] == 0
    assert config["black_reflectance"] == 0
    assert min(config["photon_budgets"]) >= 1_000_000
    for backend in ("optix", "vulkan", "embree"):
        scene = scene_from_input(config, backend)
        previous = None
        for photons in config["photon_budgets"]:
            directions = [(abs(t), 90 if t >= 0 else 270)
                          for t in config["signed_view_degrees"]]
            brf = scene.simulate(less.BRFSensor(
                bands=config["wavelengths_nm"], directions=directions,
                num_photons=photons))
            field = scene.simulate(less.ForwardRadiationFieldSensor(
                bands=config["wavelengths_nm"], mode="grid", grid=(2, 2),
                layer_heights=[0, 2], photons_per_m2=photons // 4))
            incident = field.incident_energy.astype(float)
            absorbed = (field.absorption_grid.sum(axis=(0, 1, 2), dtype=float)
                        + field.terrain_absorbed)
            albedo = field.escaped_energy / incident
            absorptance = absorbed / incident
            truncated = field.truncated_energy / incident
            brf_values = brf.brf.astype(float)
            errors = {
                "brf_absolute_error": float(np.max(np.abs(brf_values - reference_brf))),
                "albedo_absolute_error": float(np.max(np.abs(albedo - reference_albedo))),
                "absorptance_absolute_error": float(np.max(np.abs(
                    absorptance - (1 - reference_albedo)))),
                "closure_absolute_error": float(np.max(np.abs(
                    albedo + absorptance + truncated - 1))),
                "truncated_fraction": float(np.max(truncated)),
            }
            current = np.r_[brf_values.ravel(), albedo, absorptance]
            if previous is not None:
                errors["final_budget_change_absolute"] = float(
                    np.max(np.abs(current - previous)))
            checks = {key: value <= thresholds[key] for key, value in errors.items()}
            record = {
                "backend": backend, "photons": photons,
                "brf": brf_values.tolist(), "albedo": albedo.tolist(),
                "absorptance": absorptance.tolist(),
                "incident_W_nm": incident.tolist(),
                "truncated_fraction": truncated.tolist(),
                "errors": errors, "checks": checks, "passed": all(checks.values()),
                "execution": {"brf": native_identity(brf, backend),
                              "flux": native_identity(field, backend)},
            }
            records.append(record)
            previous = current
            print(backend, photons, errors, flush=True)
    return {
        "less_version": less.__version__, "input": config,
        "reference": {"brf": reference_brf.tolist(),
                      "albedo": reference_albedo.tolist(),
                      "absorptance": (1 - reference_albedo).tolist()},
        "sampling_note": "Budget refinement with " + config.get("sampling_method", "independent") + " sampling; "
                         "not independent replicates or confidence intervals.",
        "records": records,
        "passed": all(r["passed"] for r in records
                      if r["photons"] == config["photon_budgets"][-1]),
    }


def figures(result, output):
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    from mpl_toolkits.mplot3d.art3d import Poly3DCollection

    plt.rcParams.update({"font.size": 10, "svg.fonttype": "none"})
    config = result["input"]
    angles = np.asarray(config["signed_view_degrees"])
    reference = np.asarray(result["reference"]["brf"])
    colors = {"optix": "#0072B2", "vulkan": "#D55E00", "embree": "#009E73"}
    final = [r for r in result["records"]
             if r["photons"] == config["photon_budgets"][-1]]
    fig, axes = plt.subplots(1, 2, figsize=(10, 3.7), layout="constrained")
    for band, wavelength in enumerate(config["wavelengths_nm"]):
        axes[0].plot(angles, reference[:, band], color="#333333", lw=1.5,
                     label="Analytic" if band == 0 else None)
        axes[0].text(angles[-1] + 1, reference[-1, band], str(wavelength), fontsize=8)
    for row in final:
        values = np.asarray(row["brf"])
        axes[0].plot(angles, values, ".", color=colors[row["backend"]], ms=4,
                     label=[row["backend"], None, None])
        axes[1].plot(angles, np.max(np.abs(values - reference), axis=1),
                     "o-", ms=3, color=colors[row["backend"]], label=row["backend"])
    axes[0].set(xlabel="Signed view zenith (deg; + = east)", ylabel="BRF",
                title="Three wavelengths; vertical direct illumination", xlim=(-75, 85))
    axes[0].legend(fontsize=8)
    axes[1].axhline(config["thresholds"]["brf_absolute_error"], color="#555555",
                    ls="--", label="Acceptance limit")
    axes[1].set(xlabel="Signed view zenith (deg)", ylabel="Max absolute BRF error",
                title="Error across the three wavelengths")
    axes[1].set_yscale("symlog", linthresh=1e-8)
    axes[1].legend(fontsize=8)
    for ax in axes: ax.grid(alpha=0.2)
    fig.savefig(output / "brf.svg"); plt.close(fig)

    fig, axes = plt.subplots(1, 2, figsize=(10, 3.7), layout="constrained")
    rho = config["grey_reflectance"]
    axes[0].plot(rho, result["reference"]["albedo"], "k-", label="Analytic albedo")
    axes[0].plot(rho, result["reference"]["absorptance"], "k--", label="Analytic absorption")
    for row in final:
        for key, marker in (("albedo", "o"), ("absorptance", "x")):
            axes[0].plot(rho, row[key], marker, color=colors[row["backend"]],
                         label=row["backend"] if key == "albedo" else None)
    axes[0].set(xlabel="Grey slope reflectance", ylabel="Fraction of incident flux",
                title="Whole-valley energy partition")
    for backend, color in colors.items():
        rows = [r for r in result["records"] if r["backend"] == backend]
        axes[1].plot([r["photons"] for r in rows],
                     [max(r["errors"]["brf_absolute_error"],
                          r["errors"]["albedo_absolute_error"]) for r in rows],
                     "o-", color=color, label=backend)
    axes[1].axhline(0.005, color="#555555", ls="--")
    axes[1].set(xscale="log", xlabel="Photon budget per product",
                ylabel="Max absolute error (BRF, albedo)", title="Budget refinement")
    for ax in axes:
        ax.grid(alpha=0.2); ax.legend(fontsize=8)
    fig.savefig(output / "energy.svg"); plt.close(fig)

    fig = plt.figure(figsize=(7, 4.5), layout="constrained")
    ax = fig.add_subplot(projection="3d", computed_zorder=False)
    vertices = np.array([[0, 0, 1.1], [1, 0, .1], [2, 0, 1.1],
                         [0, 2, 1.1], [1, 2, .1], [2, 2, 1.1]])
    faces = [vertices[[0, 1, 4, 3]], vertices[[1, 2, 5, 4]]]
    ax.add_collection3d(Poly3DCollection(faces, facecolors=["#d2ad64", "#333c48"],
                                        edgecolors="#182535", linewidths=1, zorder=1))
    for x in (.5, 1.5):
        ax.quiver(x, 1, 1.8, 0, 0, -.5, color="#ce641c", linewidth=2, zorder=5)
    ax.set(xlim=(0, 2), ylim=(0, 2), zlim=(0, 2), xlabel="X / m", ylabel="Y / m",
           zlabel="Z / m", title="90-degree V valley; 2 m periodic cell")
    ax.view_init(elev=24, azim=-65)
    # Keep labels in the annotation layer above the 3D polygons.
    from mpl_toolkits.mplot3d import proj3d
    for text, position, label in (
            ("Lambertian", (.4, .7, .7), (.05, .84)),
            ("Black", (1.6, .7, .7), (.81, .84))):
        x, y, _ = proj3d.proj_transform(*position, ax.get_proj())
        ax.annotate(text, xy=(x, y), xytext=label, textcoords="axes fraction",
                    bbox=dict(boxstyle="round,pad=0.3", fc="white", ec="#bdcddd"),
                    arrowprops=dict(arrowstyle="-", color="#182535"), zorder=10)
    fig.savefig(output / "scene.svg"); plt.close(fig)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output", type=Path, default=Path("build/v-valley"))
    args = parser.parse_args()
    args.output.mkdir(parents=True, exist_ok=True)
    config = json.loads((DATA / "input.json").read_text(encoding="utf-8"))
    result = run(config)
    (args.output / "results.json").write_text(
        json.dumps(result, indent=2, ensure_ascii=False, allow_nan=False), encoding="utf-8")
    figures(result, args.output)
    raise SystemExit(0 if result["passed"] else 1)


if __name__ == "__main__":
    main()