A Simple Training of an AI Transformer to Play
Tic-Tac-Toe in C/C++

// player.h
// http://forejune.co/cuda/
#ifndef __player_H__
#define __player_H__

const int SOS = 9;    // start of sequence
const int EOS = 10;   // end of sequence

enum Player {EMPTY=0, X, O};

#endif

// util.h

#ifndef __UTIL_H__
#define __UTIL_H__
#include <vector>
#include <algorithm>
#include <math.h>
#include <iostream>
#include <iomanip>
#include <random>

using namespace std;

typedef vector<vector<double>>  matrixd;

// an activation function
double f(double z);

// Relu activation function with matrix input
matrixd relu(const matrixd &X);

// Derivative of Relu
matrixd relu_deriv (const matrixd &X);


// calculates entropy loss
double crossEntropy(const matrixd &P, const vector<int> &target);

// add two matrices
matrixd addMat(const matrixd &A, const matrixd &B);

// Transpose matrix n x m, B = A^T : m x n
matrixd transpose(const matrixd& A);

// Matrix multiplication   C = A X B   A: nxm, B: mxr, C: nxr
matrixd matmul(const matrixd& A, const matrixd& B);

// multiply a matrix by a scalar
matrixd mulMats (const matrixd &A, const double s);

// scale a matrix
void scaleMat(matrixd &A, const double s);

// apply masking to S
matrixd causalMask(matrixd S);

// update weight matrix W -= dL_dW * eta
void updateWeight(matrixd &W, const matrixd &dL_dW, const double eta);

// initialize an input matrix with certain random values
void initMatrix(matrixd& M, const int n, const int m);

void initMatrix(matrixd& M, const int n, const int m, double dModel);

void saveMatrix(const matrixd &W, const int n, const int m, FILE *fp);

void readMatrix(matrixd &W, const int n, const int m, FILE *fp);

// print a token dictionary
void printDictionary(unordered_map<string, int> &d);

// print a matrix
void printMatrix(const matrixd &y);

// print a vector
void printVec(const vector<int>& v);

// Add & Norm
matrixd addNnorm(const matrixd &A, const matrixd &B);

void clip(matrixd &A, double max_val);

// softmax, 1D vector input
vector<double> softmax(const vector<double> &v);

// softmax activation function, 2D input vector
matrixd softmax(const matrixd& input);

// derivative function of softmax
matrixd softmax_derivative(const matrixd &dL_dP, const matrixd &P);
#endif

// transGpt.h
#ifndef TRANSGPT_H
#define TRANSGPT_H

#include "player.h"

class Tokenizer
{
private:
    unordered_map<string, int> words;
    unordered_map<int, string> wordIndex;
    int nTokens;        // number of words

    // special tokens
    const string PAD = "<PAD>";
    const string UNK = "<UNK>";
    const string BOS = "<BOS>";
    const string EOS = "<EOS>";
    const string SEP = "<SEP>";

    void addWord(const string &w) 
    {
        if (words.find(w) == words.end()) {
            int id = nTokens;
            words[w] = id;
            wordIndex[id] = w;
            nTokens++;
        }
    }

    vector<string> split(const string &str) 
    {
        vector<string> tokens;
        stringstream ss(str);
        string word;
        while (ss >> word) {
            tokens.push_back(word);
        }
        return tokens;
    }

public:
    Tokenizer();
    Tokenizer(ifstream &fs);
    unordered_map<string, int>  getTokensWords();
    vector<int> tokenize(const string& str);
    string detokenize(const vector<int>& tokenIDs);
    int get_nTokens() const;
    unordered_map<int, string> get_wordIndex();
};

class Embedding
{
private:
    int nTokens;        // number of words
    int dim;            // embedding dimension
    matrixd em_matrix;  // embedding matrix;
public:
    Embedding(int sequence_length, int embeddingDimension);
    // create embedding matrix with tokenIDs
    matrixd embed(const vector<int>& tokenIDs);
    int get_embedDim() const;
    int get_nTokens() const;
    void backProp(const vector<int>& tokenIDs, const matrixd& dL_dX,
                         double eta);

};

// Positional Encoding
class PositionalEncoding
{
private:
    int maxTokens; // maximum sequence length
    int dModel;    // embedding dimension
    matrixd PE;    // positional_encodings matrix;
public:
    PositionalEncoding(int max_seq_length, int embedding_dimension);
    matrixd addPE(const matrixd& embeddings);
    matrixd getPE();
};

// Add & Norm
class AddnNorm
{
private:
    vector<double> gamma;   // size dModel
    vector<double> beta;    // size dModel

    matrixd Zhat;           // normalized values (nt x d)
    vector<double> mean;    // per row (nt)
    vector<double> var;     // per row (nt)

    double eps;             // small value

public:
   // constructor
   AddnNorm(int dModel);

   // forward computation
   matrixd  forward (const matrixd &A, const matrixd &B);

   // Assume forward already filled: Zhat, mean, var
   matrixd backProp(const matrixd &X, const matrixd &dL_dY, double eta);

   // Optional setters (to plug in forward results)
   void setCache(const matrixd& zhat, const vector<double>& m,
                  const vector<double>& v);

    // Accessors (optional)
    const vector<double>& getGamma();
    const vector<double>& getBeta();  
    void saveGammaBeta(FILE *fp);
    void readGammaBeta(FILE *fp);
};

// Feed-forward Network
class FeedForward
{
private:
    matrixd W1, W2;
    int dModel, d_ff;

    // cache
    matrixd F1;   // after activation function ReLU
    matrixd Z;    // before ReLU (needed for derivative)

public:
    // constructor
    FeedForward(int dModel, int d_ff_);
    // ------------------------------
    // Forward (for cache)
    // ------------------------------
    matrixd FFoutput(const matrixd &X);
 // ------------------------------
    // Backprop
    // ------------------------------
    matrixd backProp(const matrixd &X, const matrixd &dL_dY,  double eta);
    void saveWeights(FILE *fp);
    void readWeights(FILE *fp);
};

// ----------- MultiHeadAttention -----------
class MultiHeadAttention {
private:
    int dModel, nHeads, d_k;

    vector<matrixd> WQ, WK, WV; // per head
    matrixd WO;

    // cache
    vector<matrixd> Qs, Ks, Vs, Ps, As;
public:
    MultiHeadAttention(int dModel, int nHeads);
    matrixd computeAttention(const matrixd &X_Q, const matrixd &X_K, const matrixd &X_V,
                             bool mask);
     matrixd backProp( const matrixd &X_Q, const matrixd &X_K, const matrixd &X_V,
        const matrixd &dL_dH, double eta, bool mask);
     void saveWeights(FILE *fp);
     void readWeights(FILE *fp);
};

// ----------------- OutputLayer ---------------
class OutputLayer {
private:
    matrixd W;   // dModel x vocab_size
    int dModel, vocab_size;

public:
    OutputLayer(int dModel, int vocab_size);
    matrixd forward(const matrixd &X);
    matrixd backward(const matrixd &X, const matrixd &P,
                     const vector<int> &targets, double eta);
    void saveWeights(FILE *fp);
    void readWeights(FILE *fp);
};

// ------------------- GPT-style Transformer ---------------
class GPTBlock {
private:
    MultiHeadAttention mha;
    AddnNorm norm1;
    FeedForward ffn;
    AddnNorm norm2;

public:
    GPTBlock(int dModel, int nHeads, int d_ff);
    matrixd forward(const matrixd &X);
    matrixd backward(const matrixd &X, const matrixd &dL_dY, double eta);
    void saveWeights(FILE *fp);
    void readWeights(FILE *fp);
};

class MiniGPT {
private:
    Embedding embed;
    PositionalEncoding pe;
    vector<GPTBlock> layers;
    OutputLayer output;

    int seq_len, vocab_size;

public:
    MiniGPT(int vocab_size, int seq_len, int dModel, 
		    int nHeads, int d_ff, int nLayers);
    matrixd forward(const vector<int> &tokens);
    double trainStep(const vector<int> &tokens, double eta);
    int saveWeights(char fname[]);
    int readWeights(char fname[]);
};
#endif
// util.cpp: helper functions
// http://forejune.co/cuda/

#include "util.h"
#include <random>

using namespace std;

// Activation function: using ReLU (Rectified Linear Unit)
double f(double z)
{
  return max(0.0, z);
}

// Relu activation function with matrix input
matrixd relu(const matrixd &X)
{
    int rows = X.size();
    int cols = X[0].size();
    matrixd Y = X;	// same dimension of X
   
    for(int i = 0; i < rows; i++)
        for(int j = 0; j < cols; j++)
            Y[i][j] = max(0.0, X[i][j]);

    return Y;
}

// Derivative of Relu
matrixd relu_deriv(const matrixd &X)
{
    int rows = X.size();
    int cols = X[0].size();
    matrixd Y = X;	// same dimension as X
   
    for(int i = 0; i < rows; i++)
        for(int j = 0; j < cols; j++)
	  if ( X[i][j] > 0 )
            Y[i][j] = 1;
	  else
            Y[i][j] = 0;

    return Y;
}

// calculates entropy loss 
double crossEntropy(const matrixd &P, const vector<int> &target)
{
  double L = 0;
  int nt = P.size();
  for (int i = 0; i < nt; i++){
    int k = target[i];
    L -= log( P[i][k] + 1e-12 );
  }

  return L / nt;
}

// C = A + B
matrixd addMat(const matrixd &A, const matrixd &B)
{
     int n = A.size();          // rows
     int m = A[0].size();       // cols

     matrixd C = A;       // same dimension as A
     for (int i = 0; i < n; i++)
       for(int j = 0; j < m; j++)
         C[i][j] = A[i][j] + B[i][j];

     return C;
}

// Transpose matrix n x m, B = A^T : m x n
matrixd transpose(const matrixd& A)
{
   int n = A.size();    //number of rows
   int m = A[0].size(); //number of columns

   matrixd B(m, vector<double>(n));
   for (int i = 0; i < n; i++)
     for (int j = 0; j < m; j++)
       B[j][i] = A[i][j];

   return B;
}

// Matrix multiplication   C = A X B   A: nxm, B: mxr, C: nxr
matrixd matmul(const matrixd& A, const matrixd& B)
{
  int n  = A.size();
  int m = A[0].size();
  int r = B[0].size();

  matrixd C(n, vector<double>(r, 0));

  for (int i = 0; i < n; i++)
    for (int j = 0; j < r; j++)
      for (int k = 0; k < m; k++)
        C[i][j] += A[i][k] * B[k][j];

   return C;
}

// scale a matrix
void scaleMat(matrixd &A, const double s)
{
   int m = A.size();	// number of rows
   int n = A[0].size();    // number of columns

   for (int i = 0; i < m; i++)
     for (int j = 0; j < n; j++)
       A[i][j] *= s;
}

// maultiply a matrix by a scalar s
matrixd mulMats (const matrixd &A, const double s)
{
   int m = A.size();	// number of rows
   int n = A[0].size();    // number of columns
  
   matrixd B = A;

   for (int i = 0; i < m; i++)
     for (int j = 0; j < n; j++)
       B[i][j] *= s;

   return B;
}

void initMatrix(matrixd& M, const int n, const int m)
{
   // resizing 2D vector to n x m with initial values 0
   M.assign(n, vector<double>(m, 0));
   for (int i = 0; i < n; ++i)
     for (int j = 0; j < m; ++j)
        M[i][j] = ((double)rand() / RAND_MAX - 0.5);
}

void saveMatrix(const matrixd &W, const int n, const int m, FILE *fp)
{
   for (int i = 0; i < n; i++){
    for (int j = 0; j < m; j++)
      fprintf(fp, "%9.4f", W[i][j]);
    fprintf(fp, "\n");
  }
}

void readMatrix(matrixd &W, const int n, const int m, FILE *fp)
{
   for (int i = 0; i < n; i++){
    for (int j = 0; j < m; j++)
      fscanf(fp, "%lf", &W[i][j]);
  }
}

void printDictionary(unordered_map<string, int> &d)
{
  unordered_map<string, int>::iterator it = d.begin();
  int k = 0;
  while ( it != d.end() ){
    cout << right << setw(12) << it->first << ": " << setw(3) << it->second;
    it++;
    k++;
    if ( k % 5 == 0 )
      cout << endl;
  }
}

void printMatrix(const matrixd &y)
{
    int cols = y.size();
    int rows = y[0].size();

    for (int i = 0; i < cols; ++i) {
        cout << "Token " << i << ": [";
        for (int j = 0; j < rows; ++j) {
            cout << fixed << setw(7) << setprecision(3) << y[i][j];
            if (j < 4) cout << ", ";
        }
        cout << " ]\n";
    }
}

void printVec(const vector<int>& v) 
{
    for (int x : v) cout << x << " ";
    cout << endl;
}

matrixd addNnorm(const matrixd &A, const matrixd &B)
{
     auto x = addMat(A, B);     // add two matrices
     int n = x.size();          // rows
     int m = x[0].size();       // cols
     matrixd y = x;             // output same size as x 

     // normalize the result
     double epsilon = 1e-5;

     for (int i = 0; i < n; i++) {
       double sum = 0;          // sum of one row
       double std = 0;          //sigma_i square
       for (int j = 0; j < m; j++)
         sum += x[i][j];
        double mean = sum / m;  // mu_i, mean of i-th row

        sum = 0;
        for (int j = 0; j < m; j++) {
          double diff = x[i][j] - mean;
          sum += diff * diff;
        }
        std = sum / m;
        for (int j = 0; j < m; j++) {
          double xd = (x[i][j] - mean) / sqrt(std + epsilon);
          y[i][j] = xd;         // gamma = 1, beta = 0;
        }
     }

     return y;  // final output of Add & norm
}

matrixd scaledAttention(const matrixd &Q, const matrixd &K, const matrixd &V,
        bool mask)
{
  double d_k, scale;
  matrixd KT, A;	// K transpose, attention
  
  d_k = Q[0].size();
  scale = 1.0 / sqrt( d_k );
  KT = transpose ( K );
  A = matmul(Q, KT);
  scaleMat(A, scale); 

  if ( mask ) {
     for (int i = 0; i < A.size(); i++)
       for (int j = i+1; j < A[0].size(); j++) 
         A[i][j] = -1e9;	// causal mask
  }

  A = softmax( A );
  
  matrixd output = matmul(A, V);

  return output;
}

matrixd causalMask(matrixd S) 
{
    int n = S.size();
    for (int i = 0; i < n; i++) {
        for (int j = i+1; j < n; j++) {
            S[i][j] = -1e9;  // effectively -∞
        }
    }
    return S;
}

void updateWeight(matrixd &W, const matrixd &dL_dW, const double eta)
{
   for (int i = 0; i < W.size(); i++)
     for (int j = 0; j < W[0].size(); j++) {
	 double grad = dL_dW[i][j];
	 grad = max(-1.0, min(1.0, grad));	// clipping the gradient
         W[i][j] -= eta * grad;
      }
}

void clip(matrixd &A, double max_val = 5.0)
{
    for (auto &row : A)
        for (auto &x : row)
            if (x > max_val) x = max_val;
            else if (x < -max_val) x = -max_val;
}

// softmax activation function, 1D input vector
vector<double> softmax(const vector<double> &z)
{
  int m = z.size();
  vector<double> y(m, 0);

   //maximum value of vector
  double maxVal = *max_element(z.begin(), z.end());

  double sum = 0;
  // Subtract max for numerical stability
  for (int j = 0; j < m; j++) {
    y[j] = exp(z[j] - maxVal);
    sum += y[j];
  }

  // Normalize
  for (int j = 0; j < m; j++)
    y[j] /= sum;

  return y;
}

// softmax activation function, 2D input vector
matrixd softmax(const matrixd& input)
{
  int n = input.size(), m = input[0].size();    // n x m matrix 

  // 2D output vector n x m y[n][m]
  vector<vector<double>> y(n, vector<double>(m));

  for (int i = 0; i < n; i++) {
    //maximum value of the row
    double maxVal = *max_element(input[i].begin(), input[i].end());

    double sum = 0;
    // Subtract max for numerical stability
    for (int j = 0; j < m; j++) {
      y[i][j] = exp(input[i][j] - maxVal);
      sum += y[i][j];
    }

    // Normalize
    for (int j = 0; j < input[i].size(); j++)
      y[i][j] /= sum;
  }

  return y;
}

// derivative function of softmax
matrixd softmax_derivative(const matrixd &dL_dP, const matrixd &P)
{
    int nt = P.size(), ns = P[0].size();
    matrixd dL_dS(nt, vector<double>(ns,0));

    for (int i = 0; i < nt; i++)
    {
          double dot = 0.0;
          for (int j = 0; j < ns; j++)
            dot += P[i][j] * dL_dP[i][j];  // dot product of two row-vectors

          for (int j = 0; j < ns; j++)  // dot used by all columns of dL/dA
             dL_dS[i][j] = P[i][j] * (dL_dP[i][j] - dot);
    }

    return dL_dS;
}

// http://forejune.co/cuda/
// transGpt.cpp : GPT-style (decoder-only) transformer classes 

#include <iostream>
#include <fstream>
#include <vector>
#include <string>
#include <unordered_map>
#include <cmath>
#include <random>
#include <algorithm>
#include <iomanip>
#include "util.h"
#include "transGpt.h"

using namespace std;

// ------------------------ Tokenizer class --------------------------
// skip here, see earlier videos

//------------------   Embedding Class ------------------
Embedding::Embedding(int sequence_length, int embeddingDimension)
{ 
  nTokens = sequence_length;
  dim = embeddingDimension;
  initMatrix(em_matrix, nTokens, dim);	
}
    
// create embedding matrix with tokenIDs 
matrixd Embedding::embed(const vector<int>& tokenIDs) 
{
  matrixd embeddings;
  embeddings.reserve(tokenIDs.size()); // set capacity of embeddings vector
  for(int i = 0; i < tokenIDs.size(); i++) {
    if (tokenIDs[i] >= 0 && tokenIDs[i] < nTokens)
       embeddings.push_back(em_matrix[tokenIDs[i]]);
    else 
      // Out-of-dictionary tokens
      embeddings.push_back(vector<double>(dim, 0.0));
   }
    
   return embeddings;
}

int Embedding::get_embedDim() const 
{ 
       return dim; 
}

int Embedding::get_nTokens() const 
{ 
       return nTokens; 
}

void Embedding::backProp(const vector<int>& tokenIDs,
                         const matrixd& dL_dX,
                         double eta)
{
    int nt = tokenIDs.size();

    for (int i = 0; i < nt; i++) {
        int token = tokenIDs[i];

        for (int j = 0; j < dim; j++){
	   double grad = dL_dX[i][j];
	   grad = max(-1.0, min(1.0, grad));      // clipping the gradient
           //em_matrix[token][j] -= eta * dL_dX[i][j];
           em_matrix[token][j] -= eta * grad;
	 }
    }
}

// --------------------- Positional Encoding class ---------------
PositionalEncoding::PositionalEncoding(int max_seq_length, int embed_dimension) 
{
      maxTokens = max_seq_length;
      dModel = embed_dimension;

      PE.resize(maxTokens, vector<double>(dModel));
      for (int pos = 0; pos < maxTokens; pos++) {
    for (int i = 0; i < dModel/2; i++) {
    		PE[pos][2*i] = sin( pos / pow(10000, 2.0*i/dModel) );
	    	PE[pos][2*i+1] = cos( pos / pow(10000, 2.0*i/dModel) );
	}
      }
}

// add positional encodings to embeddings
matrixd PositionalEncoding::addPE(const matrixd& embeddings) 
{
    matrixd x = embeddings;
    int nTokens = embeddings.size();

    for (int pos = 0; pos < nTokens; pos++) 
      for (int i = 0; i < dModel; i++) 
         x[pos][i] += PE[pos][i];

     return x;
}

matrixd PositionalEncoding::getPE() 
{
    return PE;
}

// -------------------- Add & Norm Class---------------------------------
// constructor
AddnNorm::AddnNorm(int dModel)
{
    gamma.assign(dModel, 1.0);
    beta.assign(dModel, 0.0),
    eps = 1e-5;
}

matrixd AddnNorm::forward (const matrixd &A, const matrixd &B)
{
    matrixd x = addMat(A, B); // add two matrices
    int n = x.size();         // rows
    int m = x[0].size();      // cols
    mean.resize( n );
    var.resize( n );
    Zhat.resize(n, vector<double>(m, 0));

    // normalize the result
    for (int i = 0; i < n; i++) {
       double sum = 0;      // sum of one row
       var[i] = 0;
       for (int j = 0; j < m; j++)
          sum += x[i][j];
       mean[i] = sum / m;  // mu_i, mean of i-th row

       sum = 0;
       for (int j = 0; j < m; j++) {
          double diff = x[i][j] - mean[i];
          sum += diff * diff;
       }
       var[i] = sum / m;
       for (int j = 0; j < m; j++) {
          double xd = (x[i][j] - mean[i]) / sqrt(var[i] + eps); //epsilon);
          Zhat[i][j] = gamma[j] * xd + beta[j];
       }
    }

    return Zhat;  // final output of Add & norm
}

// Assume forward already filled: Zhat, mean, var
matrixd AddnNorm::backProp(const matrixd &X, const matrixd &dL_dY, double eta)
{
    int nt = X.size();
    int m  = X[0].size();

    matrixd dL_dZ(nt, vector<double>(m, 0.0));

    vector<double> dL_dGamma(m, 0.0);	// gradient wrt gamma
    vector<double> dL_dBeta(m, 0.0);	// gradient wrt beta

    // 1. Compute dL_dGamma and dL_dBeta for each token
    for (int i = 0; i < nt; i++)
    {
        for (int j = 0; j < m; j++){
           dL_dBeta[j]  += dL_dY[i][j];
           dL_dGamma[j] += dL_dY[i][j] * Zhat[i][j];
        }
    }


    // 3. Backprop through LayerNorm (row-wise)
    for (int i = 0; i < nt; i++)
    {
        double inv_sigma = 1.0 / sqrt(var[i] + eps);

        // Step A: compute intermediate values
        vector<double> dL_dZhat(m);
        for (int j = 0; j < m; j++)
        {
            dL_dZhat[j] = dL_dY[i][j] * gamma[j];
        }

        double sum_dL_dZhat = 0.0;
        double sum_dL_dZhat_zhat = 0.0;

        for (int j = 0; j < m; j++)
        {
            sum_dL_dZhat += dL_dZhat[j];
            sum_dL_dZhat_zhat += dL_dZhat[j] * Zhat[i][j];
        }

        // Final gradient
        for (int j = 0; j < m; j++)
        {
            dL_dZ[i][j] = inv_sigma * (dL_dZhat[j] - sum_dL_dZhat / m
                 - Zhat[i][j] * sum_dL_dZhat_zhat / m);
        }
    }

    // 2. Update parameters
    for (int j = 0; j < m; j++)
    {
        gamma[j] -= eta * dL_dGamma[j];
        beta[j]  -= eta * dL_dBeta[j];
    }

    return dL_dZ;
}

// Optional setters (to plug in forward results)
void AddnNorm::setCache(const matrixd& zhat, const vector<double>& m, 
		                 const vector<double>& v)
{
    Zhat = zhat;
    mean = m;
    var  = v;
}

void AddnNorm::saveGammaBeta(FILE *fp)
{
  int dModel = gamma.size();

  for(int i = 0; i < dModel; i++) 
    fprintf(fp, "%9.4f, %9.4f", gamma[i], beta[i] );
}

void AddnNorm::readGammaBeta(FILE *fp)
{
  int dModel = gamma.size();

  for(int i = 0; i < dModel; i++) 
    fscanf(fp, "%lf, %lf", &gamma[i], &beta[i] );
}

// ---------------------------  Feed Forward -----------------------------------
FeedForward::FeedForward(int dModel0, int d_ff0)
{
    dModel = dModel0;
    d_ff = d_ff0;

    W1.resize(dModel, vector<double>(d_ff));
    W2.resize(d_ff, vector<double>(dModel));
    initMatrix(W1, dModel, d_ff);
    initMatrix(W2, d_ff, dModel);
}

// ------------------------------
// Forward (for cache)
// ------------------------------
matrixd FeedForward::FFoutput(const matrixd &X)
{
    Z  = matmul(X, W1);   // pre-activation
    F1 = relu( Z );  //f(Z);        // F1 = f(XW1)
    matrixd F2 = matmul(F1, W2); // F2 = F1 W2

    return F2;
}

// ------------------------------
// Backprop
// ------------------------------
matrixd FeedForward::backProp(const matrixd &X, const matrixd &dL_dY,  // dL/dF2
             double eta)
{
    // ---- W2 ----
    matrixd dL_dF2 = dL_dY;

    // dL/dW2 = F1^T * dL_dF2
    matrixd F1T = transpose(F1);
    matrixd dL_dW2 = matmul(F1T, dL_dF2);

    // dL_dF1 = dL_dF2 * W2^T
    matrixd W2T = transpose(W2);
    matrixd dL_dF1 = matmul(dL_dF2, W2T);

    // ---- find dL_dZ1 ----
    matrixd fd_Z  = relu_deriv(Z); // f_deriv(Z);
    int rows = dL_dF1.size();
    int cols = dL_dF1[0].size();

    matrixd dL_dZ1 = dL_dF1;	  // same dimension
    for (int i = 0; i < rows; i++)
       for (int j = 0; j < cols; j++)
          dL_dZ1[i][j] = dL_dF1[i][j] * fd_Z[i][j];

    // ---- W1 ----
    // dL/dW1 = X^T * dL/dZ1
    matrixd XT = transpose(X);
    matrixd dL_dW1 = matmul(XT, dL_dZ1);

    // dL/dX = dL/dZ1 * W1^T
    matrixd W1T = transpose(W1);
    matrixd dL_dX = matmul(dL_dZ1, W1T);

    // ---- Update weights ----
    updateWeight(W1, dL_dW1, eta);
    updateWeight(W2, dL_dW2, eta);

    return dL_dX;	// This becomes the dL_dY of next stage
}

void FeedForward::saveWeights(FILE *fp )
{
   saveMatrix(W1, dModel, d_ff, fp);
   saveMatrix(W2, d_ff, dModel, fp);
}

void FeedForward::readWeights(FILE *fp )
{
   readMatrix(W1, dModel, d_ff, fp);
   readMatrix(W2, d_ff, dModel, fp);
}

// ------------------ MultiHeadAttention --------------------
MultiHeadAttention::MultiHeadAttention(int dModel0, int nHeads0)
{
    dModel = dModel0;
    nHeads = nHeads0;
    d_k = dModel / nHeads;

    WQ.resize(nHeads);
    WK.resize(nHeads);
    WV.resize(nHeads);

    for (int h = 0; h < nHeads; h++) {
        initMatrix(WQ[h], dModel, d_k);
        initMatrix(WK[h], dModel, d_k);
        initMatrix(WV[h], dModel, d_k);
    }

    initMatrix(WO, dModel, dModel);
}

matrixd MultiHeadAttention::computeAttention(const matrixd &X_Q,
     const matrixd &X_K, const matrixd &X_V, bool mask)
{
    Qs.clear(); Ks.clear(); Vs.clear();
    Ps.clear(); As.clear();

    vector<matrixd> heads;

    int nq = X_Q.size();
    int nk = X_K.size();

    for (int h = 0; h < nHeads; h++) {
        matrixd Qh = matmul(X_Q, WQ[h]);  // nq x d_k
        matrixd Kh = matmul(X_K, WK[h]);  // nk x d_k
        matrixd Vh = matmul(X_V, WV[h]);  // nk x d_k

        matrixd S = matmul(Qh, transpose(Kh));
        scaleMat(S, 1.0 / sqrt(d_k));


        if (mask)
          S = causalMask(S);

        matrixd Ph = softmax(S);  
        matrixd Ah = matmul(Ph, Vh);  // nq x d_k

        Qs.push_back(Qh);
        Ks.push_back(Kh);
        Vs.push_back(Vh);
        Ps.push_back(Ph);
        As.push_back(Ah);

        heads.push_back(Ah);
    }

    // Concatenate heads
    int n = heads[0].size();
    matrixd concat(n, vector<double>(dModel, 0.0));

    for (int h = 0; h < nHeads; h++) 
       for (int i = 0; i < n; i++) 
          for (int j = 0; j < d_k; j++) 
             concat[i][h * d_k + j] = heads[h][i][j];

        // Final projection
    matrixd H = matmul(concat, WO);

    return H;
}

// backpropagation
matrixd MultiHeadAttention::backProp( const matrixd &X_Q, const matrixd &X_K,
    const matrixd &X_V, const matrixd &dL_dH, double eta, bool mask)
{
   if (Ps.empty() || As.empty()) {
      cout << "ERROR: MHA cache empty in backProp!" << endl;
      exit(1);
   }

    // ---- WO gradients ----
    matrixd A_cat; // rebuild concat from cached As
 //   int nt = As[0].size();
    int nt = X_Q.size();
    double s = 1.0 / sqrt(d_k);   // for scaling 

    A_cat = matrixd(nt, vector<double>(dModel, 0.0));
    for (int h = 0; h < nHeads; h++) 
       for (int i = 0; i < nt; i++) 
          for (int j = 0; j < d_k; j++) 
             A_cat[i][h*d_k + j] = As[h][i][j];

    matrixd dL_dWO = matmul(transpose(A_cat), dL_dH);

    // Compute before updating WO
    matrixd dL_dAcat = matmul(dL_dH, transpose(WO));
 
    // ---- initialize accumulators ----
    matrixd dL_dXQ(X_Q.size(), vector<double>(dModel, 0.0));
    matrixd dL_dXK(X_K.size(), vector<double>(dModel, 0.0));
    matrixd dL_dXV(X_V.size(), vector<double>(dModel, 0.0));

    // ---- per head ----
    for (int h = 0; h < nHeads; h++) {
       // slice gradient
       matrixd dL_dA(nt, vector<double>(d_k));
       for (int i = 0; i < nt; i++)
          for (int j = 0; j < d_k; j++)
             dL_dA[i][j] = dL_dAcat[i][h*d_k + j];

       matrixd P = Ps[h];
       matrixd V = Vs[h];
       matrixd Q = Qs[h];
       matrixd K = Ks[h];

       // ---- A = P V ----
       matrixd dL_dP = matmul(dL_dA, transpose(V));
       matrixd dL_dV = matmul(transpose(P), dL_dA);

       // ---- softmax ----
       matrixd dL_dS = softmax_derivative(dL_dP, P);

         // apply mask if used
        if (mask) {
            for (int i = 0; i < nt; i++)
                for (int j = i+1; j < nt; j++)
                    dL_dS[i][j] = 0.0;
        }


       // ---- S = QK^T ----
       matrixd dL_dQ = matmul(dL_dS, K);
       matrixd dL_dK = matmul(transpose(dL_dS), Q);

       scaleMat(dL_dQ, s);
       scaleMat(dL_dK, s);

       // ---- Weight gradients ----
       matrixd dL_dWQ = matmul(transpose(X_Q), dL_dQ);
       matrixd dL_dWK = matmul(transpose(X_K), dL_dK);
       matrixd dL_dWV = matmul(transpose(X_V), dL_dV);

       // ---- input gradients ----
       matrixd dXQ = matmul(dL_dQ, transpose(WQ[h]));
       matrixd dXK = matmul(dL_dK, transpose(WK[h]));
       matrixd dXV = matmul(dL_dV, transpose(WV[h]));

       // accumulate
       dL_dXQ = addMat(dL_dXQ, dXQ);
       dL_dXK = addMat(dL_dXK, dXK);
       dL_dXV = addMat(dL_dXV, dXV);

        // =====================
        // Update weights after gradients
        // =====================
        updateWeight(WQ[h], dL_dWQ, eta);
        updateWeight(WK[h], dL_dWK, eta);
        updateWeight(WV[h], dL_dWV, eta);

    }  // for h

    updateWeight(WO, dL_dWO, eta);

    // For self-attention, these are the same
    return dL_dXQ;
}

void MultiHeadAttention::saveWeights(FILE *fp)
{
  for (int h = 0; h < nHeads; h++) {
        saveMatrix(WQ[h], dModel, d_k, fp);
        saveMatrix(WK[h], dModel, d_k, fp);
        saveMatrix(WV[h], dModel, d_k, fp);
    }

    saveMatrix(WO, dModel, dModel, fp);
}

void MultiHeadAttention::readWeights(FILE *fp)
{
  for (int h = 0; h < nHeads; h++) {
        readMatrix(WQ[h], dModel, d_k, fp);
        readMatrix(WK[h], dModel, d_k, fp);
        readMatrix(WV[h], dModel, d_k, fp);
    }

    readMatrix(WO, dModel, dModel, fp);
}

// ---------- OutputLayer ----------------------
OutputLayer::OutputLayer(int dModel0, int vocab_size0)
{
    dModel = dModel0;
    vocab_size = vocab_size0;

    initMatrix(W, dModel, vocab_size);
}

// Forward: logits → probabilities
matrixd OutputLayer::forward(const matrixd &X) {
    // X: (seq_len x dModel)
    // output: (seq_len x vocab_size)
    matrixd logits = matmul(X, W);

    return softmax(logits);
}

// Backward: returns dL/dX
matrixd OutputLayer::backward(const matrixd &X, const matrixd &P,
             const vector<int> &targets, double eta)
{
    int nt = P.size();

    // dL/dY = P - T (cross-entropy + softmax)
    matrixd dL_dY = P;

    for (int i = 0; i < nt; i++) 
        dL_dY[i][targets[i]] -= 1.0;

      // optional but recommended: normalize
    for (int i = 0; i < nt; i++)
        for (int j = 0; j < P[0].size(); j++)
            dL_dY[i][j] /= nt;

    // Gradient wrt weights
    matrixd dL_dW = matmul(transpose(X), dL_dY);

    // Gradient wrt input (back to decoder)
    matrixd dL_dX = matmul(dL_dY, transpose(W));

    //update weights after gradient computation
    updateWeight(W, dL_dW, eta);

    return dL_dX;
}

void OutputLayer::saveWeights(FILE *fp)
{
   saveMatrix(W, dModel, vocab_size, fp);
}

void OutputLayer::readWeights(FILE *fp)
{
   readMatrix(W, dModel, vocab_size, fp);
}

// ------------ GPT-style Transformer (Decoder-only) -------------------
GPTBlock::GPTBlock(int dModel, int nHeads, int d_ff)   // Decoder
    : mha(dModel, nHeads), norm1(dModel), ffn(dModel, d_ff), norm2(dModel) 
{
}

matrixd GPTBlock::forward(const matrixd &X) 
{
    matrixd H = mha.computeAttention(X, X, X, true); // masked
    matrixd X1 = norm1.forward(X, H);

    matrixd F = ffn.FFoutput(X1);
    matrixd Y = norm2.forward(X1, F);

    return Y;
}

matrixd GPTBlock::backward(const matrixd &X, const matrixd &dL_dY, double eta) 
{

    matrixd dF = norm2.backProp(X, dL_dY, eta);
    matrixd dX1 = ffn.backProp(X, dF, eta);
    matrixd dH = norm1.backProp(X, dX1, eta);

    return mha.backProp(X, X, X, dH, eta, true);
}

void GPTBlock::saveWeights(FILE *fp)
{
  fprintf(fp, "#GPTBlock: Multihead Attention mha:\n");
  mha.saveWeights( fp );
  fprintf(fp, "\n");
  fprintf(fp, "#GPTBlock: AddnNorm norm1:\n");
  norm1.saveGammaBeta( fp );
  fprintf(fp, "\n");
  fprintf(fp, "#GPTBlock: FeedForward ffn:\n");
  ffn.saveWeights( fp );
  fprintf(fp, "\n");
  fprintf(fp, "#GPTBlock: AddnNorm norm2:\n");
  norm2.saveGammaBeta( fp );
  fprintf(fp, "\n");
}

void readComment(FILE *fp)
{
  char buffer[256];
  do {
    if (fgets(buffer, sizeof(buffer), fp) == NULL) break;
  } while ( buffer[0] != '#' );
}

void GPTBlock::readWeights(FILE *fp)
{
  char buffer[256];
  readComment( fp );
  mha.readWeights( fp );
  readComment( fp );
  norm1.readGammaBeta( fp );
  readComment( fp );
  ffn.readWeights( fp );
  readComment( fp );
  norm2.readGammaBeta( fp );
}

// ----------------------- MiniGPT -----------------------
MiniGPT::MiniGPT(int vocab_size, int seq_len, int dModel, int nHeads, int d_ff, 
	int nLayers) : embed(vocab_size, dModel), pe(seq_len, dModel),
      	output(dModel, vocab_size), seq_len(seq_len), vocab_size(vocab_size)
{
    for (int i = 0; i < nLayers; i++) {
       GPTBlock gptBlock(dModel, nHeads, d_ff);
       layers.push_back(gptBlock);
    }
}

matrixd MiniGPT::forward(const vector<int> &tokens) 
{
    matrixd X = embed.embed(tokens);
    X = pe.addPE(X);

    for(int i = 0; i < layers.size(); i++)
        X = layers[i].forward(X);

    return output.forward(X);
}

double MiniGPT::trainStep(
        const vector<int> &tokens,
        double eta)
{
    // ---------------------------------
    // shifted sequences
    // ---------------------------------

    vector<int> input(tokens.begin(), tokens.end() - 1);

    vector<int> target(tokens.begin() + 1, tokens.end());

    // ---------------------------------
    // embedding
    // ---------------------------------

    matrixd X = embed.embed(input);

    X = pe.addPE(X);

    // ---------------------------------
    // transformer forward
    // ---------------------------------

    vector<matrixd> cache;

    for (int i = 0; i < layers.size(); i++)
    {
        cache.push_back(X);   // IMPORTANT, cache input

        X = layers[i].forward(X);
    }

    // ---------------------------------
    // output
    // ---------------------------------

    matrixd P = output.forward(X);

    double loss =
        crossEntropy(P, target);

    // ---------------------------------
    // output backward
    // ---------------------------------

    matrixd dL_dX = output.backward(X, P, target, eta);

    // ---------------------------------
    // transformer backward
    // ---------------------------------

    for (int i = layers.size()-1; i >= 0; i--)
    {
        dL_dX = layers[i].backward(cache[i], dL_dX, eta);
    }

    // ---------------------------------
    // embedding backward
    // ---------------------------------

    embed.backProp(input, dL_dX, eta);

    return loss;
}

int MiniGPT::saveWeights(char fname[])
{
    FILE *fp;

  if ( (fp = fopen(fname, "wt")) == NULL ) {
    printf("\nError opening file %s\n", fname);
    return -1;
  }

  fprintf(fp, "# Ouptut Layer\n");
  output.saveWeights( fp );	// Output Layer
  for (int i = 0; i < layers.size(); i++)
    layers[i].saveWeights( fp ); // GPT blocks
  fclose( fp );

  return 1;
}	

int MiniGPT::readWeights(char fname[])
{
    FILE *fp;

  if ( (fp = fopen(fname, "rt")) == NULL ) {
    printf("\nError opening file %s\n", fname);
    return -1;
  }

  readComment( fp );
  output.readWeights( fp );	// Output Layer

  for (int i = 0; i < layers.size(); i++)
    layers[i].readWeights( fp ); // GPT blocks

  fclose( fp );

  return 1;
}