//+------------------------------------------------------------------+
//|                                              KronosSelfTests.mq5 |
//|                                      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"
#property script_show_inputs
#property strict

#include <Kronos\KronosTokenizerMath.mqh>
#include <Kronos\KronosTransformerCore.mqh>
#include <Kronos\KronosSampling.mqh>

int g_pass = 0, g_fail = 0;

//+------------------------------------------------------------------+
//| Record a single assertion result.                                |
//+------------------------------------------------------------------+
void CHECK(bool c, string label) { if(c) g_pass++; else { g_fail++; PrintFormat("FAIL: %s", label); } }

//+------------------------------------------------------------------+
//| Build an n x n identity matrix.                                  |
//+------------------------------------------------------------------+
matrix Identity(ulong n)
  {
   matrix I = matrix::Zeros(n, n);
   for(ulong i = 0; i < n; i++)
      I[i][i] = 1.0;
   return I;
  }

//+------------------------------------------------------------------+
//| Unit 1: tokenizer math                                           |
//+------------------------------------------------------------------+
//| BSQ indices -> code -> indices round-trips, with unit-magnitude  |
//| code entries.                                                    |
//+------------------------------------------------------------------+
void Test_BSQ_RoundTrip()
  {
   int probes_s1[] = {0, 1, 2, 512, 1023, 341, 682};
   int probes_s2[] = {0, 1023, 256, 7, 900, 42, 511};
   double q = 1.0 / MathSqrt((double)KR_CODEBOOK_DIM);
   for(int a = 0; a < ArraySize(probes_s1); a++)
      for(int b = 0; b < ArraySize(probes_s2); b++)
        {
         int s1 = probes_s1[a], s2 = probes_s2[b];
         double code[];
         KronosBSQ_IndicesToCode(s1, s2, code);
         bool mag_ok = true;
         for(int i = 0; i < KR_CODEBOOK_DIM; i++)
            if(MathAbs(MathAbs(code[i]) - q) > 1e-12)
               mag_ok = false;
         CHECK(mag_ok, StringFormat("code magnitude for (%d,%d)", s1, s2));
         int r1, r2;
         KronosBSQ_SignsToIndices(code, r1, r2);
         CHECK(r1 == s1 && r2 == s2, StringFormat("round-trip (%d,%d)->(%d,%d)", s1, s2, r1, r2));
        }
  }

//+------------------------------------------------------------------+
//| Normalize/denormalize: stats, zero-mean column, round-trip.      |
//+------------------------------------------------------------------+
void Test_Normalize()
  {
   matrix raw = matrix::Zeros(5, 1);
   raw[0][0]=1;
   raw[1][0]=2;
   raw[2][0]=3;
   raw[3][0]=4;
   raw[4][0]=5;
   matrix norm;
   vector mean, stdv;
   KronosNormalize(raw, norm, mean, stdv);
   CHECK(MathAbs(mean[0]-3.0)<1e-12, "mean=3");
   CHECK(MathAbs(stdv[0]-MathSqrt(2.0))<1e-12, "popstd=sqrt(2)");
   double colsum=0;
   for(ulong i=0;i<5;i++)
      colsum+=norm[i][0];
   CHECK(MathAbs(colsum)<1e-9, "normalized col sum ~0");
   double e0=(1.0-3.0)/(MathSqrt(2.0)+KR_EPS);
   CHECK(MathAbs(norm[0][0]-e0)<1e-12, "first normalized value");
   matrix back;
   KronosDenormalize(norm, mean, stdv, back);
   bool rt=true;
   for(ulong i=0;i<5;i++)
      if(MathAbs(back[i][0]-raw[i][0])>1e-9)
         rt=false;
   CHECK(rt, "denormalize inverts normalize");
  }

//+------------------------------------------------------------------+
//| Timestamp features, including the weekday remap to pandas.       |
//+------------------------------------------------------------------+
void Test_Stamp()
  {
   int s[];
   KronosStamp(D'2023.01.02 13:45', s);
   CHECK(s[0]==45 && s[1]==13 && s[3]==2 && s[4]==1, "minute/hour/day/month");
   CHECK(s[2]==0, "Monday -> pandas weekday 0");
   KronosStamp(D'2023.01.01 00:00', s);
   CHECK(s[2]==6, "Sunday -> pandas weekday 6");
  }

//+------------------------------------------------------------------+
//| Unit 2: transformer core                                         |
//+------------------------------------------------------------------+
//| LinearT computes y = x @ W with the weight stored PRE-TRANSPOSED |
//| ([in,out] == the W^T that KronosLoadMatrixT produces). The same  |
//| nn.Linear [out=3,in=2] = [[1,0],[0,1],[1,1]] is stored here as   |
//| its transpose [in=2,out=3]; then x=[1,2] -> [1,2,3].             |
//+------------------------------------------------------------------+
void Test_Linear()
  {
   matrix X = matrix::Zeros(1,2);
   X[0][0]=1;
   X[0][1]=2;
   matrix W = matrix::Zeros(2,3);   // [in,out] == W^T (as stored by KronosLoadMatrixT)
   W[0][0]=1;
   W[0][1]=0;
   W[0][2]=1;
   W[1][0]=0;
   W[1][1]=1;
   W[1][2]=1;
   matrix Y = LinearT(X,W);
   CHECK(MathAbs(Y[0][0]-1)<1e-12 && MathAbs(Y[0][1]-2)<1e-12 && MathAbs(Y[0][2]-3)<1e-12,
         "LinearT pre-transposed weight convention");
  }

//+------------------------------------------------------------------+
//| SiLU activation at a few reference points.                       |
//+------------------------------------------------------------------+
void Test_SiLU()
  {
   CHECK(MathAbs(SiLU(0.0))<1e-15, "silu(0)=0");
   CHECK(MathAbs(SiLU(1.0)-0.7310585786300049)<1e-12, "silu(1)");
   CHECK(MathAbs(SiLU(-1.0)+0.2689414213699951)<1e-12, "silu(-1)");
  }

//+------------------------------------------------------------------+
//| RMSNorm output has unit root-mean-square.                        |
//+------------------------------------------------------------------+
void Test_RMSNorm()
  {
   matrix X = matrix::Zeros(1,2);
   X[0][0]=3;
   X[0][1]=4;
   vector w = vector::Ones(2);
   matrix Y = RMSNorm(X,w);
   double ms=(Y[0][0]*Y[0][0]+Y[0][1]*Y[0][1])/2.0;
   CHECK(MathAbs(ms-1.0)<1e-4, "RMSNorm output RMS ~1");
  }

//+------------------------------------------------------------------+
//| RoPE is identity at t=0 and norm-preserving elsewhere.           |
//+------------------------------------------------------------------+
void Test_RoPE()
  {
   ulong hd=4;
   matrix x = matrix::Zeros(2,hd);
   x[0][0]=0.3;
   x[0][1]=-1.2;
   x[0][2]=0.7;
   x[0][3]=2.1;
   x[1][0]=0.3;
   x[1][1]=-1.2;
   x[1][2]=0.7;
   x[1][3]=2.1;
   matrix c,s;
   RoPETables(2,hd,c,s);
   matrix y = ApplyRoPE(x,c,s);
   bool id0=true;
   for(ulong k=0;k<hd;k++)
      if(MathAbs(y[0][k]-x[0][k])>1e-12)
         id0=false;
   CHECK(id0, "RoPE t=0 identity");
   double ni=0,no=0;
   for(ulong k=0;k<hd;k++)
     {
      ni+=x[1][k]*x[1][k];
      no+=y[1][k]*y[1][k];
     }
   CHECK(MathAbs(MathSqrt(ni)-MathSqrt(no))<1e-12, "RoPE preserves norm");
  }

//+------------------------------------------------------------------+
//| Causal attention: each row is the mean of values up to that row. |
//+------------------------------------------------------------------+
void Test_Attention_Causal()
  {
   ulong d=4, T=4;
   matrix X = matrix::Zeros(T,d);
   for(ulong i=0;i<T;i++)
      X[i][0]=(double)i;
   matrix Z = matrix::Zeros(d,d), I = Identity(d);
   vector zb = vector::Zeros(d);
   matrix O = MHA(X, Z,zb, Z,zb, I,zb, I,zb, 1);
   CHECK(MathAbs(O[0][0]-0.0)<1e-12, "attn row0 = V[0]");
   CHECK(MathAbs(O[1][0]-0.5)<1e-12, "attn row1 = mean(V[0..1])");
   CHECK(MathAbs(O[3][0]-1.5)<1e-12, "attn row3 = mean(V[0..3])");
  }

//+------------------------------------------------------------------+
//| Block residual wiring -> identity when the out-projections are   |
//| zeroed.                                                          |
//+------------------------------------------------------------------+
void Test_Block_ResidualWiring()
  {
   ulong d=4, ff=8, T=3;
   matrix X = matrix::Zeros(T,d);
   for(ulong i=0;i<T;i++)
      for(ulong j=0;j<d;j++)
         X[i][j]=MathSin(0.7*(double)(i*d+j));
   vector n1=vector::Ones(d), n2=vector::Ones(d);
   matrix I=Identity(d), Z=matrix::Zeros(d,d);
   vector zb=vector::Zeros(d);
//--- weights stored PRE-TRANSPOSED for the new LinearT convention: w1/w3 are
//--- [d,ff] and w2 is [ff,d] (the W^T of the [ff,d]/[d,ff] nn.Linear weights).
//--- all zero here, but the shapes must match X.MatMul(W) dimensions
   matrix W1=matrix::Zeros(d,ff), W3=matrix::Zeros(d,ff), W2=matrix::Zeros(ff,d);
   matrix Y = TransformerBlock(X, n1, I,zb, I,zb, I,zb, Z,zb, 1, n2, W1, W3, W2);
   bool same=true;
   for(ulong i=0;i<T;i++)
      for(ulong j=0;j<d;j++)
         if(MathAbs(Y[i][j]-X[i][j])>1e-12)
            same=false;
   CHECK(same, "block residual wiring -> identity when out-projs zeroed");
  }

//+------------------------------------------------------------------+
//| Unit 3: sampling                                                 |
//+------------------------------------------------------------------+
//| Softmax sums to one and preserves the input ordering.            |
//+------------------------------------------------------------------+
void Test_Softmax()
  {
   double x[]= {1.0,2.0,3.0};
   Softmax(x);
   CHECK(MathAbs(x[0]+x[1]+x[2]-1.0)<1e-12, "softmax sums to 1");
   CHECK(x[2]>x[1] && x[1]>x[0], "softmax preserves order");
  }

//+------------------------------------------------------------------+
//| Top-k drops low logits and keeps the mass on the survivors.      |
//+------------------------------------------------------------------+
void Test_TopK()
  {
   double l[]= {1.0,2.0,3.0,4.0};
   TopKTopPFilter(l,2,1.0);
   double p[];
   ArrayCopy(p,l);
   Softmax(p);
   CHECK(MathAbs(p[0])<1e-15 && MathAbs(p[1])<1e-15, "top-k drops low logits");
   CHECK(MathAbs(p[2]+p[3]-1.0)<1e-12, "top-k mass on survivors");
  }

//+------------------------------------------------------------------+
//| Top-p removes the tail token and keeps the nucleus.              |
//+------------------------------------------------------------------+
void Test_TopP()
  {
   double target[]= {0.5,0.3,0.15,0.05};
   double l[];
   ArrayResize(l,4);
   for(int i=0;i<4;i++)
      l[i]=MathLog(target[i]);
   TopKTopPFilter(l,0,0.9);
   double p[];
   ArrayCopy(p,l);
   Softmax(p);
   CHECK(MathAbs(p[3])<1e-15, "top-p removes tail token");
   CHECK(p[2]>0.0 && p[0]>0.0, "top-p keeps nucleus");
  }

//+------------------------------------------------------------------+
//| Greedy decoding picks the argmax logit.                          |
//+------------------------------------------------------------------+
void Test_Greedy()
  {
   double l[]= {0.2,3.1,2.9,-1.0};
   CHECK(SampleFromLogits(l,1.0,0,1.0,true)==1, "greedy picks argmax");
  }

//+------------------------------------------------------------------+
//| Higher temperature flattens the distribution.                    |
//+------------------------------------------------------------------+
void Test_Temperature()
  {
   double l[]= {0.0,2.0,1.0};
   double a[];
   ArrayCopy(a,l);
   Softmax(a);
   double b[];
   ArrayCopy(b,l);
   for(int i=0;i<3;i++)
      b[i]/=4.0;
   Softmax(b);
   CHECK(b[Argmax(b)] < a[Argmax(a)], "higher T flattens distribution");
  }

//+------------------------------------------------------------------+
//| Multinomial sampling concentrates on a sharply peaked logit.     |
//+------------------------------------------------------------------+
void Test_Multinomial_PeakedSeed()
  {
   double l[]= {0.0,50.0,0.0,0.0};
   MathSrand(12345);
   int hits=0;
   for(int t=0;t<100;t++)
      if(SampleFromLogits(l,1.0,0,0.9,false)==1)
         hits++;
   CHECK(hits==100, "multinomial concentrates on peak");
  }

//+------------------------------------------------------------------+
//| Script entry point: run all units and report the tally.          |
//+------------------------------------------------------------------+
void OnStart()
  {
   g_pass=0;
   g_fail=0;
//--- Unit 1: tokenizer math
   Test_BSQ_RoundTrip();
   Test_Normalize();
   Test_Stamp();
//--- Unit 2: transformer core
   Test_Linear();
   Test_SiLU();
   Test_RMSNorm();
   Test_RoPE();
   Test_Attention_Causal();
   Test_Block_ResidualWiring();
//--- Unit 3: sampling
   Test_Softmax();
   Test_TopK();
   Test_TopP();
   Test_Greedy();
   Test_Temperature();
   Test_Multinomial_PeakedSeed();

   PrintFormat("Kronos library self-test: %d passed, %d failed", g_pass, g_fail);
   if(g_fail==0)
      Print("ALL PASS -- weight-free units 1-3 verified.");
  }
//+------------------------------------------------------------------+