"""A9: smooth-soil microwave brightness against independent Fresnel references."""
from pathlib import Path
import argparse
import cmath
import json
import math

import less
import numpy as np

DATA = Path(__file__).resolve().parents[1]/"data/microwave-brightness"
BACKENDS = ("optix", "vulkan", "embree")
SOURCES = ("soil", "crown", "trunk", "sky", "atmosphere")


def reference(config, case):
    """Hallikainen polynomial, complex Fresnel amplitudes and scalar atmosphere."""
    texture = [1, 100*config["sand_fraction"], 100*config["clay_fraction"]]
    moisture = case["soil_moisture"]
    def polynomial(coefficients):
        return sum(sum(a*b for a, b in zip(row, texture))*moisture**order
                   for order, row in enumerate(coefficients))
    real = polynomial(config["dielectric_coefficients"]["real"])
    loss = polynomial(config["dielectric_coefficients"]["loss"])
    epsilon = complex(real, -loss)
    mu = math.cos(math.radians(case["view_zenith_degrees"]))
    root = cmath.sqrt(epsilon-(1-mu*mu))
    rh = abs((mu-root)/(mu+root))**2
    rv = abs((epsilon*mu-root)/(epsilon*mu+root))**2
    reflection = np.array([rv, rh, 0., 0.])
    emission = np.array([1-rv, 1-rh, 0., 0.])
    transmission = math.exp(-case["tau"]/mu)
    sky = case["air_K"]*(1-transmission)+config["background_K"]*transmission
    sources = dict(soil=(transmission*case["soil_K"]*emission).tolist(),
        sky=(transmission*sky*reflection).tolist(),
        atmosphere=[case["air_K"]*(1-transmission)]*2+[0., 0.],
        crown=[0.]*4, trunk=[0.]*4)
    total = np.sum(list(sources.values()), axis=0)
    return dict(permittivity=[real, loss], reflection_VH=[rv, rh],
        emissivity_VH=[1-rv, 1-rh], transmission=transmission, sky_K=sky,
        sources=sources, brightness_VH=total[:2].tolist(),
        stokes=[float(total[0]+total[1]), float(total[1]-total[0]), 0., 0.])


def scene_for(config, case, backend):
    scene = less.Scene(backend=backend)
    scene.sampler = less.Sampler(config.get("sampling_method", "independent"))
    scene.size = config["scene_size_m"]
    scene.repetitive = True
    scene.terrain = less.Terrain(property=less.Lambertian(reflectance=.1))
    scene.terrain.set_property(less.MicrowaveSoil(
        sand_fraction=config["sand_fraction"], clay_fraction=config["clay_fraction"],
        soil_moisture=case["soil_moisture"], dielectric_model=config["dielectric_model"],
        surface_model=config["surface_model"], rms_height_cm=config["rms_height_cm"],
        temperature_surface=case["soil_K"], temperature_deep=case["soil_K"]))
    scene.microclimate = less.Microclimate(air_temperature=case["air_K"]-273.15)
    return scene


def execute(config, case, scene, quality, angular_samples, seed, polarization="full"):
    sensor = less.MicrowaveRadiometer(geometry=less.Orthographic(
        image_size=config["image_size"], extent=config["footprint_m"],
        altitude=config["sensor_altitude_m"],
        sampling_seed=seed if config.get("sampling_method", "independent") != "independent" else 0,
        view_zenith=case["view_zenith_degrees"], view_azimuth=case["view_azimuth_degrees"]),
        frequency=config["frequency_Hz"], polarization=polarization, quality=quality,
        scattering_samples=angular_samples, atmosphere_optical_depth=case["tau"], seed=seed)
    product = scene.simulate(sensor)
    identity = product.metadata["execution"]
    assert identity["native"] and identity["backend"] == scene.backend
    assert identity["workflow"] == "passive_microwave"
    assert identity["executor"] == f"{scene.backend}.slang.microwave_radiometer"
    exact = reference(config, case)
    values = np.stack([product.Tb_V, product.Tb_H], axis=-1).astype(float)
    expected = np.array(exact["brightness_VH"])
    sources = {name: np.asarray(product.sources[name], float) for name in SOURCES}
    means = {name: v.mean(axis=(0, 1)).tolist() for name, v in sources.items()}
    t = exact["transmission"]
    reflect = sources["sky"][..., :2]/(t*exact["sky_K"])
    emissivity = sources["soil"][..., :2]/(t*case["soil_K"])
    total = np.sum(list(sources.values()), axis=0)
    errors = dict(brightness_K=float(np.max(abs(values-expected))),
        source_K=max(float(np.max(abs(v-np.array(exact["sources"][name]))))
                     for name, v in sources.items()),
        reflection_absolute=float(np.max(abs(reflect-exact["reflection_VH"]))),
        kirchhoff_absolute=float(np.max(abs(emissivity+reflect-1))),
        source_sum_K=float(np.max(abs(total[..., :2]-values))),
        uniformity_K=float(np.max(np.ptp(values, axis=(0, 1)))),
        cross_stokes_K=float(np.max(abs(total[..., 2:]))))
    if polarization == "full":
        errors["stokes_K"] = float(np.max(abs(np.asarray(product.data)-exact["stokes"])))
    channels = {"H": expected[1:2], "V": expected[:1], "HV": expected[::-1],
                "full": np.array(exact["stokes"])}[polarization]
    errors["polarization_mode_K"] = float(np.max(abs(np.asarray(product.data)-channels)))
    limits = config["thresholds"]
    return dict(backend=scene.backend, case=case["name"], quality=quality,
        angular_samples=angular_samples, seed=seed, polarization=polarization,
        brightness_VH=values.mean(axis=(0, 1)).tolist(),
        sources=means, stokes=([float(values[..., 0].mean()+values[..., 1].mean()),
            float(values[..., 1].mean()-values[..., 0].mean()),
            float(total[..., 2].mean()), float(total[..., 3].mean())]),
        reflection_VH=reflect.mean(axis=(0, 1)).tolist(),
        emissivity_VH=emissivity.mean(axis=(0, 1)).tolist(), errors=errors,
        equilibrium_recombined_VH=(case["soil_K"]*(emissivity+reflect)).mean(axis=(0, 1)).tolist(),
        execution={k: identity[k] for k in ("backend", "workflow", "executor", "native")},
        passed=all(value <= limits[name] for name, value in errors.items()))


def summary(row):
    return {k: row[k] for k in ("quality", "angular_samples", "seed", "brightness_VH", "errors")}


def run(config, backends=BACKENDS, selected=None):
    cases = [c for c in config["cases"] if selected is None or c["name"] == selected]
    references = {c["name"]: reference(config, c) for c in cases}
    records, sampling, stability, polarization_checks, angular_checks = [], [], [], [], []
    for backend in backends:
        for case in cases:
            scene = scene_for(config, case, backend)
            levels = []
            for quality, angular in zip(config["quality_levels"], config["angular_sample_levels"]):
                row = execute(config, case, scene, quality, angular, config["seeds"][0])
                levels.append(summary(row))
                print(backend, case["name"], quality, row["passed"], row["errors"], flush=True)
            change = float(np.max(abs(np.array(levels[-1]["brightness_VH"])-levels[-2]["brightness_VH"])))
            row["sampling_change_K"] = change
            row["passed"] &= change <= config["thresholds"]["sampling_change_K"]
            records.append(row)
            sampling.append(dict(backend=backend, case=case["name"], levels=levels))
            if case["name"] in config["stability_cases"]:
                repeats = [summary(row)]
                for seed in config["seeds"][1:]:
                    repeats.append(summary(execute(config, case, scene,
                        config["quality_levels"][-1], config["angular_sample_levels"][-1], seed)))
                temperatures = np.array([r["brightness_VH"] for r in repeats])
                spread = float(np.max(np.std(temperatures, axis=0, ddof=1)))
                stability.append(dict(backend=backend, case=case["name"], repeats=repeats,
                    standard_deviation_K=spread, passed=(spread <= config["thresholds"]["seed_spread_K"]
                        and all(v <= config["thresholds"][k] for r in repeats for k, v in r["errors"].items()))))
                angular_levels = []
                for angular in config["angular_sample_levels"][:-1]:
                    angular_levels.append(summary(execute(config, case, scene,
                        config["quality_levels"][-1], angular, config["seeds"][0])))
                angular_levels.append(summary(row))
                delta = float(np.max(np.ptp([l["brightness_VH"] for l in angular_levels], axis=0)))
                angular_checks.append(dict(backend=backend, case=case["name"], levels=angular_levels,
                    maximum_change_K=delta, passed=delta <= config["thresholds"]["sampling_change_K"]))
            if case["name"] == "wet_60":
                controls = []
                for mode in config["polarization_controls"]:
                    other = execute(config, case, scene, config["quality_levels"][-1],
                        config["angular_sample_levels"][-1], config["seeds"][0], mode)
                    delta = float(np.max(abs(np.array(other["brightness_VH"])-row["brightness_VH"])))
                    controls.append(dict(mode=mode, maximum_change_K=delta,
                        channel_error_K=other["errors"]["polarization_mode_K"]))
                polarization_checks.append(dict(backend=backend, controls=controls,
                    passed=all(max(c["maximum_change_K"], c["channel_error_K"]) <=
                               config["thresholds"]["polarization_mode_K"]
                               for c in controls)))
    parity = []
    for case in cases:
        rows = [r for r in records if r["case"] == case["name"]]
        delta = float(np.max(np.ptp([r["brightness_VH"] for r in rows], axis=0)))
        parity.append(dict(case=case["name"], maximum_backend_difference_K=delta,
                           passed=delta <= config["thresholds"]["brightness_K"]))
    return dict(input=config, references=references, records=records, sampling=sampling,
        stability=stability, polarization_checks=polarization_checks, angular_checks=angular_checks,
        backend_parity=parity,
        passed=all(r["passed"] for r in records+stability+polarization_checks+angular_checks+parity))


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output", type=Path, default=Path("build/microwave-brightness-validation"))
    parser.add_argument("--backend", choices=BACKENDS)
    parser.add_argument("--case")
    parser.add_argument("--figures-only", action="store_true")
    args = parser.parse_args()
    args.output.mkdir(parents=True, exist_ok=True)
    config = json.loads((DATA/"input.json").read_text(encoding="utf-8"))
    path = args.output/"results.json"
    if args.figures_only:
        result = json.loads(path.read_text(encoding="utf-8"))
    else:
        result = run(config, [args.backend] if args.backend else BACKENDS, args.case)
        if args.case is None and "canopy" in config:
            from appendix_microwave_canopy import run as run_canopy
            result["canopy"] = run_canopy(config, [args.backend] if args.backend else BACKENDS)
            result["passed"] = result["passed"] and result["canopy"]["passed"]
        path.write_text(json.dumps(result, indent=2, allow_nan=False), encoding="utf-8")
    if args.case is None:
        figures(result, args.output)
        if "canopy" in result and args.backend is None:
            from appendix_microwave_canopy import figures as canopy_figures
            canopy_figures(config, result["canopy"], args.output)
    raise SystemExit(0 if result["passed"] else 1)


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

    plt.rcParams.update({"font.size": 10, "svg.fonttype": "none",
                         "axes.spines.top": False, "axes.spines.right": False})
    colors = ("#2563a6", "#d47713", "#27976b")
    markers = ("o", "s", "^")
    fig, axes = plt.subplots(2, 2, figsize=(12, 8), layout="constrained")
    config, refs = result["input"], result["references"]
    for prefix, linestyle in (("dry", "-"), ("wet", "--")):
        for index, polarization in ((0, "V"), (1, "H")):
            names = [f"{prefix}_{angle}" for angle in (0, 30, 60, 75)]
            axes[0, 0].plot([0, 30, 60, 75], [refs[n]["brightness_VH"][index] for n in names],
                linestyle, color="#334155" if index == 0 else "#7b536b",
                label=f"{prefix}, {polarization}: reference")
            for backend, color, marker in zip(BACKENDS, colors, markers):
                rows = {r["case"]: r for r in result["records"] if r["backend"] == backend}
                if not rows:
                    continue
                axes[0, 0].plot([0, 30, 60, 75], [rows[n]["brightness_VH"][index] for n in names],
                    ls="none", marker=marker, markerfacecolor="none", ms=8, color=color)
    axes[0, 0].set(xlabel="Incidence angle / degrees", ylabel="Brightness temperature / K",
                    title="(a) Dry/wet soil and H/V polarization")
    axes[0, 0].legend(frameon=False, fontsize=9)
    names = ("wet_60", "cool_atmosphere", "warm_atmosphere")
    positions = np.arange(6)
    bottoms = np.zeros(6)
    for name, color in zip(("soil", "sky", "atmosphere"), ("#b28d58", "#83bfd5", "#d69273")):
        values = np.array([refs[n]["sources"][name][p] for n in names for p in (0, 1)])
        axes[0, 1].bar(positions, values, bottom=bottoms, color=color, label=name)
        bottoms += values
    for backend, color, marker in zip(BACKENDS, colors, markers):
        rows = {r["case"]: r for r in result["records"] if r["backend"] == backend}
        if not rows:
            continue
        axes[0, 1].plot(positions, [rows[n]["brightness_VH"][p] for n in names for p in (0, 1)],
            ls="none", marker=marker, markerfacecolor="none", color=color, ms=8, label=backend)
        angles = [0, 30, 60, 75]
        axes[1, 0].plot(angles, [rows[f"wet_{a}"]["errors"]["brightness_K"] for a in angles],
                       marker=marker, color=color, label=backend)
        levels = next(s["levels"] for s in result["sampling"] if
                      s["backend"] == backend and s["case"] == "wet_75")
        axes[1, 1].semilogx([l["quality"] for l in levels],
            [l["errors"]["brightness_K"] for l in levels], marker=marker,
            markerfacecolor="none", color=color, label=backend)
    axes[0, 1].set(xticks=positions, xticklabels=["V", "H"]*3,
        xlabel="No atmosphere          Cool atmosphere          Warm atmosphere",
        ylabel="Source contribution / K", title="(b) Source partition: bars reference, symbols LESS")
    axes[0, 1].legend(frameon=False, ncols=3, fontsize=9,
                       loc="lower center", bbox_to_anchor=(.5, 1.16))
    axes[1, 0].set(xlabel="Incidence angle / degrees", ylabel="Maximum brightness error / K",
                    title="(c) Wet-soil accuracy")
    axes[1, 0].legend(frameon=False, fontsize=9)
    axes[1, 1].set(xlabel="Samples per pixel", ylabel="Maximum brightness error / K",
                    title="(d) Wet soil at 75°: sampling invariance")
    axes[1, 1].legend(frameon=False, fontsize=9)
    for ax in axes.flat:
        ax.grid(alpha=.18, axis="y")
    fig.savefig(output/"results.svg")
    fig.savefig(output/"results.png", dpi=150)
    plt.close(fig)


if __name__ == "__main__":
    main()