import json, math
from pathlib import Path
import numpy as np
from scipy.special import lambertw

# Resource-driven collective attention toy model.
# Agents never observe targets; a captured target leaves only a transient motion/action cue.

def lambda_star(ell, tau, r, B=0.64):
    K = math.pi/(4*r) * (1.0 - B*r/tau)
    z = ell*math.exp(K)/tau
    return ell/float(lambertw(z).real)

def formula_check():
    # Directly verify logarithmic tau dependence of the claimed collective MFPT.
    b, D, ell, lam = 10.0, 1.0, 1.0, 1.0
    taus = np.array([4., 8., 16., 32., 64., 128., 256.])
    Tc = b*b/(2*D)*np.log(taus/(lam*np.exp(-ell/lam)))
    # finite differences dTc/d(log tau) should be constant b^2/(2D)
    slopes = np.diff(Tc)/np.diff(np.log(taus))
    return {"taus": taus.tolist(), "Tc": Tc.tolist(),
            "log_tau_slopes": slopes.tolist(),
            "expected_slope": b*b/(2*D),
            "max_relative_slope_error": float(np.max(np.abs(slopes-b*b/(2*D))/(b*b/(2*D))))}

def torus_delta(a, b, L):
    d = a-b
    return (d + L/2) % L - L/2

def simulate(attention_radius, seed, N=8, L=20., steps=300, n_targets=5,
             tau=35, target_r=.42, speed=.22, cue_steps=10, social=True):
    rng = np.random.default_rng(seed)
    x = rng.uniform(0, L, (N, 2))
    ang = rng.uniform(-math.pi, math.pi, N)
    targets = rng.uniform(0, L, (n_targets, 2))
    cooldown = np.zeros(n_targets, dtype=int)
    # cue vector and expiry are internal action history, not target information
    cues = np.zeros((N, 2)); expiry = np.zeros(N, dtype=int)
    captures = 0; duplicate = 0; agg_sum = 0.; agg_count = 0
    for t in range(steps):
        # local conspecific attention: copy the most recent successful motion cue
        new_ang = ang.copy()
        if social and attention_radius > 0:
            for i in range(N):
                d = torus_delta(x, x[i], L)
                dist = np.sqrt((d*d).sum(1)); neigh = np.where((dist <= attention_radius) & (dist > 0) & (expiry > t))[0]
                if len(neigh):
                    # nearest active cue; no proximity or group reward is used
                    j = neigh[np.argmin(dist[neigh])]
                    v = cues[j]
                    if np.linalg.norm(v) > 1e-8:
                        desired = math.atan2(v[1], v[0])
                        # social influence is modest, preserving exploration
                        diff = (desired-new_ang[i]+math.pi)%(2*math.pi)-math.pi
                        new_ang[i] += .65*diff
        # persistent random exploration
        new_ang += rng.normal(0, .20, N)
        ang = (new_ang + math.pi)%(2*math.pi)-math.pi
        x = (x + speed*np.c_[np.cos(ang), np.sin(ang)]) % L
        # target replenishment
        cooldown = np.maximum(0, cooldown-1)
        for k in range(n_targets):
            if cooldown[k] == 0 and np.linalg.norm(torus_delta(x[0], targets[k], L)) < -1: # no-op, keeps target hidden
                pass
        # each target can be captured once, then replenishes after tau
        for i in range(N):
            if not np.any(cooldown == 0): break
            ds = torus_delta(targets, x[i], L)
            hit = np.where((cooldown == 0) & ((ds*ds).sum(1) <= target_r**2))[0]
            if len(hit):
                k = hit[0]; captures += 1
                # capture cue is the agent's recent action, without exposing target location/reward
                cues[i] = np.array([math.cos(ang[i]), math.sin(ang[i])])
                expiry[i] = t + cue_steps
                cooldown[k] = tau
                # target is hidden while depleted, then respawns at an unobserved location
                targets[k] = rng.uniform(0, L, 2)
        # aggregation order parameter from the stated neighbor statistic
        for i in range(N):
            d = torus_delta(x, x[i], L); dist=np.sqrt((d*d).sum(1))
            neigh=np.where((dist>0)&(dist<1.5))[0]
            agg_sum += (np.sum(dist[neigh]<1.5)/len(neigh)) if len(neigh) else 0.
            agg_count += 1
        # duplicate-search proxy: close pair during depleted-target intervals
        if np.sum(cooldown > 0) > 0:
            for i in range(N):
                d=torus_delta(x, x[i], L); duplicate += int(np.sum(((d*d).sum(1)>0)&((d*d).sum(1)<target_r**2))>0)
    return {"capture_rate": captures/(steps/1000), "captures": captures,
            "aggregation_A": agg_sum/agg_count, "duplicate_fraction": duplicate/(steps*N)}

def run():
    ell, tau, r = 1.0, 35.0, .42
    ls = lambda_star(ell, tau, r)
    multipliers = [.25,.5,.75,1.,1.5,2.,4.]
    seeds = [3, 11]
    rows=[]
    for m in multipliers:
        lam=m*ls
        vals=[simulate(lam,s) for s in seeds]
        rows.append({"multiplier":m,"lambda":lam,
          "capture_rate":float(np.mean([v['capture_rate'] for v in vals])),
          "aggregation_A":float(np.mean([v['aggregation_A'] for v in vals])),
          "duplicate_fraction":float(np.mean([v['duplicate_fraction'] for v in vals]))})
    base=[simulate(0,s,social=False) for s in seeds]
    out={"formula_check":formula_check(),"parameters":{"ell":ell,"tau":tau,"r":r,"lambda_star":ls},
         "sweep":rows,"independent_baseline":{
          "capture_rate":float(np.mean([v['capture_rate'] for v in base])),
          "aggregation_A":float(np.mean([v['aggregation_A'] for v in base])),
          "duplicate_fraction":float(np.mean([v['duplicate_fraction'] for v in base]))}}
    Path('results.json').write_text(json.dumps(out,indent=2))
    print(json.dumps(out,indent=2))

if __name__ == '__main__': run()