"""Sub-species as birth events (spec 2026-09-30 gradual adaptation §2). A group's progress toward a related race is
read from its adapted optima; births (once per sub-species) and returns (any number) happen with a chance that
rises with progress; the born group takes the new race's native curves."""
import numpy as np

from .config import CONDITIONS
from .state import convert, move


def native(race):
    return np.array([race.tolerance[c].optimum if c in race.tolerance else 0.0 for c in CONDITIONS], float)


def progress(O, r_from, r_to):
    """0..1 per group: how far its optima have moved from r_from's native curves toward r_to's (spec §2.2)."""
    num, den = np.zeros(O.shape[1]), 0.0
    for ci, c in enumerate(CONDITIONS):
        a, b = r_from.tolerance.get(c), r_to.tolerance.get(c)
        if a is None or b is None:
            continue
        gap = b.optimum - a.optimum
        w = a.side(gap > 0)
        if abs(gap) <= 0.25 * w:
            continue
        weight = abs(gap) / w
        num += weight * np.clip((O[ci] - a.optimum) / gap, 0.0, 1.0)
        den += weight
    return num / den if den > 0 else num


def smoothstep(x):
    x = np.clip(x, 0.0, 1.0)
    return x * x * (3 - 2 * x)


def chance(a, lo, hi, rate):
    a = np.asarray(a, float)
    if hi <= lo:
        return np.where(a >= hi, rate, 0.0)
    return np.where(a < lo, 0.0, rate * smoothstep((a - lo) / (hi - lo)))


def init_rules(races):
    idx = {r.id: i for i, r in enumerate(races)}
    return [{"s": s, "p": idx[r.emerge["parent"]], "origin": None, "year": None, "peak": None, "peak_t": None,
             "victims": 0.0, "backlash": False, "rate": r.emerge["ritual_rate"]}
            for s, r in enumerate(races) if r.emerge is not None]


def _crisis(rule, spec, P, t):
    """Parent groups almost eradicated: P ≤ (1 − crisis_drop) of their recent peak (memory crisis_years, decaying),
    with a peak of at least crisis_min people. Updates the rule's peak record."""
    if rule["peak"] is None:
        rule["peak"] = P.copy()
    else:
        rule["peak"] = np.maximum(P, rule["peak"] * np.exp(-(t - rule["peak_t"]) / spec["crisis_years"]))
    rule["peak_t"] = t
    return (P > 0) & (P <= (1 - spec["crisis_drop"]) * rule["peak"]) & (rule["peak"] >= spec["crisis_min"])


def _where(spec, world, regions, F):
    m = regions(spec["region"]) if spec["region"] else np.ones(world.n, bool)
    cond = spec["condition"]
    if cond:
        x = np.asarray(F[cond["field"]], float)
        if "below" in cond:
            m = m & (x < cond["below"])
        if "above" in cond:
            m = m & (x > cond["above"])
    return m


def _shelter(world, where, ok, km):
    """Per cell: for crisis cells (where), the nearest cell within km the new race can live in (ok), else −1; km 0 → the
    crisis cell itself must fit."""
    dest = np.full(world.n, -1, np.int64)
    for c in np.flatnonzero(where):
        near = world.within(c, km)
        near = near[ok[near]]
        if len(near):
            dest[c] = near[np.argmin(world.km(np.full(len(near), c), near))]
    return dest


def _shifted(st, p, s, cells, races, shift):
    """Optima the born people arrive with: `shift` of the way from their own to the new race's native curves."""
    return (1 - shift) * st.O[p][:, cells] + shift * native(races[s])[:, None]


def birth_step(st, world, races, rules, regions, rng, ok, F=None):
    F = world.fields if F is None else F
    new = []
    for rule in rules:
        if rule["origin"] is not None:
            continue
        s, p = rule["s"], rule["p"]
        spec, P = races[s].emerge, st.P[p]
        crisis = _crisis(rule, spec, P, st.t) if spec["crisis_drop"] > 0 else True
        if st.t < spec["after"]:
            continue
        ritual = spec["mode"] == "ritual"
        where = (P > 0) & _where(spec, world, regions, F) & crisis
        if not ritual:
            where &= ok[s]
        if not where.any():
            continue
        a = progress(st.O[p], races[p], races[s])
        if ritual:
            dest = _shelter(world, where, ok[s], spec["crisis_km"])
            where &= dest >= 0
            if not where.any() or rng.random() >= spec["trigger"]:
                continue
            h = np.where(where, P, 0.0)
        else:
            h = np.where(where, chance(a, spec["birth_min"], spec["birth_sure"], spec["birth_rate"])
                         * P / (P + races[p].founder), 0.0)
            if spec["settlement_density"] > 0 and h.any():
                dens = np.where(h > 0, P / world.area, -1.0)
                best = int(np.argmax(dens))
                h = np.where(np.arange(world.n) == best, h, 0.0) * (dens[best] >= spec["settlement_density"])
            if not h.any() or rng.random() >= 1 - np.prod(1 - np.clip(h, 0, 1)):
                continue
        c = np.flatnonzero(h > 0)
        wgt = h[c]
        if spec["isolated"]:
            around = world.nb_sum(st.P.sum(0))[c]
            wgt = wgt * (P[c] / (P[c] + around)) ** 4
        cell = int(rng.choice(c, p=wgt / wgt.sum()))
        convert(st, p, s, [cell], [spec["convert"] * P[cell]], O_new=_shifted(st, p, s, [cell], races,
                                                                               spec["birth_shift"]))
        ev = {"race": races[s].id, "cell": cell, "year": int(st.t), "a": float(a[cell])}
        if ritual:                                  # survivors of the crisis flee to the nearest land they can live in
            ev["cell"], ev["crisis"] = int(dest[cell]), cell
            if ev["cell"] != cell:
                move(st, s, [cell], [ev["cell"]], [st.P[s, cell]])
        rule["origin"], rule["year"] = ev["cell"], int(st.t)
        new.append(ev)
    return new


def return_step(st, world, races, rules, rng, ok):
    """Sub-species groups adapted back toward the parent give birth to parent-race people (spec §2.6)."""
    out = []
    for rule in rules:
        s, p = rule["s"], rule["p"]
        spec = races[s].emerge
        if rule["origin"] is None or spec["mode"] == "ritual":
            continue
        P = st.P[s]
        a = progress(st.O[s], races[s], races[p])
        h = np.where((P > 0) & ok[p], chance(a, spec["return_min"], spec["return_sure"], spec["return_rate"])
                     * P / (P + races[s].founder), 0.0)
        cells = np.flatnonzero(rng.random(world.n) < h)
        if not len(cells):
            continue
        amt = spec["convert"] * P[cells]
        convert(st, s, p, cells, amt, O_new=_shifted(st, s, p, cells, races, spec["birth_shift"]))
        out += [{"race": races[p].id, "from": races[s].id, "cell": int(c), "year": int(st.t), "people": float(x)}
                for c, x in zip(cells, amt)]
    return out


def _born(rules):
    return [r for r in rules if r["origin"] is not None]


def displacement(K, P, fit, rules, races):
    """In cells a race shares with its parent / sub-species, the less fit one keeps (1 − displace·other's share)
    of its capacity (spec §2.4)."""
    for rule in _born(rules):
        s, p = rule["s"], rule["p"]
        d = races[s].emerge["displace"]
        tot = P[s] + P[p]
        share_s = np.where(tot > 0, P[s] / np.maximum(tot, 1e-12), 0.0)
        both = (P[s] > 0) & (P[p] > 0)
        K[p] *= np.where(both & (fit[p] < fit[s]), 1 - d * share_s, 1.0)
        K[s] *= np.where(both & (fit[s] < fit[p]), 1 - d * (1 - share_s), 1.0)


def inherit(st, world, races, rules, births, regions, ok):
    """Dominant trait: a share of the parent groups' new births in or next to the sub-species is born as it, gated
    by the parents' own adaptation (spec §2.5; ritual races: no gate, §2.7)."""
    for rule in _born(rules):
        s, p = rule["s"], rule["p"]
        spec = races[s].emerge
        Ps, Pp = st.P[s], st.P[p]
        c = np.flatnonzero((Pp > 0) & (births[p] > 0) & ok[s])          # only where parents are born and the new
        if not len(c):                                                   # curves fit: compute there
            continue
        sub, cols = world.nb_local(c)
        near_s = Ps[c] + world.nb_apply(sub, Ps[cols])
        near = near_s + Pp[c] + world.nb_apply(sub, Pp[cols])
        share = np.where(near > 0, near_s / np.maximum(near, 1e-12), 0.0)
        where = near_s > 0
        if spec["spread"] == "region" and spec["region"]:
            where &= regions(spec["region"])[c]
        if spec["mode"] == "ritual":
            gate = np.ones(len(c))
        else:
            a = progress(st.O[p][:, c], races[p], races[s])
            gate = smoothstep((a - spec["mix_min"]) / max(spec["birth_sure"] - spec["mix_min"], 1e-12))
        amt = np.where(where, spec["dominance"] * gate * share * births[p, c], 0.0)
        keep = amt > 0
        cells, amt = c[keep], amt[keep]
        if len(cells):
            convert(st, p, s, cells, amt, O_new=np.repeat(native(races[s])[:, None], len(cells), 1))


def raids(st, world, races, rules, rng, ok):
    """Before the smackdown, ritual races send out cells: groups of at least raid_min send ~Poisson(raid_rate) raids of
    raid_size people each to occupied non-ritual land within raid_km (weighted by its people) where the race can
    live; the rituals there do the killing. Each group keeps at least raid_min / 2."""
    out = []
    for rule in _born(rules):
        s = rule["s"]
        spec = races[s].emerge
        if spec["mode"] != "ritual" or spec["raid_rate"] <= 0 or rule["backlash"]:
            continue
        others = np.sum(np.delete(st.P, s, axis=0), axis=0)
        for c in np.flatnonzero(st.P[s] >= spec["raid_min"]):
            k = min(rng.poisson(spec["raid_rate"]), int((st.P[s, c] - spec["raid_min"] / 2) // spec["raid_size"]))
            if k <= 0:
                continue
            near = world.within(c, spec["raid_km"])
            near = near[(near != c) & (others[near] >= 1) & ok[s][near]]
            if not len(near):
                continue
            to = rng.choice(near, size=k, p=others[near] / others[near].sum())
            move(st, s, np.full(k, c), to, np.full(k, float(spec["raid_size"])))
            out.append({"race": races[s].id, "cell": int(c), "year": int(st.t), "raids": k})
    return out


def rituals(st, world, races, rules, rng):
    """Ritual races win converts at a blood price: rituals per cell ~ Poisson(rate·√people); each kills ritual_cost
    and converts ritual_converts, taken from the neighbour cell (or own) with most parents. victims = "all": the dead
    come from every nearby non-ritual people by presence. After backlash_victims dead in all, the neighbours strike
    back once: the race loses backlash_loss of its people and holds rituals at backlash_calm × the rate."""
    out = []
    for rule in _born(rules):
        s, p = rule["s"], rule["p"]
        spec = races[s].emerge
        if spec["mode"] != "ritual":
            continue
        others = np.array([r for r in range(len(races)) if r != s]) if spec["victims"] == "all" else np.array([p])
        per = spec["ritual_cost"] + spec["ritual_converts"]
        for c in np.flatnonzero(st.P[s] >= 1):
            k = rng.poisson(rule["rate"] * st.P[s, c] ** spec["ritual_power"])
            if k == 0:
                continue
            nb = world.nb[c]
            opts = np.r_[c, nb[nb >= 0]]
            v = int(opts[np.argmax(st.P[p, opts])])
            pool = st.P[np.ix_(others, opts)]
            k = min(k, int(st.P[p, v] // per), int(pool.sum() // per))
            if k == 0:
                continue
            dead = k * spec["ritual_cost"]
            taken = dead * pool / pool.sum()
            st.P[np.ix_(others, opts)] = pool - taken
            conv = k * spec["ritual_converts"]
            st.P[p, v] -= conv
            st.P[s, c] += conv                   # the converts join the ritual group, taking its curves (unchanged O)
            rule["victims"] += dead
            out.append({"race": races[s].id, "cell": int(c), "year": int(st.t), "victims": float(dead),
                        "converts": float(conv),
                        "dead": {races[o].id: float(x) for o, x in zip(others, taken.sum(1)) if x > 0}})
        by_count = spec["backlash_victims"] > 0 and rule["victims"] >= spec["backlash_victims"]
        by_time = spec["backlash_years"] > 0 and st.t >= rule["year"] + spec["backlash_years"]
        if not rule["backlash"] and (by_count or by_time):
            rule["backlash"] = True
            lost = _smackdown(st, world, s, others, spec["backlash_loss"])
            rule["rate"] = spec["ritual_rate"] * spec["backlash_calm"]
            out.append({"race": races[s].id, "cell": -1, "year": int(st.t), "victims": 0.0, "converts": 0.0,
                        "backlash": lost})
    return out


def _smackdown(st, world, s, others, loss):
    """Kill `loss` of race s; each group's share ∝ its exposure (own people + non-ritual people in and next to its
    cell), capped at all of it: big, exposed groups are wiped out, small remote ones survive. Returns people lost."""
    P = st.P[s]
    enemies = st.P[others].sum(0)
    h = np.where(P > 0, P + enemies + world.nb_sum(enemies), 0.0)
    target = loss * P.sum()
    if target <= 0 or not h.any():
        return 0.0
    lo, hi = 0.0, 1.0 / h[h > 0].min()                      # at hi every group is wiped out
    for _ in range(200):
        k = (lo + hi) / 2
        lo, hi = (k, hi) if (P * np.minimum(1.0, k * h)).sum() < target else (lo, k)
    frac = np.minimum(1.0, hi * h)
    st.P[s] = P * (1 - frac)
    return float(target)


def summarise(events, sums):
    """Per-cell events of one step → one record per race (and source race, if given): groups and summed `sums`
    fields; dict fields are summed per key (keeps stats small)."""
    out = {}
    for e in events:
        k = (e["race"], e.get("from"), e["year"])
        head = {"race": e["race"], **({"from": e["from"]} if "from" in e else {}), "year": e["year"], "groups": 0}
        rec = out.setdefault(k, {**head, **{f: 0.0 for f in sums}})
        rec["groups"] += e["cell"] >= 0                 # cell −1: a race-wide record (the backlash), not a group
        for f in sums:
            if f not in e:
                continue
            x = e[f]
            if isinstance(x, dict):
                rec[f] = rec[f] or {}
                for key, v in x.items():
                    rec[f][key] = rec[f].get(key, 0.0) + v
            else:
                rec[f] += x
    return list(out.values())