//+------------------------------------------------------------------+
//|                                               KronosSampling.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_SAMPLING_MQH
#define KRONOS_SAMPLING_MQH

#define KR_NEG_INF  (-1.0e308)

//+------------------------------------------------------------------+
//| In-place softmax over an array, with the -inf entries mapped to  |
//| zero probability and the usual max-subtraction for stability.    |
//+------------------------------------------------------------------+
void Softmax(double &x[])
  {
   int n = ArraySize(x);
   double mx = KR_NEG_INF;
   for(int i = 0; i < n; i++)
      if(x[i] > mx)
         mx = x[i];
   double s = 0.0;
   for(int i = 0; i < n; i++)
     {
      x[i] = (x[i] <= KR_NEG_INF * 0.5) ? 0.0 : MathExp(x[i] - mx);
      s += x[i];
     }
   for(int i = 0; i < n; i++)
      x[i] /= s;
  }

//+------------------------------------------------------------------+
//| Return the index order that sorts key[] in descending value.     |
//+------------------------------------------------------------------+
void ArgsortDesc(const double &key[], int &order[])
  {
   int n = ArraySize(key);
   ArrayResize(order, n);
   for(int i = 0; i < n; i++)
      order[i] = i;
   for(int i = 0; i < n - 1; i++)
      for(int j = 0; j < n - 1 - i; j++)
         if(key[order[j]] < key[order[j + 1]])
           { int t = order[j]; order[j] = order[j + 1]; order[j + 1] = t; }
  }

//+------------------------------------------------------------------+
//| Top-k then top-p filtering. Removed logits are set to -inf, at   |
//| least one token is kept, and the first token that crosses the    |
//| cumulative top-p threshold is itself kept.                       |
//+------------------------------------------------------------------+
void TopKTopPFilter(double &logits[], int top_k, double top_p)
  {
   int n = ArraySize(logits);

   if(top_k > 0)
     {
      int k = (int)MathMin(MathMax(top_k, 1), n);
      double tmp[];
      ArrayCopy(tmp, logits);
      int order[];
      ArgsortDesc(tmp, order);
      double kth = logits[order[k - 1]];
      for(int i = 0; i < n; i++)
         if(logits[i] < kth)
            logits[i] = KR_NEG_INF;
     }

   if(top_p < 1.0)
     {
      int order[];
      ArgsortDesc(logits, order);
      double probs[];
      ArrayResize(probs, n);
      for(int i = 0; i < n; i++)
         probs[i] = logits[order[i]];
      Softmax(probs);

      bool remove_sorted[];
      ArrayResize(remove_sorted, n);
      double cum = 0.0;
      for(int i = 0; i < n; i++)
        {
         cum += probs[i];
         remove_sorted[i] = (cum > top_p);
        }
      for(int i = n - 1; i > 0; i--)
         remove_sorted[i] = remove_sorted[i - 1];
      remove_sorted[0] = false;

      for(int i = 0; i < n; i++)
         if(remove_sorted[i])
            logits[order[i]] = KR_NEG_INF;
     }
  }

//+------------------------------------------------------------------+
//| Index of the largest element (greedy pick).                      |
//+------------------------------------------------------------------+
int Argmax(const double &x[])
  {
   int n = ArraySize(x), best = 0;
   double bv = x[0];
   for(int i = 1; i < n; i++)
      if(x[i] > bv)
        {
         bv = x[i];
         best = i;
        }
   return best;
  }

//+------------------------------------------------------------------+
//| Draw one index from a probability array (inverse-CDF sampling).  |
//+------------------------------------------------------------------+
int MultinomialDraw(const double &probs[])
  {
   double r = (double)MathRand() / 32767.0;
   double c = 0.0;
   int n = ArraySize(probs);
   for(int i = 0; i < n; i++)
     {
      c += probs[i];
      if(r <= c)
         return i;
     }
   return n - 1;
  }

//+------------------------------------------------------------------+
//| Full decode step: temperature scale, top-k/top-p filter,         |
//| softmax, then sample. greedy=true returns the argmax instead.    |
//+------------------------------------------------------------------+
int SampleFromLogits(const double &logits_in[], double T, int top_k, double top_p, bool greedy)
  {
   double l[];
   ArrayCopy(l, logits_in);
   int n = ArraySize(l);
   for(int i = 0; i < n; i++)
      l[i] /= T;
   if(top_k > 0 || top_p < 1.0)
      TopKTopPFilter(l, top_k, top_p);
   double probs[];
   ArrayCopy(probs, l);
   Softmax(probs);
   return greedy ? Argmax(probs) : MultinomialDraw(probs);
  }

#endif // KRONOS_SAMPLING_MQH
//+------------------------------------------------------------------+