"""
Lagrangian particle trajectory simulator – compact block output.

10 000 particles seeded along the Catalan coast, advected by a synthetic
Western Mediterranean current field for 72 h (3 days).
Output: 3 JSON blocks of 24 h each + a metadata file.
Each block stores per-step arrays of [lon, lat, u, v, beached] per particle.
Trails are reconstructed client-side for zero data duplication.
"""

import json
import math
import os
import numpy as np

DATA_DIR = os.path.join(os.path.dirname(__file__), "data")

NUM_STEPS = 72
DT_SECONDS = 3600
DEG_PER_M_LAT = 1.0 / 111_320
BLOCK_SIZE = 24  # steps per block
PARTICLES_PER_SEED = 333

JITTER_DEG = 0.008  # small spread so particles start in water near coast
OFFSHORE_DEG = 0.012  # push seaward if initial position is on land

SEED_POINTS = [
    (3.21, 41.39), (2.82, 41.62), (2.17, 41.45), (1.82, 41.22),
    (1.25, 41.12), (1.00, 40.95), (0.63, 40.62), (0.85, 40.72),
    (3.12, 41.98), (3.19, 42.30), (2.82, 41.97), (2.07, 41.38),
    (3.10, 41.77), (3.17, 41.70), (1.52, 41.08), (3.04, 41.87),
    (0.75, 40.78), (2.93, 41.69), (1.60, 41.15), (2.40, 41.53),
    (0.50, 40.55), (2.60, 41.57), (3.16, 42.12), (2.50, 41.55),
    (1.40, 41.05), (3.05, 42.05), (2.92, 41.83), (1.15, 40.90),
    (2.70, 41.64), (3.22, 42.20),
]

SEED_NAMES = [
    "Barcelona", "Badalona", "Castelldefels", "Vilanova",
    "Tarragona", "Cambrils", "Delta Ebre S", "Delta Ebre N",
    "Roses", "Portbou", "L'Estartit", "Sitges",
    "Blanes", "Lloret de Mar", "Salou", "Tossa de Mar",
    "L'Ampolla", "St. Feliu Guíxols", "Torredembarra", "El Masnou",
    "Alcanar", "Premià de Mar", "Cadaqués", "Vilassar de Mar",
    "Altafulla", "L'Escala", "Palamós", "Miami Platja",
    "Arenys de Mar", "Llançà",
]

# ── Land polygons ────────────────────────────────────────────────────────────

IBERIAN_PENINSULA = [
    (-2.0, 36.7), (-1.5, 36.8), (-0.5, 37.4), (0.0, 38.0),
    (0.2, 38.8), (-0.2, 39.5), (0.0, 39.9), (0.2, 40.4),
    (0.45, 40.5), (0.5, 40.65), (0.7, 40.75), (0.8, 40.85),
    (0.9, 40.9), (1.0, 41.0), (1.2, 41.1), (1.5, 41.1),
    (1.8, 41.2), (2.0, 41.3), (2.1, 41.35), (2.2, 41.4),
    (2.5, 41.5), (2.8, 41.6), (2.9, 41.65), (3.0, 41.7),
    (3.1, 41.8), (3.15, 41.9), (3.15, 42.0), (3.2, 42.1),
    (3.2, 42.3), (3.25, 42.45), (3.05, 42.5), (2.5, 42.6),
    (1.8, 42.7), (1.0, 42.75), (0.0, 42.7), (-1.0, 42.8),
    (-2.0, 42.5), (-2.0, 36.7),
]
FRANCE_SOUTH = [
    (3.05, 42.5), (3.25, 42.45), (3.3, 42.5), (3.5, 42.6),
    (3.6, 43.0), (3.5, 43.2), (3.9, 43.3), (4.1, 43.4),
    (4.5, 43.3), (4.8, 43.4), (5.0, 43.3), (5.5, 43.2),
    (5.8, 43.1), (6.0, 43.1), (6.3, 43.15), (6.6, 43.0),
    (7.0, 43.5), (7.5, 43.8), (7.5, 44.5), (6.0, 44.5),
    (4.0, 44.0), (3.0, 43.5), (3.0, 42.8), (3.05, 42.5),
]
ITALY_WEST = [
    (7.5, 43.8), (7.7, 43.8), (8.0, 43.9), (8.5, 44.0),
    (9.0, 44.1), (9.5, 44.3), (10.0, 44.0), (10.0, 43.5),
    (10.5, 43.0), (10.5, 42.5), (11.0, 42.0), (11.5, 41.5),
    (12.0, 41.2), (12.5, 41.0), (13.0, 40.5), (13.5, 40.0),
    (14.0, 39.5), (14.0, 38.5), (14.0, 44.5), (7.5, 44.5),
    (7.5, 43.8),
]
SARDINIA = [
    (8.1, 39.0), (8.3, 38.8), (8.5, 38.85), (8.8, 38.9),
    (9.2, 39.0), (9.6, 39.2), (9.7, 39.8), (9.6, 40.2),
    (9.5, 40.6), (9.2, 40.9), (8.8, 41.0), (8.3, 41.0),
    (8.1, 40.8), (8.1, 40.4), (8.2, 39.8), (8.1, 39.0),
]
CORSICA = [
    (8.6, 41.4), (8.8, 41.4), (9.2, 41.6), (9.4, 42.0),
    (9.5, 42.4), (9.4, 42.8), (9.2, 43.0), (8.8, 42.7),
    (8.6, 42.3), (8.6, 41.8), (8.6, 41.4),
]
MALLORCA = [
    (2.3, 39.25), (2.7, 39.3), (3.1, 39.5), (3.4, 39.7),
    (3.5, 39.85), (3.4, 39.95), (3.1, 39.95), (2.7, 39.9),
    (2.4, 39.75), (2.3, 39.55), (2.3, 39.25),
]
MENORCA = [
    (3.8, 39.8), (3.9, 39.82), (4.1, 39.95), (4.25, 40.05),
    (4.1, 40.1), (3.85, 40.0), (3.8, 39.9), (3.8, 39.8),
]
IBIZA = [
    (1.2, 38.8), (1.3, 38.82), (1.45, 38.9), (1.55, 39.0),
    (1.5, 39.1), (1.35, 39.1), (1.2, 39.0), (1.15, 38.9),
    (1.2, 38.8),
]
NORTH_AFRICA = [
    (-2.0, 36.7), (-2.0, 35.0), (10.0, 35.0), (10.0, 37.0),
    (9.5, 37.2), (9.0, 37.0), (8.5, 37.0), (8.0, 36.8),
    (7.5, 36.9), (7.0, 37.0), (6.0, 37.0), (5.0, 36.8),
    (4.0, 36.7), (3.0, 36.7), (2.0, 36.5), (1.0, 36.3),
    (0.0, 35.8), (-1.0, 35.5), (-2.0, 36.7),
]

LAND_POLYGONS = [
    IBERIAN_PENINSULA, FRANCE_SOUTH, ITALY_WEST,
    SARDINIA, CORSICA, MALLORCA, MENORCA, IBIZA, NORTH_AFRICA,
]


def _pip(px, py, poly):
    n = len(poly)
    inside = False
    j = n - 1
    for i in range(n):
        xi, yi = poly[i]
        xj, yj = poly[j]
        if ((yi > py) != (yj > py)) and (px < (xj - xi) * (py - yi) / (yj - yi) + xi):
            inside = not inside
        j = i
    return inside


def is_land(lon, lat):
    for poly in LAND_POLYGONS:
        if _pip(lon, lat, poly):
            return True
    return False


# ── Current field (NW Mediterranean / Catalan coast, realistic) ───────────────

# Typical surface currents: Northern Current ~0.1–0.25 m/s along coast (SW), weaker offshore
BASE_SPEED_MS = 0.18  # m/s, main flow magnitude
DIFFUSION_MS = 0.025  # turbulent fluctuation (m/s)

def current_field(lon, lat, rng, t_frac):
    # Main flow: SW along Catalan coast (negative v = south, slight negative u = west)
    angle_deg = 215 + 8 * math.sin(2 * math.pi * t_frac)  # slow seasonal wobble
    rad = math.radians(angle_deg)
    u_base = BASE_SPEED_MS * math.cos(rad)
    v_base = BASE_SPEED_MS * math.sin(rad)

    # Coastal weakening: reduce speed when close to Iberian coast (lon < 2.5, lat 40–42.5)
    dist_coast = 1.0
    if lon < 3.2 and 40.4 < lat < 42.5:
        dist_coast = max(0.15, (lon - 0.5) * 0.4 + (42.2 - lat) * 0.15)
    factor = min(1.0, dist_coast)
    u_main = u_base * factor
    v_main = v_base * factor

    # Weak cyclonic gyre east of Barcelona (3.5, 41.5), radius ~0.8 deg
    cx, cy = 3.6, 41.4
    dx, dy = lon - cx, lat - cy
    r_deg = math.sqrt(dx * dx + dy * dy)
    if r_deg < 0.9 and r_deg > 0.05:
        gyre_speed = 0.06 * math.exp(-r_deg / 0.5)  # m/s
        # tangential: (-dy, dx) for anticlockwise
        u_gyre = -gyre_speed * (dy / r_deg)
        v_gyre = gyre_speed * (dx / r_deg)
    else:
        u_gyre, v_gyre = 0.0, 0.0

    # Small-scale turbulence (Gaussian, uncorrelated)
    u_noise = rng.normal(0, DIFFUSION_MS)
    v_noise = rng.normal(0, DIFFUSION_MS)

    return (u_main + u_gyre + u_noise, v_main + v_gyre + v_noise)


def advect_rk4(lon, lat, rng, t_frac):
    cos_lat = math.cos(math.radians(lat))
    dpm_lon = DEG_PER_M_LAT / max(cos_lat, 1e-6)

    def vel(lo, la):
        u, v = current_field(lo, la, rng, t_frac)
        return u * DT_SECONDS * dpm_lon, v * DT_SECONDS * DEG_PER_M_LAT

    k1 = vel(lon, lat)
    k2 = vel(lon + k1[0] / 2, lat + k1[1] / 2)
    k3 = vel(lon + k2[0] / 2, lat + k2[1] / 2)
    k4 = vel(lon + k3[0], lat + k3[1])

    new_lon = lon + (k1[0] + 2 * k2[0] + 2 * k3[0] + k4[0]) / 6
    new_lat = lat + (k1[1] + 2 * k2[1] + 2 * k3[1] + k4[1]) / 6
    u_f, v_f = current_field(new_lon, new_lat, rng, t_frac)
    return new_lon, new_lat, u_f, v_f


# ── Accumulation grid ────────────────────────────────────────────────────────

ACCUM_RES = 0.1

def build_accumulation(b_lons, b_lats):
    if len(b_lons) == 0:
        return []
    cells: dict[tuple[int, int], int] = {}
    for lo, la in zip(b_lons, b_lats):
        key = (round(lo / ACCUM_RES), round(la / ACCUM_RES))
        cells[key] = cells.get(key, 0) + 1
    return [[k[0] * ACCUM_RES, k[1] * ACCUM_RES, v] for k, v in cells.items()]


# ── Main ─────────────────────────────────────────────────────────────────────

def generate():
    os.makedirs(DATA_DIR, exist_ok=True)
    rng = np.random.default_rng(42)

    init_lons, init_lats, origins = [], [], []
    for si, (slon, slat) in enumerate(SEED_POINTS):
        for _ in range(PARTICLES_PER_SEED):
            lon = slon + rng.uniform(-JITTER_DEG, JITTER_DEG)
            lat = slat + rng.uniform(-JITTER_DEG, JITTER_DEG)
            init_lons.append(lon)
            init_lats.append(lat)
            origins.append(si)

    n = len(init_lons)
    lons = np.array(init_lons, dtype=float)
    lats = np.array(init_lats, dtype=float)

    # Ensure every particle starts in water: push seaward (east/northeast) if on land
    for i in range(n):
        for _ in range(50):
            if not is_land(lons[i], lats[i]):
                break
            lons[i] += rng.uniform(0, OFFSHORE_DEG)
            lats[i] += rng.uniform(0, OFFSHORE_DEG * 0.5)
    init_lons = lons.tolist()
    init_lats = lats.tolist()

    # Staggered release over first 24h for more realistic dispersion
    release_steps = rng.integers(0, min(24, NUM_STEPS), size=n).tolist()
    init_lons_arr = np.array(init_lons, dtype=float)
    init_lats_arr = np.array(init_lats, dtype=float)

    beached = np.zeros(n, dtype=bool)

    # Collect ALL frames: frames[step] = [[lon, lat, u, v, beached], ...]
    all_frames = []
    all_accum = []

    for step in range(NUM_STEPS):
        t_frac = step / NUM_STEPS
        frame = np.zeros((n, 5))

        for i in range(n):
            if step < release_steps[i]:
                frame[i] = [init_lons_arr[i], init_lats_arr[i], 0.0, 0.0, 0]
                continue
            if step == release_steps[i]:
                lons[i], lats[i] = init_lons_arr[i], init_lats_arr[i]

            if not beached[i]:
                nl, nla, u, v = advect_rk4(float(lons[i]), float(lats[i]), rng, t_frac)
                if is_land(nl, nla):
                    beached[i] = True
                    u, v = 0.0, 0.0
                else:
                    lons[i], lats[i] = nl, nla
                    frame[i] = [lons[i], lats[i], u, v, 0]
                    continue

            frame[i] = [lons[i], lats[i], 0.0, 0.0, 1 if beached[i] else 0]

        all_frames.append(np.round(frame, 5).tolist())

        b_mask = beached
        accum = build_accumulation(lons[b_mask].tolist(), lats[b_mask].tolist())
        all_accum.append(accum)

        act = int(np.sum(~beached))
        bch = int(np.sum(beached))
        print(f"  Step {step:>2d}/{NUM_STEPS - 1}  |  active {act:>5d}  beached {bch:>5d}")

    # Write blocks
    num_blocks = NUM_STEPS // BLOCK_SIZE
    for b in range(num_blocks):
        s0 = b * BLOCK_SIZE
        s1 = s0 + BLOCK_SIZE
        block = {
            "block": b,
            "step_start": s0,
            "step_end": s1,
            "frames": all_frames[s0:s1],
            "accumulation": all_accum[s0:s1],
        }
        path = os.path.join(DATA_DIR, f"block_{b}.json")
        with open(path, "w") as f:
            json.dump(block, f, separators=(",", ":"))
        size_mb = os.path.getsize(path) / 1_048_576
        print(f"  Block {b} ({s0}h-{s1 - 1}h) → {path}  [{size_mb:.1f} MB]")

    # Metadata (release_steps: step at which each particle is released, 0..num_steps-1)
    meta = {
        "num_particles": n,
        "num_steps": NUM_STEPS,
        "num_blocks": num_blocks,
        "block_size": BLOCK_SIZE,
        "dt_hours": 1,
        "seed_names": SEED_NAMES,
        "origins": origins,
        "release_steps": release_steps,
    }
    meta_path = os.path.join(DATA_DIR, "metadata.json")
    with open(meta_path, "w") as f:
        json.dump(meta, f, separators=(",", ":"))
    print(f"  Metadata → {meta_path}")


if __name__ == "__main__":
    n = len(SEED_POINTS) * PARTICLES_PER_SEED
    print(f"Simulating {n} particles for {NUM_STEPS} hours …")
    generate()
    print("Done.")
