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

#ifndef KRONOS_TRANSFORMER_CORE_MQH
#define KRONOS_TRANSFORMER_CORE_MQH

//--- transformer core constants
#define KR_ROPE_BASE   10000.0
#define KR_NORM_EPS    1e-5

//+------------------------------------------------------------------+
//| Linear (PyTorch convention): y = x @ W^T. The weight is loaded   |
//| pre-transposed (KronosLoadMatrixT stores W as [in,out] == W^T),  |
//| so this is a plain MatMul with NO per-call transpose. Profiling  |
//| showed the old per-call W.Transpose() was ~60% of the forward    |
//| pass; the transpose now happens once at load. Math is identical. |
//+------------------------------------------------------------------+
matrix LinearT(const matrix &X, const matrix &W) { return X.MatMul(W); }

//+------------------------------------------------------------------+
//| Add a per-output-column bias to every row of Y.                  |
//+------------------------------------------------------------------+
void AddRowBias(matrix &Y, const vector &b)
  {
   ulong T = Y.Rows(), O = Y.Cols();
   for(ulong i = 0; i < T; i++)
      for(ulong j = 0; j < O; j++)
         Y[i][j] += b[j];
  }

//+------------------------------------------------------------------+
//| RMSNorm (row-wise): out = x / sqrt(mean(x^2) + eps) * weight.    |
//+------------------------------------------------------------------+
matrix RMSNorm(const matrix &X, const vector &w, double eps = KR_NORM_EPS)
  {
   ulong T = X.Rows(), D = X.Cols();
   matrix out = matrix::Zeros(T, D);
   for(ulong i = 0; i < T; i++)
     {
      double ms = 0.0;
      for(ulong j = 0; j < D; j++)
         ms += X[i][j] * X[i][j];
      ms /= (double)D;
      double scale = 1.0 / MathSqrt(ms + eps);
      for(ulong j = 0; j < D; j++)
         out[i][j] = X[i][j] * scale * w[j];
     }
   return out;
  }

//+------------------------------------------------------------------+
//| SiLU activation: x * sigmoid(x).                                 |
//+------------------------------------------------------------------+
double SiLU(double x) { return x / (1.0 + MathExp(-x)); }

//+------------------------------------------------------------------+
//| SwiGLU feed-forward: w2( SiLU(w1 x) * w3 x ), all bias-free.     |
//+------------------------------------------------------------------+
matrix SwiGLU(const matrix &X, const matrix &W1, const matrix &W3, const matrix &W2)
  {
   matrix H = LinearT(X, W1);
   matrix G = LinearT(X, W3);
   ulong T = H.Rows(), FF = H.Cols();
   matrix gated = matrix::Zeros(T, FF);
   for(ulong i = 0; i < T; i++)
      for(ulong j = 0; j < FF; j++)
         gated[i][j] = SiLU(H[i][j]) * G[i][j];
   return LinearT(gated, W2);
  }

//+------------------------------------------------------------------+
//| Build the RoPE cos/sin tables (T, hd). emb = cat(freqs, freqs),  |
//| so the two halves of each row are duplicated.                    |
//+------------------------------------------------------------------+
void RoPETables(ulong T, ulong hd, matrix &cosT, matrix &sinT)
  {
   ulong half = hd / 2;
   cosT = matrix::Zeros(T, hd);
   sinT = matrix::Zeros(T, hd);
   for(ulong t = 0; t < T; t++)
      for(ulong k = 0; k < half; k++)
        {
         double inv_freq = 1.0 / MathPow(KR_ROPE_BASE, (2.0 * (double)k) / (double)hd);
         double ang = (double)t * inv_freq;
         double c = MathCos(ang), s = MathSin(ang);
         cosT[t][k] = c;
         cosT[t][k + half] = c;
         sinT[t][k] = s;
         sinT[t][k + half] = s;
        }
  }

//+------------------------------------------------------------------+
//| Apply RoPE: out = x*cos + rotate_half(x)*sin, where              |
//| rotate_half(x) = cat(-x2, x1).                                   |
//+------------------------------------------------------------------+
matrix ApplyRoPE(const matrix &X, const matrix &cosT, const matrix &sinT)
  {
   ulong T = X.Rows(), hd = X.Cols(), half = hd / 2;
   matrix out = matrix::Zeros(T, hd);
   for(ulong t = 0; t < T; t++)
      for(ulong k = 0; k < hd; k++)
        {
         double rh = (k < half) ? -X[t][k + half] : X[t][k - half];
         out[t][k] = X[t][k] * cosT[t][k] + rh * sinT[t][k];
        }
   return out;
  }

//+------------------------------------------------------------------+
//| Extract one head's columns [c0, c0+hd) into a (T, hd) matrix.    |
//+------------------------------------------------------------------+
matrix SliceCols(const matrix &M, ulong c0, ulong hd)
  {
   ulong T = M.Rows();
   matrix s = matrix::Zeros(T, hd);
   for(ulong i = 0; i < T; i++)
      for(ulong k = 0; k < hd; k++)
         s[i][k] = M[i][c0 + k];
   return s;
  }

//+------------------------------------------------------------------+
//| Write one head's (T, hd) output back into columns [c0, c0+hd).   |
//+------------------------------------------------------------------+
void WriteCols(matrix &M, const matrix &S, ulong c0)
  {
   ulong T = S.Rows(), hd = S.Cols();
   for(ulong i = 0; i < T; i++)
      for(ulong k = 0; k < hd; k++)
         M[i][c0 + k] = S[i][k];
  }

//+------------------------------------------------------------------+
//| Scaled dot-product attention for one head.                       |
//| Qh(Tq,hd), Kh(Tk,hd), Vh(Tk,hd). causal=true masks j>i and so    |
//| needs Tq==Tk. scores = Q.Kt * inv_sqrt -> stable row-softmax     |
//| -> weights . V. Returns Oh(Tq,hd).                               |
//+------------------------------------------------------------------+
matrix SDPA(const matrix &Qh, const matrix &Kh, const matrix &Vh, bool causal)
  {
   ulong Tq = Qh.Rows(), Tk = Kh.Rows(), hd = Qh.Cols();
   double inv_sqrt = 1.0 / MathSqrt((double)hd);
//--- scaled scores
   matrix scores = Qh.MatMul(Kh.Transpose());   // (Tq,Tk)
   scores *= inv_sqrt;
//--- row-wise softmax with optional causal mask
   for(ulong i = 0; i < Tq; i++)
     {
      //--- causal mask: query i attends to keys 0..i only
      ulong jmax = causal ? i : (Tk - 1);
      double smax = -DBL_MAX;
      for(ulong j = 0; j <= jmax; j++)
         if(scores[i][j] > smax)
            smax = scores[i][j];
      double denom = 0.0;
      for(ulong j = 0; j < Tk; j++)
        {
         if(causal && j > i)
           {
            scores[i][j] = 0.0;
            continue;
           }
         double e = MathExp(scores[i][j] - smax);
         scores[i][j] = e;
         denom += e;
        }
      double inv = 1.0 / denom;
      for(ulong j = 0; j < Tk; j++)
         scores[i][j] *= inv;   // masked entries already 0
     }
   return scores.MatMul(Vh);   // (Tq,hd)
  }

//+------------------------------------------------------------------+
//| Causal multi-head self-attention with RoPE.                      |
//| q/k/v/out projections carry a bias.                              |
//+------------------------------------------------------------------+
matrix MHA(const matrix &X,
           const matrix &Wq, const vector &bq,
           const matrix &Wk, const vector &bk,
           const matrix &Wv, const vector &bv,
           const matrix &Wo, const vector &bo,
           int n_heads)
  {
   ulong T = X.Rows();
   ulong d_model = X.Cols();
   ulong hd = d_model / (ulong)n_heads;
//--- q/k/v projections (with bias)
   matrix Q = LinearT(X, Wq);
   AddRowBias(Q, bq);
   matrix K = LinearT(X, Wk);
   AddRowBias(K, bk);
   matrix V = LinearT(X, Wv);
   AddRowBias(V, bv);
//--- shared RoPE tables for all heads
   matrix cosT, sinT;
   RoPETables(T, hd, cosT, sinT);
   matrix ctx = matrix::Zeros(T, d_model);
//--- per-head causal attention
   for(int h = 0; h < n_heads; h++)
     {
      ulong c0 = (ulong)h * hd;
      matrix Qs = SliceCols(Q, c0, hd);   // named locals: MQL5 passes matrices by reference only
      matrix Ks = SliceCols(K, c0, hd);
      matrix Qh = ApplyRoPE(Qs, cosT, sinT);
      matrix Kh = ApplyRoPE(Ks, cosT, sinT);
      matrix Vh = SliceCols(V, c0, hd);

      matrix Oh = SDPA(Qh, Kh, Vh, true);   // causal
      WriteCols(ctx, Oh, c0);
     }
//--- output projection (with bias)
   matrix out = LinearT(ctx, Wo);
   AddRowBias(out, bo);
   return out;
  }

//+------------------------------------------------------------------+
//| Non-causal multi-head CROSS-attention with RoPE (inference).     |
//| q from Xq (sibling embed), k/v from Xkv (context); each query    |
//| attends to ALL key positions. q/k/v/out carry a bias. The        |
//| predictor's dep_layer uses n_heads=4 (head_dim=128), NOT 8.      |
//|                                                                  |
//| RoPE quirk (matches RotaryPositionalEmbedding.forward): the      |
//| cos/sin cache is sized to the QUERY length and the same cache is |
//| applied to the keys. So when Tq==1 (a single broadcast s1 pick)  |
//| every key is rotated at position 0; when Tq==Tk keys rotate by   |
//| their own positions. We index the key rotation by (Tq==1 ? 0:j), |
//| reusing the query's table.                                       |
//+------------------------------------------------------------------+
matrix CrossMHA(const matrix &Xq, const matrix &Xkv,
                const matrix &Wq, const vector &bq,
                const matrix &Wk, const vector &bk,
                const matrix &Wv, const vector &bv,
                const matrix &Wo, const vector &bo,
                int n_heads)
  {
   ulong Tq = Xq.Rows();
   ulong Tk = Xkv.Rows();
   ulong d_model = Xq.Cols();
   ulong hd = d_model / (ulong)n_heads;   // scaling handled inside SDPA
//--- q from Xq, k/v from Xkv (all with bias)
   matrix Q = LinearT(Xq,  Wq);
   AddRowBias(Q, bq);
   matrix K = LinearT(Xkv, Wk);
   AddRowBias(K, bk);
   matrix V = LinearT(Xkv, Wv);
   AddRowBias(V, bv);
//--- single cache sized to the query length, reused for keys (PyTorch quirk).
//--- valid only when Tq==1 (broadcast) or Tq==Tk (position-matched); any other
//--- mix would have failed PyTorch's broadcast, so reject it loudly.
   if(!(Tq == 1 || Tq == Tk))
     { PrintFormat("CrossMHA: unsupported Tq=%I64u, Tk=%I64u (need Tq==1 or Tq==Tk)", Tq, Tk); }
   matrix cosT, sinT;
   RoPETables(Tq, hd, cosT, sinT);
//--- per-key rotation table: row j uses position (Tq==1 ? 0 : j)
   matrix cosK = matrix::Zeros(Tk, hd), sinK = matrix::Zeros(Tk, hd);
   for(ulong j = 0; j < Tk; j++)
     {
      ulong p = (Tq == 1) ? 0 : j;            // broadcast when single query
      for(ulong k = 0; k < hd; k++)
        {
         cosK[j][k] = cosT[p][k];
         sinK[j][k] = sinT[p][k];
        }
     }
   matrix ctx = matrix::Zeros(Tq, d_model);
//--- per-head non-causal cross-attention
   for(int h = 0; h < n_heads; h++)
     {
      ulong c0 = (ulong)h * hd;
      matrix Qs = SliceCols(Q, c0, hd);
      matrix Ks = SliceCols(K, c0, hd);
      matrix Qh = ApplyRoPE(Qs, cosT, sinT);
      matrix Kh = ApplyRoPE(Ks, cosK, sinK);
      matrix Vh = SliceCols(V, c0, hd);

      matrix Oh = SDPA(Qh, Kh, Vh, false);   // non-causal: all keys
      WriteCols(ctx, Oh, c0);
     }
//--- output projection (with bias)
   matrix out = LinearT(ctx, Wo);
   AddRowBias(out, bo);
   return out;
  }

//+------------------------------------------------------------------+
//| Pre-norm block:                                                  |
//|   x += MHA(RMSNorm1(x)); x += SwiGLU(RMSNorm2(x)).               |
//+------------------------------------------------------------------+
matrix TransformerBlock(const matrix &X,
                        const vector &norm1_w,
                        const matrix &Wq, const vector &bq,
                        const matrix &Wk, const vector &bk,
                        const matrix &Wv, const vector &bv,
                        const matrix &Wo, const vector &bo,
                        int n_heads,
                        const vector &norm2_w,
                        const matrix &W1, const matrix &W3, const matrix &W2)
  {
//--- attention sub-block: x1 = x + MHA(RMSNorm1(x))
   matrix n1 = RMSNorm(X, norm1_w);                 // named locals: no temporaries by reference
   matrix a  = MHA(n1, Wq, bq, Wk, bk, Wv, bv, Wo, bo, n_heads);
   matrix x1 = X + a;
//--- feed-forward sub-block: out = x1 + SwiGLU(RMSNorm2(x1))
   matrix n2 = RMSNorm(x1, norm2_w);
   matrix f  = SwiGLU(n2, W1, W3, W2);
   return x1 + f;
  }

#endif // KRONOS_TRANSFORMER_CORE_MQH
//+------------------------------------------------------------------+