#include <iostream>
#include <conio.h>
#include <windows.h>
#include <vector>
#include <cmath>

bool gameOver;
const int width = 20;
const int height = 20;
int x, y, fruitX, fruitY, score;
int tailX[100], tailY[100];
int nTail;
enum eDirection { STOP = 0, LEFT, RIGHT, UP, DOWN };
eDirection dir;

const int inputSize = (width * height) + 1 + (2 * nTail);
const int hiddenSize = 16;
const int outputSize = 4;

struct NeuralNetwork
{
    std::vector<std::vector<double>> weights1;
    std::vector<std::vector<double>> weights2;
};

NeuralNetwork neuralNetwork;

void InitializeNeuralNetwork()
{
    neuralNetwork.weights1.resize(hiddenSize, std::vector<double>(inputSize, 0.0));
    neuralNetwork.weights2.resize(outputSize, std::vector<double>(hiddenSize, 0.0));

    // Randomly initialize the weights
    for (int i = 0; i < hiddenSize; i++)
    {
        for (int j = 0; j < inputSize; j++)
        {
            neuralNetwork.weights1[i][j] = (2.0 * rand() / RAND_MAX) - 1.0;
        }
    }

    for (int i = 0; i < outputSize; i++)
    {
        for (int j = 0; j < hiddenSize; j++)
        {
            neuralNetwork.weights2[i][j] = (2.0 * rand() / RAND_MAX) - 1.0;
        }
    }
}

std::vector<double> FlattenInput()
{
    std::vector<double> input(inputSize, 0.0);

    // Set the positions of the apples
    input[(fruitY * width) + fruitX] = 1.0;

    // Set the length of the snake's body
    input[width * height] = static_cast<double>(nTail) / (width * height);

    // Set the positions of the snake's body segments
    for (int i = 0; i < nTail; i++)
    {
        input[(tailY[i] * width) + tailX[i]] = 1.0;
    }

    return input;
}

std::vector<double> GetOutput(std::vector<double>& input)
{
    std::vector<double> hidden(hiddenSize, 0.0);
    std::vector<double> output(outputSize, 0.0);

    // Calculate hidden layer values
    for (int i = 0; i < hiddenSize; i++)
    {
        for (int j = 0; j < inputSize; j++)
        {
            hidden[i] += neuralNetwork.weights1[i][j] * input[j];
        }
        hidden[i] = 1.0 / (1.0 + exp(-hidden[i]));
    }

    // Calculate output layer values
    for (int i = 0; i < outputSize; i++)
    {
        for (int j = 0; j < hiddenSize; j++)
        {
            output[i] += neuralNetwork.weights2[i][j] * hidden[j];
        }
        output[i] = 1.0 / (1.0 + exp(-output[i]));
    }

    return output;
}

int GetMaxIndex(const std::vector<double>& values)
{
    int maxIndex = 0;
    double maxValue = values[0];
    for (int i = 1; i < values.size(); i++)
    {
        if (values[i] > maxValue)
        {
            maxValue = values[i];
            maxIndex = i;
        }
    }
    return maxIndex;
}

void Setup()
{
    gameOver = false;
    dir = STOP;
    x = width / 2;
    y = height / 2;
    fruitX = rand() % width;
    fruitY = rand() % height;
    score = 0;
}

void Draw()
{
    system("cls");
    for (int i = 0; i < width + 2; i++)
        std::cout << "#";
    std::cout << std::endl;

    for (int i = 0; i < height; i++)
    {
        for (int j = 0; j < width; j++)
        {
            if (j == 0)
                std::cout << "#";
            if (i == y && j == x)
                std::cout << "O";
            else if (i == fruitY && j == fruitX)
                std::cout << "F";
            else
            {
                bool printTail = false;
                for (int k = 0; k < nTail; k++)
                {
                    if (tailX[k] == j && tailY[k] == i)
                    {
                        std::cout << "o";
                        printTail = true;
                    }
                }
                if (!printTail)
                    std::cout << " ";
            }
            if (j == width - 1)
                std::cout << "#";
        }
        std::cout << std::endl;
    }

    for (int i = 0; i < width + 2; i++)
        std::cout << "#";
    std::cout << std::endl;
    std::cout << "Score:" << score << std::endl;
}

void Input()
{
    if (_kbhit())
    {
        switch (_getch())
        {
        case 'a':
            dir = LEFT;
            break;
        case 'd':
            dir = RIGHT;
            break;
        case 'w':
            dir = UP;
            break;
        case 's':
            dir = DOWN;
            break;
        case 'x':
            gameOver = true;
            break;
        }
    }
}

void Logic()
{
    int prevX = tailX[0];
    int prevY = tailY[0];
    int prev2X, prev2Y;
    tailX[0] = x;
    tailY[0] = y;
    for (int i = 1; i < nTail; i++)
    {
        prev2X = tailX[i];
        prev2Y = tailY[i];
        tailX[i] = prevX;
        tailY[i] = prevY;
        prevX = prev2X;
        prevY = prev2Y;
    }

    switch (dir)
    {
    case LEFT:
        x--;
        break;
    case RIGHT:
        x++;
        break;
    case UP:
        y--;
        break;
    case DOWN:
        y++;
        break;
    }

    if (x >= width)
        x = 0;
    else if (x < 0)
        x = width - 1;

    if (y >= height)
        y = 0;
    else if (y < 0)
        y = height - 1;

    for (int i = 0; i < nTail; i++)
    {
        if (tailX[i] == x && tailY[i] == y)
            gameOver = true;
    }

   // Update the input vector with the flattened game state
    std::vector<double> input = FlattenInput();

    // Get the output from the neural network
    std::vector<double> output = GetOutput(input);

    // Determine the direction to move based on the output
    int maxIndex = GetMaxIndex(output);
    switch (maxIndex)
    {
        case 0:
            dir = LEFT;
            break;
        case 1:
            dir = RIGHT;
            break;
        case 2:
            dir = UP;
            break;
        case 3:
            dir = DOWN;
            break;
    }

    if (x == fruitX && y == fruitY)
    {
        score += 10;
        fruitX = rand() % width;
        fruitY = rand() % height;
        nTail++;
    }
}

int main()
{
    InitializeNeuralNetwork();
    Setup();
    while (!gameOver)
    {
        Draw();
        Input();
        Logic();
        Sleep(10);  // Pause for a short time to control game speed
    }
    return 0;
}