
#include <cassert>
#include <cmath>
#include "render.h"
#include <unordered_map>
Board board;
const char* shaderSource =
        R"(
#version 430
layout(local_size_x = 1) in;
layout(std430, binding = 1) buffer shaderBuffer {
  float data[];
};

void main() {
  int size = 4;
  uint blockIdx = gl_GlobalInvocationID.x; 

  uint start = blockIdx * size;
  if (start + size <= data.length()) {
    data[start] = 1;
    data[start + 1] = 2;
    data[start + 2] = 3;
    data[start + 3] = 4;
  }
}
)";
int turn=0;
Board createStartingBoard() {
  Board board{};
  board.black = (1ULL << index(3,3)) | (1ULL << index(4,4));
board.white = (1ULL << index(3,4)) | (1ULL << index(4,3));
  return board;
}
std::vector<Move> legal;
bool botIsBlack=false;
bool botTurn=false;

std::vector<int> options(64);

std::vector<Move> cSquares = {
  {0,1}, {1,0}, {6,0}, {7,1}, {0,6}, {1,7}, {6,7}, {7,6}
};
std::vector<Move> xSquares = {
  {1,1}, {6,1}, {1,6}, {6,6}
};
int maxDepth=3;

void reward_fn(const std::vector<Move>& legalMoves, int root, Board state, bool botTurn){
  int &reward = options[root];

  //Corner bonus
  std::vector<Move> corners = {{0,0},{0,7},{7,0},{7,7}};
  for(auto &c : corners){
      if(hasPiece(state.black,c.x,c.y)) reward += botIsBlack ? 5000 : -4000;
      if(hasPiece(state.white,c.x,c.y)) reward += botIsBlack ? -4000 : 5000;
  }

  //C-squares penalty
  for(auto &c : cSquares){
      if(hasPiece(state.black,c.x,c.y)) reward += botIsBlack ? -50 : 50;
      if(hasPiece(state.white,c.x,c.y)) reward += botIsBlack ? 50 : -200;
  }

  //X-squares penalty
  for(auto &x : xSquares){
      if(hasPiece(state.black,x.x,x.y)) reward += botIsBlack ? -300 : 300;
      if(hasPiece(state.white,x.x,x.y)) reward += botIsBlack ? 300 : -300;
  }

  //Mobility bonus
  
  int blackCount = __builtin_popcountll(state.black);
  int whiteCount = __builtin_popcountll(state.white);
  int totalPieces = blackCount+whiteCount;
  float progress = static_cast<float>(totalPieces) / 64.0f;
  if(botTurn) reward += legalMoves.size()*(1.0f-progress);
  else        reward -= legalMoves.size()*(1.0f-progress);

 
  if(progress > 0.5f){

    int diff = blackCount - whiteCount;

    if(botIsBlack){
        reward += diff * (progress*progress);
    } else{
        reward -= diff * (progress*progress);
    }
  }

}

int explore(int layer,bool botTurn,Board state,Move id={-1,-1}){ 
  if(layer==maxDepth){ return 0; } 
  bool isBlack=!(botTurn ^ botIsBlack); 
  std::vector<Move> legalMoves=getLegalMoves(state,isBlack); 
  for(auto& move:legalMoves){ 
    Move ID=id; 
    if(layer==0){ 
      options[index(move.x,move.y)]=0; 
      ID=move;
    } 
    int root=index(ID.x,ID.y); 
    Board boardCopy=state; 
    playMove(boardCopy,isBlack,move.x,move.y);
    
    reward_fn(legalMoves, root, boardCopy, botTurn); 
    explore(layer+1,!botTurn,boardCopy,ID); 
  } 
  return 0; 
}

float timer=-1;
float timerDuration=1.5f;
float timerX=-1;
float timerY=-1;
void makeMove(int x, int y) {
  std::fill(options.begin(), options.end(), INT_MIN);
  playMove(board, !botIsBlack, x,y);
  explore(0, true, board);
  turn++;
  
  int bestScore = INT_MIN;
  std::vector<int> bestMove;
  std::vector<Move> legalMoves = getLegalMoves(board, botIsBlack);

  for (auto &m : legalMoves) {
      int id = index(m.x, m.y);

      std::cout<<m.x<<" "<<m.y<<" "<<options[id]<<std::endl;
      if (options[id] > bestScore) {
          bestScore = options[id];
          bestMove = {m.x, m.y};
      }
  }
  turn++;
  if (bestMove.size() >= 2) {
    timer=0;
    timerX=bestMove[0];
    timerY=bestMove[1];
  } else {
    std::cout << "No valid bestMove found." << std::endl;
  }
}
int main() {
  
    SetConfigFlags(FLAG_MSAA_4X_HINT); 
    InitWindow(500, 500, "Othello Engine");
    board=createStartingBoard();

    unsigned int shaderID = rlCompileShader(shaderSource, RL_COMPUTE_SHADER);
    unsigned int shaderProgram = rlLoadComputeShaderProgram(shaderID);
    


    std::vector<float> data{1,2,3,4,5,6,7,8,9,10,11,12};
    unsigned int ssboA = rlLoadShaderBuffer(data.size() * sizeof(float), data.data(), RL_DYNAMIC_DRAW);
    

    legal=getLegalMoves(board,!botIsBlack);
    std::cout<<legal.size()<<std::endl<<std::endl;
    while (!WindowShouldClose()) {
        if(timer>-1){
          timer+=GetFrameTime();
          if(timer>=timerDuration){
            timer=-1;
            playMove(board, botIsBlack, timerX,timerY);
            legal = getLegalMoves(board, !botIsBlack);
          }
        }
        bool clicked=IsMouseButtonPressed(0);  
        Vector2 mousePos=GetMousePosition();
        rlEnableShader(shaderProgram);
        rlBindShaderBuffer(ssboA, 1);
        rlComputeShaderDispatch(data.size(), 1, 1);
        rlReadShaderBuffer(ssboA, data.data(), data.size() * sizeof(float), 0);
        
        
        //for (int i = 0; i < data.size(); ++i) std::cout << data[i] << " ";
        //std::cout << std::endl;
        
        
        BeginDrawing();
        ClearBackground(ORANGE);
        renderBoard(board, mousePos, clicked, legal, makeMove,timer!=-1);
        EndDrawing();
    }
    
    puts("##unloading##");
    rlUnloadShaderProgram(shaderProgram);
    rlUnloadShaderBuffer(ssboA);
    CloseWindow();
}
