"""Eras: the base world plus local events. build_era applies an era's events
to the built base, re-runs the stages they affect inside an influence mask (changed cells + mask_km) and keeps every
cell outside it exactly as the base. Results: out/r<res>/eras/<era>/ (as out/r<res>/), previews in
previews/r<res>/eras/<era>/. An era's cache key is the base world's inputs key plus its events (config/eras.toml). Each era's events happen at the end of the era before it; its `years` erode what the events so far
changed; zone events follow the land at the start of their own era."""
from __future__ import annotations

import hashlib
import importlib
import json
import re
import shutil
from pathlib import Path

import numpy as np

from . import config as C, environment as EN, events as EV, fields as FL, pipeline as P
from .config import params
from .erosion import DEFAULTS as EROSION_DEFAULTS, K_MULT, erode
from .graph import distance_to, ocean_mask

DEFAULTS = {"mask_km": 1500.0, "climate_blend_km": 300.0, "erosion_steps": 1, "years": 10000.0}
RERUN = ("climate", "hydrology", "seabed", "environment", "ice", "fields")


def era_dir(root: Path, res: int, name: str) -> Path:
    return Path(root) / "out" / f"r{res}" / "eras" / name


def steps(tect: dict, name: str) -> list:
    """[(era, its own events, years)] for every era up to and including `name` ([eras] order): an era's events happen
    at the end of the era before it; `years` is the time from them to the era's map (None: the default)."""
    C.era_events(tect, name)                                  # an unknown era: ConfigError
    eras = tect.get("eras") or {}
    order = eras["order"]
    by_name = {e["name"]: e for e in tect.get("event", [])}
    return [(n, [by_name[e] for e in eras.get(n, {}).get("events", [])], eras.get(n, {}).get("years"))
            for n in order[: order.index(name) + 1]]


def fingerprint(base_key: str, steps_: list) -> str:
    content = [[own, yrs] for _, own, yrs in steps_]
    return hashlib.sha256((base_key + json.dumps(content, sort_keys=True)).encode()).hexdigest()[:16]


def _same(a, b) -> bool:
    """Equal arrays; NaN equals NaN in float arrays (a field may hold NaN where it has no value)."""
    a, b = np.asarray(a), np.asarray(b)
    return np.array_equal(a, b, equal_nan=a.dtype.kind in "fc" and b.dtype.kind in "fc")


def merge_lake_ids(base, new, mask):
    """Lake ids of an era: the base's outside the mask; inside it a lake the era left as it was keeps its base id and
    every other lake gets a fresh id after the base's, so one id never names two lakes."""
    base, new, mask = np.asarray(base), np.asarray(new), np.asarray(mask, bool)
    out = base.copy()
    nxt = max(int(base.max()) + 1, 0) if len(base) else 0
    counts = np.bincount(base[base >= 0]) if (base >= 0).any() else np.zeros(0, np.int64)
    ks = np.unique(new[mask & (new >= 0)])
    idx = np.flatnonzero(np.isin(new, ks))
    idx = idx[np.argsort(new[idx], kind="stable")]
    starts = np.searchsorted(new[idx], ks)
    ends = np.append(starts[1:], len(idx))
    for s, e in zip(starts.tolist(), ends.tolist()):
        cells = idx[s:e]
        b = base[cells]
        kept = b[0] >= 0 and bool(np.all(b == b[0])) and counts[b[0]] == len(cells)
        into = cells[mask[cells]]
        if kept:
            out[into] = b[0]
        else:
            out[into] = nxt
            nxt += 1
    out[mask & (new < 0)] = -1
    return out.astype(base.dtype)


def zone_key(name: str) -> str:
    return "zone_" + re.sub(r"[^A-Za-z0-9]", "_", name)


def _heights(g, z, events, min_sea, log, name, uncut=lambda z: z):
    for ev in events:                                        # heights first, in era order
        if ev["kind"] == "disintegrate":
            z, z_ref = EV.disintegrate(g.xyz, z, ~ocean_mask(g, uncut(z), min_sea), ev, g.radius_km)
            if z_ref is None:
                log(f"era {name}: warning: {ev['name']} has no land in its rim ring (0.9–1.0 × radius): "
                    f"its bowl hangs from 0 m")
            else:
                log(f"era {name}: {ev['name']} (ground at its rim {z_ref:.0f} m)")
        elif ev["kind"] == "volcano":
            z = EV.volcano(g.xyz, z, ev, g.radius_km)
    return z


def build_era(root: Path, res: int, name: str, log=print, base=None, low_memory: bool = False) -> Path:
    root = Path(root)
    cfg, tect = C.load(root)
    events = C.era_events(tect, name)
    if not events:
        log(f"era {name}: the base world (out/r{res})")
        return root / "out" / f"r{res}"
    out = era_dir(root, res, name)
    base_key = P.inputs_key(root, res)
    plan = steps(tect, name)
    key = fingerprint(base_key, plan)
    label = tect["eras"][name].get("label", name)
    meta_path = out / "cells_meta.json"
    if meta_path.exists() and (out / "cells.npz").exists():
        meta = json.loads(meta_path.read_text())
        era = meta.get("era", {})
        if era.get("fingerprint") == key:
            if (era.get("label"), era.get("name")) != (label, name):   # a new label or name needs no rebuild
                era["label"], era["name"] = label, name
                meta_path.write_text(json.dumps(meta, indent=1))
            log(f"era {name}: cached")
            return out
    ctx = base if base is not None else P.build(root, res, stop="fields", log=lambda m: None, low_memory=low_memory)
    low_memory = low_memory or ctx.low_memory
    g, d = ctx.grid, ctx.data
    E = params(ctx.cfg, "eras", DEFAULTS)
    EP = {**params(ctx.cfg, "erosion", EROSION_DEFAULTS), "steps": E["erosion_steps"]}
    kmult = np.vectorize(K_MULT.get)(np.asarray(d["age_class"])).astype(np.float64)
    z0 = np.asarray(d["elevation_eroded_m"], dtype=np.float64)
    z, ocean = z0.copy(), np.asarray(d["ocean"])
    water = np.asarray(d.get("open_water", ocean))
    cut = np.asarray(d.get("lake_cut_m", np.zeros(g.n)), dtype=np.float64)
    uncut = lambda zz: np.where(zz == z0, zz + cut, zz)       # sea masks: the ground before lake beds were carved
    zones, lands = [], []
    for _, own, yrs in plan:                                 # era by era: its events at the end of the era before
        for ev in own:                                       # zones follow the land before their own era's events
            if ev["kind"] == "zone":
                zones.append(ev)
                lands.append(EV.landmass(g, ~water, ev["seed"], ev["name"]) if ev["shape"] == "landmass" else None)
        z = _heights(g, z, own, EP["min_sea_km2"], log, name, uncut)
        hit = z != z0
        if hit.any():                                        # the changed ground weathers for the era's years
            years = float(E["years"] if yrs is None else yrs)
            ze = erode(g, z, kmult, {**EP, "dt_myr": years / 1.0e6 / E["erosion_steps"]})
            z = np.where(hit, np.maximum(np.minimum(z, ze), -11000.0), z)
            ocean, water = ocean_mask(g, uncut(z), np.inf), ocean_mask(g, uncut(z), EP["min_sea_km2"])
    hit = z != z0
    weights = [EV.zone_weight(g.xyz, ev, g.radius_km, ctx.seed, None if land is None else g.xyz[land])
               for ev, land in zip(zones, lands)]
    changed = hit | (ocean != np.asarray(d["ocean"])) | (water != np.asarray(d.get("open_water", d["ocean"])))
    for w in weights:
        changed |= w > 0
    mask = distance_to(g, changed) <= E["mask_km"] if changed.any() else np.zeros(g.n, bool)

    ec = P.Ctx(ctx.root, ctx.cfg, ctx.tect, ctx.res, dict(d), g, low_memory=low_memory)
    ec.data.update({"elevation_eroded_m": z.astype(np.asarray(d["elevation_eroded_m"]).dtype), "ocean": ocean,
                    "open_water": water, "lake_carve": hit})   # lake beds: geology, carved where the ground moved
    new_keys = {}
    for s in RERUN:
        o = importlib.import_module(f"mapgen.{s}").run(ec)
        ec.data.update(o)
        new_keys[s] = list(o)
    blend = np.clip(distance_to(g, ~mask) / E["climate_blend_km"], 0.0, 1.0) if (~mask).any() else np.ones(g.n)
    shape = lambda a, v: v.reshape(-1, *([1] * (a.ndim - 1)))
    merged = dict(d)
    for s, keys in new_keys.items():
        for k in keys:
            nv, bv = np.asarray(ec.data[k]), np.asarray(d[k])
            if s == "climate" and nv.dtype.kind == "f":      # solved on the whole world, blended in over 300 km
                nv = (bv + shape(nv, blend) * (nv - bv)).astype(nv.dtype)
            merged[k] = np.where(shape(nv, mask), nv, bv)
    for k in ("deposits", "deposit_main", "iron_potential"):  # ore is the base's geology: events move ground, the
        if k in d:                                           # rank-by-share placement would shift it world-wide
            merged[k] = d[k]
    if "lake_id" in merged:                                  # one id never names two lakes
        merged["lake_id"] = merge_lake_ids(d["lake_id"], ec.data["lake_id"], mask)
    merged["elevation_eroded_m"] = np.where(mask, ec.data["elevation_eroded_m"], d["elevation_eroded_m"])   # events and
    merged["ocean"] = ocean                                                                  # lake beds: inside only
    if "open_water" in d:
        merged["open_water"] = water
    H = float(ctx.cfg["planet"]["scale_height_km"]) * 1000.0
    for ev, w, land in zip(zones, weights, lands):          # zone events write their fields last
        merged.update(EV.apply_zone(merged, w, ev, merged["z_surface_m"], H))
        merged[zone_key(ev["name"])] = w.astype(np.float32)
        if land is not None:
            merged[zone_key(ev["name"]) + "_land"] = land
    if zones:                                                # plants settle into the zone's air and gravity
        pl = ctx.cfg["planet"]
        merged["plant_height_x"] = FL.plant_height(pl["gravity_g"], merged["gravity_g"]).astype(
            np.asarray(d["plant_height_x"]).dtype)
        sea_p = np.asarray(merged["pressure_bar"], dtype=np.float64) / FL.pressure(1.0, merged["z_surface_m"], H)
        wet = np.sqrt(sea_p / pl["sea_level_pressure_bar"])  # more CO₂ per breath: less water lost per growth
        hz, hr = EN.holdridge(merged["biotemp"], np.asarray(merged["P_ann"], dtype=np.float64) * wet, merged["T_min"])
        inside = np.logical_or.reduce([w > 0 for w in weights])
        merged["holdridge"] = np.where(inside, hz, merged["holdridge"]).astype(np.asarray(d["holdridge"]).dtype)
        merged["hold_region"] = np.where(inside, hr, merged["hold_region"]).astype(np.asarray(d["hold_region"]).dtype)
    merged["era_mask"] = mask
    for k, a in d.items():                                   # guard: nothing outside the mask may differ
        a, b = np.asarray(a), np.asarray(merged[k])
        if a.shape[:1] == (g.n,) and not _same(a[~mask], b[~mask]):
            raise RuntimeError(f"era {name}: {k} changed outside the influence mask")

    rc = P.Ctx(ctx.root, ctx.cfg, ctx.tect, ctx.res, merged, g, out_dir=out,
               preview_dir=root / "previews" / f"r{res}" / "eras" / name, low_memory=low_memory)
    importlib.import_module("mapgen.render").run(rc)
    meta = json.loads(meta_path.read_text())
    meta["era"] = {"name": name, "label": label, "events": events, "fingerprint": key, "base_key": base_key,
                   "steps": [[n, [e["name"] for e in own], yrs] for n, own, yrs in plan],
                   "mask_km": E["mask_km"], "mask_cells": int(mask.sum())}
    meta_path.write_text(json.dumps(meta, indent=1))
    log(f"era {name}: {int(mask.sum())} of {g.n} cells inside the influence mask")
    return out


def build_all(root: Path, res: int, log=print, base=None, low_memory: bool = False) -> list:
    """Every configured era with events (the base era is the base build); era folders no longer configured go."""
    _, tect = C.load(Path(root))
    names = [n for n in (tect.get("eras") or {}).get("order", []) if C.era_events(tect, n)]
    built = [build_era(root, res, n, log, base, low_memory) for n in names]
    top = Path(root) / "out" / f"r{res}" / "eras"
    for p in (top.iterdir() if top.is_dir() else []):
        if p.is_dir() and p.name not in names:
            shutil.rmtree(p, ignore_errors=True)
    return built