// physarum-cpu: one step of a Physarum (slime-mould) transport-network sim.
// Synchronous update: all agents sense the start-of-step field along three
// sensors, steer up-gradient and move (parallel over agents, each agent's jitter
// from an index-seeded RNG); deposits are then applied in a serial index-order
// pass, and the field is diffused (separable 3-tap) and decayed. On a cadence
// the agents are counting-sorted into tile order so that consecutive (and
// same-thread) agents sense a small hot window instead of scattering three
// gathers across the whole field -- the gathers dominate and otherwise miss.
// Read-only sensing + ordered deposit + serial sort make the step deterministic
// for any thread count. Agent trig/arithmetic and the diffuse stencil vectorise
// on NEON (aarch64); a portable scalar path covers everything else.
#include <torch/torch.h>
#include <ATen/Parallel.h>
#include <cstdint>
#include <cstring>
#include <cmath>
#include <thread>
#include <vector>

static const int TILE = 32;        // sort granularity (px); ~8-step drift stays inside
static const int RESORT = 8;       // re-sort every RESORT steps (amortises the sort)

namespace {

// Fan a range [0,total) across at::get_num_threads() worker threads, each given
// a contiguous chunk; f(begin,end) must touch only its own slice. std::thread
// (not OpenMP) so it parallelises under the stock kernel-builder compile, which
// exposes no way to pass -fopenmp to a CPU kernel. Falls back to serial for a
// single thread or a small range.
template<class F>
inline void parallel_chunks(int total, F f){
    int T = at::get_num_threads(); if (T < 1) T = 1;
    if (T == 1 || total < 4096){ f(0, total); return; }
    if (T > total) T = total;
    int chunk = (total + T - 1) / T;
    std::vector<std::thread> pool; pool.reserve(T - 1);
    for (int t=1; t<T; t++){
        int b=t*chunk, e=b+chunk>total?total:b+chunk;
        if (b>=e) break;
        pool.emplace_back([&,b,e]{ f(b,e); });
    }
    f(0, chunk>total?total:chunk);                  // calling thread runs chunk 0
    for (auto& th : pool) th.join();
}

inline uint32_t xr(uint32_t &s){ s ^= s<<13; s ^= s>>17; s ^= s<<5; return s; }
inline float frand(uint32_t &s){ return (xr(s) >> 8) * (1.0f/16777216.0f); }

// per-agent, per-step independent seed: the step's `seed` argument already
// encodes base_seed+t, so hashing it with the agent index gives a stream that
// varies per agent and per step yet does not depend on iteration order.
inline uint32_t agent_seed(uint32_t base, int i){
    uint32_t s = (uint32_t)i*2654435761u ^ base*2246822519u;
    s ^= s>>15; s *= 2246822519u; s ^= s>>13; s *= 3266489917u; s ^= s>>16;
    return s ? s : 1u;
}

inline float fsin(float x){                         // parabola approx, |err|<1e-3
    float k = x * 0.15915494f;
    k = x - 6.2831853f * (float)((int)(k + (k < 0 ? -0.5f : 0.5f)));
    float y = 1.27323954f*k - 0.405284735f*k*(k<0?-k:k);
    return 0.225f*(y*(y<0?-y:y) - y) + y;
}
inline float fcos(float x){ return fsin(x + 1.5707963f); }

inline float samp(const float* t, float x, float y, int W, int H){
    int xi = ((int)x % W + W) % W, yi = ((int)y % H + H) % H;
    return t[(size_t)yi*W + xi];
}

void agents(const float* t, float* ax, float* ay, float* ah, int W, int H, int N,
            float sd, float sa, float turn, float speed, uint32_t base);
void deposit(float* t, const float* ax, const float* ay, int W, int H, int N, float dep);
void diffuse(float* t, float* tmp, int W, int H, float decay);

#if defined(__ARM_NEON)
#include <arm_neon.h>
inline float32x4_t vsin(float32x4_t x){
    float32x4_t k = vrndnq_f32(vmulq_n_f32(x, 0.15915494f));
    x = vmlsq_n_f32(x, k, 6.2831853f);
    float32x4_t ax = vabsq_f32(x);
    float32x4_t y = vsubq_f32(vmulq_n_f32(x,1.27323954f),
                              vmulq_n_f32(vmulq_f32(x,ax),0.405284735f));
    float32x4_t ay = vabsq_f32(y);
    return vmlaq_n_f32(y, vsubq_f32(vmulq_f32(y,ay), y), 0.225f);
}
inline float32x4_t vcos(float32x4_t x){ return vsin(vaddq_f32(x, vdupq_n_f32(1.5707963f))); }

// sense + steer + move for the agent quads [q0,q1); reads the field only, writes
// ax/ay/ah. No field writes, so any set of quads runs on any thread.
inline void move_quads(const float* t, float* ax, float* ay, float* ah,
                       int q0, int q1, int W, int H, float sd, float sa,
                       float turn, float speed, uint32_t base){
    float32x4_t vsd=vdupq_n_f32(sd), vsp=vdupq_n_f32(speed), vsa=vdupq_n_f32(sa);
    float32x4_t vW=vdupq_n_f32((float)W), vH=vdupq_n_f32((float)H);
    float32x4_t iW=vdupq_n_f32(1.0f/W), iH=vdupq_n_f32(1.0f/H);
    for (int q=q0; q<q1; q++){
        int i=q*4;
        float32x4_t x=vld1q_f32(ax+i), y=vld1q_f32(ay+i), h=vld1q_f32(ah+i);
        float32x4_t hl=vsubq_f32(h,vsa), hr=vaddq_f32(h,vsa);
        float cx[4],cy[4],lx[4],ly[4],rx[4],ry[4];
        vst1q_f32(cx, vmlaq_f32(x, vcos(h), vsd));  vst1q_f32(cy, vmlaq_f32(y, vsin(h), vsd));
        vst1q_f32(lx, vmlaq_f32(x, vcos(hl),vsd));  vst1q_f32(ly, vmlaq_f32(y, vsin(hl),vsd));
        vst1q_f32(rx, vmlaq_f32(x, vcos(hr),vsd));  vst1q_f32(ry, vmlaq_f32(y, vsin(hr),vsd));
        float FC[4],FL[4],FR[4];
        for (int l=0;l<4;l++){ FC[l]=samp(t,cx[l],cy[l],W,H);
            FL[l]=samp(t,lx[l],ly[l],W,H); FR[l]=samp(t,rx[l],ry[l],W,H); }
        float32x4_t fc=vld1q_f32(FC), fl=vld1q_f32(FL), fr=vld1q_f32(FR);
        float jit[4]; for(int l=0;l<4;l++){ uint32_t s=agent_seed(base,i+l); jit[l]=0.5f+0.5f*frand(s); }
        float32x4_t vj=vld1q_f32(jit), vturn=vmulq_n_f32(vj,turn);
        uint32x4_t straight=vandq_u32(vcgtq_f32(fc,fl), vcgtq_f32(fc,fr));
        float32x4_t dh=vbslq_f32(vcgtq_f32(fl,fr), vnegq_f32(vturn),
                        vbslq_f32(vcgtq_f32(fr,fl), vturn,
                          vmulq_n_f32(vsubq_f32(vj,vdupq_n_f32(0.5f)), 2.0f*turn)));
        dh=vbslq_f32(straight, vdupq_n_f32(0.0f), dh);
        h=vaddq_f32(h,dh);
        x=vmlaq_f32(x, vcos(h), vsp); y=vmlaq_f32(y, vsin(h), vsp);
        x=vsubq_f32(x, vmulq_f32(vW, vrndmq_f32(vmulq_f32(x,iW))));
        y=vsubq_f32(y, vmulq_f32(vH, vrndmq_f32(vmulq_f32(y,iH))));
        vst1q_f32(ax+i,x); vst1q_f32(ay+i,y); vst1q_f32(ah+i,h);
    }
}
void agents(const float* t, float* ax, float* ay, float* ah, int W, int H, int N,
            float sd, float sa, float turn, float speed, uint32_t base){
    int nq = N/4;
    parallel_chunks(nq, [&](int q0, int q1){
        move_quads(t, ax, ay, ah, q0, q1, W, H, sd, sa, turn, speed, base);
    });
    for (int i=nq*4; i<N; i++){                         // scalar tail, same sync rule
        float x=ax[i], y=ay[i], h=ah[i];
        float fc=samp(t,x+fcos(h)*sd,y+fsin(h)*sd,W,H);
        float fl=samp(t,x+fcos(h-sa)*sd,y+fsin(h-sa)*sd,W,H);
        float fr=samp(t,x+fcos(h+sa)*sd,y+fsin(h+sa)*sd,W,H);
        uint32_t s=agent_seed(base,i); float j=0.5f+0.5f*frand(s);
        if(fc>fl&&fc>fr){} else if(fl>fr) h-=turn*j; else if(fr>fl) h+=turn*j; else h+=(j-0.5f)*2.0f*turn;
        x=fmodf(x+fcos(h)*speed+W,W); y=fmodf(y+fsin(h)*speed+H,H);
        ax[i]=x; ay[i]=y; ah[i]=h;
    }
}
void diffuse(float* t, float* tmp, int W, int H, float dec){
    for (int y=0;y<H;y++){ const float* r=t+(size_t)y*W; float* o=tmp+(size_t)y*W;
        o[0]=0.25f*r[W-1]+0.5f*r[0]+0.25f*r[1]; int x=1;
        for (; x+4<=W-1; x+=4){ float32x4_t L=vld1q_f32(r+x-1),C=vld1q_f32(r+x),R=vld1q_f32(r+x+1);
            vst1q_f32(o+x, vmlaq_n_f32(vmlaq_n_f32(vmulq_n_f32(C,0.5f),L,0.25f),R,0.25f)); }
        for (; x<W-1; x++) o[x]=0.25f*r[x-1]+0.5f*r[x]+0.25f*r[x+1];
        o[W-1]=0.25f*r[W-2]+0.5f*r[W-1]+0.25f*r[0]; }
    float32x4_t vd=vdupq_n_f32(dec);
    for (int y=0;y<H;y++){ int yu=y?y-1:H-1, yd=y+1<H?y+1:0;
        const float* U=tmp+(size_t)yu*W,*C=tmp+(size_t)y*W,*D=tmp+(size_t)yd*W;
        float* o=t+(size_t)y*W; int x=0;
        for (; x+4<=W; x+=4){ float32x4_t s=vmlaq_n_f32(vmlaq_n_f32(vmulq_n_f32(vld1q_f32(C+x),0.5f),
            vld1q_f32(U+x),0.25f), vld1q_f32(D+x),0.25f); vst1q_f32(o+x, vmulq_f32(s,vd)); }
        for (; x<W; x++) o[x]=(0.25f*U[x]+0.5f*C[x]+0.25f*D[x])*dec; }
}
#else   // ---------- portable scalar ----------
void agents(const float* t, float* ax, float* ay, float* ah, int W, int H, int N,
            float sd, float sa, float turn, float speed, uint32_t base){
    parallel_chunks(N, [&](int i0, int i1){
        for (int i=i0;i<i1;i++){
            float x=ax[i], y=ay[i], h=ah[i];
            float fc=samp(t,x+fcos(h)*sd,y+fsin(h)*sd,W,H);
            float fl=samp(t,x+fcos(h-sa)*sd,y+fsin(h-sa)*sd,W,H);
            float fr=samp(t,x+fcos(h+sa)*sd,y+fsin(h+sa)*sd,W,H);
            uint32_t s=agent_seed(base,i); float j=0.5f+0.5f*frand(s);
            if(fc>fl&&fc>fr){} else if(fl>fr) h-=turn*j; else if(fr>fl) h+=turn*j; else h+=(j-0.5f)*2.0f*turn;
            x=fmodf(x+fcos(h)*speed+W,W); y=fmodf(y+fsin(h)*speed+H,H);
            ax[i]=x; ay[i]=y; ah[i]=h;
        }
    });
}
void diffuse(float* t, float* tmp, int W, int H, float dec){
    for (int y=0;y<H;y++){ const float* r=t+(size_t)y*W; float* o=tmp+(size_t)y*W;
        for (int x=0;x<W;x++){ int xl=x?x-1:W-1, xr=x+1<W?x+1:0;
            o[x]=0.25f*r[xl]+0.5f*r[x]+0.25f*r[xr]; } }
    for (int y=0;y<H;y++){ int yu=y?y-1:H-1, yd=y+1<H?y+1:0;
        const float* U=tmp+(size_t)yu*W,*C=tmp+(size_t)y*W,*D=tmp+(size_t)yd*W;
        float* o=t+(size_t)y*W;
        for (int x=0;x<W;x++) o[x]=(0.25f*U[x]+0.5f*C[x]+0.25f*D[x])*dec; }
}
#endif

// deposit is a serial index-order scatter so the field sum is identical for any
// thread count; positions are already wrapped into range by the move.
void deposit(float* t, const float* ax, const float* ay, int W, int H, int N, float dep){
    for (int i=0;i<N;i++){ int xi=(int)ax[i], yi=(int)ay[i];
        if(xi>=0&&xi<W&&yi>=0&&yi<H) t[(size_t)yi*W+xi]+=dep; }
}

// counting-sort agents into row-major tile order (serial, deterministic). Tile
// coords are clamped because the position wrap can round to exactly x==W.
void bin_sort(float* ax, float* ay, float* ah, int W, int H, int N, int tile){
    int TW=(W+tile-1)/tile, TH=(H+tile-1)/tile, nt=TW*TH;
    std::vector<int> off(nt+1, 0), cur(nt);
    auto tid=[&](int i){ int xi=(int)ax[i]; xi=xi<0?0:(xi>=W?W-1:xi);
                         int yi=(int)ay[i]; yi=yi<0?0:(yi>=H?H-1:yi);
                         return (yi/tile)*TW + xi/tile; };
    for (int i=0;i<N;i++) off[tid(i)+1]++;
    for (int t=1;t<=nt;t++) off[t]+=off[t-1];
    for (int t=0;t<nt;t++) cur[t]=off[t];
    std::vector<float> sx(N), sy(N), sh(N);
    for (int i=0;i<N;i++){ int p=cur[tid(i)]++; sx[p]=ax[i]; sy[p]=ay[i]; sh[p]=ah[i]; }
    memcpy(ax, sx.data(), (size_t)N*4);
    memcpy(ay, sy.data(), (size_t)N*4);
    memcpy(ah, sh.data(), (size_t)N*4);
}

} // namespace

void physarum_step(torch::Tensor trail, torch::Tensor tmp,
                   torch::Tensor ax, torch::Tensor ay, torch::Tensor ah,
                   int64_t seed, double sense_dist, double sense_ang,
                   double turn, double speed, double deposit_, double decay){
    TORCH_CHECK(trail.dtype()==torch::kFloat32 && trail.is_contiguous(), "trail must be f32 contiguous");
    TORCH_CHECK(trail.dim()==2 && tmp.sizes()==trail.sizes(), "trail/tmp must be [H,W]");
    int H=(int)trail.size(0), W=(int)trail.size(1), N=(int)ax.size(0);
    uint32_t base = (uint32_t)(seed ^ 0x9e3779b9u); if(!base) base=1;
    float* t = trail.data_ptr<float>();
    // seed increments per step (base_seed + t), so this re-sorts every RESORT steps
    if ((uint64_t)seed % (uint64_t)RESORT == 0)
        bin_sort(ax.data_ptr<float>(), ay.data_ptr<float>(), ah.data_ptr<float>(), W,H,N, TILE);
    agents(t, ax.data_ptr<float>(), ay.data_ptr<float>(), ah.data_ptr<float>(),
           W,H,N, (float)sense_dist,(float)sense_ang,(float)turn,(float)speed, base);
    deposit(t, ax.data_ptr<float>(), ay.data_ptr<float>(), W,H,N, (float)deposit_);
    diffuse(t, tmp.data_ptr<float>(), W,H,(float)decay);
}