"""Short-hop settlement simulator in a periodic box of linearly moving stars.

Rule.  A settled system becomes *active* t_delay after it was settled.  An
active system launches to any unsettled, unclaimed star that is (or comes)
within d_hop of it, provided their relative speed is below dv_max (the
velocity change a ship needs to rendezvous).  The ship covers the separation
at launch at speed v_ship, so the target becomes settled sep/v_ship later.
The seed is active at t = 0.

Every star is a valid target and a future launcher.  A settleable fraction
f < 1 is modelled by populating the box at density f*n (the other stars are
inert and play no role).

Stars move on straight lines (no Galactic potential), the box is periodic.
Close approaches are detected exactly within each time step for all pairs
whose relative speed is below v_cap: candidate pairs are those within
d_hop + v_cap*dt/2 at mid-step, and the time of first entry inside d_hop is
solved analytically for linear relative motion.  A star that becomes active
part-way through a step launches from its activation time onward (revision:
in the first pass it waited for the next step, which added a lag of up to one
step).  The step is limited to dt_max so that a star settled and activated
within one step (which does wait for the next step) is rare.
"""
import numpy as np
from scipy.spatial import cKDTree

KMS = 1.0227121650537077  # km/s -> pc/Myr


def make_box(n_density, box, vel_sample, weights, rng, axes=(0, 1, 2)):
    """Uniform random positions at number density n_density [pc^-3] in a
    periodic box (Lx, Ly, Lz) [pc]; velocities [pc/Myr] resampled with
    replacement from the weighted empirical sample (km/s), in the frame in
    which the weighted mean velocity vanishes.  `axes` permutes the velocity
    components: axes[0] is the Galactic component (0 = U, 1 = V, 2 = W) that
    is laid along the box x axis (the long axis of a slab)."""
    box = np.asarray(box, float)
    N = int(rng.poisson(n_density * np.prod(box)))
    pos = rng.random((N, 3)) * box
    mean = np.sum(weights[:, None] * vel_sample, axis=0)
    idx = rng.choice(len(vel_sample), size=N, p=weights)
    vel = ((vel_sample[idx] - mean) * KMS)[:, list(axes)]
    return pos, vel


def simulate(pos0, vel, box, d_hop, t_delay, v_ship_kms, seed_idx, t_end,
             dv_max_kms=np.inf, v_cap_kms=200.0, margin_frac=None, stop_n=None,
             stop_front=None, front_origin=None, record_every=1, bin_pc=20.0,
             dt_max=None):
    """Run the settlement process.  Returns dict with arrays t, n_settled and
    (slab runs) front positions along x: f50 (half-settled) and lead (farthest
    settled star), each for +x and -x."""
    box = np.asarray(box, float)
    N = len(pos0)
    v_ship = v_ship_kms * KMS
    v_cap = v_cap_kms * KMS
    dv_max = dv_max_kms * KMS
    if margin_frac is None:
        margin_frac = 2.0 if d_hop <= 1.0 else (1.0 if d_hop <= 2.0 else 0.5)
    margin_frac = min(margin_frac, 0.49 * box.min() / d_hop - 1.0)
    margin = margin_frac * d_hop
    if dt_max is not None:
        margin = min(margin, 0.5 * v_cap * dt_max)
    r_s = d_hop + margin
    assert r_s < 0.5 * box.min(), "search radius must be < half the box"
    dt = 2.0 * margin / v_cap

    t_arr = np.full(N, np.inf)     # time at which the star becomes settled
    t_act = np.full(N, np.inf)     # time at which it becomes active
    t_arr[seed_idx] = 0.0
    t_act[seed_idx] = 0.0
    hist = dict(t=[], n=[], f50p=[], f50m=[], leadp=[], leadm=[])
    if front_origin is None:
        front_origin = pos0[np.atleast_1d(seed_idx)[0], 0]
    Lx = box[0]
    edges = np.arange(0.0, Lx / 2 + 1e-9, bin_pc)
    nxb = int(np.ceil(Lx / bin_pc))
    # slices to look at on either side (2 if the last slice is a partial one)
    reach = 1 if abs(nxb * bin_pc - Lx) < 1e-9 else 2
    assert r_s < bin_pc, "search radius must be smaller than the x-slice width"
    t = 0.0
    step = 0
    d2 = d_hop**2
    while t < t_end:
        active = np.flatnonzero(t_act < t + dt)     # active now or during this step
        free = np.flatnonzero(np.isinf(t_arr))
        if len(free) == 0 and t > t_arr.max():
            break                              # every star has been reached
        search = len(active) > 0 and len(free) > 0
        if search:
            pm = (pos0 + vel * (t + 0.5 * dt)) % box
            # keep only launchers and targets in x-slices that hold (or adjoin
            # a slice that holds) the other kind: far behind the front nothing
            # is free, far ahead nothing is active
            xb = np.minimum((pm[:, 0] / bin_pc).astype(int), nxb - 1)
            fc = np.bincount(xb[free], minlength=nxb)
            ac = np.bincount(xb[active], minlength=nxb)
            near_f, near_a = fc.copy(), ac.copy()
            for s_ in range(1, reach + 1):
                near_f += np.roll(fc, s_) + np.roll(fc, -s_)
                near_a += np.roll(ac, s_) + np.roll(ac, -s_)
            active = active[near_f[xb[active]] > 0]
            free = free[near_a[xb[free]] > 0]
            search = len(active) > 0 and len(free) > 0
        if search:
            if len(active) <= 8:
                # seed phase: brute force is much cheaper than building trees
                pf = pm[free]
                rws, cls = [], []
                for k_, a_ in enumerate(active):
                    dp = pf - pm[a_]
                    dp -= box * np.round(dp / box)
                    hit = np.flatnonzero(np.einsum("ij,ij->i", dp, dp) < r_s**2)
                    rws.append(np.full(len(hit), k_))
                    cls.append(hit)
                row, col = np.concatenate(rws), np.concatenate(cls)
            else:
                ta = cKDTree(pm[active], boxsize=box)
                tf = cKDTree(pm[free], boxsize=box)
                coo = ta.sparse_distance_matrix(tf, r_s, output_type="coo_matrix")
                row, col = coo.row, coo.col
            if len(row):
                ia = active[row]
                jf = free[col]
                w = vel[jf] - vel[ia]
                pmid = pm[jf] - pm[ia]
                pmid -= box * np.round(pmid / box)
                # the launcher can act only from its activation time a0 (measured
                # from the step start) to the end of the step
                a0 = np.clip(t_act[ia] - t, 0.0, dt)
                p0 = pmid - w * (0.5 * dt - a0)[:, None]      # separation at t + a0
                a = np.einsum("ij,ij->i", w, w)
                b = np.einsum("ij,ij->i", p0, w)
                c = np.einsum("ij,ij->i", p0, p0)
                ts = np.clip(-b / np.maximum(a, 1e-30), 0.0, dt - a0)
                dmin2 = c + 2 * b * ts + a * ts**2
                ok = (dmin2 < d2) & (a < dv_max**2)
                if ok.any():
                    a, b, c, jf, a0 = a[ok], b[ok], c[ok], jf[ok], a0[ok]
                    inside = c < d2
                    disc = np.maximum(b**2 - a * (c - d2), 0.0)
                    tau = np.where(inside, 0.0, (-b - np.sqrt(disc)) / np.maximum(a, 1e-30))
                    sep = np.where(inside, np.sqrt(c), d_hop)
                    arrive = t + a0 + tau + sep / v_ship
                    np.minimum.at(t_arr, jf, arrive)
                    newly = np.unique(jf)
                    t_act[newly] = t_arr[newly] + t_delay
        t += dt
        step += 1
        if step % record_every == 0:
            settled = t_arr <= t
            ns = int(settled.sum())
            hist["t"].append(t)
            hist["n"].append(ns)
            x = (pos0[:, 0] + vel[:, 0] * t - front_origin + Lx / 2) % Lx - Lx / 2
            for sgn, kf, kl in ((1, "f50p", "leadp"), (-1, "f50m", "leadm")):
                xs = sgn * x
                side = xs >= 0
                h_all, _ = np.histogram(xs[side], edges)
                h_set, _ = np.histogram(xs[side & settled], edges)
                frac = h_set / np.maximum(h_all, 1)
                below = np.flatnonzero(frac < 0.5)
                if len(below) == 0:
                    f50 = edges[-1]
                elif below[0] == 0:
                    f50 = 0.0
                else:
                    k = below[0]
                    f1, f2 = frac[k - 1], frac[k]
                    f50 = edges[k] - 0.5 * bin_pc + bin_pc * (f1 - 0.5) / max(f1 - f2, 1e-9)
                hist[kf].append(f50)
                hist[kl].append(xs[side & settled].max() if (side & settled).any() else 0.0)
            if stop_n is not None and ns >= stop_n:
                break
            if stop_front is not None and min(hist["f50p"][-1], hist["f50m"][-1]) >= stop_front:
                break
    out = {k: np.asarray(v) for k, v in hist.items()}
    out["dt"] = dt
    out["N"] = N
    out["t_arr"] = t_arr
    return out


def time_to_reach(t, n, target):
    """First time the settled count reaches `target` (linear interpolation in log n)."""
    n = np.asarray(n, float)
    k = np.flatnonzero(n >= target)
    if len(k) == 0:
        return np.nan
    k = k[0]
    if k == 0:
        return t[0]
    n0, n1 = max(n[k - 1], 1.0), n[k]
    if n1 == n0:
        return t[k]
    f = (np.log(target) - np.log(n0)) / (np.log(n1) - np.log(n0))
    return t[k - 1] + f * (t[k] - t[k - 1])


def arrival_time(t, f, x):
    """First time the half-settled front position f(t) reaches distance x."""
    t, f = np.asarray(t), np.asarray(f)
    k = np.flatnonzero(f >= x)
    if len(k) == 0:
        return np.nan
    k = k[0]
    if k == 0 or f[k] == f[k - 1]:
        return float(t[k])
    return float(t[k - 1] + (x - f[k - 1]) / (f[k] - f[k - 1]) * (t[k] - t[k - 1]))