// Code Predictor (Qwen3-TTS) inference, extrait de
// /opt/Kazeia/kazeia-tts-decoder-ggml/src/cp_runner.cpp et exposé en API
// utilisable depuis tts_orchestrate. Architecture : transformer qwen3 5 couches,
// hidden 1024, GQA 16/8, head_dim 128, RMS eps 1e-6, RoPE NEOX theta=1e6.
// q_norm/k_norm AVANT RoPE (gotcha critique).
//
// Deux variantes de forward :
//   - cp_predict          : approche "recompute" historique (15 graphes, L=2..16). Reste
//                           comme oracle bit-exact référence.
//   - cp_predict_cached   : KV cache f32 persistant entre étapes. 1 graphe prefill
//                           (L=2) + 14 graphes decode (L=1). Speedup ~8x mesuré, math
//                           équivalente mais ordre des sommes f32 peut différer.
#pragma once
#include "sampler.h"
#include <cstdint>
#include <string>
#include <vector>
#include <unordered_map>

struct ggml_context;
struct ggml_tensor;
struct ggml_backend;
typedef struct ggml_backend * ggml_backend_t;
struct ggml_backend_sched;
struct ggml_backend_buffer;

// Nombre de couches du CP — fixe par l'archi. Déclaré ici car taille tableaux KV cache.
#define CP_N_LAYER 5

struct CPState {
    ggml_context * weights_ctx = nullptr;
    std::unordered_map<std::string, ggml_tensor*> tensors;  // weights par nom
    std::vector<float> heads;        // [15, 2048, 1024] f32
    std::vector<float> codec_embs;   // [15, 2048, 1024] f32
    int n_threads = 4;

    // --- Backend multi-device (Phase F : CP forward sur HMX V79) ---
    // Si backends est non vide après cp_load_with_backends(), les poids sont alloués
    // sur backends[0] (HTP) au lieu du weights_ctx CPU pur. Un sched est créé pour
    // splitter les ops auto entre HTP et CPU. cp_forward_cached_step détecte la
    // présence du sched et utilise la variante sched-friendly (no_alloc=true +
    // ggml_set_input/output + tensor_set/get).
    //
    // CP est un transformer 5L 1024-hidden, batches très petits (n_new=1 décode,
    // n_new=2 prefill), donc passe sans accroc la limite hexagon nrows(src1) ≤ 1024
    // qui bloquait BigVGAN.
    std::vector<ggml_backend_t> backends;
    ggml_backend_sched * sched = nullptr;
    ggml_backend_buffer * weights_buf = nullptr;
    ggml_backend_buffer * cache_buf   = nullptr;  // pour les K/V cache, sur le même backend
    bool owns_backends = false;
    // Sampler HF-style (rep_penalty + top_k + top_p + temperature). Par défaut = greedy
    // (top_k=1 -> argmax). Pour TTS prod, configurer en {temp=0.9, top_k=50, rep_penalty=1.05}
    // ce qui matche Python subtalker_*.
    Sampler sampler = { /*temp=*/0.0f, /*top_k=*/1, /*top_p=*/1.0f, /*rep_penalty=*/1.0f,
                        /*rep_window=*/16, /*rng_state=*/1u, {} };

    // --- KV cache (utilisé uniquement par cp_predict_cached). Alloué paresseusement au
    // premier appel via cp_cache_init(). Layout per-layer : [HEAD_DIM=128, N_KV=8, T_MAX=16]
    // f32 contigu, même layout que K/V juste après reshape+norm+RoPE (avant permute).
    // Pas d'état conservé entre frames : à chaque cp_predict_cached() on réécrit pos 0..15.
    ggml_context * cache_ctx = nullptr;
    ggml_tensor * K_cache[CP_N_LAYER] = { nullptr };
    ggml_tensor * V_cache[CP_N_LAYER] = { nullptr };
};

// Charge cp_f16.gguf + cp_heads.bin + cp_codec_embs.bin. Retourne false en cas d'échec.
// Path CPU pur : poids dans weights_ctx alloué par ggml_init.
bool cp_load(CPState& s, const char* gguf_path, const char* heads_path, const char* embs_path);

// Charge le même contenu mais alloue les poids sur le buft du 1er backend (HTP si
// présent) et crée un sched pour le forward. devs[end] doit être CPU (assertion sched).
// owns_backends=true -> cp_free libère aussi les backends.
bool cp_load_with_backends(CPState& s, const char* gguf_path,
                            const char* heads_path, const char* embs_path,
                            const std::vector<ggml_backend_t>& devs,
                            bool owns_backends);

// Libère les ressources.
void cp_free(CPState& s);

// Prédit CB1..CB15 (15 codes int32) depuis le hidden state du Talker et l'embed de CB0.
//   hidden : float[1024]
//   cb0_emb: float[1024]
//   out_codes : int32_t[15] (CB1..CB15)
void cp_predict(CPState& s, const float* hidden, const float* cb0_emb, int32_t* out_codes);

// Même contrat que cp_predict, mais avec KV cache : 1 prefill (n_new=2) + 14 decode (n_new=1).
// Speedup ~8x (CP 233ms -> ~28ms par frame mesuré sur SM8750), permet de viser RTF <1.
// La math est équivalente (mêmes éléments contractés, attention causale identique) mais l'ordre
// f32 des sommations dans mul_mat dépend de la shape, donc bit-exactness non garantie.
void cp_predict_cached(CPState& s, const float* hidden, const float* cb0_emb, int32_t* out_codes);