//+------------------------------------------------------------------+
//|                                            KronosPredictorS2.mqh |
//|                                      MMQ — Muhammad Minhas Qamar |
//|                                                     www.mql5.com |
//+------------------------------------------------------------------+
#property copyright "MMQ — Muhammad Minhas Qamar"
#property link      "https://www.mql5.com"
#property version   "1.00"

#ifndef KRONOS_PREDICTOR_S2_MQH
#define KRONOS_PREDICTOR_S2_MQH

#include <Kronos\KronosTokenizerMath.mqh>     // loaders
#include <Kronos\KronosTransformerCore.mqh>   // CrossMHA, RMSNorm, LinearT, AddRowBias

#ifndef KR_PRED_VOCAB
#define KR_PRED_VOCAB   1024
#endif

#define KR_DEP_N_HEADS  4              // dep_layer cross-attn heads (source default)

//+------------------------------------------------------------------+
//| Predictor, decode_s2 stage.                                      |
//+------------------------------------------------------------------+
class CKronosPredictorS2
  {
protected:
   ulong             m_dm;           // 512
   string            m_dir;

   matrix            m_embS1;        // emb_s1.weight : [vocab, d_model] (shared table)

   //--- dep_layer cross-attention (n_heads = 4)
   matrix            m_cWq;
   vector            m_cBq;
   matrix            m_cWk;
   vector            m_cBk;
   matrix            m_cWv;
   vector            m_cBv;
   matrix            m_cWo;
   vector            m_cBo;
   vector            m_depNormW;     // dep_layer.norm.weight

   //--- s2 head
   matrix            m_projS2W;
   vector            m_projS2B;   // head.proj_s2 : [vocab, d_model]

   string            F(const string name) { return m_dir + name + ".bin"; }

public:
   //+---------------------------------------------------------------+
   //| Load the shared emb_s1 table, the dep_layer cross-attention   |
   //| weights (n_heads=4) and the s2 head.                          |
   //+---------------------------------------------------------------+
   bool              Init(const string weight_dir, ulong d_model)
     {
      m_dir = weight_dir;
      m_dm = d_model;
      if(d_model == 0)
        {
         Print("CKronosPredictorS2::Init: d_model=0");
         return false;
        }

      bool ok = true;
      //--- emb_s1 is a lookup table (sibling rows by id) -> keep row-major
      ok &= KronosLoadMatrix(F("embedding_emb_s1_weight"), KR_PRED_VOCAB, m_dm, m_embS1);
      //--- cross-attn q/k/v/out and the s2 head are LinearT weights -> transposed
      ok &= KronosLoadMatrixT(F("dep_layer_cross_attn_q_proj_weight"),   m_dm, m_dm, m_cWq);
      ok &= KronosLoadVector(F("dep_layer_cross_attn_q_proj_bias"),     m_dm,       m_cBq);
      ok &= KronosLoadMatrixT(F("dep_layer_cross_attn_k_proj_weight"),   m_dm, m_dm, m_cWk);
      ok &= KronosLoadVector(F("dep_layer_cross_attn_k_proj_bias"),     m_dm,       m_cBk);
      ok &= KronosLoadMatrixT(F("dep_layer_cross_attn_v_proj_weight"),   m_dm, m_dm, m_cWv);
      ok &= KronosLoadVector(F("dep_layer_cross_attn_v_proj_bias"),     m_dm,       m_cBv);
      ok &= KronosLoadMatrixT(F("dep_layer_cross_attn_out_proj_weight"), m_dm, m_dm, m_cWo);
      ok &= KronosLoadVector(F("dep_layer_cross_attn_out_proj_bias"),   m_dm,       m_cBo);
      ok &= KronosLoadVector(F("dep_layer_norm_weight"),                m_dm,       m_depNormW);

      ok &= KronosLoadMatrixT(F("head_proj_s2_weight"), KR_PRED_VOCAB, m_dm, m_projS2W);
      ok &= KronosLoadVector(F("head_proj_s2_bias"),   KR_PRED_VOCAB,        m_projS2B);

      if(!ok)
        {
         Print("CKronosPredictorS2::Init: load failed");
         return false;
        }
      return true;
     }

   //+---------------------------------------------------------------+
   //| decode_s2: context (L, d_model) from decode_s1 plus the chosen|
   //| s1 ids. Returns s2_logits (L, 1024).                          |
   //|                                                               |
   //| s1_ids may have length L (one sibling per position) or length |
   //| 1 (a single sibling broadcast across all L context rows). The |
   //| length-1 case mirrors the reference capture (s1_pick = argmax |
   //| of the last step), where PyTorch broadcasts (L,d)+(1,d) and   |
   //| the single query sits at RoPE position 0.                     |
   //+---------------------------------------------------------------+
   bool              DecodeS2(const matrix &context, const int &s1_ids[], matrix &s2_logits)
     {
      ulong L = context.Rows();
      int   Q = ArraySize(s1_ids);
      if(L == 0 || (Q != (int)L && Q != 1))
        { PrintFormat("DecodeS2: length mismatch (ids %d, ctx %I64u)", Q, L); return false; }

      //--- sibling_embed = raw emb_s1 table rows (NO sqrt(d) scale here).
      //--- Q rows: either L (per-position) or 1 (broadcast).
      matrix sib = matrix::Zeros((ulong)Q, m_dm);
      for(int t = 0; t < Q; t++)
        {
         ulong id = (ulong)s1_ids[t];
         for(ulong j = 0; j < m_dm; j++)
            sib[t][j] = m_embS1[id][j];
        }
      //--- cross-attention: q = sibling_embed (Q rows), k/v = context (L rows),
      //--- n_heads=4, non-causal. attn has Q rows.
      matrix attn = CrossMHA(sib, context,
                             m_cWq, m_cBq, m_cWk, m_cBk,
                             m_cWv, m_cBv, m_cWo, m_cBo,
                             KR_DEP_N_HEADS);
      //--- dep_layer: RMSNorm(context + attn). Broadcast attn row 0 across
      //--- all L context rows when Q == 1 (matches PyTorch (L,d)+(1,d)).
      matrix sum = matrix::Zeros(L, m_dm);
      for(ulong i = 0; i < L; i++)
        {
         ulong ar = (Q == 1) ? 0 : i;
         for(ulong j = 0; j < m_dm; j++)
            sum[i][j] = context[i][j] + attn[ar][j];
        }
      matrix x2 = RMSNorm(sum, m_depNormW);
      //--- s2 head: d_model -> 1024
      s2_logits = LinearT(x2, m_projS2W);
      AddRowBias(s2_logits, m_projS2B);
      return true;
     }
  };

#endif // KRONOS_PREDICTOR_S2_MQH
//+------------------------------------------------------------------+