"""Hazards: static mortality layers and roaming units (few, very powerful creatures that raid)."""
import numpy as np


class Static:
    def __init__(self, spec):
        self.spec = spec
        self.danger = None          # cells drift avoids (push hazards)

    def step(self, st, world, races, regions, t, rng):
        s = self.spec
        y0, y1 = s["years"]
        if not (y0 <= t < y1):
            self.danger = np.zeros(world.n) if s["push"] > 0 else None
            return {"name": s["name"]}
        f = s["decay"] ** ((t - y0) / 10.0)
        m = s["mortality"] * f
        mask = regions(s["region"])
        st.P[:, mask] *= 1 - m
        if s["push"] > 0:               # on land, a share flees each step (e.g. megafauna driving people into the sea)
            land = mask & ~np.asarray(world.fields["ocean"], bool)
            st.push[:, land] += s["push"] * f * st.P[:, land]
            self.danger = land.astype(float)
        return {"name": s["name"], "mortality": m}


class Roaming:
    def __init__(self, spec):
        self.spec = spec
        self.home = self.cur = self.kind = None

    def _start(self, st, world, races, regions, rng):
        s = self.spec
        region = regions(s["home_region"])
        ids = [r.id for r in races]
        w = np.zeros(world.n)
        if s["home_near"] and s["home_near"] in ids:
            w = st.P[ids.index(s["home_near"])] * region
        if w.sum() <= 0:
            w = (region & ~np.asarray(world.fields["ocean"], bool)).astype(float)
        if w.sum() <= 0:
            w = region.astype(float)
        n = s["local"] + s["wanderers"]
        self.home = rng.choice(world.n, size=n, p=w / w.sum()).astype(np.int64)
        self.cur = self.home.copy()
        self.kind = np.r_[np.zeros(s["local"], np.int8), np.ones(s["wanderers"], np.int8)]

    def _burst(self, st, world, cell, share):
        cells = world.within(int(cell), self.spec["radius_km"])
        st.P[:, cells] *= 1 - share

    def _walk(self, world, cell, km, rng):
        v = world.xyz[cell]
        rnd = rng.normal(size=3)
        t = rnd - rnd.dot(v) * v
        t /= np.linalg.norm(t)
        d = km / world.radius_km
        return int(world.nearest(v * np.cos(d) + t * np.sin(d)))

    def step(self, st, world, races, regions, t, rng):
        s = self.spec
        if t < s["start"]:
            return {"name": s["name"], "units": 0, "raids": 0, "births": 0, "deaths": 0}
        if self.home is None:
            self._start(st, world, races, regions, rng)
        tot = st.P.sum(0)
        occ = np.flatnonzero(regions(s["home_region"]) & (tot >= 1))
        raids, dead = 0, []
        for i in range(len(self.home)):
            if self.kind[i] == 0:
                self.cur[i] = self.home[i]
                if len(occ) and rng.random() < s["raid_chance"]:
                    tgt = int(rng.choice(occ, p=tot[occ] / tot[occ].sum()))
                    self.cur[i] = tgt
                    self._burst(st, world, tgt, s["raid_burst"])
                    raids += 1
                    dens = tot[tgt] / world.area[tgt]
                    if rng.random() < s["raid_death"] * min(1.0, dens / s["density_ref"]):
                        dead.append(i)
                else:
                    self._burst(st, world, self.home[i], s["local_mortality"])
            else:
                self.cur[i] = self._walk(world, self.cur[i], s["wander_step_km"], rng)
                if rng.random() < s["wander_stop"]:
                    self._burst(st, world, self.cur[i], s["wander_burst"])
        n = len(self.home)
        dead = set(dead) | {i for i in range(n) if rng.random() < s["infight"]}
        keep = np.array([i not in dead for i in range(n)], bool)
        self.home, self.cur, self.kind = self.home[keep], self.cur[keep], self.kind[keep]
        births = 0
        lo, hi = s["clutch"]
        local = np.flatnonzero(self.kind == 0)
        if len(self.home) and s["growth"] > 0:
            n = len(self.home)
            room = max(0.0, 1.0 - n / s["cap"]) if s["cap"] > 0 else 1.0  # logistic when capped
            for _ in range(rng.poisson(s["growth"] * n * room / ((lo + hi) / 2))):
                size = int(rng.integers(lo, hi + 1))
                at = self.home[rng.choice(local)] if len(local) else self.home[rng.integers(len(self.home))]
                self.home = np.r_[self.home, np.full(size, at)]
                self.cur = np.r_[self.cur, np.full(size, at)]
                self.kind = np.r_[self.kind, np.zeros(size, np.int8)]
                births += size
        return {"name": s["name"], "units": int(len(self.home)), "raids": raids, "births": births,
                "deaths": int(len(dead))}


def make_hazard(spec):
    return Static(spec) if spec["kind"] == "static" else Roaming(spec)