/* Experimental pixel-owned reduction; intentionally not a production backend.
 * Sharing this TU reuses the exact production cache indexing and row bounds. */
#include "../src/hip_psf.hip"
#include <omp.h>
#include <vector>
#include <algorithm>
#include <cstdlib>
#include <climits>
#include <atomic>
#include <thread>
#include <sys/resource.h>

struct TileTask { unsigned tile, first, count, partial; };
struct Prepared {
  int bx, by, support;
  size_t phase_base;
  double tx, ty, r, g, b;
};
__global__ static void prepare_tiles(const PsfCachedEvent *events, size_t count,
    int phases, int radius, double max_radius, Prepared *prepared, int *bounds) {
  const size_t thread=(size_t)blockIdx.x*blockDim.x+threadIdx.x;
  const size_t i=thread/32;const int lane=thread%32;
  if(i>=count)return;
  const PsfCachedEvent e=events[i];
  const int bx=(int)floor(e.x),by=(int)floor(e.y);
  const double fx=e.x-bx,fy=e.y-by,ax=fx*phases,ay=fy*phases;
  const int x0=(int)floor(ax),y0=(int)floor(ay);
  const int support=(int)fmin(ceil(e.support_radius),max_radius);
  if(!lane)prepared[i]={bx,by,support,weight_index(phases,radius,x0,y0,0,0),
                         ax-x0,ay-y0,e.color.r*e.flux,e.color.g*e.flux,e.color.b*e.flux};
  const int side=2*radius+1;
  for(int row=lane;row<side;row+=32) {
    int first=1,last=0;
    const int dy=row-radius;
    if(dy>=-support && dy<=support)row_range(e.support_radius,fx,fy,support,dy,&first,&last);
    bounds[2*(i*side+row)]=first;bounds[2*(i*side+row)+1]=last;
  }
}
__global__ static void tile_partial(const Prepared *events, const int *bounds, const unsigned *refs,
    const TileTask *tasks, int tile_size, int tiles_x, int width, int height,
    const float *weights, int phases, int radius, double *partial, double *hdr) {
  const TileTask task=tasks[blockIdx.x];
  const int local=threadIdx.x;
  const int px=(task.tile%tiles_x)*tile_size+local%tile_size;
  const int py=(task.tile/tiles_x)*tile_size+local/tile_size;
  const int side=2*radius+1;
  const size_t plane=(size_t)side*side;
  double r=0,g=0,b=0;
  if(px<width && py<height) for(unsigned k=0;k<task.count;++k) {
    const unsigned id=refs[task.first+k];
    const Prepared e=events[id];
    const int dx=px-e.bx,dy=py-e.by;
    if(dy < -e.support || dy > e.support) continue;
    const size_t row=2*((size_t)id*side+dy+radius);
    if(dx<bounds[row] || dx>bounds[row+1])continue;
    const size_t p=e.phase_base+(ptrdiff_t)dy*side+dx;
    const double w00=weights[p],w10=weights[p+plane];
    const double w01=weights[p+(phases+1)*plane],w11=weights[p+(phases+2)*plane];
    const double w=(1-e.ty)*((1-e.tx)*w00+e.tx*w10)+e.ty*((1-e.tx)*w01+e.tx*w11);
    r+=e.r*w;g+=e.g*w;b+=e.b*w;
  }
  if(task.partial==UINT_MAX) {
    if(px<width && py<height) {
      const size_t p=3*((size_t)py*width+px);
      hdr[p]+=r;hdr[p+1]+=g;hdr[p+2]+=b;
    }
  } else {
    const size_t p=3*((size_t)task.partial*tile_size*tile_size+local);
    partial[p]=r;partial[p+1]=g;partial[p+2]=b;
  }
}
__global__ static void tile_merge(const unsigned *starts, const TileTask *tasks, int tile_size,int tiles_x,
    int width,int height,const double *partial,double *hdr) {
  const unsigned tile=blockIdx.x; const int local=threadIdx.x;
  const int px=(tile%tiles_x)*tile_size+local%tile_size;
  const int py=(tile/tiles_x)*tile_size+local/tile_size;
  if(px>=width || py>=height || starts[tile+1]-starts[tile]<=1) return;
  double r=0,g=0,b=0;
  for(unsigned k=starts[tile];k<starts[tile+1];++k) {
    const size_t p=3*((size_t)tasks[k].partial*tile_size*tile_size+local);
    r+=partial[p];g+=partial[p+1];b+=partial[p+2];
  }
  const size_t p=3*((size_t)py*width+px);
  hdr[p]+=r;hdr[p+1]+=g;hdr[p+2]+=b;
}
static void check(hipError_t e) {if(e!=hipSuccess){fprintf(stderr,"HIP: %s\n",hipGetErrorString(e));exit(1);}}
static double sync_time(HipPsfSink *s,double start) {check(hipStreamSynchronize(s->stream));return omp_get_wtime()-start;}
static size_t contributions(const std::vector<PsfCachedEvent>& events,int width,int height,const PsfKernelCache& cache) {
  size_t n=0;
  for(const auto&e:events) {
    const int bx=(int)floor(e.x),by=(int)floor(e.y);
    const double fx=e.x-bx,fy=e.y-by;
    const int support=(int)fmin(ceil(e.support_radius),cache.max_radius_pixels);
    for(int dy=-support;dy<=support;++dy) {
      if(by+dy<0 || by+dy>=height) continue;
      const double left=e.support_radius*e.support_radius-(dy+0.5-fy)*(dy+0.5-fy);
      if(left<0) continue;
      int lo=std::max(-support,(int)ceil(fx-0.5-sqrt(left)));
      int hi=std::min(support,(int)floor(fx-0.5+sqrt(left)));
      lo=std::max(lo,-bx);hi=std::min(hi,width-bx-1);
      if(hi>=lo)n+=(size_t)(hi-lo+1);
    }
  }
  return n;
}
int main(int argc,char **argv) {
  if(argc<4 || argc>6) {puts("Usage: replay_psf INPUT.events EVENTS(1..65536) atomic|16|32|adaptive|production|production-prepared|production-parallel|production-busy [disperse|mixed] [rounds=1..512]\nExternal timeout <=45s; sequential processes only. production-busy adds 15 CPU spin workers to test host contention. Repeated rounds compare against a scaled CPU double reference.");return 2;}
  char *end=nullptr;const long requested=strtol(argv[2],&end,10);
  if(!*argv[2] || *end || requested<1 || requested>65536) return 2;
  const bool adaptive=!strcmp(argv[3],"adaptive");
  const bool parallel_production=!strcmp(argv[3],"production-parallel");
  const bool busy_production=!strcmp(argv[3],"production-busy");
  const bool prepared_production=!strcmp(argv[3],"production-prepared");
  const bool production=!strcmp(argv[3],"production") || prepared_production || parallel_production || busy_production;
  const int requested_tile=(!strcmp(argv[3],"atomic") || production) ? 0 : !strcmp(argv[3],"16") ? 16 :
                           !strcmp(argv[3],"32") ? 32 : adaptive ? 16 : -1;
  const bool distribution_arg=argc>=5 && (!strcmp(argv[4],"disperse") || !strcmp(argv[4],"mixed"));
  if(requested_tile<0 || (argc==6 && !distribution_arg))return 2;
  const char *rounds_arg=argc==6 ? argv[5] : argc==5 && !distribution_arg ? argv[4] : nullptr;
  long rounds=1;
  if(rounds_arg) {
    char *rounds_end=nullptr;
    rounds=strtol(rounds_arg,&rounds_end,10);
    if(!*rounds_arg || *rounds_end || rounds<1 || rounds>512)return 2;
  }
  if(parallel_production && (rounds<16 || rounds%16 || distribution_arg)) return 2;
  FILE *f=fopen(argv[1],"r");if(!f){perror(argv[1]);return 2;}
  int width,height;size_t count;PointSpreadFunction psf;double tail;
  if(fscanf(f,"PSFEVENTS1 %d %d %zu %lf %lf %lf",&width,&height,&count,&psf.fwhm_pixels,&psf.moffat_beta,&tail)!=6 ||
      width<1 || height<1 || width>3840 || height>2160 || count>65536 || count<(size_t)requested ||
      !std::isfinite(psf.fwhm_pixels) || psf.fwhm_pixels<=0 || psf.fwhm_pixels>2.7 ||
      !std::isfinite(psf.moffat_beta) || psf.moffat_beta<4.5 || !std::isfinite(tail) || tail<1e-8 || tail>=1) return 2;
  PsfKernelCache cache={};if(psf_kernel_cache_init(&cache,&psf,tail))return 1;
  psf_kernel_cache_report_ready(&cache,stdout);
  std::vector<PsfCachedEvent> events(requested);
  for(auto &e:events) {
    if(fscanf(f,"%lf %lf %lf %lf %lf %lf %lf",&e.x,&e.y,&e.color.r,&e.color.g,&e.color.b,&e.flux,&e.support_radius)!=7) return 2;
    const double fields[]={e.x,e.y,e.color.r,e.color.g,e.color.b,e.flux,e.support_radius};
    for(double v:fields)if(!std::isfinite(v))return 2;
    if(fabs(e.x)>width+cache.max_radius_pixels || fabs(e.y)>height+cache.max_radius_pixels || e.support_radius<0 || e.support_radius>1e6)return 2;
  }
  fclose(f);
  // Wing-clipped events retain their physical radius; traversal alone is cache-limited.
  const size_t original_contributions=contributions(events,width,height,cache);
  const bool mixed=distribution_arg && !strcmp(argv[4],"mixed");
  if(distribution_arg && !mixed) {
    unsigned state=12345;
    for(auto&e:events) {
      // Preserve fractional phase and clipping: move only fully interior events.
      const int s=(int)ceil(e.support_radius)+1;
      if(e.x<s || e.y<s || e.x>=width-s || e.y>=height-s)continue;
      state=1664525u*state+1013904223u; e.x=s+state%(width-2*s)+(e.x-floor(e.x));
      state=1664525u*state+1013904223u; e.y=s+state%(height-2*s)+(e.y-floor(e.y));
    }
  }
  const size_t work=contributions(events,width,height,cache);
  if(work!=original_contributions)return 1;
  const size_t values=(size_t)width*height*3;
  std::vector<double> cpu(values,0),gpu(values,0);
  const LinearRgb direct_color={0.7,0.2,0.5};
  auto direct=[&](double *hdr){splat_moffat_direct(hdr,width,height,12.25,20.75,direct_color,0.2,&psf,tail,0);};
  double start=omp_get_wtime();
  for(size_t i=0;i<events.size();++i) {
    splat_prepared_cached_event(cpu.data(),width,height,&events[i],&cache);
    if(mixed && i+1==std::min(events.size(),(size_t)16384))direct(cpu.data());
  }
  const double cpu_time=omp_get_wtime()-start;
  char message[256]={};HipPsfSink*s=nullptr;
  start=omp_get_wtime();if(hip_psf_sink_create(&s,width,height,&cache,16384,message,sizeof message)){fprintf(stderr,"%s\n",message);return 1;}
  const double create=omp_get_wtime()-start;
  hipDeviceProp_t prop;check(hipGetDeviceProperties(&prop,0));
  printf("device=%s arch=%s events=%zu rounds=%ld frame=%dx%d mode=%s distribution=%s contributions_per_round=%zu CPU_reference=%.9f create=%.9f\n",prop.name,prop.gcnArchName,events.size(),rounds,width,height,argv[3],mixed?"fixture-mixed-direct":distribution_arg?"synthetic-phase-preserving-dispersal":"captured",work,cpu_time,create);fflush(stdout);
  double select=0,bin=0,upload=0,prepare=0,kernel=0,merge=0,clear=0;
  size_t peak_scratch=0,ref_total=0,task_total=0,center_tiles=0,tile_chunks=0,atomic_chunks=0;
  Prepared *prepared=nullptr;int *bounds=nullptr;
  PsfCachedEvent *tile_events=nullptr;
  unsigned *drefs=nullptr,*dstarts=nullptr;TileTask *dtasks=nullptr;double *partial=nullptr;
  auto direct_boundary=[&]() {
    if(hip_psf_sink_finish(s,gpu.data(),message,sizeof message)) {fprintf(stderr,"%s\n",message);exit(1);}
    direct(gpu.data());
    if(hip_psf_sink_load_hdr(s,gpu.data(),message,sizeof message)) {fprintf(stderr,"%s\n",message);exit(1);}
  };
  const double replay_start=omp_get_wtime();
  std::atomic<bool> stop_busy(false);
  std::vector<std::thread> busy_workers;
  if(busy_production) for(unsigned worker=0;worker<15;++worker)
    busy_workers.emplace_back([&,worker] {
      volatile uint64_t state=worker+1;
      while(!stop_busy.load(std::memory_order_relaxed))
        for(int i=0;i<4096;++i)state=state*1664525u+1013904223u;
    });
  HipPsfPreparedChunk *serial_chunk=nullptr;
  if(prepared_production && hip_psf_prepared_chunk_create(
      &serial_chunk,s,message,sizeof message)) {fprintf(stderr,"%s\n",message);return 1;}
  if(requested_tile) {
    // Fixed hard bounds; no full-frame event list or per-chunk allocation on GPU.
    check(hipMalloc(&tile_events,16384*sizeof(PsfCachedEvent)));
    check(hipMalloc(&prepared,16384*sizeof(Prepared)));
    check(hipMalloc(&bounds,16384*(size_t)(2*cache.radius_pixels+1)*2*sizeof(int)));
    check(hipMalloc(&drefs,32*1024*1024));check(hipMalloc(&dtasks,2*1024*1024));
    check(hipMalloc(&dstarts,131072));check(hipMalloc(&partial,128*1024*1024));
  }
  if(parallel_production) {
    omp_lock_t lock; omp_init_lock(&lock);
    int failed=0;
#pragma omp parallel num_threads(16) reduction(+:failed)
    {
      char local_message[256]={};
      HipPsfPreparedChunk *chunk=nullptr;
      if(hip_psf_prepared_chunk_create(&chunk,s,local_message,sizeof local_message)) {
        fprintf(stderr,"parallel chunk create: %s\n",local_message);failed=1;
      }
      for(long round=0;round<rounds/omp_get_num_threads() && !failed;++round)
        for(size_t offset=0;offset<events.size();offset+=16384) {
          const size_t n=std::min((size_t)16384,events.size()-offset);
          if(hip_psf_prepared_chunk_prepare(chunk,events.data()+offset,n,
                                             local_message,sizeof local_message)) {
            fprintf(stderr,"parallel prepare: %s\n",local_message);failed=1;break;
          }
          omp_set_lock(&lock);
          const int result=hip_psf_sink_submit_prepared(s,chunk,events.data()+offset,n,
                                                        local_message,sizeof local_message);
          omp_unset_lock(&lock);
          if(result) {fprintf(stderr,"parallel submit: %s\n",local_message);failed=1;break;}
        }
      hip_psf_prepared_chunk_destroy(chunk);
    }
    omp_destroy_lock(&lock);
    if(failed)return 1;
  } else for(long round=0;round<rounds;++round) for(size_t offset=0;offset<events.size();offset+=16384) {
    const size_t n=std::min((size_t)16384,events.size()-offset);
    int tile=requested_tile;
    if(adaptive) {
      start=omp_get_wtime();
      const int selector_tile=32,nx=(width+selector_tile-1)/selector_tile;
      const int ny=(height+selector_tile-1)/selector_tile;
      std::vector<unsigned char> occupied((size_t)nx*ny,0);
      size_t count_occupied=0;
      for(size_t i=0;i<n;++i) {
        const int x=(int)floor(events[offset+i].x),y=(int)floor(events[offset+i].y);
        if(x<0 || x>=width || y<0 || y>=height)continue;
        unsigned char &mark=occupied[(size_t)(y/selector_tile)*nx+x/selector_tile];
        if(!mark) {mark=1;++count_occupied;}
      }
      /* This threshold is deliberately experimental. The captured dense and
       * lensed chunks have >=45 events per occupied center tile, while the
       * equal-work dispersed control has about 2.4. Small chunks retain the
       * lower-overhead production atomic path. */
      tile=n>=8192 && count_occupied && n/count_occupied>=32 ? 16 : 0;
      center_tiles+=count_occupied;select+=omp_get_wtime()-start;
    }
    if(!tile) {
      if(!production)++atomic_chunks;
      if(prepared_production && hip_psf_prepared_chunk_prepare(
          serial_chunk,events.data()+offset,n,message,sizeof message)) {
        fprintf(stderr,"%s\n",message);return 1;
      }
      const int submit_result=prepared_production
          ? hip_psf_sink_submit_prepared(s,serial_chunk,events.data()+offset,n,
                                         message,sizeof message)
          : hip_psf_sink_submit(s,events.data()+offset,n,message,sizeof message);
      if(submit_result) {fprintf(stderr,"%s\n",message);return 1;}
      if(mixed && offset==0)direct_boundary();continue;
    }
    ++tile_chunks;
    start=omp_get_wtime();
    const int nx=(width+tile-1)/tile,ny=(height+tile-1)/tile;
    std::vector<std::vector<unsigned>> lists(nx*ny);
    size_t refs_count=0;
    for(size_t i=0;i<n;++i) {
      const auto&e=events[offset+i];const int support=(int)fmin(ceil(e.support_radius),cache.max_radius_pixels);
      const int bx=(int)floor(e.x),by=(int)floor(e.y);
      const int x0=std::max(0,bx-support),x1=std::min(width-1,bx+support);
      const int y0=std::max(0,by-support),y1=std::min(height-1,by+support);
      if(x0>x1 || y0>y1)continue;
      for(int y=y0/tile;y<=y1/tile;++y)for(int x=x0/tile;x<=x1/tile;++x) {
        if(++refs_count>32*1024*1024/sizeof(unsigned)){fputs("reference limit exceeded\n",stderr);return 2;}
        lists[y*nx+x].push_back((unsigned)i);
      }
    }
    std::vector<unsigned> refs,starts(nx*ny+1);refs.reserve(refs_count);
    std::vector<TileTask> tasks;size_t partial_count=0;
    for(unsigned t=0;t<lists.size();++t) {
      starts[t]=tasks.size();
      for(size_t k=0;k<lists[t].size();k+=256) {
        const unsigned length=std::min((size_t)256,lists[t].size()-k);
        const unsigned part=lists[t].size()<=256 ? UINT_MAX : (unsigned)partial_count++;
        if(partial_count*(size_t)tile*tile*3*sizeof(double)>128*1024*1024 ||
           (tasks.size()+1)*sizeof(TileTask)>2*1024*1024) {
          fputs("partial/task limit exceeded\n",stderr);return 2;
        }
        tasks.push_back({t,(unsigned)refs.size(),length,part});
        refs.insert(refs.end(),lists[t].begin()+k,lists[t].begin()+k+length);
      }
    }
    starts.back()=tasks.size();
    const size_t bytes=partial_count*(size_t)tile*tile*3*sizeof(double);
    if(bytes>128*1024*1024 || tasks.size()*sizeof(TileTask)>2*1024*1024 || starts.size()*sizeof(unsigned)>131072){fputs("partial/task limit exceeded\n",stderr);return 2;}
    peak_scratch=std::max(peak_scratch,bytes+refs.size()*sizeof(unsigned)+tasks.size()*sizeof(TileTask)+starts.size()*sizeof(unsigned)+n*(sizeof(Prepared)+(size_t)(2*cache.radius_pixels+1)*2*sizeof(int)));
    ref_total+=refs.size();task_total+=tasks.size();bin+=omp_get_wtime()-start;
    start=omp_get_wtime();
    check(hipMemcpyAsync(tile_events,events.data()+offset,n*sizeof(PsfCachedEvent),hipMemcpyHostToDevice,s->stream));
    if(!refs.empty())check(hipMemcpyAsync(drefs,refs.data(),refs.size()*sizeof(unsigned),hipMemcpyHostToDevice,s->stream));
    if(!tasks.empty())check(hipMemcpyAsync(dtasks,tasks.data(),tasks.size()*sizeof(TileTask),hipMemcpyHostToDevice,s->stream));
    check(hipMemcpyAsync(dstarts,starts.data(),starts.size()*sizeof(unsigned),hipMemcpyHostToDevice,s->stream));
    upload+=sync_time(s,start);
    start=omp_get_wtime();
    hipLaunchKernelGGL(prepare_tiles,dim3((n*32+127)/128),dim3(128),0,s->stream,tile_events,n,s->phase_resolution,s->radius_pixels,s->max_radius_pixels,prepared,bounds);
    check(hipGetLastError());prepare+=sync_time(s,start);
    start=omp_get_wtime();
    if(!tasks.empty())hipLaunchKernelGGL(tile_partial,dim3(tasks.size()),dim3(tile*tile),0,s->stream,prepared,bounds,drefs,dtasks,tile,nx,width,height,s->weights,s->phase_resolution,s->radius_pixels,partial,s->hdr);
    check(hipGetLastError());kernel+=sync_time(s,start);
    start=omp_get_wtime();hipLaunchKernelGGL(tile_merge,dim3(nx*ny),dim3(tile*tile),0,s->stream,dstarts,dtasks,tile,nx,width,height,partial,s->hdr);
    check(hipGetLastError());merge+=sync_time(s,start);
    if(mixed && offset==0)direct_boundary();
    // Every partial pixel is written by exactly one thread: no memset required.
  }
  if(hip_psf_sink_finish(s,gpu.data(),message,sizeof message)){fprintf(stderr,"%s\n",message);return 1;}
  stop_busy.store(true,std::memory_order_relaxed);
  for(auto &worker:busy_workers)worker.join();
  hip_psf_prepared_chunk_destroy(serial_chunk);
  if(requested_tile) {start=omp_get_wtime();check(hipFree(tile_events));check(hipFree(prepared));check(hipFree(bounds));check(hipFree(drefs));check(hipFree(dtasks));check(hipFree(dstarts));check(hipFree(partial));clear=omp_get_wtime()-start;}
  const double wall=omp_get_wtime()-replay_start;
  HipPsfTiming timing={};hip_psf_sink_get_timing(s,&timing);
  upload+=timing.upload_seconds;kernel+=timing.kernel_seconds;
  if(production) {
    select+=timing.selection_seconds;bin+=timing.bin_seconds;
    atomic_chunks+=timing.atomic_batch_count;tile_chunks+=timing.tile_batch_count;
  }
  double max_abs=0,max_rel=0;long double sums[3]={},diffs[3]={};bool pass=true;
  for(size_t i=0;i<values;++i) {
    if(!std::isfinite(cpu[i]) || !std::isfinite(gpu[i]))pass=false;
    const double reference=cpu[i]*rounds;
    const double d=fabs(gpu[i]-reference);max_abs=fmax(max_abs,d);
    if(reference!=0)max_rel=fmax(max_rel,d/fabs(reference));
    sums[i%3]+=reference;diffs[i%3]+=(long double)gpu[i]-reference;
  }
  long double ysum=.2126L*sums[0]+.7152L*sums[1]+.0722L*sums[2],ydiff=.2126L*diffs[0]+.7152L*diffs[1]+.0722L*diffs[2];
  struct rusage usage;getrusage(RUSAGE_SELF,&usage);
  printf("complete_replay=%.9f replay_wall=%.9f select=%.9f bin=%.9f upload=%.9f prepare=%.9f accumulation=%.9f merge=%.9f download=%.9f cleanup=%.9f events_s=%.3f contributions_s=%.3f center_tiles=%zu tile_chunks=%zu atomic_chunks=%zu refs=%zu tasks=%zu scratch_used_peak=%zu scratch_device_reserved=%zu RSS_KiB=%ld max_abs=%.17g max_rel=%.17g flux_R=%.17Lg flux_G=%.17Lg flux_B=%.17Lg flux_Y=%.17Lg validation=%s\n",create+wall,wall,select,bin,upload,prepare,kernel,merge,timing.download_seconds,clear,events.size()*rounds/wall,work*rounds/wall,center_tiles,tile_chunks,atomic_chunks,ref_total,task_total,peak_scratch,requested_tile?163*1024*1024+131072+16384*(sizeof(Prepared)+(size_t)(2*cache.radius_pixels+1)*2*sizeof(int)):0,usage.ru_maxrss,max_abs,max_rel,sums[0]?diffs[0]/sums[0]:0,sums[1]?diffs[1]/sums[1]:0,sums[2]?diffs[2]/sums[2]:0,ysum?ydiff/ysum:0,"scaled-cpu-double-reference");
  for(int c=0;c<3;++c)if(fabsl(diffs[c])>1e-10L*fmaxl(fabsl(sums[c]),1e-30L))pass=false;
  if(fabsl(ydiff)>1e-10L*fmaxl(fabsl(ysum),1e-30L))pass=false;
  hip_psf_sink_destroy(s);psf_kernel_cache_destroy(&cache);
  return pass && max_rel<1e-10 && (rounds>1 || max_abs<1e-10) ? 0 : 1;
}