"""Analytic geometry and radiance checks across distance and triangle scales."""
from pathlib import Path
import argparse
import json

import less
import numpy as np

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


def reference(config, case):
    slope, gap = config["slope_dz_dx"], config["strip_height_gap_m"]
    angle = np.deg2rad(case["view_zenith_degrees"])
    sine, cosine, tangent = np.sin(angle), np.cos(angle), np.tan(angle)
    normalizer = np.sqrt(1+slope*slope)
    factor = config["sun_normal_irradiance"]/np.pi/normalizer
    radiance = config["plate_reflectance"]*factor
    strip_radiance = config["strip_reflectance"]*factor
    half = config["strip_width_m"]/2
    projected = (1-slope*tangent)*half
    shadow = (-projected, projected)
    visible_strip = (-projected-tangent*gap, projected-tangent*gap)
    width = config["image_size"][0]
    edges = np.linspace(-config["image_extent_m"][0]/2, config["image_extent_m"][0]/2, width+1)
    def overlap(interval):
        return np.maximum(0, np.minimum(edges[1:], interval[1])-
                          np.maximum(edges[:-1], interval[0]))/np.diff(edges)
    intersection = (max(shadow[0], visible_strip[0]), min(shadow[1], visible_strip[1]))
    union = overlap(shadow)+overlap(visible_strip)
    if intersection[1]>intersection[0]:
        union -= overlap(intersection)
    profile = radiance*(1-union)+strip_radiance*overlap(visible_strip)
    centre = np.full(2, case["plate_side_m"]/2)
    axis = np.array([sine, 0, cosine])
    shift = sine*gap/(cosine-slope*sine)
    targets = [-1.5, -half-shift-.01, -half-shift+.01, 0,
               half-shift-.01, half-shift+.01, 1.5]
    rays = []
    for x in targets:
        target = np.array([centre[0]+x, centre[1]+.13, config["plate_height_m"]+slope*x])
        origin = target+case["distance_m"]*axis
        upper_range = case["distance_m"]-gap/(cosine-slope*sine)
        upper = origin-upper_range*axis
        on_strip = abs(upper[0]-centre[0])<=half
        point = upper if on_strip else target
        rays.append(dict(origin=origin.tolist(), point=point.tolist(),
            range_m=float(upper_range if on_strip else case["distance_m"]),
            surface="strip" if on_strip else "plate"))
    return dict(radiance=float(radiance), strip_radiance=float(strip_radiance),
        profile=profile.tolist(), shadow=list(shadow), visible_strip=list(visible_strip), rays=rays)


def scene_for(config, case, variant, backend):
    scene = less.Scene(backend=backend)
    scene.size = case["plate_side_m"]
    scene.terrain = None
    scene.repetitive = False
    scene.measurement_plane = config["plate_height_m"]
    scene.illumination = less.Illumination(source=less.Sun(zenith=0,
        irradiance=config["sun_normal_irradiance"]), atmosphere=less.NoAtmosphere())
    centre = np.full(2, case["plate_side_m"]/2)
    for name, path, height, rho in (
            ("plate", f'plate-{case["plate_side_m"]}.obj', config["plate_height_m"],
             config["plate_reflectance"]),
            ("strip", "strip.obj", config["plate_height_m"]+config["strip_height_gap_m"],
             config["strip_reflectance"])):
        obj = less.Object(name, mesh=str(DATA/path))
        obj.set_property(less.Lambertian(reflectance=rho))
        scene.add(obj, positions=[(*centre, height)])
    if variant=="refined":
        scene.subdivide_large_primitives(max_area=case["plate_side_m"]**2/256,
                                         max_total_triangles=4096)
    return scene


def identity(product, backend, workflow):
    execution = product.metadata["execution"]
    assert execution["backend"]==backend and execution["native"] is True
    assert execution["workflow"]==workflow
    expected = "lidar" if workflow=="lidar_point_cloud" else "optical_backward"
    assert execution["executor"]==f"{backend}.slang.{expected}"
    return {key:execution[key] for key in ("backend", "workflow", "executor", "native")}


def run(config):
    references = {case["name"]:reference(config, case) for case in config["cases"]}
    records, convergence = [], []
    limits = config["thresholds"]
    for backend in BACKENDS:
        for case in config["cases"]:
            exact = references[case["name"]]
            centre = np.full(2, case["plate_side_m"]/2)
            angle = np.deg2rad(case["view_zenith_degrees"])
            profiles = {}
            for variant in config["mesh_variants"]:
                scene = scene_for(config, case, variant, backend)
                survey = less.LiDARSurvey(scanner=less.Scanner(beam_samples=1, beam_divergence=0,
                    wavelengths=[config["wavelength_nm"]]), platform=less.Tripod(height=1),
                    legs=[less.ScanPosition(*ray["origin"], azimuth_range=(270,270.01),
                        zenith_range=(180-case["view_zenith_degrees"],180-case["view_zenith_degrees"]+.01))
                        for ray in exact["rays"]],
                    range_min=0, range_max=case["distance_m"]+2, range_resolution=1)
                cloud = scene.simulate(survey, product="fast_point_cloud")
                point_identity = identity(cloud, backend, "lidar_point_cloud")
                primitive_count = scene.n_primitives()
                assert primitive_count==4 if variant=="coarse" else primitive_count>4
                expected = np.array([ray["point"] for ray in exact["rays"]])
                origins = np.array([ray["origin"] for ray in exact["rays"]])
                points = np.asarray(cloud.points, float)
                valid = points.shape==expected.shape
                position_error = float(np.max(np.linalg.norm(points-expected, axis=1))) if valid else None
                range_error = (float(np.max(abs(np.linalg.norm(points-origins,axis=1)-
                    np.array([ray["range_m"] for ray in exact["rays"]])))) if valid else None)
                relative_z = points[:,2]-config["plate_height_m"]-config["slope_dz_dx"]*(points[:,0]-centre[0])
                plane_error = float(np.max(np.minimum(abs(relative_z),abs(relative_z-
                    config["strip_height_gap_m"]))/np.sqrt(1+config["slope_dz_dx"]**2))) if len(points) else None
                previous = None
                profile_ref = np.array(exact["profile"])
                lit = np.isclose(profile_ref, exact["radiance"], rtol=0,atol=1e-12)
                for quality in config["qualities"]:
                    product = scene.simulate(less.OpticalImager(less.Orthographic(
                        image_size=tuple(config["image_size"]), extent=tuple(config["image_extent_m"]),
                        pixel_sampling=config.get("image_sampling_method", "independent"),
                        center=tuple(centre), altitude=case["distance_m"]*np.cos(angle),
                        view_zenith=case["view_zenith_degrees"],view_azimuth=90),
                        bands=[config["wavelength_nm"]], quality=quality), scattering_order=1)
                    execution = identity(product, backend, "optical_backward")
                    image = np.asarray(product.data, float)[:,:,0]
                    profile = image.mean(axis=0)
                    error = float(np.mean(abs(image-profile_ref))/exact["radiance"])
                    lit_error = float(np.max(abs(image[:,lit]/exact["radiance"]-1)))
                    change = None if previous is None else float(np.mean(abs(profile-previous))/exact["radiance"])
                    convergence.append(dict(backend=backend,case=case["name"],mesh=variant,
                        quality=quality,image_normalized_l1=error,lit_relative_error=lit_error,
                        budget_change=change))
                    previous = profile
                profiles[variant] = profile
                passed = (valid and position_error<=limits["position_error_m"]
                    and range_error<=limits["range_error_m"] and plane_error<=limits["plane_residual_m"]
                    and error<=limits["image_normalized_l1"] and lit_error<=limits["lit_relative_error"]
                    and change<=limits["last_budget_image_change"])
                records.append(dict(backend=backend,case=case["name"],mesh=variant,
                    primitive_count=primitive_count,
                    position_error_m=position_error,range_error_m=range_error,plane_residual_m=plane_error,
                    points=points.tolist(),point_execution=point_identity,
                    image=image.tolist(),execution=execution,
                    image_normalized_l1=error,lit_relative_error=lit_error,budget_change=change,
                    passed=bool(passed)))
                print(backend,case["name"],variant,position_error,error,lit_error,passed,flush=True)
            difference = float(np.mean(abs(profiles["coarse"]-profiles["refined"]))/exact["radiance"])
            for row in records[-2:]:
                row["mesh_image_difference"] = difference
                row["passed"] = bool(row["passed"] and difference<=limits["mesh_image_difference"])
    return dict(input=config,references=references,records=records,convergence=convergence,
                passed=all(row["passed"] for row in records))


def plots(result, output):
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    plt.rcParams.update({"font.size":10,"svg.fonttype":"none"})
    config = result["input"]
    colors = dict(optix="#0072b2",vulkan="#d55e00",embree="#009e73")
    markers = dict(optix="o",vulkan="x",embree="+")
    fig, axes = plt.subplots(2,2,figsize=(11,7),layout="constrained")
    x = np.linspace(-config["image_extent_m"][0]/2,config["image_extent_m"][0]/2,
                    config["image_size"][0],endpoint=False)+config["image_extent_m"][0]/config["image_size"][0]/2
    ref = result["references"]["grazing"]
    axes[0,0].step(x,np.array(ref["profile"])/ref["radiance"],where="mid",color="k",label="Analytic pixel means")
    for backend in BACKENDS:
        row = next(r for r in result["records"] if r["backend"]==backend and r["case"]=="grazing" and r["mesh"]=="coarse")
        axes[0,0].plot(x,np.array(row["image"]).mean(axis=0)/ref["radiance"],ls="none",marker=markers[backend],ms=4,color=colors[backend],label=backend)
    axes[0,0].set(title="75° view: occluder and cast shadow",xlabel="Image-plane X relative to centre (m)",ylabel="Radiance / sunlit plane radiance")
    for ax, names, values, title, xlabel in (
            (axes[0,1],["baseline","range_1km","range_10km"],[10,1000,10000],"Ray distance","Distance (m)"),
            (axes[1,0],["baseline","triangle_1km","triangle_10km"],[20,1000,10000],"Triangle scale; distance 10 m","Square side (m)")):
        for backend in BACKENDS:
            for variant,style in (("coarse","-"),("refined","--")):
                rows=[next(r for r in result["records"] if r["backend"]==backend and r["case"]==name and r["mesh"]==variant) for name in names]
                ax.plot(values,[r["position_error_m"] for r in rows],style,marker=markers[backend],color=colors[backend],label=f"{backend} {variant}")
        ax.axhline(config["thresholds"]["position_error_m"],color="k",ls=":")
        ax.set(xscale="log",yscale="log",title=title,xlabel=xlabel,ylabel="Maximum point error (m)")
    ax = axes[1,1]
    for backend in BACKENDS:
        for variant,style in (("coarse","-"),("refined","--")):
            rows=[r for r in result["convergence"] if r["backend"]==backend and r["mesh"]==variant]
            ax.plot(config["qualities"],[max(r["image_normalized_l1"] for r in rows if r["quality"]==n)*100 for n in config["qualities"]],style,color=colors[backend],marker=markers[backend],label=f"{backend} {variant}")
    ax.axhline(config["thresholds"]["image_normalized_l1"]*100,color="k",ls=":")
    ax.set(xscale="log",yscale="log",title="All seven cases",xlabel="Samples per pixel",ylabel="Maximum normalized image L1 (%)")
    for ax in axes.flat:
        ax.grid(alpha=.2);ax.legend(frameon=False,fontsize=7)
    fig.savefig(output/"results.svg");fig.savefig(output/"results.png",dpi=140);plt.close(fig)
    fig, axes = plt.subplots(4,4,figsize=(11,6),layout="constrained")
    for index,name in enumerate(("nadir","range_10km","triangle_10km","grazing")):
        ref=result["references"][name]
        images=[np.tile(ref["profile"],(config["image_size"][1],1))]
        images += [next(r["image"] for r in result["records"] if r["case"]==name and r["backend"]==backend and r["mesh"]=="coarse") for backend in BACKENDS]
        for col,im in enumerate(images):
            ax=axes[index,col];artist=ax.imshow(np.asarray(im)/ref["radiance"],cmap="viridis",vmin=0,vmax=1,aspect="auto",extent=(-2,2,1,-1),interpolation="nearest")
            if index==0:ax.set_title(["Analytic",*BACKENDS][col])
            if col==0:ax.set_ylabel(name.replace("_"," "))
            ax.set_xticks([-2,0,2]);ax.set_yticks([])
    fig.colorbar(artist,ax=axes,label="Radiance / sunlit plane radiance",shrink=.8)
    fig.savefig(output/"images.svg");fig.savefig(output/"images.png",dpi=140);plt.close(fig)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output",type=Path,default=Path("build/geometry-scale"))
    parser.add_argument("--plot-only",type=Path)
    args = parser.parse_args()
    args.output.mkdir(parents=True,exist_ok=True)
    if args.plot_only:
        plots(json.loads(args.plot_only.read_text()),args.output)
        return
    config = json.loads((DATA/"input.json").read_text())
    result = run(config)
    (args.output/"results.json").write_text(json.dumps(result,indent=2,allow_nan=False))
    raise SystemExit(0 if result["passed"] else 1)


if __name__=="__main__":
    main()