"""The engine: one step = events → hazards → capacity → growth → conflict → migration → adaptation →
sub-species → tech. Later sub-projects add phases to Engine.phases."""
import resource
import time
from pathlib import Path

import numpy as np

from . import tech as techm
from .adaptation import adapt, prospective
from .config import dumps_toml, resolved
from .conflict import apply_conflict, apply_curse, conflict
from .demography import capacity, crowding, grow, harvest_shocks, overlap_matrix, stochastic_round, veins, weather
from .events import apply_event
from .habitat import comfort, environment, fitness, native_optima, quality, suitability
from .hazards import make_hazard
from .migration import drift_and_bud, long_jumps, travel_cost
from .nudges import nudge_multipliers
from .regions import Regions
from .snapshot import engine_version, write_json, write_snapshot
from .state import new_state
from .lineage import (birth_step, displacement, inherit, init_rules, native, progress, raids, return_step,
                      rituals, summarise)


class Engine:
    def __init__(self, world, history, races, seed=None):
        self.w, self.h, self.races = world, history, races
        self.R, self.n = len(races), world.n
        self.seed = history.seed if seed is None else seed
        self.regions = Regions(world, history.regions)
        self.micro = np.exp(0.3 * world.noise(self.seed + 1000))
        self.state = new_state(self.R, self.n, native_optima(races), self.seed)
        self.rules = init_rules(races)
        self.returns, self.ritual_log, self.raid_log = [], [], []
        self.births = np.zeros((self.R, self.n))
        self.hazards = [make_hazard(s) for s in history.hazards]
        self.events_log, self.emergence, self.stats_steps = [], [], []
        self.hazard_stats = []
        self.years_per_step = history.step
        self._fc = {}                                  # per race: inputs and values of _fit_comf
        self._nudge_memo = None
        self._world_changed()
        self._capacity()
        self.hostile = np.zeros((self.R, self.n))
        self.crowd = np.zeros((self.R, self.n))
        self.phases = [("events", self.phase_events), ("hazards", self.phase_hazards),
                       ("capacity", self._capacity), ("growth", self.phase_growth), ("conflict", self.phase_conflict),
                       ("migration", self.phase_migration), ("adaptation", self.phase_adaptation),
                       ("subspecies", self.phase_subspecies), ("tech", self.phase_tech)]

    # ---- derived state ----
    def _world_changed(self):
        self.env = environment(self.w)
        self.suit = np.stack([suitability(self.w, r) for r in self.races])
        self.q = np.stack([quality(self.suit[i], r, self.micro) for i, r in enumerate(self.races)])
        self.alpha = overlap_matrix(self.suit, self.races)
        self.native_ok = np.stack([(self.suit[i] > 0) & (fitness(self.env, np.repeat(native(r)[:, None], self.n, 1), r)
                                                         >= self.h.f_min) for i, r in enumerate(self.races)])
        self.tcost0 = np.stack([travel_cost(self.w, r, 0.0) for r in self.races])

    def _capacity(self):
        st, tp = self.state, self.h.tech
        prospective(st, self.w)
        self.nudge_g, nudge_c, self.nudge_e, self.nudge_a = self._nudges(st.t)
        e = weather(self.w, st.rng)
        K = np.zeros((self.R, self.n))
        self.q_eff = np.zeros((self.R, self.n))
        self.eff = []
        self.OK = np.zeros((self.R, self.n), bool)
        self.fit = np.zeros((self.R, self.n))
        self.nudge_share = {}
        for r, race in enumerate(self.races):
            eff = techm.effective(st.T[r], tp)
            self.eff.append(eff)
            q = self.q[r]
            q_eff = np.where(q > 0, q + st.improve[r] * (race.qmax - q), 0.0)
            fit, comf = self._fit_comf(r, race, st.O[r], techm.reach(eff, race, tp))
            self.fit[r] = fit * comf                  # displacement compares the whole position on the curve
            shock = harvest_shocks(self.w, q_eff, techm.shock_scale(eff, tp), e)
            k = capacity(self.w, race, q_eff, fit, comf, techm.k_gain(eff, tp), shock, nudge_c[r])
            ok = (self.suit[r] > 0) & (fit >= self.h.f_min)
            K[r] = np.where(ok, k, 0.0)
            self.q_eff[r], self.OK[r] = q_eff, ok
            tot = K[r].sum()
            self.nudge_share[race.id] = float((K[r] * (1 - 1 / nudge_c[r])).sum() / tot) if tot > 0 else 0.0
        displacement(K, st.P, self.fit, self.rules, self.races)
        self.K = K

    def _nudges(self, t):
        """nudge_multipliers, rebuilt only when the set of active nudges or the era changes (read-only arrays)."""
        key = (tuple(i for i, nd in enumerate(self.h.nudges) if nd["years"][0] <= t < nd["years"][1]), self.w.era)
        if self._nudge_memo is None or self._nudge_memo[0] != key:
            arrs = nudge_multipliers(self.h.nudges, self.races, self.regions, t, self.n)
            for a in arrs:
                a.flags.writeable = False
            self._nudge_memo = (key, arrs)
        return self._nudge_memo[1]

    def _fit_comf(self, r, race, O, rch):
        """fitness and comfort of race r, recomputed only in cells whose optima or reach changed since the last
        step (both are per-cell functions of those and the environment): the same values, much less work."""
        rch = np.broadcast_to(np.asarray(rch), (self.n,))       # keeps its dtype (float32 from tech): same rounding
        c = self._fc.get(r)
        if c is None or c[0] is not self.env:
            fit, comf = fitness(self.env, O, race, rch), comfort(O, race)
        else:
            _, O0, r0, fit, comf = c
            ch = np.flatnonzero((O != O0).any(0) | (rch != r0))
            if len(ch):
                fit[ch] = fitness(self.env, O[:, ch], race, rch[ch], cells=ch)
                comf[ch] = comfort(O[:, ch], race)
        self._fc[r] = (self.env, O.copy(), rch.copy(), fit, comf)
        return fit, comf

    # ---- phases ----
    def phase_events(self):
        for i, ev in enumerate(self.h.events):
            if ev["year"] != self.state.t:
                continue
            log = apply_event(self.state, self.w, self.races, ev, self.regions, self.q, self.OK, self.h, i)
            self.events_log.append(log)
            if ev["kind"] == "era_switch":
                self._world_changed()

    def phase_hazards(self):
        self.hazard_stats = [hz.step(self.state, self.w, self.races, self.regions, self.state.t, self.state.rng)
                             for hz in self.hazards]

    def phase_growth(self):
        st = self.state
        self.crowd = crowding(st.P, self.K, self.alpha, self.races, self.h.competition["intolerance"], world=self.w)
        self.births = grow(st.P, self.K, self.crowd, self.q_eff, self.races, self.nudge_g)
        for r, race in enumerate(self.races):
            if race.veins > 0:
                st.P[r] = veins(st.P[r], self.K[r], race.veins, race.allee)

    def phase_conflict(self):
        blame = np.zeros_like(self.state.P)
        loss, self.hostile = conflict(self.state.P, self.q_eff, self.crowd, self.races, blame=blame)
        apply_conflict(self.state, loss, self.races)
        apply_curse(self.state, blame, self.h.competition["curse_decay"])

    def phase_migration(self):
        st, tp, mp = self.state, self.h.tech, self.h.migration
        tcost = self.tcost0 - np.stack([techm.travel_bonus(self.eff[i], tp) for i in range(self.R)])
        danger = sum((hz.danger for hz in self.hazards if getattr(hz, "danger", None) is not None), np.zeros(self.n))
        drift_and_bud(st, self.w, self.races, self.K, self.OK, tcost, self.hostile + danger[None], self.nudge_e, st.rng,
                      mp)
        boost = np.stack([techm.jump_boost(self.eff[i], tp) for i in range(self.R)]) * self.nudge_e
        long_jumps(st, self.w, self.races, self.K, self.OK, st.rng, mp, boost, env=self.env, suit=self.suit,
                   f_min=self.h.f_min)

    def phase_adaptation(self):
        adapt(self.state, self.env, self.races, self.q_eff, self.years_per_step, self.nudge_a)

    def phase_subspecies(self):
        st, rng = self.state, self.state.rng
        for e in birth_step(st, self.w, self.races, self.rules, self.regions, rng, self.native_ok):
            e["lat"], e["lon"] = float(self.w.lat[e["cell"]]), float(self.w.lon[e["cell"]])
            self.emergence.append(e)
        self.returns += summarise(return_step(st, self.w, self.races, self.rules, rng, self.native_ok), ("people",))
        inherit(st, self.w, self.races, self.rules, self.births, self.regions, self.native_ok)
        self.raid_log += summarise(raids(st, self.w, self.races, self.rules, rng, self.native_ok), ("raids",))
        self.ritual_log += summarise(rituals(st, self.w, self.races, self.rules, rng), ("victims", "converts", "backlash", "dead"))

    def phase_tech(self):
        steps = self.years_per_step / 10.0
        techm.step_tech(self.state, self.w, self.races, self.h.tech, steps)
        techm.step_land(self.state, self.races, self.h.tech, steps)
        stochastic_round(self.state.P, self.state.rng)

    # ---- driving ----
    def step(self):
        """Process year state.t, then advance."""
        for _, fn in self.phases:
            fn()
        self.state.t += self.h.step

    def _stats(self, t):
        P = self.state.P
        ids = [r.id for r in self.races]
        fam = {}
        for i, r in enumerate(self.races):
            fam[r.family] = fam.get(r.family, 0.0) + float(P[i].sum())
        return {"year": t, "pop": {ids[i]: float(P[i].sum()) for i in range(self.R)},
                "family": fam,
                "occupied": {ids[i]: int((P[i] >= 1).sum()) for i in range(self.R)},
                "regions": {name: {ids[i]: float(P[i, self.regions(name)].sum()) for i in range(self.R)}
                            for name in self.h.stats_regions},
                "hazards": self.hazard_stats, "nudge_capacity": self.nudge_share,
                "progress": {f"{self.races[r['p']].id}>{self.races[r['s']].id}": self._mean_progress(r)
                             for r in self.rules}}

    def _mean_progress(self, rule):
        p, s = rule["p"], rule["s"]
        P = self.state.P[p]
        tot = P.sum()
        return float((progress(self.state.O[p], self.races[p], self.races[s]) * P).sum() / tot) if tot > 0 else 0.0

    def run(self, out_dir, years=None, progress=None):
        years = self.h.years if years is None else years
        out = Path(out_dir)
        (out / "snap").mkdir(parents=True, exist_ok=True)
        meta = {"seed": self.seed, "engine": engine_version(), "world": str(self.w.meta.get("dir", "")),
                "years": years}
        (out / "run.toml").write_text(dumps_toml(resolved(self.h, self.races, **meta)))
        t0 = time.time()
        while self.state.t <= years:
            t = self.state.t
            self.step()
            self.stats_steps.append(self._stats(t))
            if t % self.h.snapshot_every == 0 or t == years:
                write_snapshot(out / "snap" / f"y{t:04d}.npz", self.state, t, self.races)
            if progress and t % 500 == 0:
                progress(f"year {t}: " + ", ".join(f"{k} {v:,.0f}" for k, v in self.stats_steps[-1]["pop"].items()))
        stats = {"steps": self.stats_steps, "events": self.events_log, "emergence": self.emergence,
                 "returns": self.returns, "rituals": self.ritual_log, "raids": self.raid_log,
                 "timing": {"total_s": time.time() - t0,
                            "peak_rss_mb": resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024},
                 "meta": meta}
        write_json(out / "stats.json", stats)
        return stats