/*
  StarForth — Steady-State Virtual Machine Runtime

  Copyright (c) 2023–2025 Robert A. James
  All rights reserved.

  This file is part of the StarForth project.

  Licensed under the StarForth License, Version 1.0 (the "License");
  you may not use this file except in compliance with the License.

  You may obtain a copy of the License at:
      https://github.com/star.4th@proton.me/StarForth/LICENSE.txt

  This software is provided "AS IS", WITHOUT WARRANTY OF ANY KIND,
  express or implied, including but not limited to the warranties of
  merchantability, fitness for a particular purpose, and noninfringement.

 See the License for the specific language governing permissions and
 limitations under the License.

 */

#include "compudynamics.h"
#include <string.h>
#include <stdlib.h>

/* ============================================================================
 * Per-Level Tuning Accessors
 * ============================================================================
 *
 * Returned by value, not exposed as addressable globals -- see the
 * "Tuning knobs are returned by value" comment in compudynamics.h for why.
 */

/* Matches the constants ssm_jacquard.c/.h used before this module existed
 * (SSM_SCORE_MAX, SSM_SCORE_DECAY_FACTOR_Q16, SSM_UCB_K, SSM_REWARD_GAIN,
 * SSM_JOINT_WEIGHT_Q16, SSM_ANOVA_WEIGHT_Q16, SSM_ENTROPY/CV/TEMPORAL_DECAY
 * thresholds) -- a lift, not a retune. id_window_depth=32 and
 * locality_lookback=8 are new: they size the entity-ID history the
 * deterministic locality signal replaces temporal_decay with. */
CDTuning cd_tuning_word(void)
{
    CDTuning t;
    /* classifier */
    t.id_window_depth              = 32u;
    t.locality_lookback            = 8u;
    t.diversity_high_threshold_q16  = 49152u; /* 0.75 * 65536 */
    t.volatility_high_threshold_q16 = 9830u;  /* 0.15 * 65536 */
    t.locality_high_threshold_q16   = 32768u; /* 0.5  * 65536 */
    /* bandit */
    t.num_configs             = 128u; /* 7-bit L1-L7 */
    t.num_regimes              = 8u;   /* 3-bit diversity/volatility/locality */
    t.score_max                = 65535u;
    t.score_min                = 0u;
    t.score_decay_factor_q16   = 64881u; /* 0.99 * 65536 */
    t.ucb_k                    = 6554u;  /* 10% of score_max */
    t.reward_gain               = 16384u; /* 25% of score_max */
    t.joint_weight_q16          = 49152u; /* 0.75 */
    t.anova_weight_q16          = 16384u; /* 0.25 */
    return t;
}

/* VM level: classifier-side parameters only, matching capsule_vm_physics.c's
 * existing VM_FLEET_WINDOW_DEPTH so a future migration has a ready-made
 * starting point. Bandit fields are zero and unused -- see compudynamics.h. */
CDTuning cd_tuning_vm(void)
{
    CDTuning t;
    /* classifier */
    t.id_window_depth               = 64u; /* matches VM_FLEET_WINDOW_DEPTH */
    t.locality_lookback             = 16u;
    t.diversity_high_threshold_q16  = 49152u;
    t.volatility_high_threshold_q16 = 9830u;
    t.locality_high_threshold_q16   = 32768u;
    /* bandit -- unused at VM level today */
    t.num_configs           = 0u;
    t.num_regimes            = 0u;
    t.score_max               = 0u;
    t.score_min               = 0u;
    t.score_decay_factor_q16  = 0u;
    t.ucb_k                   = 0u;
    t.reward_gain              = 0u;
    t.joint_weight_q16         = 0u;
    t.anova_weight_q16         = 0u;
    return t;
}

/* ============================================================================
 * Integer Helpers (isqrt, ln approximation)
 * ============================================================================
 */

static uint32_t cd_isqrt32(uint32_t n)
{
    uint32_t x, y;
    if (n == 0u) return 0u;
    x = n;
    y = (x + 1u) / 2u;
    while (y < x) {
        x = y;
        y = (x + n / x) / 2u;
    }
    return x;
}

static uint64_t cd_isqrt64(uint64_t n)
{
    uint64_t x, y;
    if (n == 0u) return 0u;
    x = n;
    y = (x + 1u) / 2u;
    while (y < x) {
        x = y;
        y = (x + n / x) / 2u;
    }
    return x;
}

/* floor(log2(n)) * ln(2) * 65536, integer shifts only. See ssm_jacquard.c's
 * former ln_approx_q16 -- identical algorithm, moved here so the generic
 * bandit doesn't depend on ssm_jacquard.c. */
static uint32_t cd_ln_approx_q16(uint32_t n)
{
    uint32_t bits, tmp;
    if (n <= 1u) return 0u;
    bits = 0u;
    tmp = n >> 1u;
    while (tmp > 0u) { tmp >>= 1u; bits++; }
    return bits * 45426u;
}

/* ============================================================================
 * Deterministic Regime Classification
 * ============================================================================
 */

CDRegimeMetrics cd_classify_ids(const uint32_t *ids, uint32_t count,
                                 uint32_t locality_lookback)
{
    CDRegimeMetrics m;
    uint32_t n, i, j, distinct, lookback;
    uint8_t  counted[CD_MAX_CLASSIFY_DEPTH];
    uint32_t per_id_count[CD_MAX_CLASSIFY_DEPTH];

    m.diversity_q16  = 0u;
    m.volatility_q16 = 0u;
    m.locality_q16   = 0u;

    if (!ids || count == 0u) return m;

    n = (count > CD_MAX_CLASSIFY_DEPTH) ? CD_MAX_CLASSIFY_DEPTH : count;
    memset(counted, 0, n * sizeof(counted[0]));

    /* --- Diversity: distinct IDs / n, and per-ID touch counts for the
     * volatility pass below (single O(n^2) sweep, n <= 64) --- */
    distinct = 0u;
    for (i = 0u; i < n; i++) {
        uint32_t c;
        if (counted[i]) continue;
        c = 1u;
        counted[i] = 1u;
        for (j = i + 1u; j < n; j++) {
            if (ids[j] == ids[i]) { c++; counted[j] = 1u; }
        }
        per_id_count[distinct] = c;
        distinct++;
    }
    m.diversity_q16 = (uint32_t)(((uint64_t)distinct * 65536u) / n);

    /* --- Volatility: coefficient of variation of per-ID touch counts,
     * computed in Q8 to avoid overflow, result in Q16 --- */
    if (distinct > 0u) {
        uint32_t mean_q8 = (n << 8u) / distinct;
        uint64_t sum_var_q16 = 0u;
        uint64_t stddev_q8, cv_q16;

        for (i = 0u; i < distinct; i++) {
            int64_t diff_q8 = (int64_t)(per_id_count[i] << 8u) - (int64_t)mean_q8;
            uint64_t diff2_q16 = (uint64_t)(diff_q8 * diff_q8);
            sum_var_q16 += diff2_q16;
        }
        stddev_q8 = cd_isqrt64(sum_var_q16 / distinct);
        cv_q16 = (mean_q8 > 0u) ? ((stddev_q8 * 65536u) / mean_q8) : 0u;
        m.volatility_q16 = (cv_q16 > 0xFFFFFFFFu) ? 0xFFFFFFFFu : (uint32_t)cv_q16;
    }

    /* --- Locality: fraction of touches (from lookback..n-1) that repeat
     * an ID seen within the preceding lookback touches --- */
    lookback = (locality_lookback > n) ? n : locality_lookback;
    if (n > lookback && lookback > 0u) {
        uint32_t matches = 0u;
        uint32_t denom = n - lookback;
        for (i = lookback; i < n; i++) {
            for (j = i - lookback; j < i; j++) {
                if (ids[j] == ids[i]) { matches++; break; }
            }
        }
        m.locality_q16 = (uint32_t)(((uint64_t)matches * 65536u) / denom);
    }

    return m;
}

/* ============================================================================
 * Generic UCB1 Config-Space Bandit
 * ============================================================================
 */

int cd_config_table_init(CDConfigTable *table, const CDTuning *tuning)
{
    if (!table || !tuning) return -1;
    memset(table, 0, sizeof(*table));

    if (tuning->num_configs == 0u || tuning->num_regimes == 0u) return -1;

    table->entries = (CDConfigEntry *)malloc(tuning->num_configs * sizeof(CDConfigEntry));
    table->regime_scores = (uint32_t *)malloc(
        (size_t)tuning->num_regimes * tuning->num_configs * sizeof(uint32_t));
    table->regime_ticks = (uint32_t *)malloc(
        (size_t)tuning->num_regimes * tuning->num_configs * sizeof(uint32_t));
    table->total_regime_ticks = (uint32_t *)malloc(tuning->num_regimes * sizeof(uint32_t));

    if (!table->entries || !table->regime_scores || !table->regime_ticks ||
        !table->total_regime_ticks) {
        cd_config_table_free(table);
        return -1;
    }

    table->num_configs = tuning->num_configs;
    table->num_regimes = tuning->num_regimes;
    return 0;
}

void cd_config_table_free(CDConfigTable *table)
{
    if (!table) return;
    free(table->entries);
    free(table->regime_scores);
    free(table->regime_ticks);
    free(table->total_regime_ticks);
    memset(table, 0, sizeof(*table));
}

void cd_config_table_seed(CDConfigTable *table, const CDTuning *tuning,
                           const uint32_t *initial_scores, uint8_t initial_config)
{
    uint32_t c, r;
    if (!table || !table->entries || !tuning || !initial_scores) return;

    for (c = 0u; c < table->num_configs; c++) {
        CDConfigEntry *e = &table->entries[c];
        e->config_bits       = (uint8_t)c;
        e->has_prev           = 0u;
        e->_pad[0]            = 0u;
        e->_pad[1]            = 0u;
        e->benefit_score      = initial_scores[c];
        e->tick_count         = 0u;
        e->prev_window        = 0u;
        e->prev_locality_q16  = 0u;

        for (r = 0u; r < table->num_regimes; r++) {
            table->regime_scores[r * table->num_configs + c] = initial_scores[c];
            table->regime_ticks[r * table->num_configs + c]  = 0u;
        }
    }

    table->current_config = initial_config;
    table->current_regime = 0u;
    for (r = 0u; r < table->num_regimes; r++) {
        table->total_regime_ticks[r] = 0u;
    }
    (void)tuning;
}

uint8_t cd_config_table_tick(CDConfigTable *table, const CDTuning *tuning,
                              uint8_t regime, uint32_t anova_stability_q16,
                              int use_joint_convergence,
                              uint32_t current_window, uint32_t current_locality_q16)
{
    uint32_t r, c;
    CDConfigEntry *entry;
    uint32_t combined_stability;
    uint32_t *score_cell;
    uint32_t old_score, decayed, new_score;
    int64_t  delta_score, new_s;
    uint32_t total, ln_total, best_c, best_ucb, ci;

    if (!table || !table->entries || !tuning) return 0u;

    table->current_regime = regime;
    r = (uint32_t)regime;
    c = (uint32_t)table->current_config;
    entry = &table->entries[c];

    if (use_joint_convergence && entry->has_prev) {
        uint32_t dw_raw = (current_window > entry->prev_window)
                          ? (current_window - entry->prev_window)
                          : (entry->prev_window - current_window);
        uint32_t dl_raw = (current_locality_q16 > entry->prev_locality_q16)
                          ? (current_locality_q16 - entry->prev_locality_q16)
                          : (entry->prev_locality_q16 - current_locality_q16);
        uint64_t dw_q16 = (entry->prev_window > 0u)
                          ? (((uint64_t)dw_raw * 65536u) / entry->prev_window)
                          : 0u;
        uint64_t dl_q16 = dl_raw; /* already Q16, bounded to [0,65536] */
        uint64_t sum_sq, joint_err;
        uint32_t stability;

        if (dw_q16 > 65536u) dw_q16 = 65536u;
        if (dl_q16 > 65536u) dl_q16 = 65536u;

        /* joint_error = sqrt(dw^2 + dl^2) / sqrt(2), normalised to [0,65536] */
        sum_sq = dw_q16 * dw_q16 + dl_q16 * dl_q16;
        joint_err = (cd_isqrt64(sum_sq) * 65536u) / 92682u; /* sqrt(2)*65536 ~= 92682 */
        if (joint_err > 65536u) joint_err = 65536u;

        stability = (uint32_t)(65536u - (uint32_t)joint_err);

        combined_stability = (uint32_t)(
            ((uint64_t)tuning->joint_weight_q16 * stability +
             (uint64_t)tuning->anova_weight_q16 * anova_stability_q16) / 65536u
        );
    } else {
        combined_stability = anova_stability_q16;
    }

    delta_score = ((int64_t)combined_stability - 32768) *
                  (int64_t)tuning->reward_gain / 32768;

    score_cell = &table->regime_scores[r * table->num_configs + c];
    old_score  = *score_cell;
    decayed    = (uint32_t)(((uint64_t)old_score * tuning->score_decay_factor_q16) / 65536u);
    new_s      = (int64_t)decayed + delta_score;

    if      (new_s < (int64_t)tuning->score_min) new_score = tuning->score_min;
    else if (new_s > (int64_t)tuning->score_max) new_score = tuning->score_max;
    else                                          new_score = (uint32_t)new_s;

    *score_cell = new_score;
    table->regime_ticks[r * table->num_configs + c]++;
    if (table->total_regime_ticks[r] < 0xFFFFFFFFu) table->total_regime_ticks[r]++;
    entry->tick_count++;

    entry->prev_window       = current_window;
    entry->prev_locality_q16 = current_locality_q16;
    entry->has_prev          = 1u;

    /* --- UCB1: select best config for the next tick --- */
    total    = table->total_regime_ticks[r];
    ln_total = cd_ln_approx_q16(total > 0u ? total : 1u);
    best_c   = 0u;
    best_ucb = 0u;

    for (ci = 0u; ci < table->num_configs; ci++) {
        uint32_t nc = table->regime_ticks[r * table->num_configs + ci];
        uint32_t ucb;

        if (nc == 0u) {
            ucb = tuning->score_max;
        } else {
            uint32_t ratio  = ln_total / nc;
            uint32_t sqrt_r = cd_isqrt32(ratio);
            uint32_t bonus  = (uint32_t)(((uint64_t)tuning->ucb_k * sqrt_r) / 256u);
            uint32_t sc     = table->regime_scores[r * table->num_configs + ci];
            ucb = sc + bonus;
            if (ucb > tuning->score_max) ucb = tuning->score_max;
        }

        if (ucb > best_ucb) { best_ucb = ucb; best_c = ci; }
    }

    table->current_config = (uint8_t)best_c;
    return table->entries[best_c].config_bits;
}

void cd_config_table_force(CDConfigTable *table, uint8_t config_idx)
{
    if (!table || !table->entries) return;
    if ((uint32_t)config_idx >= table->num_configs) return;
    table->current_config = config_idx;
}