"""Tutorial 11b — IrradianceMapSensor (up and down irradiance diagram).

Flat-Lambertian scene: downwelling / upwelling per pixel + analysis verification
(I·cos(SZA) and ρ·I·cos(SZA)), and demonstrate `resolution` parameters and three
Surface= Differences in normalization options."""
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import numpy as np

import less

SUN_ZENITH = 30.0
BANDS = [550, 850]
HORIZ = np.array([1.6203, 0.7190], dtype=np.float32)  # Total horizontal irradiance W/m²/nm
REFL = 0.15

# ── Scene ──────────────────────────────────────────────────
scene = less.Scene()                  # Automatic selection; optix/vulkan/embree can be specified explicitly
scene.size = 40.0
scene.terrain = less.Terrain(property=less.Lambertian(reflectance=REFL))
scene.illumination = less.Illumination(
    source=less.Sun(
        zenith=SUN_ZENITH, azimuth=180, wavelengths=BANDS,
        irradiance=HORIZ / np.cos(np.radians(SUN_ZENITH))),
    atmosphere=less.NoAtmosphere(),
)

# ── Sensor: Use resolution to specify meters/pixel ────────────────────────
sensor = less.IrradianceMapSensor(
    bands=BANDS,
    resolution=1.0,            # 1 m/pixel → grid automatically (40, 40)
    photons_per_m2=400,
    max_bounces=None,
)
p = scene.simulate(sensor)

# ── Original vs normalized ───────────────────────────────────────
print("Original downwelling (W/nm per pixel column):",
      p.downwelling.mean(axis=(0, 1)))

down_h = p.downwelling_irradiance('horizontal')      # W/m²/nm
up_h   = p.upwelling_irradiance  ('horizontal')

# TOA obtains HORIZ after beam normal irradiance projection; analytical expectation down = HORIZ, up = REFL*HORIZ
print(f"\n=== Flat Lambertian surface (rho={REFL}, SZA={SUN_ZENITH} deg) ===")
for i, wl in enumerate(BANDS):
    print(f"  {wl} nm  down={down_h[:,:,i].mean():.4e}  "
          f"up={up_h[:,:,i].mean():.4e}  "
          f"(predicted down={HORIZ[i]:.4e}, up={REFL * HORIZ[i]:.4e}) [W/m^2/nm]")

# ── On flat ground, slope / surface is equal to horizontal (cos β = 1)──────
# Flat terrain does not have a DEM set, so 'slope' / 'surface' will cause ValueError.
# Only horizontal is demonstrated here. The DEM scene is available in val16_terrain_albedo.py.

# ── Visualization ────────────────────────────────────────────
for i, wl in enumerate(BANDS):
    fig, ax = plt.subplots(1, 2, figsize=(10, 4))
    im0 = ax[0].imshow(down_h[:, :, i], origin='lower', cmap='viridis')
    ax[0].set_title(f'Downwelling {wl} nm')
    plt.colorbar(im0, ax=ax[0], label='W / m² / nm')
    im1 = ax[1].imshow(up_h[:, :, i], origin='lower', cmap='magma')
    ax[1].set_title(f'Upwelling {wl} nm')
    plt.colorbar(im1, ax=ax[1], label='W / m² / nm')
    fig.tight_layout()
    fig.savefig(f'11b_irradiance_{wl}nm.png', dpi=120)
    plt.close(fig)

print("\nThe output has been written to the current directory: 11b_irradiance_{550,850}nm.png")