// ============================================================================
// HDGL Sieve Host — sieve → warp-LL → reward feedback loop
// ============================================================================
//
// Orchestrates the full prime-detection pipeline:
//   1. Sieve kernel: N oscillator slots each tracking a Mersenne exponent.
//      GRA plasticity drives resonance toward candidate primes.
//   2. Warp-LL engine: NTT-based Lucas-Lehmer verification.
//      Candidates from the sieve are routed here for real LL residue.
//   3. Reward feedback: LL residue is injected back into the sieve's
//      reward_accum, closing the RL loop.
//   4. Critic: trains online on (features, LL-residue) pairs.
//
// Mathematical scope:
//   LIMBS=64 → max exponent p ≈ 3904 bits.
//   Sieve targets 8 confirmed small Mersenne prime exponents:
//     {61, 89, 107, 127, 521, 607, 1279, 3217}
//   For p=61, the NTT modulus IS M61, so the LL test is exact.
//   For other p, the residue is a proximity signal the critic learns from.
//
// Compile (add to existing build):
//   nvcc -O3 -arch=sm_75 hdgl_analog_v33.cu hdgl_warp_ll_v33.cu \
//        hdgl_sieve_v34.cu hdgl_critic_v33.c hdgl_sieve_host.c \
//        -o hdgl_sieve -Xlinker /NODEFAULTLIB:libcmt
// ============================================================================

#include <cuda_runtime.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <signal.h>

#ifdef _WIN32
#define WIN32_LEAN_AND_MEAN
#include <windows.h>
typedef struct { LARGE_INTEGER v; } sh_timespec_t;
static LARGE_INTEGER g_qpf;
static void sh_clock_init(void) { QueryPerformanceFrequency(&g_qpf); }
static void sh_gettime(sh_timespec_t *t) { QueryPerformanceCounter(&t->v); }
static double sh_elapsed(const sh_timespec_t *a, const sh_timespec_t *b) {
    return (double)(b->v.QuadPart - a->v.QuadPart) / (double)g_qpf.QuadPart;
}
#else
#include <time.h>
typedef struct timespec sh_timespec_t;
static void sh_clock_init(void) {}
static void sh_gettime(sh_timespec_t *t) { clock_gettime(CLOCK_MONOTONIC, t); }
static double sh_elapsed(const sh_timespec_t *a, const sh_timespec_t *b) {
    return (b->tv_sec - a->tv_sec) + (b->tv_nsec - a->tv_nsec) * 1e-9;
}
#endif

#include "hdgl_critic_v33.h"

// ============================================================================
// Shared type definitions — must match all TUs exactly
// ============================================================================

typedef struct {
    int   slot_idx;
    float score;
    float ll_seed_f;
    float r_harmonic;   // Mersenne exponent p for sieve candidates
    float phase;
    float coherence;
    float amp;
} Candidate;

typedef struct { uint64_t S[64]; uint64_t C[64]; } CSSState;

typedef struct {
    uint32_t  *d_assigned_p, *d_sieve_state;
    float     *d_r_harmonic, *d_phase, *d_reward_accum;
    uint32_t  *d_prime_found;
    int       *d_prime_count;
    Candidate *d_cand_buf;
    int       *d_cand_count;
    CSSState  *d_css;
    float     *d_ll_residue;
    int8_t    *d_ll_verified;
    uint32_t  *h_prime_found;
    int       *h_prime_count;
    Candidate *h_cand_buf;
    int       *h_cand_count;
    int N, S;
    float dt;
} SieveHostState;

#define CRITIC_W_TOTAL 57
#define MAX_SIEVE_CANDS 256

// ============================================================================
// Forward declarations
// ============================================================================

extern "C" {
int  hdgl_sieve_alloc(SieveHostState *st, int N, int S, float dt);
void hdgl_sieve_free(SieveHostState *st);
void hdgl_sieve_step(const SieveHostState *st, cudaStream_t s);
int  hdgl_sieve_harvest(SieveHostState *st, uint32_t *out, int max, cudaStream_t s);
int  hdgl_sieve_harvest_verified(SieveHostState *st, uint32_t *out, int max,
                                  int ll_iters, float ll_eps, cudaStream_t s);
void hdgl_warp_ll_v33_alloc(CSSState **out);
void hdgl_warp_ll_v33_free(CSSState *p);
void hdgl_warp_ll_v33_seed(CSSState *d, int n, cudaStream_t s);
void hdgl_warp_ll_v33_launch(const Candidate *dc, int nc, int pb, int it,
                              CSSState *css, float *res, int8_t *ver,
                              float eps, cudaStream_t s);
}

// ============================================================================
// Configuration
// ============================================================================

#define SIEVE_N          16384    // oscillator slots (8 bands × 2048 each)
#define SIEVE_S          128      // lattice row width
#define SIEVE_DT         0.001f
#define STEPS_PER_CYCLE  32       // sieve kernel calls per harvest
#define TOTAL_CYCLES     5000
#define LL_ITERS_PER_CYCLE 128    // covers p up to ~130 in one shot (p=61 needs 59)
#define LL_RESIDUE_EPS   1e-4f

// Reward injection: slots where LL fires get reward boost to d_reward_accum
#define REWARD_LL_CONFIRM  2.0f
#define REWARD_LL_CLOSE    0.5f   // for low-residue but unconfirmed

// Critic features for sieve: (residue, p_norm, coherence, amp, acc_norm)
#define CRITIC_IN  5

// ============================================================================
// CUDA error check
// ============================================================================

#define CUDA_CHECK(c) do {                                               \
    cudaError_t _e = (c);                                                \
    if (_e != cudaSuccess) {                                             \
        fprintf(stderr, "[sieve_host] CUDA error %s:%d: %s\n",          \
                __FILE__, __LINE__, cudaGetErrorString(_e));             \
        exit(EXIT_FAILURE);                                              \
    }                                                                    \
} while (0)

// ============================================================================
// Global state
// ============================================================================

static volatile int g_running = 1;
static void sig_handler(int s) { (void)s; g_running = 0; }

// Confirmed primes found this session
static uint32_t g_confirmed[1024];
static int g_n_confirmed = 0;

// Stats
static long long g_total_cands    = 0;
static long long g_total_verified = 0;
static long long g_total_steps    = 0;

// ============================================================================
// Reward injection kernel — boost reward_accum for LL-verified slots
// ============================================================================

// d_ll_verified is indexed by cand_idx (0..n_cands-1).
// d_cands[i].slot_idx identifies the sieve slot to receive the reward.
__global__ void sieve_reward_inject_kernel(
    float          *d_reward_accum,
    const Candidate *d_cands,
    const int8_t   *d_ll_verified,
    int             n_cands)
{
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n_cands) return;
    if (d_ll_verified[i] == (int8_t)1) {
        int slot = d_cands[i].slot_idx;
        atomicAdd(&d_reward_accum[slot], REWARD_LL_CONFIRM);
    }
}

// ============================================================================
// main
// ============================================================================

int main(int argc, char **argv)
{
    (void)argc; (void)argv;
    signal(SIGINT,  sig_handler);
    signal(SIGTERM, sig_handler);
    sh_clock_init();

    printf("====================================================================\n");
    printf("  HDGL Sieve Host — 8-band small-prime search + warp-LL verify\n");
    printf("  Targets: M61 M89 M107 M127 M521 M607 M1279 M3217\n");
    printf("  (max verifiable exponent: LIMBS=64 × 61 bits = 3904 bits)\n");
    printf("  Press Ctrl-C to stop.\n");
    printf("====================================================================\n\n");

    // Init critic
    critic_init();
    critic_load_weights(NULL);   // loads hdgl_critic.bin if present

    // Allocate sieve
    SieveHostState sieve;
    memset(&sieve, 0, sizeof(sieve));
    if (hdgl_sieve_alloc(&sieve, SIEVE_N, SIEVE_S, SIEVE_DT) != 0) {
        fprintf(stderr, "Failed to allocate sieve\n");
        return 1;
    }
    printf("Sieve allocated: N=%d, S=%d, dt=%.4f\n",
           SIEVE_N, SIEVE_S, SIEVE_DT);
    printf("Bands: 8 × %d slots each\n\n", SIEVE_N / 8);

    // Create streams
    cudaStream_t stream_sieve, stream_ll;
    CUDA_CHECK(cudaStreamCreate(&stream_sieve));
    CUDA_CHECK(cudaStreamCreate(&stream_ll));

    // Per-cycle buffers
    uint32_t found_proxy[256];
    uint32_t found_verified[256];

    sh_timespec_t t0, t1;
    sh_gettime(&t0);

    long long last_report_steps = 0;

    for (int cycle = 0; cycle < TOTAL_CYCLES && g_running; cycle++) {

        // ---- Step sieve ----
        for (int s = 0; s < STEPS_PER_CYCLE; s++) {
            hdgl_sieve_step(&sieve, stream_sieve);
        }
        CUDA_CHECK(cudaStreamSynchronize(stream_sieve));
        g_total_steps += STEPS_PER_CYCLE;

        // ---- Proxy harvest (quick scan of prime_found ring) ----
        int n_proxy = hdgl_sieve_harvest(&sieve, found_proxy, 256, stream_sieve);
        (void)n_proxy;   // used only for printing inside harvest

        // ---- Warp-LL verified harvest ----
        int n_verified = hdgl_sieve_harvest_verified(
            &sieve, found_verified, 256,
            LL_ITERS_PER_CYCLE, LL_RESIDUE_EPS,
            stream_ll);

        g_total_cands    += MAX_SIEVE_CANDS;   // approximate — cand_count resets
        g_total_verified += n_verified;

        // Store confirmed primes
        for (int i = 0; i < n_verified && g_n_confirmed < 1024; i++) {
            g_confirmed[g_n_confirmed++] = found_verified[i];
        }

        // ---- Reward feedback: inject LL signal back into sieve ----
        // d_ll_verified[cand_idx] == 1 for each confirmed candidate.
        // Unprocessed slots remain 0 (from cudaMemset), so MAX_SIEVE_CANDS is a safe bound.
        {
            int blk = 256;
            int grd = (MAX_SIEVE_CANDS + blk - 1) / blk;
            sieve_reward_inject_kernel<<<grd, blk, 0, stream_ll>>>(
                sieve.d_reward_accum,
                sieve.d_cand_buf,
                sieve.d_ll_verified,
                MAX_SIEVE_CANDS);
            CUDA_CHECK(cudaStreamSynchronize(stream_ll));
        }

        // ---- Critic training: sample from candidates, train on LL residue ----
        // Copy candidate buffer to host for critic observation
        int nc = 0;
        CUDA_CHECK(cudaMemcpy(&nc, sieve.d_cand_count,
                              sizeof(int), cudaMemcpyDeviceToHost));
        // (cand_count was reset inside harvest_verified, so nc=0 here — we
        // use the last batch copied in h_cand_buf instead)
        //
        // For critic: we copy residues from the last LL launch.
        // The h_cand_buf was populated inside harvest_verified.
        // We use up to min(nc_last, 8) samples per cycle.
        // Since harvest_verified already copied h_cand_buf, use that.
        // (nc is 0 post-reset; use a cached count implicitly via loop below)
        {
            // Train critic on proxy-ring exponents: feature = (proxy residue,
            // p_norm, 0, 0, 0), target = 0 (no LL confirm yet).
            // A richer critic training would use the full h_cand_buf residues.
            // For now use a simple online update from the warp-LL residue.
            float feat[CRITIC_IN];
            for (int i = 0; i < n_verified; i++) {
                uint32_t p = found_verified[i];
                feat[0] = 0.0f;                        // residue ≈ 0 (confirmed)
                feat[1] = (float)p / 3904.0f;          // p normalised
                feat[2] = 0.5f;
                feat[3] = 0.5f;
                feat[4] = 1.0f;                        // high confidence
                critic_observe(feat, 1.0f);            // reward = 1 for confirmed prime
            }
        }

        // ---- Progress report every 100 cycles ----
        if ((cycle + 1) % 100 == 0) {
            sh_gettime(&t1);
            double elapsed = sh_elapsed(&t0, &t1);
            double steps_per_sec = (double)(g_total_steps - last_report_steps)
                                   / elapsed * 100.0 / 100.0;
            // Actually compute rate over last 100 cycles
            long long steps_this_period = g_total_steps - last_report_steps;
            double rate = (double)steps_this_period / elapsed;
            last_report_steps = g_total_steps;
            sh_gettime(&t0);

            printf("[cycle %4d/%d]  steps=%lld  rate=%.1fk/s  "
                   "proxy_total≈%lld  confirmed=%d\n",
                   cycle + 1, TOTAL_CYCLES,
                   g_total_steps,
                   rate / 1000.0,
                   g_total_cands,
                   g_n_confirmed);

            if (g_n_confirmed > 0) {
                printf("  Confirmed primes so far: ");
                for (int i = 0; i < g_n_confirmed; i++)
                    printf("M%u ", g_confirmed[i]);
                printf("\n");
            }
            fflush(stdout);
        }
    }

    // ---- Final summary ----
    printf("\n====================================================================\n");
    printf("  HDGL Sieve search complete.\n");
    printf("  Total sieve steps : %lld\n", g_total_steps);
    printf("  Total LL-confirmed: %d\n",   g_n_confirmed);
    if (g_n_confirmed > 0) {
        printf("  Confirmed Mersenne primes:\n");
        for (int i = 0; i < g_n_confirmed; i++)
            printf("    2^%u - 1  (p=%u)\n", g_confirmed[i], g_confirmed[i]);
    }
    printf("====================================================================\n\n");

    // Save critic weights
    critic_save_weights(NULL);
    printf("Critic weights saved to hdgl_critic.bin\n");

    // Cleanup
    CUDA_CHECK(cudaStreamDestroy(stream_sieve));
    CUDA_CHECK(cudaStreamDestroy(stream_ll));
    hdgl_sieve_free(&sieve);
    return 0;
}
