// physarum-flow: the Tero adaptive-conductivity Physarum model, whose compute
// core is a weighted-grid-Laplacian solve. Given per-edge conductivities on a
// masked grid and a balanced source/sink injection, flow_cg solves L(cE,cS) p = b
// for the node pressures p by Jacobi-preconditioned conjugate gradients. One node
// is grounded with a diagonal penalty so the Laplacian is positive-definite (no
// singular nullspace to project out); the gauge choice does not affect edge flux.
// The caller iterates: solve -> flux Q = c*grad(p) -> reinforce conductive edges,
// which converges onto the shortest path (single source/sink) or an efficient
// network (multiple terminals). The 5-point weighted mat-vec is in gather form
// (each node reads its neighbours, no scatter) and vectorises with NEON; CG on a
// modest grid is latency-bound and sequential, so it is not threaded.
#include <torch/torch.h>
#include <cstdint>
#include <cmath>
#include <vector>

namespace {

const float GROUND = 1e3f;      // diagonal penalty pinning the ground node's gauge
                                // (kept modest: too large floors the f32 CG residual)

#if defined(__ARM_NEON)
#include <arm_neon.h>
#endif

// Ap = (L + GROUND e_g e_g^T) p. cE[y*W+x] is the edge (y,x)-(y,x+1); cS the edge
// (y,x)-(y+1,x); both zero on non-open edges, so wall nodes yield 0.
void matvec(float* Ap, const float* p, const float* cE, const float* cS,
            int H, int W, int g){
    for (int y=0; y<H; y++){
        const float* pr = p + (size_t)y*W;
        const float* pu = (y>0)   ? p + (size_t)(y-1)*W : nullptr;
        const float* pd = (y<H-1) ? p + (size_t)(y+1)*W : nullptr;
        const float* cEr = cE + (size_t)y*W;
        const float* cSr = cS + (size_t)y*W;
        const float* cSu = (y>0) ? cS + (size_t)(y-1)*W : nullptr;
        float* o = Ap + (size_t)y*W;
        int x = 0;
#if defined(__ARM_NEON)
        if (y>0 && y<H-1){                              // interior rows: all 4 edges exist
            { float pc=pr[0], s=cEr[0]*(pc-pr[1])+cSr[0]*(pc-pd[0])+cSu[0]*(pc-pu[0]); o[0]=s; }
            for (x=1; x+4<=W-1; x+=4){                  // vectorise interior cols [1, W-2]
                float32x4_t pc=vld1q_f32(pr+x);
                float32x4_t ce=vld1q_f32(cEr+x),  cw=vld1q_f32(cEr+x-1);
                float32x4_t cs=vld1q_f32(cSr+x),  cn=vld1q_f32(cSu+x);
                float32x4_t s=vmulq_f32(ce, vsubq_f32(pc, vld1q_f32(pr+x+1)));
                s=vmlaq_f32(s, cw, vsubq_f32(pc, vld1q_f32(pr+x-1)));
                s=vmlaq_f32(s, cs, vsubq_f32(pc, vld1q_f32(pd+x)));
                s=vmlaq_f32(s, cn, vsubq_f32(pc, vld1q_f32(pu+x)));
                vst1q_f32(o+x, s);
            }
            for (; x<W; x++){                           // scalar tail (incl. col W-1)
                float pc=pr[x], s=cEr[x-1]*(pc-pr[x-1])+cSr[x]*(pc-pd[x])+cSu[x]*(pc-pu[x]);
                if (x<W-1) s+=cEr[x]*(pc-pr[x+1]);
                o[x]=s;
            }
            continue;
        }
#endif
        for (x=0; x<W; x++){                            // boundary rows / scalar fallback
            float pc=pr[x], s=0.0f;
            if (x<W-1) s+=cEr[x]*(pc-pr[x+1]);
            if (x>0)   s+=cEr[x-1]*(pc-pr[x-1]);
            if (y<H-1) s+=cSr[x]*(pc-pd[x]);
            if (y>0)   s+=cSu[x]*(pc-pu[x]);
            o[x]=s;
        }
    }
    Ap[g] += GROUND * p[g];
}

#if defined(__ARM_NEON)
inline double hadd(float32x4_t v){ float32x2_t s=vadd_f32(vget_low_f32(v),vget_high_f32(v)); return (double)vget_lane_f32(vpadd_f32(s,s),0); }
#endif
double ddot(const float* a, const float* b, int N){
#if defined(__ARM_NEON)
    float32x4_t acc=vdupq_n_f32(0); int i=0;
    for (; i+4<=N; i+=4) acc=vmlaq_f32(acc, vld1q_f32(a+i), vld1q_f32(b+i));
    double s=hadd(acc); for (; i<N; i++) s+=(double)a[i]*b[i]; return s;
#else
    double s=0; for (int i=0;i<N;i++) s += (double)a[i]*(double)b[i]; return s;
#endif
}
// p += a*d ; r -= a*Ad ; return sum r^2
double axpy_rr(float* p, float* r, const float* d, const float* Ad, float a, int N){
#if defined(__ARM_NEON)
    float32x4_t va=vdupq_n_f32(a), acc=vdupq_n_f32(0); int i=0;
    for (; i+4<=N; i+=4){
        float32x4_t di=vld1q_f32(d+i), adi=vld1q_f32(Ad+i);
        float32x4_t pi=vmlaq_f32(vld1q_f32(p+i), va, di); vst1q_f32(p+i, pi);
        float32x4_t ri=vmlsq_f32(vld1q_f32(r+i), va, adi); vst1q_f32(r+i, ri);
        acc=vmlaq_f32(acc, ri, ri);
    }
    double s=hadd(acc);
    for (; i<N; i++){ p[i]+=a*d[i]; r[i]-=a*Ad[i]; s+=(double)r[i]*r[i]; }
    return s;
#else
    double s=0; for (int i=0;i<N;i++){ p[i]+=a*d[i]; r[i]-=a*Ad[i]; s+=(double)r[i]*r[i]; } return s;
#endif
}
void precond(float* z, const float* r, const float* diag, int N){   // z = r/diag (0 where diag~0)
#if defined(__ARM_NEON)
    float32x4_t eps=vdupq_n_f32(1e-20f); int i=0;
    for (; i+4<=N; i+=4){
        float32x4_t dg=vld1q_f32(diag+i);
        float32x4_t zz=vdivq_f32(vld1q_f32(r+i), dg);
        vst1q_f32(z+i, vbslq_f32(vcgtq_f32(dg,eps), zz, vdupq_n_f32(0)));
    }
    for (; i<N; i++) z[i]= diag[i]>1e-20f ? r[i]/diag[i] : 0.f;
#else
    for (int i=0;i<N;i++) z[i]= diag[i]>1e-20f ? r[i]/diag[i] : 0.f;
#endif
}
void xpby(float* d, const float* z, float beta, int N){             // d = z + beta*d
#if defined(__ARM_NEON)
    float32x4_t vb=vdupq_n_f32(beta); int i=0;
    for (; i+4<=N; i+=4) vst1q_f32(d+i, vmlaq_f32(vld1q_f32(z+i), vb, vld1q_f32(d+i)));
    for (; i<N; i++) d[i]=z[i]+beta*d[i];
#else
    for (int i=0;i<N;i++) d[i]=z[i]+beta*d[i];
#endif
}

} // namespace

// Solve L p = b with node `ground` pinned, by Jacobi-preconditioned CG; p is
// warm-started and overwritten with the solution.
void flow_cg(torch::Tensor cE, torch::Tensor cS, torch::Tensor b,
             torch::Tensor p, int64_t ground, int64_t max_iters, double tol){
    TORCH_CHECK(p.dtype()==torch::kFloat32 && p.is_contiguous(), "p must be f32 contiguous");
    int H=(int)p.size(0), W=(int)p.size(1), N=H*W, g=(int)ground;
    const float* cEp=cE.data_ptr<float>(); const float* cSp=cS.data_ptr<float>();
    const float* bp=b.data_ptr<float>();
    float* pp=p.data_ptr<float>();
    std::vector<float> diag(N,0.f), r(N), z(N), d(N), Ad(N);
    for (int y=0;y<H;y++) for (int x=0;x<W;x++){        // diag = sum of incident conductivities
        int i=y*W+x; float s=0.f;
        if (x<W-1) s+=cEp[i]; if (x>0) s+=cEp[i-1];
        if (y<H-1) s+=cSp[i]; if (y>0) s+=cSp[i-W];
        diag[i]=s;
    }
    diag[g]+=GROUND;

    matvec(Ad.data(), pp, cEp, cSp, H, W, g);
    for (int i=0;i<N;i++) r[i]=bp[i]-Ad[i];
    precond(z.data(), r.data(), diag.data(), N);
    for (int i=0;i<N;i++) d[i]=z[i];
    double rz = ddot(r.data(), z.data(), N), tol2 = tol*tol;
    for (int it=0; it<max_iters; it++){
        matvec(Ad.data(), d.data(), cEp, cSp, H, W, g);
        double dAd = ddot(d.data(), Ad.data(), N);
        if (dAd<=1e-30) break;
        float alpha=(float)(rz/dAd);
        double rr = axpy_rr(pp, r.data(), d.data(), Ad.data(), alpha, N);
        if (rr < tol2) break;
        precond(z.data(), r.data(), diag.data(), N);
        double rzn = ddot(r.data(), z.data(), N);
        float beta=(float)(rzn/rz);
        xpby(d.data(), z.data(), beta, N);
        rz=rzn;
    }
}