// ============================================================================
// HDGL Learned Critic v33 — Tiny Host-Side Policy Reward Network
// ============================================================================
//
// [ROADMAP ITEM 10] Replaces the hand-tuned reward formula in v32:
//   OLD: reward = 1/(1+residue) + 0.25*coherence + 0.15*amp
//   NEW: reward = critic_forward(state_features)
//
// Architecture: 2-layer MLP, pure C, no external dependencies
//   Input  dim = 5: [residue, coherence, amp, r_h_norm, acc_norm]
//   Hidden dim = 8: fully-connected, ReLU activation
//   Output dim = 1: linear, scalar reward estimate
//
// Training: online TD(0) with SGD
//   Target: r_observed + GAMMA_TD * V(s')
//   Loss:   MSE (critic_forward(s) - target)^2
//   Update: every CRITIC_UPDATE_INTERVAL reward injections
//
// Thread safety: single-threaded (called from host loop only).
// The critic weights are pushed to the GPU as a __constant__ symbol
// by hdgl_host_v33.c each time they are updated.
//
// Compile: included as part of nvcc compilation of hdgl_host_v33.c
// ============================================================================

#include <math.h>
#include <string.h>
#include <stdio.h>
#include <stdlib.h>
#include "hdgl_critic_v33.h"

// ============================================================================
// Constants
// ============================================================================

#define CRITIC_IN    5
#define CRITIC_HIDE  8
#define CRITIC_OUT   1

#define CRITIC_LR         3e-4f    // Adam works well at 3e-4
#define CRITIC_GAMMA_TD   0.99f
#define CRITIC_CLIP       1.0f     // gradient clipping (tighter for Adam)
#define CRITIC_UPDATE_INTERVAL 32  // reward samples before a weight update

#define ADAM_BETA1  0.9f
#define ADAM_BETA2  0.999f
#define ADAM_EPS    1e-8f

// ============================================================================
// Weight storage (single global instance, host-only)
// ============================================================================

typedef struct {
    // Layer 1: [CRITIC_HIDE × CRITIC_IN] weights + [CRITIC_HIDE] bias
    float w1[CRITIC_HIDE][CRITIC_IN];
    float b1[CRITIC_HIDE];
    // Layer 2: [CRITIC_OUT × CRITIC_HIDE] weights + [CRITIC_OUT] bias
    float w2[CRITIC_OUT][CRITIC_HIDE];
    float b2[CRITIC_OUT];
    // Gradient accumulators
    float dw1[CRITIC_HIDE][CRITIC_IN];
    float db1[CRITIC_HIDE];
    float dw2[CRITIC_OUT][CRITIC_HIDE];
    float db2[CRITIC_OUT];
    // Running normalisation stats (online Welford)
    float feat_mean[CRITIC_IN];
    float feat_M2[CRITIC_IN];
    long long feat_n;
    // Replay mini-buffer for batch updates
    float  replay_s[CRITIC_UPDATE_INTERVAL][CRITIC_IN];
    float  replay_target[CRITIC_UPDATE_INTERVAL];
    int    replay_head;
    int    replay_count;
    // Training step counter
    long long steps;
    // Adam optimizer first and second moment estimates (zero-inited by memset)
    float adam_mw1[CRITIC_HIDE][CRITIC_IN];
    float adam_vw1[CRITIC_HIDE][CRITIC_IN];
    float adam_mb1[CRITIC_HIDE];
    float adam_vb1[CRITIC_HIDE];
    float adam_mw2[CRITIC_OUT][CRITIC_HIDE];
    float adam_vw2[CRITIC_OUT][CRITIC_HIDE];
    float adam_mb2[CRITIC_OUT];
    float adam_vb2[CRITIC_OUT];
    // Exponential moving average of MSE loss (for monitoring)
    float loss_ema;
} Critic;

static Critic g_critic;

// ============================================================================
// Activation
// ============================================================================

static inline float relu(float x) { return x > 0.0f ? x : 0.0f; }
static inline float relu_d(float x) { return x > 0.0f ? 1.0f : 0.0f; }

// ============================================================================
// Forward pass — stores activations in caller's buffers for backprop
// ============================================================================

static float critic_forward_full(
    const float s[CRITIC_IN],
    float h1[CRITIC_HIDE])   // hidden activations (output for backprop)
{
    for (int j = 0; j < CRITIC_HIDE; j++) {
        float z = g_critic.b1[j];
        for (int i = 0; i < CRITIC_IN; i++) z += g_critic.w1[j][i] * s[i];
        h1[j] = relu(z);
    }
    float out = g_critic.b2[0];
    for (int j = 0; j < CRITIC_HIDE; j++) out += g_critic.w2[0][j] * h1[j];
    return out;
}

// ============================================================================
// Public: forward-only (no stored activations)
// ============================================================================

float critic_forward(const float s[CRITIC_IN]) {
    float h1[CRITIC_HIDE];
    return critic_forward_full(s, h1);
}

// ============================================================================
// Online feature normalisation (Welford one-pass)
// ============================================================================

static void normalise_inplace(float s[CRITIC_IN]) {
    g_critic.feat_n++;
    for (int i = 0; i < CRITIC_IN; i++) {
        float delta  = s[i] - g_critic.feat_mean[i];
        g_critic.feat_mean[i] += delta / (float)g_critic.feat_n;
        float delta2 = s[i] - g_critic.feat_mean[i];
        g_critic.feat_M2[i] += delta * delta2;
        float var = (g_critic.feat_n > 1)
                    ? g_critic.feat_M2[i] / (float)(g_critic.feat_n - 1)
                    : 1.0f;
        float std = sqrtf(var + 1e-8f);
        s[i] = (s[i] - g_critic.feat_mean[i]) / std;
        // Clamp to [-3, 3] to prevent extreme inputs
        if (s[i] >  3.0f) s[i] =  3.0f;
        if (s[i] < -3.0f) s[i] = -3.0f;
    }
}

// ============================================================================
// Push a (state, target) sample into the replay mini-buffer
// ============================================================================

void critic_observe(const float raw_s[CRITIC_IN], float target) {
    // Normalise features
    float s[CRITIC_IN];
    memcpy(s, raw_s, CRITIC_IN * sizeof(float));
    normalise_inplace(s);

    int pos = g_critic.replay_head % CRITIC_UPDATE_INTERVAL;
    memcpy(g_critic.replay_s[pos], s, CRITIC_IN * sizeof(float));
    g_critic.replay_target[pos] = target;
    g_critic.replay_head++;
    if (g_critic.replay_count < CRITIC_UPDATE_INTERVAL) g_critic.replay_count++;

    // Trigger a mini-batch update when buffer is full
    if (g_critic.replay_count == CRITIC_UPDATE_INTERVAL &&
        g_critic.replay_head % CRITIC_UPDATE_INTERVAL == 0) {
        critic_update();
    }
}

// ============================================================================
// Mini-batch gradient update
// ============================================================================

void critic_update(void) {
    if (g_critic.replay_count == 0) return;

    // Zero accumulators
    memset(g_critic.dw1, 0, sizeof(g_critic.dw1));
    memset(g_critic.db1, 0, sizeof(g_critic.db1));
    memset(g_critic.dw2, 0, sizeof(g_critic.dw2));
    memset(g_critic.db2, 0, sizeof(g_critic.db2));

    int n = g_critic.replay_count;
    float inv_n = 1.0f / (float)n;
    float total_loss = 0.0f;

    for (int b = 0; b < n; b++) {
        const float *s     = g_critic.replay_s[b];
        float        tgt   = g_critic.replay_target[b];
        float h1[CRITIC_HIDE];

        float pred  = critic_forward_full(s, h1);
        float err   = pred - tgt;
        total_loss += err * err;
        float dout  = err * inv_n;

        // Clip
        if (dout >  CRITIC_CLIP) dout =  CRITIC_CLIP;
        if (dout < -CRITIC_CLIP) dout = -CRITIC_CLIP;

        // Layer 2 gradients
        for (int j = 0; j < CRITIC_HIDE; j++) {
            g_critic.dw2[0][j] += dout * h1[j];
        }
        g_critic.db2[0] += dout;

        // Layer 1 gradients (backprop through ReLU)
        for (int j = 0; j < CRITIC_HIDE; j++) {
            float dh = dout * g_critic.w2[0][j];
            float dz = dh * (h1[j] > 0.0f ? 1.0f : 0.0f);
            for (int i = 0; i < CRITIC_IN; i++) {
                g_critic.dw1[j][i] += dz * s[i];
            }
            g_critic.db1[j] += dz;
        }
    }

    // Update loss EMA
    float loss = total_loss * inv_n;
    g_critic.loss_ema = (g_critic.steps == 0)
                        ? loss
                        : 0.9f * g_critic.loss_ema + 0.1f * loss;

    // Adam step (bias-corrected)
    g_critic.steps++;
    float bc1 = 1.0f - powf(ADAM_BETA1, (float)g_critic.steps);
    float bc2 = 1.0f - powf(ADAM_BETA2, (float)g_critic.steps);

    for (int j = 0; j < CRITIC_HIDE; j++) {
        for (int i = 0; i < CRITIC_IN; i++) {
            float g = g_critic.dw1[j][i];
            g_critic.adam_mw1[j][i] = ADAM_BETA1 * g_critic.adam_mw1[j][i] + (1.0f - ADAM_BETA1) * g;
            g_critic.adam_vw1[j][i] = ADAM_BETA2 * g_critic.adam_vw1[j][i] + (1.0f - ADAM_BETA2) * g * g;
            g_critic.w1[j][i] -= CRITIC_LR * (g_critic.adam_mw1[j][i] / bc1)
                                            / (sqrtf(g_critic.adam_vw1[j][i] / bc2) + ADAM_EPS);
        }
        float gb = g_critic.db1[j];
        g_critic.adam_mb1[j] = ADAM_BETA1 * g_critic.adam_mb1[j] + (1.0f - ADAM_BETA1) * gb;
        g_critic.adam_vb1[j] = ADAM_BETA2 * g_critic.adam_vb1[j] + (1.0f - ADAM_BETA2) * gb * gb;
        g_critic.b1[j] -= CRITIC_LR * (g_critic.adam_mb1[j] / bc1)
                                     / (sqrtf(g_critic.adam_vb1[j] / bc2) + ADAM_EPS);
    }
    for (int j = 0; j < CRITIC_HIDE; j++) {
        float g = g_critic.dw2[0][j];
        g_critic.adam_mw2[0][j] = ADAM_BETA1 * g_critic.adam_mw2[0][j] + (1.0f - ADAM_BETA1) * g;
        g_critic.adam_vw2[0][j] = ADAM_BETA2 * g_critic.adam_vw2[0][j] + (1.0f - ADAM_BETA2) * g * g;
        g_critic.w2[0][j] -= CRITIC_LR * (g_critic.adam_mw2[0][j] / bc1)
                                        / (sqrtf(g_critic.adam_vw2[0][j] / bc2) + ADAM_EPS);
    }
    {
        float g = g_critic.db2[0];
        g_critic.adam_mb2[0] = ADAM_BETA1 * g_critic.adam_mb2[0] + (1.0f - ADAM_BETA1) * g;
        g_critic.adam_vb2[0] = ADAM_BETA2 * g_critic.adam_vb2[0] + (1.0f - ADAM_BETA2) * g * g;
        g_critic.b2[0] -= CRITIC_LR * (g_critic.adam_mb2[0] / bc1)
                                     / (sqrtf(g_critic.adam_vb2[0] / bc2) + ADAM_EPS);
    }

    g_critic.replay_count = 0;
    g_critic.replay_head  = 0;
}

// ============================================================================
// TD target construction
//   observed_reward: immediate reward from the environment (LL result etc.)
//   s_next:          normalised features at s'
//   Returns TD(0) target: r + γ * V(s')
// ============================================================================

// Normalise using current stats WITHOUT updating them (read-only)
static void normalise_readonly(float s[CRITIC_IN]) {
    for (int i = 0; i < CRITIC_IN; i++) {
        float var = (g_critic.feat_n > 1)
                    ? g_critic.feat_M2[i] / (float)(g_critic.feat_n - 1)
                    : 1.0f;
        float std = sqrtf(var + 1e-8f);
        s[i] = (s[i] - g_critic.feat_mean[i]) / std;
        if (s[i] >  3.0f) s[i] =  3.0f;
        if (s[i] < -3.0f) s[i] = -3.0f;
    }
}

float critic_td_target(float observed_reward, const float raw_s_next[CRITIC_IN]) {
    float s_next[CRITIC_IN];
    memcpy(s_next, raw_s_next, CRITIC_IN * sizeof(float));
    normalise_readonly(s_next);   // do NOT update Welford stats for s_next
    float h1[CRITIC_HIDE];
    float v_next = critic_forward_full(s_next, h1);
    return observed_reward + CRITIC_GAMMA_TD * v_next;
}

// ============================================================================
// Serialise weights for GPU upload (packed float array)
// Layout: w1 (HIDE×IN), b1 (HIDE), w2 (OUT×HIDE), b2 (OUT)
//   = 5*8 + 8 + 8*1 + 1 = 57 floats
// ============================================================================

int critic_weight_count(void) {
    return CRITIC_HIDE * CRITIC_IN  // w1
         + CRITIC_HIDE              // b1
         + CRITIC_OUT * CRITIC_HIDE // w2
         + CRITIC_OUT;              // b2
}

void critic_pack_weights(float *out) {
    // Fold Welford normalisation into Layer-1 so the GPU can use raw inputs.
    // Effective transform: W1'[j][i] = W1[j][i] / std[i]
    //                      b1'[j]    = b1[j] - sum_i( W1'[j][i] * mean[i] )
    // This keeps GPU inference identical to CPU forward on normalised data.
    float feat_std[CRITIC_IN];
    for (int i = 0; i < CRITIC_IN; i++) {
        float var = (g_critic.feat_n > 1)
                    ? g_critic.feat_M2[i] / (float)(g_critic.feat_n - 1)
                    : 1.0f;
        feat_std[i] = sqrtf(var + 1e-8f);
    }
    int p = 0;
    for (int j = 0; j < CRITIC_HIDE; j++)
        for (int i = 0; i < CRITIC_IN; i++)
            out[p++] = g_critic.w1[j][i] / feat_std[i];
    for (int j = 0; j < CRITIC_HIDE; j++) {
        float b = g_critic.b1[j];
        for (int i = 0; i < CRITIC_IN; i++)
            b -= (g_critic.w1[j][i] / feat_std[i]) * g_critic.feat_mean[i];
        out[p++] = b;
    }
    for (int j = 0; j < CRITIC_HIDE; j++) out[p++] = g_critic.w2[0][j];
    out[p++] = g_critic.b2[0];
}

// ============================================================================
// Init — Xavier/Glorot init
// ============================================================================

void critic_init(void) {
    memset(&g_critic, 0, sizeof(g_critic));
    srand(42);
    // Xavier: std = sqrt(2 / (fan_in + fan_out))
    float std1 = sqrtf(2.0f / (float)(CRITIC_IN + CRITIC_HIDE));
    float std2 = sqrtf(2.0f / (float)(CRITIC_HIDE + CRITIC_OUT));
    for (int j = 0; j < CRITIC_HIDE; j++)
        for (int i = 0; i < CRITIC_IN; i++)
            g_critic.w1[j][i] = std1 * ((float)rand()/(float)RAND_MAX * 2.0f - 1.0f);
    for (int j = 0; j < CRITIC_HIDE; j++)
        g_critic.w2[0][j]  = std2 * ((float)rand()/(float)RAND_MAX * 2.0f - 1.0f);
    // biases initialised to 0 (memset above)
    for (int i = 0; i < CRITIC_IN; i++) g_critic.feat_M2[i] = 1.0f;  // avoid div-by-zero
}

// ============================================================================
// Debug print
// ============================================================================

void critic_print_stats(void) {
    printf("[critic] steps=%lld  feat_n=%lld  loss_ema=%.5f\n",
           g_critic.steps, g_critic.feat_n, (double)g_critic.loss_ema);
    printf("[critic] feat_mean:");
    for (int i = 0; i < CRITIC_IN; i++) printf(" %.3f", (double)g_critic.feat_mean[i]);
    printf("\n[critic] w2:");
    for (int j = 0; j < CRITIC_HIDE; j++) printf(" %.4f", (double)g_critic.w2[0][j]);
    printf("  b2=%.4f\n", (double)g_critic.b2[0]);
}

// ============================================================================
// Weight persistence
// ============================================================================

#define CRITIC_SAVE_FILE "hdgl_critic.bin"

void critic_save_weights(const char *path) {
    const char *p = path ? path : CRITIC_SAVE_FILE;
    FILE *f = fopen(p, "wb");
    if (!f) { fprintf(stderr, "[critic] save failed: %s\n", p); return; }
    fwrite(&g_critic, sizeof(Critic), 1, f);
    fclose(f);
    printf("[critic] saved to %s  (steps=%lld)\n", p, g_critic.steps);
}

int critic_load_weights(const char *path) {
    const char *p = path ? path : CRITIC_SAVE_FILE;
    FILE *f = fopen(p, "rb");
    if (!f) return 0;
    Critic loaded;
    if (fread(&loaded, sizeof(Critic), 1, f) != 1) { fclose(f); return 0; }
    fclose(f);
    g_critic = loaded;
    printf("[critic] resumed from %s  (steps=%lld  loss_ema=%.5f)\n",
           p, g_critic.steps, (double)g_critic.loss_ema);
    return 1;
}

// ============================================================================
// Pre-training from scanner-calibrated resonance nodes
//
// The BigG+Fudge10 scanner confirmed 14 prime candidates at lam=41..67
// (all passing S<0.25).  Three of those nodes project into the analog's
// r_harmonic range [1, R_MAX=1000]:
//
//   lam=9  → r_h ≈  61.2  (M61 confirmed Mersenne prime, BigG S=0.026)
//   lam=11 → r_h ≈ 160    (scanner-geometry resonance node)
//   lam=13 → r_h ≈ 420    (scanner-geometry resonance node)
//
// All 14 scanner candidates have p >> R_MAX, so their r_h saturates to 1000.
//
// Inject 32 positive + 32 negative examples (= 2 × CRITIC_UPDATE_INTERVAL)
// so two full Adam updates fire before the first simulation step.
//
// Features passed RAW — Welford normaliser in critic_observe() handles scaling.
// Feature layout (matches host usage in hdgl_host_v33.c):
//   [residue, coherence, amp, r_harmonic, candidate_score]
// ============================================================================

void critic_pretrain_scan_candidates(void) {
    // ── Positive: near-resonance-node field states ───────────────────────────
    // Vary residue and coherence slightly so the Welford stats see a spread,
    // but keep all variants clearly in the "prime-like" regime.
    static const struct { float r_h; float target; } pos_nodes[] = {
        {  61.2f, 2.0f },   // lam=9:  M61 confirmed (S=0.025687) — gold target
        { 160.0f, 1.8f },   // lam=11: scanner-geometry resonance node
        { 420.0f, 1.8f },   // lam=13: scanner-geometry resonance node
        {1000.0f, 2.0f },   // lam=41..67: all 14 scanner candidates saturate here
    };
    int n_pos_nodes = (int)(sizeof(pos_nodes)/sizeof(pos_nodes[0]));
    for (int ni = 0; ni < n_pos_nodes; ni++) {
        float rh  = pos_nodes[ni].r_h;
        float tgt = pos_nodes[ni].target;
        for (int rep = 0; rep < 8; rep++) {
            float res = 0.005f + (float)rep * 0.002f;   // 0.005..0.019
            float coh = 3.9f   - (float)rep * 0.15f;   // 3.90..2.85
            float s[CRITIC_IN] = { res, coh, 0.95f, rh, 2.0f };
            critic_observe(s, tgt);
        }
    }

    // ── Negative: off-node, incoherent field states ──────────────────────────
    // Mid-range r_h values that sit between resonance nodes.
    static const float neg_rh[] = { 20.0f, 50.0f, 120.0f, 200.0f,
                                     300.0f, 500.0f, 700.0f, 850.0f };
    for (int ni = 0; ni < 8; ni++) {
        float rh = neg_rh[ni];
        float s1[CRITIC_IN] = { 0.75f, 0.3f, 0.15f, rh, 0.6f };
        critic_observe(s1, 0.05f);
        float s2[CRITIC_IN] = { 0.50f, 1.0f, 0.40f, rh, 0.9f };
        critic_observe(s2, 0.10f);
        float s3[CRITIC_IN] = { 0.30f, 1.5f, 0.60f, rh, 1.2f };
        critic_observe(s3, 0.15f);
        float s4[CRITIC_IN] = { 0.90f, 0.5f, 0.20f, rh, 0.4f };
        critic_observe(s4, 0.05f);
    }

    // Flush any remaining partial replay buffer.
    // (64 samples = 2 × CRITIC_UPDATE_INTERVAL so two batch updates already fired;
    // this is a no-op if replay_count==0, but kept for safety.)
    critic_update();

    printf("[critic] pretrain_scan_candidates: 32 positive + 32 negative"
           "  steps=%lld\n", (long long)g_critic.steps);
}

