/*************************************************************************************

    Grid physics library, www.github.com/paboyle/Grid

    Source file: ./tests/debug/Test_ring_allreduce.cc

    Copyright (C) 2026

Author: Peter Boyle <pboyle@bnl.gov>

    This program is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 2 of the License, or
    (at your option) any later version.

    See the full license in the file "LICENSE" in the top level distribution
    directory
*************************************************************************************/
/*  END LEGAL */

//////////////////////////////////////////////////////////////////////////////
// Gate for RingAllReduce / CartesianRingAllReduce (P2P-only vector
// all-reduce) against GlobalSumVector (MPI_Allreduce):
//   T1 flat ring      == GlobalSumVector, RealD/ComplexD/ComplexF/RealF,
//                        n in {1, P-1, P, P+1, 1000003, 2^20} (n<P, n%P!=0)
//   T2 cartesian ring == GlobalSumVector, same sweep
//   T3 bitwise repeatable (deterministic order)
//   T4 timing at 16 MB, both rings vs GlobalSumVector
//   T5 CartesianRingAllGather bitwise == padded GlobalSumVector; timing at the dense-apply shape
//
//   mpirun -n 4 ./Test_ring_allreduce --grid 16.16.16.32 --mpi 1.1.2.2
// (2D process grid so the cartesian variant exercises more than one ring)
//////////////////////////////////////////////////////////////////////////////

#include <Grid/Grid.h>

using namespace Grid;

static int failures = 0;
static void Report(const std::string &name, bool pass, const std::string &detail="")
{
  std::cout << GridLogMessage << "  " << name << (pass ? "  PASS" : "  ** FAIL **");
  if ( detail.size() ) std::cout << "   " << detail;
  std::cout << std::endl;
  if ( !pass ) failures++;
}

template<class T> double Mag(const T &x){ return std::fabs((double)x); }
template<> double Mag<ComplexD>(const ComplexD &z){ return std::sqrt(z.real()*z.real()+z.imag()*z.imag()); }
template<> double Mag<ComplexF>(const ComplexF &z){ return std::sqrt((double)z.real()*z.real()+(double)z.imag()*z.imag()); }

template<class T> T Fill(uint64_t i, int rank){ return T( 0.5*std::sin(0.01*i + 0.7*rank) + 1.0e-3*rank ); }
template<> ComplexD Fill<ComplexD>(uint64_t i,int rank){ return ComplexD(0.5*std::sin(0.01*i+0.7*rank), 0.3*std::cos(0.02*i-0.1*rank)); }
template<> ComplexF Fill<ComplexF>(uint64_t i,int rank){ return ComplexF(0.5*std::sin(0.01*i+0.7*rank), 0.3*std::cos(0.02*i-0.1*rank)); }

template<class T>
void Check(const std::string &tname, GridCartesian *grid, double tol)
{
  int P  = grid->ProcessorCount();
  int me = grid->ThisRank();
  std::vector<uint64_t> sizes({1,(uint64_t)std::max(P-1,1),(uint64_t)P,(uint64_t)P+1,1000003,1<<20});
  for(uint64_t n : sizes){
    std::vector<T> h(n); for(uint64_t i=0;i<n;i++) h[i]=Fill<T>(i,me);
    std::vector<T> ref(h); grid->GlobalSumVector(&ref[0],(int)n);

    deviceVector<T> d(n);
    for(int variant=0;variant<2;variant++){
      acceleratorCopyToDevice(&h[0],&d[0],n*sizeof(T));
      if(variant==0) RingAllReduce(grid,&d[0],n);
      else           CartesianRingAllReduce(grid,&d[0],n);
      std::vector<T> out(n); acceleratorCopyFromDevice(&d[0],&out[0],n*sizeof(T));
      double worst=0.0, scale=0.0;
      for(uint64_t i=0;i<n;i++){ worst=std::max(worst,Mag<T>(out[i]-ref[i])); scale=std::max(scale,Mag<T>(ref[i])); }
      RealD w=worst; grid->GlobalMax(w);
      std::ostringstream os; os<<"n="<<n<<" worst abs "<<w<<" (scale "<<scale<<")";
      Report(std::string(variant?"T2 cartesian ":"T1 flat      ")+tname, w<tol*std::max(scale,1.0), os.str());
    }
  }
  // T3 determinism
  {
    uint64_t n=1<<18;
    std::vector<T> h(n); for(uint64_t i=0;i<n;i++) h[i]=Fill<T>(i,me);
    deviceVector<T> d(n); std::vector<T> a(n),b(n);
    acceleratorCopyToDevice(&h[0],&d[0],n*sizeof(T)); RingAllReduce(grid,&d[0],n); acceleratorCopyFromDevice(&d[0],&a[0],n*sizeof(T));
    acceleratorCopyToDevice(&h[0],&d[0],n*sizeof(T)); RingAllReduce(grid,&d[0],n); acceleratorCopyFromDevice(&d[0],&b[0],n*sizeof(T));
    RealD diff = (memcmp(&a[0],&b[0],n*sizeof(T))!=0) ? 1.0 : 0.0;
    grid->GlobalSum(diff);
    Report("T3 bitwise repeat "+tname, diff==0.0);
  }
}

int main(int argc, char **argv)
{
  Grid_init(&argc, &argv);
  GridCartesian *grid = SpaceTimeGrid::makeFourDimGrid(GridDefaultLatt(),
                             GridDefaultSimd(Nd, vComplexD::Nsimd()), GridDefaultMpi());
  std::cout << GridLogMessage << "Ring allreduce test: P=" << grid->ProcessorCount()
            << " processor grid " << grid->_processors << std::endl;

  Check<RealD>   ("RealD   ", grid, 1.0e-13);
  Check<ComplexD>("ComplexD", grid, 1.0e-13);
  Check<RealF>   ("RealF   ", grid, 1.0e-5);
  Check<ComplexF>("ComplexF", grid, 1.0e-5);

  // T5: CartesianRingAllGather delivers blocks in LEX-coordinate order; the
  // reference is the zero-padded GlobalSumVector in RANK order, compared
  // through the lex->rank table (bitwise: no arithmetic on either path).
  // On a machine where ranks are relabelled (Frontier OptimalCommunicator)
  // this is the test that catches a rank/coordinate confusion.
  {
    int P=grid->ProcessorCount(), me=grid->ThisRank(), mylex=CartesianLexIndex(grid);
    std::vector<int> l2r(P);
    for(int lp=0; lp<P; lp++){ Coordinate pc(grid->_ndimension); Lexicographic::CoorFromIndex(pc, lp, grid->_processors); l2r[lp]=grid->RankFromProcessorCoor(pc); }
    Report("T5 lex->rank table consistent for my rank", l2r[mylex]==me);
    int perm=0; for(int lp=0;lp<P;lp++) if(l2r[lp]!=lp) perm=1;
    std::cout << GridLogMessage << "  (ranks " << (perm ? "ARE" : "are not") << " permuted relative to lex coordinates on this run)" << std::endl;
    for(uint64_t chunk : std::vector<uint64_t>({1,3,64,1000,65537})){
      uint64_t n=chunk*P;
      std::vector<ComplexD> h(n,ComplexD(0.0,0.0)), ref;
      for(uint64_t i=0;i<chunk;i++) h[me*chunk+i]=Fill<ComplexD>(me*chunk+i,me);   // rank-order reference
      ref=h; grid->GlobalSumVector(&ref[0],(int)n);
      std::vector<ComplexD> hin(n,ComplexD(0.0,0.0));
      for(uint64_t i=0;i<chunk;i++) hin[mylex*chunk+i]=h[me*chunk+i];              // my slot is my lex index
      deviceVector<ComplexD> d(n); acceleratorCopyToDevice(&hin[0],&d[0],n*sizeof(ComplexD));
      CartesianRingAllGather(grid,&d[0],chunk);
      std::vector<ComplexD> out(n); acceleratorCopyFromDevice(&d[0],&out[0],n*sizeof(ComplexD));
      int bad=0; for(int lp=0;lp<P;lp++) if(memcmp(&out[lp*chunk],&ref[(uint64_t)l2r[lp]*chunk],chunk*sizeof(ComplexD))!=0) bad=1;
      RealD diff=bad; grid->GlobalSum(diff);
      Report("T5 CartesianRingAllGather (lex blocks) bitwise == rank-order padded GlobalSumVector, ComplexD chunk="+std::to_string(chunk), diff==0.0);
    }
    { uint64_t chunk=1001, n=chunk*P;
      std::vector<ComplexF> h(n,ComplexF(0.0,0.0)), ref;
      for(uint64_t i=0;i<chunk;i++) h[me*chunk+i]=Fill<ComplexF>(me*chunk+i,me);
      ref=h; grid->GlobalSumVector(&ref[0],(int)n);
      std::vector<ComplexF> hin(n,ComplexF(0.0,0.0)); for(uint64_t i=0;i<chunk;i++) hin[mylex*chunk+i]=h[me*chunk+i];
      deviceVector<ComplexF> d(n); acceleratorCopyToDevice(&hin[0],&d[0],n*sizeof(ComplexF));
      CartesianRingAllGather(grid,&d[0],chunk);
      std::vector<ComplexF> out(n); acceleratorCopyFromDevice(&d[0],&out[0],n*sizeof(ComplexF));
      int bad=0; for(int lp=0;lp<P;lp++) if(memcmp(&out[lp*chunk],&ref[(uint64_t)l2r[lp]*chunk],chunk*sizeof(ComplexF))!=0) bad=1;
      RealD diff=bad; grid->GlobalSum(diff);
      Report("T5 CartesianRingAllGather bitwise, ComplexF chunk=1001", diff==0.0);
    }
    // timing: the dense-apply shape, N=138240 x 4 rhs of ComplexF, chunk = N*4/P
    { uint64_t chunk=(uint64_t)138240*4/P, n=chunk*P;
      deviceVector<ComplexF> d(n); std::vector<ComplexF> h(n,ComplexF(1.0,0.0)); acceleratorCopyToDevice(&h[0],&d[0],n*sizeof(ComplexF));
      double t0=usecond(); CartesianRingAllGather(grid,&d[0],chunk); double t1=usecond();
      acceleratorCopyToDevice(&h[0],&d[0],n*sizeof(ComplexF));
      double t2=usecond(); CartesianRingAllReduce(grid,&d[0],n); double t3=usecond();
      std::cout << GridLogMessage << "T5 timing N=138240 x 4 ComplexF (" << n*8/1.0e6 << " MB): allgather " << (t1-t0)/1000. << " ms, cartesian allreduce " << (t3-t2)/1000. << " ms" << std::endl;
    }
  }

  // T6: orthogDim -- CartesianRingAllReduce(.., orthogDim=d) must equal the sum
  // over ranks sharing my coordinate in d.  Reference: for each value c of that
  // coordinate, GlobalSumVector of (my data if my coord==c else 0); keep c=mine.
  {
    int me=grid->ThisRank(); int Nd=grid->_ndimension;
    uint64_t n=4099;
    for(int d=-1; d<Nd; d++){
      std::vector<ComplexD> h(n); for(uint64_t i=0;i<n;i++) h[i]=Fill<ComplexD>(i,me);
      std::vector<ComplexD> ref(n,ComplexD(0.0,0.0));
      int Pd = (d<0) ? 1 : grid->_processors[d];
      int myc = (d<0) ? 0 : grid->_processor_coor[d];
      for(int c=0;c<Pd;c++){
        std::vector<ComplexD> m(n); for(uint64_t i=0;i<n;i++) m[i] = (myc==c) ? h[i] : ComplexD(0.0,0.0);
        grid->GlobalSumVector(&m[0],(int)n);
        if ( myc==c ) ref=m;
      }
      deviceVector<ComplexD> dv(n); acceleratorCopyToDevice(&h[0],&dv[0],n*sizeof(ComplexD));
      CartesianRingAllReduce(grid,&dv[0],n,d);
      std::vector<ComplexD> out(n); acceleratorCopyFromDevice(&dv[0],&out[0],n*sizeof(ComplexD));
      double worst=0.0; for(uint64_t i=0;i<n;i++) worst=std::max(worst,Mag<ComplexD>(out[i]-ref[i]));
      RealD w=worst; grid->GlobalMax(w);
      std::ostringstream os; os<<"orthogDim="<<d<<" (P_d="<<Pd<<") worst abs "<<w;
      Report("T6 CartesianRingAllReduce orthogDim == masked GlobalSumVector", w<1.0e-12, os.str());
    }
  }

  // T7: single-dimension gather.  Reference for my line along d: pad my chunk
  // at slot coor[d] of a P_d*chunk vector, mask by "other coordinates == this
  // line", GlobalSumVector; loop over all lines so every rank gets its own.
  {
    int me=grid->ThisRank(); int Nd=grid->_ndimension;
    for(int d=0; d<Nd; d++){
      int Pd=grid->_processors[d]; if ( Pd==1 ) continue;
      uint64_t chunk=1013, n=chunk*Pd;
      int myc=grid->_processor_coor[d];
      std::vector<ComplexD> pad(n,ComplexD(0.0,0.0));
      for(uint64_t i=0;i<chunk;i++) pad[myc*chunk+i]=Fill<ComplexD>(i,me);
      // line index = lex index of my coordinates with dim d removed
      auto lineIndex=[&](void){ int idx=0,stride=1; for(int e=0;e<Nd;e++){ if(e==d) continue; idx+=grid->_processor_coor[e]*stride; stride*=grid->_processors[e]; } return idx; };
      int nlines=grid->ProcessorCount()/Pd, myline=lineIndex();
      std::vector<ComplexD> ref(n);
      for(int L=0;L<nlines;L++){
        std::vector<ComplexD> m(n); for(uint64_t i=0;i<n;i++) m[i]=(myline==L)?pad[i]:ComplexD(0.0,0.0);
        grid->GlobalSumVector(&m[0],(int)n);
        if ( myline==L ) ref=m;
      }
      deviceVector<ComplexD> dv(n); acceleratorCopyToDevice(&pad[0],&dv[0],n*sizeof(ComplexD));
      CartesianRingAllGather(grid,&dv[0],chunk,d);
      std::vector<ComplexD> out(n); acceleratorCopyFromDevice(&dv[0],&out[0],n*sizeof(ComplexD));
      RealD diff=(memcmp(&out[0],&ref[0],n*sizeof(ComplexD))!=0)?1.0:0.0; grid->GlobalSum(diff);
      Report("T7 CartesianRingAllGather(dim="+std::to_string(d)+") bitwise == masked reference, P_d="+std::to_string(Pd), diff==0.0);
    }
  }

  // T8: single-dimension ALL-TO-ALL.  Fingerprint block j (destined for coord j)
  // on the rank at coord m with a tag encoding the (source,dest) coordinate pair,
  // src*Pd+dst; after the exchange block i (from coord i) must carry the tag the
  // coord-i rank prepared for me, i.e. Fill(e, i*Pd + m).  Analytic reference, so
  // a coordinate<->rank relabelling bug (wrong ShiftedRanks partner) shows up as a
  // mismatched source coordinate rather than silently passing.
  {
    int Nd=grid->_ndimension;
    for(int d=0; d<Nd; d++){
      int Pd=grid->_processors[d]; if ( Pd==1 ) continue;
      int m=grid->_processor_coor[d];
      uint64_t chunk=1009, n=chunk*(uint64_t)Pd;
      std::vector<ComplexD> sh(n), rh(n), ref(n);
      for(int j=0;j<Pd;j++) for(uint64_t e=0;e<chunk;e++) sh[j*chunk+e]=Fill<ComplexD>(e, m*Pd+j);
      for(int i=0;i<Pd;i++) for(uint64_t e=0;e<chunk;e++) ref[i*chunk+e]=Fill<ComplexD>(e, i*Pd+m);
      deviceVector<ComplexD> sd(n), rd(n);
      acceleratorCopyToDevice(&sh[0],&sd[0],n*sizeof(ComplexD));
      CartesianRingAllToAll(grid,&sd[0],&rd[0],chunk,d);
      acceleratorCopyFromDevice(&rd[0],&rh[0],n*sizeof(ComplexD));
      RealD diff=(memcmp(&rh[0],&ref[0],n*sizeof(ComplexD))!=0)?1.0:0.0; grid->GlobalSum(diff);
      Report("T8 CartesianRingAllToAll(dim="+std::to_string(d)+") bitwise == analytic reference, P_d="+std::to_string(Pd), diff==0.0);
    }
  }

  // T4 timing at 16 MB of ComplexF (the dense-apply size at 12 RHS is 13.3 MB)
  {
    uint64_t n = 2*1024*1024;
    deviceVector<ComplexF> d(n); std::vector<ComplexF> h(n,ComplexF(1.0,0.0));
    for(int rep=0;rep<2;rep++){
      acceleratorCopyToDevice(&h[0],&d[0],n*sizeof(ComplexF));
      double t0=usecond(); RingAllReduce(grid,&d[0],n);          double t1=usecond();
      acceleratorCopyToDevice(&h[0],&d[0],n*sizeof(ComplexF));
      double t2=usecond(); CartesianRingAllReduce(grid,&d[0],n); double t3=usecond();
      double t4=usecond(); grid->GlobalSumVector(&h[0],(int)n);   double t5=usecond();
      if(rep) std::cout << GridLogMessage << "T4 timing 16 MB ComplexF: flat ring " << (t1-t0)/1000.
                        << " ms, cartesian ring " << (t3-t2)/1000. << " ms, GlobalSumVector(host) " << (t5-t4)/1000. << " ms" << std::endl;
    }
  }

  std::cout << GridLogMessage << (failures ? "Test_ring_allreduce: FAILURES" : "Test_ring_allreduce: ALL PASS") << std::endl;
  Grid_finalize();
  return failures ? 1 : 0;
}
