"""
Mars (GJ0e) terrain builder.

Data source:
  - Elevation: MOLA MEGDR (Mars Orbiter Laser Altimeter)
    PDS format, big-endian int16, metres relative to areoid.
    Available at multiple resolutions. We use 4ppd (1440×720)
    or 16ppd (5760×2880) — both small enough to download quickly.

Mars properties:
  - Min elevation: ~-8200 m (Hellas Basin)
  - Max elevation: ~21229 m (Olympus Mons)
  - Polar ice caps: CO2 + water ice
  - Thin atmosphere (6 mbar) — classified as "thin" in body_def
  - Almost no liquid water (hydrosphere: "ice")
"""

import numpy as np

from sol_data.download import ensure_cached
from sol_data.shared import (
    GRID_W, GRID_H,
    load_raw_binary, resample_to_grid, normalize_01,
    compute_hillshade, assemble_terrain,
    temperature_grid_analytical,
)

# MOLA MEGDR — 4 pixels per degree (1440 × 720), big-endian int16
# Each pixel = metres relative to Mars areoid
# PDS binary with no header (data starts at byte 0 for .img files)
MOLA_4PPD_URL = "https://pds-geosciences.wustl.edu/mgs/mgs-m-mola-5-megdr-l3-v1/mgsl_300x/meg004/megt90n000cb.img"
MOLA_4PPD_FILE = "mola_megdr_4ppd.img"
MOLA_4PPD_W = 1440
MOLA_4PPD_H = 720

# Alternative: 16ppd (5760 × 2880) for higher quality
MOLA_16PPD_URL = "https://pds-geosciences.wustl.edu/mgs/mgs-m-mola-5-megdr-l3-v1/mgsl_300x/meg016/megt90n000eb.img"
MOLA_16PPD_FILE = "mola_megdr_16ppd.img"
MOLA_16PPD_W = 5760
MOLA_16PPD_H = 2880

# Mars physical constants
MARS_MIN_ELEV_M = -8200.0   # Hellas Basin
MARS_MAX_ELEV_M = 21229.0   # Olympus Mons summit
# Real Mars temperatures — we don't fudge these. Mars colour comes from
# ferric biome classes (34/35/36) applied based on iron oxide substrate.
MARS_EQUATORIAL_TEMP_K = 215.0  # daytime average near equator
MARS_POLAR_TEMP_K = 150.0
MARS_OCEAN_FRACTION = 0.0    # no liquid water (ice only)

# Ferric biome class IDs (from biomes.toml)
FERRIC_DUST = 34
FERRIC_HIGHLAND = 35
FERRIC_LOWLAND = 36


def _load_mola(use_16ppd: bool = False) -> np.ndarray:
    """Load MOLA DEM, return elevation in metres."""
    if use_16ppd:
        url, filename, w, h = MOLA_16PPD_URL, MOLA_16PPD_FILE, MOLA_16PPD_W, MOLA_16PPD_H
    else:
        url, filename, w, h = MOLA_4PPD_URL, MOLA_4PPD_FILE, MOLA_4PPD_W, MOLA_4PPD_H

    path = ensure_cached(url, filename)
    print(f"    loading MOLA: {path} ({w}x{h})")

    # MOLA MEGDR: big-endian int16, metres, no header
    arr = load_raw_binary(str(path), w, h, dtype=">i2", offset=0)

    # MOLA nodata is typically 32767 or -32768
    arr[arr > 30000] = 0.0
    arr[arr < -30000] = 0.0

    print(f"    MOLA range: [{arr.min():.0f}, {arr.max():.0f}] m")
    return arr


def build_terrain(body_def: dict) -> dict:
    """Build Mars terrain dict from MOLA data."""
    print("  Mars: loading MOLA data...")

    # ── 1. Elevation ────────────────────────────────────────────────────
    mola_raw = _load_mola(use_16ppd=False)

    # MOLA is col 0 = 0° longitude — shift to col 0 = 180°W
    from sol_data.shared import greenwich_to_dateline
    mola_shifted = greenwich_to_dateline(mola_raw)

    # Resample to grid
    elevation_m = resample_to_grid(mola_shifted, GRID_H, GRID_W, order=1)

    # Normalise to [0, 1]
    elevation = normalize_01(elevation_m, MARS_MIN_ELEV_M, MARS_MAX_ELEV_M)

    print("    elevation normalised")

    # ── 2. Temperature ──────────────────────────────────────────────────
    # Analytical: equatorial ~210K, polar ~150K, elevation lapse
    temperature_K = temperature_grid_analytical(
        base_T_K=MARS_EQUATORIAL_TEMP_K,
        elevation=elevation,
        lapse_rate_K_per_unit=30.0,
        lat_gradient_K=60.0,
    )

    # Polar ice caps: very cold at high latitudes
    v = np.linspace(0, 1, GRID_H, dtype=np.float32)
    lat_abs = np.abs(v - 0.5) * 2.0
    polar_rows = lat_abs > 0.75
    temperature_K[polar_rows, :] = np.minimum(temperature_K[polar_rows, :], 155.0)

    print(f"    temperature: [{temperature_K.min():.0f}, {temperature_K.max():.0f}] K")

    # ── 3. Moisture ─────────────────────────────────────────────────────
    # Mars has almost no moisture — thin atmosphere
    moisture = np.zeros((GRID_H, GRID_W), dtype=np.float32)
    # Slight moisture near polar caps (water ice)
    moisture[polar_rows, :] = 0.1

    # ── 4. Terraformed water bodies ─────────────────────────────────────
    # Lore: 800 years of partial terraforming. Water pools in the deepest
    # basins (Hellas, Utopia, Isidis). ~2% of surface is now liquid water.
    from sol_data.shared import compute_sea_level as _compute_sl
    from scipy.ndimage import binary_dilation

    TERRAFORM_OCEAN_FRAC = 0.02  # 2% water coverage
    sea_level = _compute_sl(elevation, TERRAFORM_OCEAN_FRAC)
    surface_water = elevation < sea_level

    # Don't flood polar regions — those stay as ice caps, not lakes
    surface_water[polar_rows, :] = False

    n_water = int(surface_water.sum())
    print(f"    terraformed water: {n_water} cells "
          f"(sea_level={sea_level:.4f})")

    # ── 5. Biome classification ─────────────────────────────────────────
    # Mars biome is built directly — compute_biome() would classify
    # everything as ice at these temperatures.
    biome = np.full((GRID_H, GRID_W), FERRIC_DUST, dtype=np.int8)

    # Elevation-based ferric variation
    biome[elevation > 0.55] = FERRIC_HIGHLAND  # volcanic highlands
    biome[elevation < 0.25] = FERRIC_LOWLAND   # basin floors

    # Polar ice caps
    biome[polar_rows, :] = 17  # ice/snow

    # Terraformed green fringe around water bodies — vegetation band
    # where the thicker local atmosphere and water access allow plants.
    # ~5 cell band around each water body.
    veg_ring = binary_dilation(surface_water, iterations=5) & ~surface_water
    # Don't put vegetation at poles
    veg_ring[polar_rows, :] = False
    biome[veg_ring] = 12  # shrubland (olive green — sparse terraformed vegetation)

    # Inner vegetation ring (closer to water = lusher)
    inner_ring = binary_dilation(surface_water, iterations=2) & ~surface_water
    inner_ring[polar_rows, :] = False
    biome[inner_ring] = 8  # temperate grassland (greener)

    # Ocean depth bands for water bodies
    if surface_water.any():
        depth = np.clip((sea_level - elevation) / (sea_level + 1e-9), 0, 1)
        biome[surface_water & (depth < 0.15)] = 2   # shallow
        biome[surface_water & (depth >= 0.15) & (depth < 0.50)] = 1  # mid
        biome[surface_water & (depth >= 0.50)] = 0   # deep

    n_ice = int((biome == 17).sum())
    n_ferric = int(((biome >= 34) & (biome <= 36)).sum())
    n_veg = int(((biome == 8) | (biome == 12)).sum())
    n_ocean = int(((biome >= 0) & (biome <= 2)).sum())
    print(f"    biomes: {len(np.unique(biome))} classes "
          f"(ferric={n_ferric}, ice={n_ice}, veg={n_veg}, water={n_ocean})")

    # ── 5. Hillshade ────────────────────────────────────────────────────
    hillshade = compute_hillshade(elevation)

    # ── 6. Assemble ─────────────────────────────────────────────────────
    return assemble_terrain(
        elevation=elevation,
        temperature_K=temperature_K,
        moisture=moisture,
        biome=biome,
        surface_water=surface_water,
        hillshade=hillshade,
        rivers=[],  # no rivers on Mars
        sea_level=sea_level,
    )