"""A9 absorbing canopy: independent dipoles and piecewise-exponential image integrals."""
from pathlib import Path
import argparse
import json
import math

import less
import numpy as np
from less.distributions import Fixed
import appendix_microwave_brightness as ground


def absorption(config, case):
    """Dual-dispersion dielectric and electrostatic oblate-dipole absorption."""
    s = config["canopy"]
    if s["volume_scattering"] is not False:
        raise ValueError("this analytic reference requires volume_scattering=False")
    if s["leaf_dry_density"] != 0 or s["leaf_orientation"] != "horizontal":
        raise ValueError("this reference requires gravimetric dielectric inputs and horizontal leaves")
    f, m, t = config["frequency_Hz"]*1e-9, s["leaf_moisture"], case["vegetation_K"]-273.15
    relaxation = 1.1109e-10+t*(-3.824e-12+t*(6.938e-14-t*5.096e-16))
    static = 88.045+t*(-.4147+t*(6.295e-4+t*1.075e-5))
    free = 4.9+(static-4.9)/(1+1j*f*1e9*relaxation)-1j*22.74/f
    bound = 2.9+55/(1+(1+1j)*math.sqrt(f/.18)/math.sqrt(2))
    epsilon = (1.7-.74*m+6.16*m*m+m*(.55*m-.076)*free+
               4.64*m*m/(7.36*m*m+1)*bound)
    ratio = s["leaf_thickness_m"]/(2*s["leaf_radius_m"])
    eccentricity = math.sqrt(1-ratio*ratio)
    transverse = ratio*(math.atan(eccentricity/ratio)-eccentricity*ratio)/(2*eccentricity**3)
    axial = 1-2*transverse
    volume = math.pi*s["leaf_radius_m"]**2*s["leaf_thickness_m"]
    chi = epsilon-1
    alpha_t, alpha_z = volume*chi/(1+chi*transverse), volume*chi/(1+chi*axial)
    theta = math.radians(case["view_zenith_degrees"])
    k = 2*math.pi*config["frequency_Hz"]/299792458
    alpha_v = alpha_t*math.cos(theta)**2+alpha_z*math.sin(theta)**2
    rates = -case["density_m3"]*k*np.array([alpha_v.imag, alpha_t.imag])
    return dict(permittivity=[epsilon.real, -epsilon.imag], absorption_VH_m1=rates.tolist(),
                depolarization=[transverse, axial])


def ground_reference(config, case):
    return ground.reference(config, dict(case, soil_moisture=.3, tau=0, air_K=290))


def path_length(config, case, x, sign):
    """Closed ray/vertical-strip overlap as a function of ground coordinate x."""
    s = config["canopy"]
    theta = math.radians(case["view_zenith_degrees"])
    mu, slope = math.cos(theta), sign*math.tan(theta)
    low, high = s["bottom_m"], s["top_m"]
    if case["kind"] == "uniform":
        return (high-low)/mu
    result = 0.
    for left, right in case["intervals_m"]:
        if slope == 0:
            result += (high-low)/mu if left <= x < right else 0
        else:
            entry, exit = sorted(((left-x)/slope, (right-x)/slope))
            result += max(0., min(high, exit)-max(low, entry))/mu
    return result


def breakpoints(config, case, left, right):
    s = config["canopy"]
    tangent = math.tan(math.radians(case["view_zenith_degrees"]))
    points = [left, right]
    for interval in case.get("intervals_m", []):
        for edge in interval:
            for height in (s["bottom_m"], s["top_m"]):
                for sign in (-1, 1):
                    x = edge+sign*tangent*height
                    if left < x < right:
                        points.append(x)
    return sorted(set(points))


def exponential_average(rate, first, last):
    difference = rate*(last-first)
    return (math.exp(-rate*first)*(-math.expm1(-difference))/difference
            if abs(difference) > 1e-12 else math.exp(-rate*(first+last)/2))


def pixel_transmissions(config, case, left, right, rate):
    up, both = 0., 0.
    points = breakpoints(config, case, left, right)
    for a, b in zip(points[:-1], points[1:]):
        length = b-a
        # Interior samples determine each affine segment, including nadir
        # jumps. No endpoint convention changes a pixel's area integral.
        x1, x2 = a+.25*length, a+.75*length
        lu1, lu2 = [path_length(config, case, x, 1) for x in (x1, x2)]
        ld1, ld2 = [path_length(config, case, x, -1) for x in (x1, x2)]
        def integral(v1, v2):
            return exponential_average(rate, 1.5*v1-.5*v2, 1.5*v2-.5*v1)*length
        up += integral(lu1, lu2)
        both += integral(lu1+ld1, lu2+ld2)
    return up/(right-left), both/(right-left)


def reference(config, case):
    s = config["canopy"]
    particle, soil = absorption(config, case), ground_reference(config, case)
    width = 1 if case["kind"] == "uniform" else s["image_size"][0]
    extent = s["image_extent_m"][0]
    edges = np.linspace(s["image_center_m"][0]-extent/2,
                        s["image_center_m"][0]+extent/2, width+1)
    sources = {name: np.zeros((width, 2)) for name in ("soil", "crown", "sky")}
    for i, (left, right) in enumerate(zip(edges[:-1], edges[1:])):
        for p, rate in enumerate(particle["absorption_VH_m1"]):
            u, ud = pixel_transmissions(config, case, left, right, rate)
            r = soil["reflection_VH"][p]
            sources["soil"][i, p] = u*(1-r)*case["soil_K"]
            sources["crown"][i, p] = ((1-u)+r*(u-ud))*case["vegetation_K"]
            sources["sky"][i, p] = r*ud*config["background_K"]
    return dict(**particle, profile_VH_K=np.sum(list(sources.values()), axis=0).tolist(),
                sources_profile_VH_K={k: v.tolist() for k, v in sources.items()})


def reference_check(config, case, exact):
    """Check exponential antiderivatives with independent Gauss integration."""
    settings = config["canopy"]
    width = len(exact["profile_VH_K"])
    extent = settings["image_extent_m"][0]
    center = settings["image_center_m"][0]
    edges = np.linspace(center-extent/2, center+extent/2, width+1)
    nodes, weights = np.polynomial.legendre.leggauss(32)
    reflection = ground_reference(config, case)["reflection_VH"]
    values = np.zeros((width, 2))
    for index, (left, right) in enumerate(zip(edges[:-1], edges[1:])):
        points = breakpoints(config, case, left, right)
        for a, b in zip(points[:-1], points[1:]):
            x = (a+b)/2+nodes*(b-a)/2
            up = np.array([path_length(config, case, p, 1) for p in x])
            down = np.array([path_length(config, case, p, -1) for p in x])
            for polarization, rate in enumerate(exact["absorption_VH_m1"]):
                u, d = np.exp(-rate*up), np.exp(-rate*down)
                r = reflection[polarization]
                temperatures = (u*(1-r)*case["soil_K"]+
                    ((1-u)+u*r*(1-d))*case["vegetation_K"]+u*r*d*config["background_K"])
                values[index, polarization] += float(weights@temperatures)*(b-a)/(2*(right-left))
    error = float(np.max(abs(values-exact["profile_VH_K"])))
    return dict(gauss_order=32, maximum_difference_K=error,
                passed=error <= settings["thresholds"]["reference_quadrature_K"])


def scene_for(config, case, backend):
    scene = ground.scene_for(config, dict(case, soil_moisture=.3, tau=0, air_K=290), backend)
    scene.measurement_plane = 0
    scene.repetitive = case["kind"] == "uniform"
    plant = less.Object("canopy")
    plant.add_component("canopy", less.TurbidBoundary(str(ground.DATA/(
        "layer.obj" if case["kind"] == "uniform" else "block.obj"))))
    plant.set_property("canopy", less.MicrowaveVegetation(components=[less.MicrowaveLeaf(
        radius=config["canopy"]["leaf_radius_m"], thickness=config["canopy"]["leaf_thickness_m"],
        density=case["density_m3"], moisture=config["canopy"]["leaf_moisture"],
        dry_density=config["canopy"]["leaf_dry_density"], temperature=case["vegetation_K"],
        orientation=Fixed(0), scattering_model="rayleigh")]))
    positions = [(0, 0, 0), (5, 0, 0)] if case["kind"] == "gap" else [(0, 0, 0)]
    scene.add(plant, positions=positions)
    return scene


def execute(config, case, scene, quality, exact, seed=None):
    s = config["canopy"]
    uniform = case["kind"] == "uniform"
    run_seed = config["seeds"][0] if seed is None else seed
    product = scene.simulate(less.MicrowaveRadiometer(
        geometry=less.Orthographic(image_size=2 if uniform else tuple(s["image_size"]),
            sampling_seed=run_seed if config.get("sampling_method", "independent") != "independent" else 0,
            extent=tuple(s["image_extent_m"]), center=tuple(s["image_center_m"]),
            altitude=config["sensor_altitude_m"], view_zenith=case["view_zenith_degrees"],
            view_azimuth=s["view_azimuth_degrees"]),
        frequency=config["frequency_Hz"], quality=quality, scattering_samples=s["angular_samples"],
        atmosphere_optical_depth=0, volume_scattering=s["volume_scattering"],
        seed=run_seed))
    identity = product.metadata["execution"]
    assert identity["native"] and identity["backend"] == scene.backend
    assert identity["executor"] == f"{scene.backend}.slang.microwave_radiometer"
    assert identity["workflow"] == "passive_microwave"
    assert product.metadata["volume_scattering"] is False
    image = np.stack([product.Tb_V, product.Tb_H], axis=-1).astype(float)
    expected = np.array(exact["profile_VH_K"])[None, :, :]
    difference = image-expected
    means = {name: np.asarray(value, float)[..., :2].mean(axis=(0, 1)).tolist()
             for name, value in product.sources.items()}
    source_error = max(float(np.max(abs(np.array(means[name])-
        np.mean(exact["sources_profile_VH_K"][name], axis=0)))) for name in ("soil", "crown", "sky"))
    response = (np.asarray(product.sources["soil"], float)[..., :2]/case["soil_K"]+
        np.asarray(product.sources["crown"], float)[..., :2]/case["vegetation_K"]+
        np.asarray(product.sources["sky"], float)[..., :2]/config["background_K"])
    row = dict(backend=scene.backend, case=case["name"], quality=quality,
        volume_scattering=product.metadata["volume_scattering"],
        seed=config["seeds"][0] if seed is None else seed,
        brightness_mean_VH_K=image.mean(axis=(0, 1)).tolist(),
        brightness_image_VH_K=image.tolist(), sources_mean_VH_K=means,
        errors=dict(mean_K=float(np.max(abs(np.mean(difference, axis=(0, 1))))),
            image_rmse_K=float(np.sqrt(np.mean(difference*difference))),
            image_max_K=float(np.max(abs(difference))), source_mean_K=source_error,
            source_sum_K=float(np.max(abs(sum(np.asarray(v, float)[..., :2]
                for v in product.sources.values())-image))),
            thermal_response_closure=float(np.max(abs(response-1)))),
        execution={k: identity[k] for k in ("backend", "workflow", "executor", "native")})
    e, limits = row["errors"], s["thresholds"]
    row["passed"] = bool(e["image_max_K"] <= limits["uniform_brightness_K"] and
        e["source_mean_K"] <= limits["uniform_source_K"]) if uniform else bool(
        e["mean_K"] <= limits["image_mean_K"] and e["image_rmse_K"] <= limits["image_rmse_K"] and
        e["image_max_K"] <= limits["image_max_K"] and e["source_mean_K"] <= limits["image_source_mean_K"])
    row["passed"] = bool(row["passed"] and e["source_sum_K"] <= limits["source_sum_K"] and
        e["thermal_response_closure"] <= limits["thermal_response_closure"])
    return row


def compact(row):
    return {k: v for k, v in row.items() if k not in
            ("brightness_image_VH_K", "sources_mean_VH_K", "execution", "passed")}


def run(config, backends=ground.BACKENDS, selected=None):
    settings = config["canopy"]
    cases = [c for c in settings["cases"] if selected is None or c["name"] == selected]
    refs = {c["name"]: reference(config, c) for c in cases}
    rows, sampling, stability = [], [], []
    for backend in backends:
        for case in cases:
            scene = scene_for(config, case, backend)
            budgets = settings["uniform_quality" if case["kind"] == "uniform" else "image_quality"]
            levels = []
            previous = None
            for quality in budgets:
                row = execute(config, case, scene, quality, refs[case["name"]])
                current = np.array(row["brightness_image_VH_K"])
                if previous is not None:
                    row["sampling_mean_change_K"] = float(np.max(abs(
                        current.mean(axis=(0, 1))-previous.mean(axis=(0, 1)))))
                    row["sampling_image_rmse_change_K"] = float(np.sqrt(np.mean((current-previous)**2)))
                previous = current
                levels.append(compact(row))
                print(backend, case["name"], quality, row["passed"], row["errors"], flush=True)
            row["passed"] = bool(row["passed"] and
                row["sampling_mean_change_K"] <= settings["thresholds"]["sampling_mean_change_K"] and
                row["sampling_image_rmse_change_K"] <= settings["thresholds"]["sampling_image_rmse_change_K"])
            rows.append(row)
            sampling.append(dict(backend=backend, case=case["name"], levels=levels))
            if case["name"] in settings.get("stability_cases", []):
                repeats = [compact(row)]
                for seed in config["seeds"][1:]:
                    repeats.append(compact(execute(config, case, scene, budgets[-1], refs[case["name"]], seed)))
                means = np.array([r["brightness_mean_VH_K"] for r in repeats])
                deviation = means.std(axis=0, ddof=1)
                expected = np.mean(refs[case["name"]]["profile_VH_K"], axis=0)
                ci = 4.302652729911275*deviation/np.sqrt(len(repeats))
                stability.append(dict(backend=backend, case=case["name"], repeats=repeats,
                    mean_VH_K=means.mean(axis=0).tolist(), standard_deviation_VH_K=deviation.tolist(),
                    confidence_95_half_width_VH_K=ci.tolist(),
                    passed=bool(np.max(abs(means.mean(axis=0)-expected)) <= settings["thresholds"]["image_mean_K"])))
    reference_checks = {c["name"]: reference_check(config, c, refs[c["name"]]) for c in cases}
    return dict(references=refs, reference_checks=reference_checks, records=rows,
        sampling=sampling, stability=stability,
        passed=all(r["passed"] for r in rows+stability+list(reference_checks.values())))


def figures(config, 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")
    names = [f"uniform_{a}" for a in (0, 40, 60)]
    for p, style in ((0, "-"), (1, "--")):
        axes[0, 0].plot([0, 40, 60], [result["references"][n]["profile_VH_K"][0][p] for n in names],
            style, color="#334155", label=("V" if p == 0 else "H")+" reference")
    x = np.linspace(4.25, 15.75, config["canopy"]["image_size"][0])
    exact = np.array(result["references"]["gap_60"]["profile_VH_K"])
    axes[0, 1].plot(x, exact[:, 0], "-", color="#334155", label="V reference")
    axes[0, 1].plot(x, exact[:, 1], "--", color="#334155", label="H reference")
    for backend, color, marker in zip(ground.BACKENDS, colors, markers):
        rows = {r["case"]: r for r in result["records"] if r["backend"] == backend}
        if not rows:
            continue
        for p in (0, 1):
            axes[0, 0].plot([0, 40, 60], [rows[n]["brightness_mean_VH_K"][p] for n in names],
                ls="none", marker=marker, markerfacecolor="none", ms=8, color=color,
                label=backend if p == 0 else None)
            axes[0, 1].plot(x, np.array(rows["gap_60"]["brightness_image_VH_K"]).mean(axis=0)[:, p],
                ls="none", marker=marker, markerfacecolor="none", ms=5, color=color)
        for name, axis in (("block_0", axes[1, 0]), ("gap_60", axes[1, 1])):
            levels = next(s["levels"] for s in result["sampling"] if
                          s["backend"] == backend and s["case"] == name)
            axis.loglog([r["quality"] for r in levels], [r["errors"]["image_rmse_K"] for r in levels],
                        marker=marker, markerfacecolor="none", color=color, label=backend)
    axes[0, 0].set(xlabel="Incidence angle / degrees", ylabel="Brightness temperature / K",
                    title="(a) Uniform absorbing layer")
    axes[0, 0].legend(frameon=False, fontsize=9, ncols=2)
    axes[0, 1].set(xlabel="Ground coordinate x / m", ylabel="Brightness temperature / K",
                    title="(b) Two blocks with a gap, view 60°")
    axes[0, 1].legend(frameon=False, fontsize=9)
    axes[1, 0].set(title="(c) Single block, nadir mixed pixels")
    axes[1, 1].set(title="(d) Two blocks with a gap, oblique view")
    for axis in axes[1]:
        axis.set(xlabel="Samples per pixel", ylabel="Image RMSE / K")
        axis.legend(frameon=False, fontsize=9)
    for axis in axes.flat:
        axis.grid(alpha=.18, axis="y")
    fig.savefig(output/"canopy-results.svg")
    fig.savefig(output/"canopy-results.png", dpi=150)
    plt.close(fig)

    fig, axes = plt.subplots(2, 5, figsize=(14, 5), layout="constrained")
    image_plot, error_plot = None, None
    error_limit = max(r["errors"]["image_max_K"] for r in result["records"]
                      if r["case"] in ("block_0", "gap_60"))
    for index, name in enumerate(("block_0", "gap_60")):
        target = np.broadcast_to(np.array(result["references"][name]["profile_VH_K"])[None, :, 1], (4, 24))
        image_plot = axes[index, 0].imshow(target, origin="upper", aspect="auto", vmin=80, vmax=280,
                                          cmap="viridis", extent=[4, 16, 11, 9])
        errors = []
        for column, backend in enumerate(ground.BACKENDS, 1):
            row = next(r for r in result["records"] if r["backend"] == backend and r["case"] == name)
            image = np.array(row["brightness_image_VH_K"])[..., 1]
            axes[index, column].imshow(image, origin="upper", aspect="auto", vmin=80, vmax=280,
                                      cmap="viridis", extent=[4, 16, 11, 9])
            errors.append(abs(image-target))
        error_plot = axes[index, 4].imshow(np.max(errors, axis=0), origin="upper", aspect="auto",
            vmin=0, vmax=max(error_limit, 1e-12), cmap="magma", extent=[4, 16, 11, 9])
        for column, title in enumerate(("Reference H", "optix", "vulkan", "embree", "Maximum |error|")):
            axes[index, column].set_title(title, fontsize=10)
            axes[index, column].set_xlabel("x / m")
        axes[index, 0].set_ylabel(("Single block, 0°" if index == 0 else "Two blocks, 60°")+"\ny / m")
    fig.colorbar(image_plot, ax=axes[:, :4], label="H brightness / K", shrink=.8)
    fig.colorbar(error_plot, ax=axes[:, 4], label="Error / K", shrink=.8)
    fig.savefig(output/"canopy-images.svg")
    fig.savefig(output/"canopy-images.png", dpi=150)
    plt.close(fig)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--backend", choices=ground.BACKENDS)
    parser.add_argument("--case")
    parser.add_argument("--output", type=Path, default=Path("build/microwave-canopy-validation"))
    parser.add_argument("--figures-only", action="store_true")
    args = parser.parse_args()
    config = json.loads((ground.DATA/"input.json").read_text(encoding="utf-8"))
    args.output.mkdir(parents=True, exist_ok=True)
    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 ground.BACKENDS, args.case)
        path.write_text(json.dumps(result, indent=2, allow_nan=False), encoding="utf-8")
    if args.case is None and args.backend is None:
        figures(config, result, args.output)
    raise SystemExit(0 if result["passed"] else 1)