"""Parallel-leaf fluorescence and exact slit-shadow validation (public LESS API)."""
from pathlib import Path
import argparse
import json

import less
import numpy as np

DATA = Path(__file__).resolve().parents[1] / "data/fluorescence"
BACKENDS = ("optix", "vulkan", "embree")


def layer_solution(reflectance, transmittance, upward_source, downward_source,
                   top_incident=0):
    """Exact diffuse-plane boundary system, with a black lower boundary.

    Outgoing quantities alternate upward/downward by layer, ordered top down.
    The same equations apply to irradiance or radiance, with consistent sources.
    """
    r, t, su, sd = np.broadcast_arrays(
        reflectance, transmittance, upward_source, downward_source)
    count, bands = r.shape
    result = np.empty((count, 2, bands))
    for band in range(bands):
        matrix = np.eye(2*count)
        rhs = np.stack([su[:, band], sd[:, band]], axis=1).ravel()
        for i in range(count):
            if i:
                matrix[2*i, 2*(i-1)+1] -= r[i, band]
                matrix[2*i+1, 2*(i-1)+1] -= t[i, band]
            else:
                boundary = np.broadcast_to(top_incident, (bands,))[band]
                rhs[0] += r[i, band]*boundary
                rhs[1] += t[i, band]*boundary
            if i+1 < count:
                matrix[2*i, 2*(i+1)] -= t[i, band]
                matrix[2*i+1, 2*(i+1)] -= r[i, band]
        result[:, :, band] = np.linalg.solve(matrix, rhs).reshape(count, 2)
    return result


def reference(config, case, coefficients):
    count = case["layers"]
    leaves = coefficients[:count]
    r = np.array([c["excitation_reflectance"] for c in leaves])
    t = np.array([c["excitation_transmittance"] for c in leaves])
    incident = (np.asarray(config["normal_irradiance_W_m2_nm"])*
                config["excitation_width_nm"]*
                np.cos(np.deg2rad(config["sun_zenith_degrees"])))
    outgoing = layer_solution(r, t, np.zeros_like(r), np.zeros_like(r), incident)
    down = np.vstack([incident, outgoing[:-1, 1]])
    up = np.vstack([outgoing[1:, 0], np.zeros_like(incident)])
    assert config["absorption_driver"] == "total"
    top, bottom = (1-r-t)*down, (1-r-t)*up
    su, sd = [], []
    for i, c in enumerate(leaves):
        f, b = np.array(c["forward_source_matrix"]), np.array(c["backward_source_matrix"])
        phi = config["leaf_parameters"][i]["fqe"]*case.get("yield_scale", 1)
        su.append(phi*(b@top[i]+f@bottom[i]))
        sd.append(phi*(f@top[i]+b@bottom[i]))
    er = np.array([c["emission_reflectance"] for c in leaves])
    et = np.array([c["emission_transmittance"] for c in leaves])
    emission = layer_solution(er, et, su, sd)
    shape = (config["image_size"][1], config["image_size"][0], len(config["emission_nm"]))
    observed = emission[-1, 1] if case.get("bottom_view") else emission[0, 0]
    image = np.broadcast_to(observed, shape).copy()
    visible = np.ones(shape[1])
    if case.get("slit"):
        half = config["slit_width_m"]/2
        height = config["slit_height_gap_m"]
        # Sun and observer are east of the leaf. Both see the aperture
        # projected westward on the leaf; their intersection emits and is seen.
        sun_shift = height*np.tan(np.deg2rad(config["sun_zenith_degrees"]))
        view_shift = height*np.tan(np.deg2rad(case["view_zenith"]))
        lower = max(-half-sun_shift, -half-view_shift)
        upper = min(half-sun_shift, half-view_shift)
        edges = np.linspace(-config["image_extent_m"][0]/2,
                            config["image_extent_m"][0]/2, shape[1]+1)
        visible = np.maximum(0, np.minimum(edges[1:], upper)-
                             np.maximum(edges[:-1], lower))/np.diff(edges)
        image *= visible[None, :, None]
    return dict(apar_top=top.tolist(), apar_bottom=bottom.tolist(),
                source_up=su, source_down=sd, outgoing=emission.tolist(),
                visible_fraction=visible.tolist(), image=image.tolist())


def scene_for(config, case, backend):
    scene = less.Scene(backend=backend)
    scene.size = config["size_m"]
    scene.repetitive = False if case.get("slit") else config["periodic_wrap_limit"]
    scene.terrain = less.Terrain(property=less.Lambertian(reflectance=0))
    scene.measurement_plane = config["layer_heights_m"][0]
    scene.illumination = less.Illumination(source=less.Sun(
        zenith=config["sun_zenith_degrees"], azimuth=config["sun_azimuth_degrees"],
        wavelengths=config["excitation_nm"],
        irradiance=config["normal_irradiance_W_m2_nm"]), atmosphere=less.NoAtmosphere())
    for i in range(case["layers"]):
        params = dict(config["leaf_parameters"][i])
        params["fqe"] *= case.get("yield_scale", 1)
        leaf = less.Object(f"leaf_{i}", mesh=str(DATA/"sheet.obj"))
        leaf.set_property(less.Fluspect(**params))
        scene.add(leaf, positions=[(0, 0, config["layer_heights_m"][i])])
    if case.get("slit"):
        mask = less.Object("black_slit", mesh=str(DATA/"slit.obj"))
        mask.set_property(less.Lambertian(reflectance=0))
        scene.add(mask, positions=[(0, 0, config["layer_heights_m"][0]+
                                      config["slit_height_gap_m"])])
    return scene


def identity(product, backend):
    execution = product.metadata["execution"]
    assert execution["backend"] == backend and execution["native"] is True
    return {k: execution[k] for k in ("backend", "workflow", "executor", "native")}


def run(config, coefficients, backends=BACKENDS):
    refs = {c["name"]: reference(config, c, coefficients) for c in config["cases"]}
    rows, convergence = [], []
    profile = less.SpectralProfile([less.SpectralBand(w, config["excitation_width_nm"], "par")
                                    for w in config["excitation_nm"]])
    limits = config["thresholds"]
    for backend in backends:
        scenes, excitation = {}, {}
        for case in config["cases"]:
            geometry = (case["layers"], case.get("yield_scale", 1), bool(case.get("slit")))
            if geometry not in scenes:
                scenes[geometry] = scene_for(config, case, backend)
            scene = scenes[geometry]
            exact = refs[case["name"]]
            target = np.asarray(exact["image"])
            scale = max(float(abs(target).sum()), 1e-30)
            previous = None
            for samples, quality in zip(config["source_samples"], config["image_samples"]):
                key = (geometry, samples)
                if key not in excitation:
                    radiation = scene.simulate(less.RadiationFieldProcess(
                        spectral_profile=profile, n_samples=samples,
                        max_depth=config["interaction_limit"]), seed=17)
                    state = scene.simulate(less.SIFProcess(rad=radiation, mode="static"), seed=17)
                    excitation[key] = radiation, state
                radiation, state = excitation[key]
                projection = less.Orthographic(
                    image_size=tuple(config["image_size"]), extent=tuple(config["image_extent_m"]),
                    altitude=10, view_zenith=case["view_zenith"], view_azimuth=90)
                if case.get("bottom_view"):
                    projection = less.Perspective(image_size=tuple(config["image_size"]),
                        position=(2, 2, .5), target=(2, 2, 3), up=(0, 1, 0), fov=30)
                projection.pixel_sampling = config.get("image_sampling_method", "independent")
                image = scene.simulate(less.SIFImager(projection,
                    bands=config["emission_nm"], quality=quality, sif_result=state),
                    scattering_order=config["interaction_limit"])
                actual = np.asarray(image.data, float)
                zero = case.get("yield_scale", 1) == 0
                error = float(np.max(abs(actual))) if zero else float(abs(actual-target).sum()/scale)
                mean_error = (error if zero else float(np.max(abs(actual.mean((0, 1))-
                    target.mean((0, 1)))/np.maximum(target.mean((0, 1)), 1e-30))))
                change = None if previous is None else (float(np.max(abs(actual-previous)))
                    if zero else float(np.max(abs(actual.mean((0, 1))-previous.mean((0, 1)))/
                                             np.maximum(target.mean((0, 1)), 1e-30))))
                source_error = None
                observed = None
                if not case.get("slit"):
                    indices = [np.flatnonzero(np.isclose(state.element_height, height))
                               for height in config["layer_heights_m"][:case["layers"]]]
                    observed = np.array([[state.apar_spectral_top[ids].mean(0),
                                          state.apar_spectral_bot[ids].mean(0)] for ids in indices])
                    expected = np.stack([exact["apar_top"], exact["apar_bottom"]], axis=1)
                    # Absolute error scaled by the strongest face in each band
                    # also covers analytically dark bottom faces.
                    source_error = float(np.max(abs(observed-expected)/
                        np.maximum(np.max(expected, axis=(0, 1)), 1e-30)))
                check = dict(backend=backend, case=case["name"], source_samples=samples,
                    image_samples=quality, image_error=error, mean_error=mean_error,
                    source_error=source_error, budget_change=change)
                convergence.append(check)
                print(json.dumps(check), flush=True)
                previous = actual
            passed = (error <= (limits["zero_absolute_radiance"] if zero else
                                limits["image_normalized_l1"]) and
                      mean_error <= (limits["zero_absolute_radiance"] if zero else
                                     limits["mean_radiance_relative_error"]) and
                      (source_error is None or source_error <= limits["source_relative_error"]) and
                      change <= (limits["zero_absolute_radiance"] if zero else
                                 limits["final_budget_change"]))
            rows.append(dict(check, image=actual.tolist(), passed=bool(passed),
                layer_excitation=None if observed is None else observed.tolist(),
                apar_top=state.apar_spectral_top.tolist(), apar_bottom=state.apar_spectral_bot.tolist(),
                executions=[identity(p, backend) for p in (radiation, state, image)]))
    controls = []
    for backend in backends:
        selected = {r["case"]: r for r in rows if r["backend"] == backend}
        if "single" in selected and "single_double" in selected:
            single = np.asarray(selected["single"]["image"])
            doubled = np.asarray(selected["single_double"]["image"])
            error = float(np.max(abs(doubled-2*single)/np.maximum(2*abs(single), 1e-30)))
            controls.append(dict(backend=backend, linearity_error=error,
                                 passed=error <= limits["linearity_relative_error"]))
    return dict(input=config, coefficients=coefficients, references=refs, records=rows,
                convergence=convergence, controls=controls,
                passed=all(r["passed"] for r in rows+controls))


def serializable(value):
    if isinstance(value, np.ndarray):
        return value.tolist()
    raise TypeError(type(value).__name__)


def figures(result, output):
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    plt.rcParams.update({"svg.fonttype": "none", "font.size": 10})
    colors = dict(optix="#0072b2", vulkan="#d55e00", embree="#009e73")
    config = result["input"]
    labels = dict(single="Single layer, upper face", single_bottom="Single layer, lower face",
                  two_nadir="Two layers", three_nadir="Three layers")
    fig, axes = plt.subplots(1, 2, figsize=(10, 4.8), layout="constrained")
    wavelengths = config["emission_nm"]
    for name, style in (("single", "-"), ("single_bottom", "-."),
                        ("two_nadir", "--"), ("three_nadir", ":")):
        reference_image = np.asarray(result["references"][name]["image"])
        axes[0].plot(wavelengths, reference_image.mean((0, 1))*1000,
                     style, color="#222", label=labels[name]+" (analytic)")
        for backend, color in colors.items():
            row = next(r for r in result["records"] if r["case"] == name and r["backend"] == backend)
            axes[0].plot(wavelengths, np.asarray(row["image"]).mean((0, 1))*1000,
                         ".", color=color, label=backend if name == "single" else None)
    axes[0].set(xlabel="Emission wavelength / nm", ylabel="SIF radiance / mW m$^{-2}$ sr$^{-1}$ nm$^{-1}$",
                title="Single and multiple leaf layers")
    for backend, color in colors.items():
        levels = config["image_samples"]
        error = [max(r["image_error"] for r in result["convergence"] if
                 r["backend"] == backend and r["image_samples"] == q) for q in levels]
        axes[1].plot(levels, np.asarray(error)*100, "o-", color=color, label=backend)
    axes[1].axhline(config["thresholds"]["image_normalized_l1"]*100,
                    color="#555", ls="--", label="Error limit")
    axes[1].set(xscale="log", xlabel="Samples per image pixel", ylabel="Maximum normalized image L1 / %",
                title="Joint source and image sampling refinement")
    for ax in axes:
        ax.grid(alpha=.2)
    handles, names = axes[0].get_legend_handles_labels()
    fig.legend(handles, names, loc="outside lower center", ncol=3, fontsize=8)
    axes[1].legend(fontsize=7)
    fig.savefig(output/"results.svg")
    fig.savefig(output/"results.png", dpi=160)
    plt.close(fig)

    names = ("slit_nadir", "slit_aligned", "slit_oblique")
    fig, axes = plt.subplots(4, 3, figsize=(10, 6), layout="constrained")
    band = config["emission_nm"].index(740)
    vmax = max(np.asarray(result["references"][name]["image"])[..., band].max() for name in names)*1000
    for column, name in enumerate(names):
        exact = np.asarray(result["references"][name]["image"])[..., band]*1000
        images = [exact]+[np.asarray(next(r for r in result["records"] if
                   r["case"] == name and r["backend"] == b)["image"])[..., band]*1000 for b in BACKENDS]
        for row, values in enumerate(images):
            picture = axes[row, column].imshow(values, origin="upper", aspect="auto", vmin=0,
                vmax=vmax, extent=(-1.5, 1.5, -.5, .5), cmap="viridis")
            axes[row, column].set_yticks([])
            if not column:
                axes[row, column].set_ylabel(("Analytic", *BACKENDS)[row])
            if not row:
                angle = next(c["view_zenith"] for c in config["cases"] if c["name"] == name)
                axes[row, column].set_title(f"740 nm · view {angle}°")
            if row == 3:
                axes[row, column].set_xlabel("X relative to leaf centre / m")
    fig.colorbar(picture, ax=axes, label="SIF radiance / mW m$^{-2}$ sr$^{-1}$ nm$^{-1}$", shrink=.75)
    fig.savefig(output/"images.svg")
    fig.savefig(output/"images.png", dpi=160)
    plt.close(fig)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output", type=Path, default=Path("build/fluorescence"))
    parser.add_argument("--backend", choices=BACKENDS)
    parser.add_argument("--case")
    parser.add_argument("--plot-only", type=Path)
    args = parser.parse_args()
    config = json.loads((DATA/"input.json").read_text())
    coefficients = json.loads((DATA/"leaf-coefficients.json").read_text())
    if args.case:
        config["cases"] = [c for c in config["cases"] if c["name"] == args.case]
    args.output.mkdir(parents=True, exist_ok=True)
    if args.plot_only:
        figures(json.loads(args.plot_only.read_text()), args.output)
        return
    result = run(config, coefficients, [args.backend] if args.backend else BACKENDS)
    (args.output/"results.json").write_text(json.dumps(result, default=serializable,
        indent=2, allow_nan=False), encoding="utf-8")
    raise SystemExit(0 if result["passed"] else 1)


if __name__ == "__main__":
    main()