"""Episode runner with persistent pheromone across churn."""
import numpy as np, sim, schedulers as sch

SCHEDULERS = {
    "RoundRobin": lambda ctx, st, **k: sch.round_robin(ctx),
    "LeastConn":  lambda ctx, st, **k: sch.least_connection(ctx),
    "ACO":        lambda ctx, st, **k: sch.standalone_aco(ctx, iters=k["iters"], state=st),
    "PSO":        lambda ctx, st, **k: sch.standalone_pso(ctx, iters=k["iters"]),
    "WOA":        lambda ctx, st, **k: sch.woa(ctx, iters=k["iters"]),
    "AdaptiSwarm": lambda ctx, st, **k: sch.adaptiswarm_edge(ctx, iters=k["iters"], sigma_u=k.get("sigma_u",0.0), state=st),
    "Abl_ACOonly": lambda ctx, st, **k: sch.adaptiswarm_edge(ctx, iters=k["iters"], use_pso=False, sigma_u=k.get("sigma_u",0.0), state=st),
    "Abl_PSOrand": lambda ctx, st, **k: sch.adaptiswarm_edge(ctx, iters=k["iters"], use_aco_seed=False, use_adaptive=False, sigma_u=k.get("sigma_u",0.0), state=st),
    "Abl_PSOseed": lambda ctx, st, **k: sch.adaptiswarm_edge(ctx, iters=k["iters"], use_adaptive=False, sigma_u=k.get("sigma_u",0.0), state=st),
}

def run_episode(name, n_nodes, lam, seed, n_steps=15, iters=40,
                alpha=0.5, beta=0.3, gamma=0.2, churn=True, probe_step=8):
    rng = np.random.default_rng(seed)
    nodes = sim.make_nodes(n_nodes, rng)
    cap, mem, p_idle, p_peak = sim.node_arrays(nodes)
    link_lat = sim.make_link_latency(n_nodes, rng)
    active = np.ones(n_nodes, dtype=bool); backlog = np.zeros(n_nodes); state = {}
    agg = dict(latency=[], energy=[], util=[], jain=[], A_met=[], all_met=[], sigma=[])
    conv_curve = None; prev_util = np.zeros(n_nodes)
    for step in range(n_steps):
        if churn and n_nodes >= 5:
            if rng.random() < 0.3:
                idx = rng.integers(n_nodes); active[idx] = not active[idx]
            if active.sum() < max(2, n_nodes // 3):
                active[:] = True
        act_idx = np.where(active)[0]
        a_cap = cap[act_idx]; a_mem = mem[act_idx]; a_link = link_lat[act_idx]
        a_pidle = p_idle[act_idx]; a_ppeak = p_peak[act_idx]; a_backlog = backlog[act_idx]
        n_arr = rng.poisson(lam * sim.DT)
        demand, tmem, deadline, cls = sim.make_tasks(n_arr, rng)
        sigma_u_live = float(np.std(prev_util[act_idx])) if act_idx.size else 0.0
        ctx = sch.Ctx(demand, tmem, deadline, cls, a_backlog, a_cap, a_link,
                      a_pidle, a_ppeak, alpha, beta, gamma, rng, node_mem=a_mem)
        local_state = {}
        for gkey in ("aco", "adapti"):
            gtau = state.get("global_" + gkey)
            if gtau is not None: local_state[gkey] = gtau[act_idx].copy()
        assign, curve = SCHEDULERS[name](ctx, local_state, iters=iters, sigma_u=sigma_u_live)
        for gkey in ("aco", "adapti"):
            if gkey in local_state:
                gtau = state.get("global_" + gkey)
                if gtau is None: gtau = np.ones(n_nodes)
                gtau[act_idx] = local_state[gkey]; state["global_" + gkey] = gtau
        if n_arr > 0:
            mtr = sim.evaluate(assign, demand, tmem, deadline, cls, a_backlog,
                               a_cap, a_link, a_pidle, a_ppeak, alpha, beta,
                               gamma, return_metrics=True, node_mem=a_mem)
            for k,mk in [("latency","avg_latency"),("energy","energy"),("util","utilization"),
                         ("jain","jain"),("A_met","classA_deadline_met"),("all_met","all_deadline_met"),("sigma","sigma_u")]:
                agg[k].append(mtr[mk])
            added = np.zeros(len(act_idx))
            for j,n in enumerate(assign): added[n] += demand[j]
            backlog[act_idx] = np.maximum(a_backlog + added - a_cap * sim.DT, 0.0)
            prev_util = np.zeros(n_nodes); prev_util[act_idx] = mtr["served"] / (a_cap * sim.DT)
        else:
            backlog[act_idx] = np.maximum(a_backlog - a_cap * sim.DT, 0.0)
        if step == probe_step and len(curve) > 1: conv_curve = curve
    def stat(key):
        v = np.array(agg[key]) if agg[key] else np.array([0.0]); return v.mean()
    conv_iter = iters
    if conv_curve is not None and len(conv_curve) > 1:
        final = conv_curve[-1]; rs = abs(conv_curve[0] - final) + 1e-9
        for i,f in enumerate(conv_curve):
            if abs(f - final) <= 0.02 * rs: conv_iter = i + 1; break
    return dict(latency_ms=stat("latency")*1000.0, utilization=stat("util")*100.0,
                energy=stat("energy"), jain=stat("jain"),
                A_deadline=stat("A_met")*100.0, all_deadline=stat("all_met")*100.0,
                conv_iter=conv_iter, conv_curve=conv_curve)