# main.py

import math
import random
import time
import traceback

class CartPole:
    def __init__(self):
        # Initialize constants for the CartPole physics
        self.gravity = 9.8
        self.masscart = 1.0
        self.masspole = 0.1
        self.total_mass = self.masspole + self.masscart
        self.length = 0.5  # Half the pole's length
        self.polemass_length = self.masspole * self.length
        self.force_mag = 10.0
        self.tau = 0.02  # Time interval for updating the values
        self.x_threshold = 2.4
        self.theta_threshold_radians = 12 * 2 * math.pi / 360  # 12 degrees in radians
        self.reset()

    def reset(self):
        # Initialize state with small random values
        while True:
            self.state = [
                random.uniform(-0.05, 0.05),  # x (cart position)
                random.uniform(-0.05, 0.05),  # x_dot (cart velocity)
                random.uniform(-0.05, 0.05),  # theta (pole angle)
                random.uniform(-0.05, 0.05)   # theta_dot (pole angular velocity)
            ]
            if not self.is_done():
                break
        return self.state

    def is_done(self):
        x, x_dot, theta, theta_dot = self.state
        done = x < -self.x_threshold \
               or x > self.x_threshold \
               or theta < -self.theta_threshold_radians \
               or theta > self.theta_threshold_radians
        return done

    def step(self, action):
        # Update the state of the environment based on the action
        x, x_dot, theta, theta_dot = self.state
        force = self.force_mag if action == 1 else -self.force_mag
        costheta = math.cos(theta)
        sintheta = math.sin(theta)

        # Equations for the dynamics of the cart-pole system
        temp = (force + self.polemass_length * theta_dot ** 2 * sintheta) / self.total_mass
        thetaacc = (self.gravity * sintheta - costheta * temp) / \
                   (self.length * (4.0 / 3.0 - self.masspole * costheta ** 2 / self.total_mass))
        xacc = temp - self.polemass_length * thetaacc * costheta / self.total_mass

        # Update the four state variables using Euler's method
        x += self.tau * x_dot
        x_dot += self.tau * xacc
        theta += self.tau * theta_dot
        theta_dot += self.tau * thetaacc

        self.state = [x, x_dot, theta, theta_dot]

        # Check if the episode is over
        done = self.is_done()

        reward = 1.0 if not done else 0.0  # Reward is 1 for every step, 0 if done

        return self.state, reward, done

class NeuralNetwork:
    def __init__(self, input_size, hidden_size, output_size):
        # Initialize weights and biases for a simple neural network
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.output_size = output_size

        # Weights between input and hidden layer
        self.W1 = [[random.uniform(-1, 1) for _ in range(hidden_size)] for _ in range(input_size)]
        self.b1 = [0.0 for _ in range(hidden_size)]

        # Weights between hidden and output layer
        self.W2 = [[random.uniform(-1, 1) for _ in range(output_size)] for _ in range(hidden_size)]
        self.b2 = [0.0 for _ in range(output_size)]

    def relu(self, x):
        # ReLU activation function
        return [max(0.0, xi) for xi in x]

    def softmax(self, x):
        # Softmax activation function for probabilities
        max_x = max(x)
        exp_x = [math.exp(xi - max_x) for xi in x]
        sum_exp_x = sum(exp_x)
        return [xi / sum_exp_x for xi in exp_x]

    def forward(self, x):
        # Forward pass through the network
        # Input to hidden layer
        h = []
        for j in range(self.hidden_size):
            hj = self.b1[j]
            for i in range(self.input_size):
                hj += x[i] * self.W1[i][j]
            h.append(hj)
        h = self.relu(h)

        # Hidden to output layer
        y_raw = []
        for k in range(self.output_size):
            yk = self.b2[k]
            for j in range(self.hidden_size):
                yk += h[j] * self.W2[j][k]
            y_raw.append(yk)
        # Output activation
        policy = self.softmax(y_raw[:2])  # First two outputs are policy probabilities
        value = math.tanh(y_raw[2])       # Third output is the value estimation
        return policy, value

    def backpropagate(self, x, target_policy, target_value, learning_rate):
        # Simplified backpropagation for updating the network
        # Forward pass (recompute for gradients)
        h_raw = []
        h = []
        for j in range(self.hidden_size):
            hj = self.b1[j]
            for i in range(self.input_size):
                hj += x[i] * self.W1[i][j]
            h_raw.append(hj)
            h.append(max(0.0, hj))  # ReLU activation
        # Hidden to output layer
        y_raw = []
        for k in range(self.output_size):
            yk = self.b2[k]
            for j in range(self.hidden_size):
                yk += h[j] * self.W2[j][k]
            y_raw.append(yk)
        # Output activation
        policy, value = self.forward(x)

        # Compute gradients
        # Value loss gradient
        dL_dvalue_output = 2.0 * (value - target_value) * (1 - value ** 2)  # Derivative of tanh

        # Policy loss gradient (cross-entropy)
        dL_dpolicy_output = [policy[k] - target_policy[k] for k in range(len(policy))]

        # Gradients for weights between hidden and output layers
        dL_dW2 = [[0.0 for _ in range(self.output_size)] for _ in range(self.hidden_size)]
        for j in range(self.hidden_size):
            for k in range(len(policy)):
                dL_dW2[j][k] = h[j] * dL_dpolicy_output[k]
            # Gradient for value output
            dL_dW2[j][2] = h[j] * dL_dvalue_output

        # Update weights and biases
        for j in range(self.hidden_size):
            for k in range(self.output_size):
                self.W2[j][k] -= learning_rate * dL_dW2[j][k]
        for k in range(len(policy)):
            self.b2[k] -= learning_rate * dL_dpolicy_output[k]
        self.b2[2] -= learning_rate * dL_dvalue_output

        # Note: For simplicity, gradients for W1 and b1 are not computed

class Node:
    def __init__(self, state, parent=None):
        self.state = state.copy()
        self.parent = parent
        self.children = {}
        self.visits = 0
        self.value = 0.0
        self.prior = 0.0

def mcts_search(root_state, neural_net, simulations=50):
    root = Node(state=root_state)
    # Check if root state is terminal
    env = CartPole()
    env.state = root_state.copy()
    if env.is_done():
        # Return a random action if the state is terminal
        return random.choice([0, 1])

    # Run MCTS simulations
    for _ in range(simulations):
        node = root
        path = [node]
        # Selection
        while node.children:
            # UCB formula to select the best child
            total_visits = sum(child.visits for child in node.children.values())
            log_total = math.log(total_visits + 1)
            best_score = -float('inf')
            best_action = None
            for action, child in node.children.items():
                ucb_score = child.value / (child.visits + 1e-5) + \
                            1.0 * child.prior * math.sqrt(log_total / (child.visits + 1e-5))
                if ucb_score > best_score:
                    best_score = ucb_score
                    best_action = action
            node = node.children[best_action]
            path.append(node)
        # Expansion
        env.state = node.state.copy()
        if env.is_done():
            value = 0.0  # Terminal node
        else:
            # Evaluate leaf node with the neural network
            policy, value = neural_net.forward(node.state)
            # Expand node with possible actions
            for a in [0, 1]:
                env.state = node.state.copy()
                next_state, reward, done = env.step(a)
                child = Node(state=next_state, parent=node)
                child.prior = policy[a]
                node.children[a] = child
        # Backpropagation
        for node in reversed(path):
            node.visits += 1
            node.value += value  # Cumulative value

    # Choose the action with the highest visit count
    if not root.children:
        # Handle the case where root has no children
        best_action = random.choice([0, 1])
    else:
        best_action = max(root.children.items(), key=lambda item: item[1].visits)[0]

    return best_action

def train_agent():
    try:
        env = CartPole()
        neural_net = NeuralNetwork(input_size=4, hidden_size=16, output_size=3)  # 2 policy outputs + 1 value output
        episode = 0
        while True:
            episode += 1
            state = env.reset()
            done = False
            total_reward = 0
            steps = 0
            while not done and steps < 200:
                # Run MCTS to select the action
                action = mcts_search(state, neural_net)
                # Take the action in the environment
                next_state, reward, done = env.step(action)
                total_reward += reward
                # Prepare targets for training
                target_policy = [0.0, 0.0]
                target_policy[action] = 1.0
                target_value = reward
                # Update the neural network
                neural_net.backpropagate(state, target_policy, target_value, learning_rate=0.01)
                # Move to the next state
                state = next_state
                steps += 1
            # Reporting statistics for the episode
            print(f"Episode: {episode} | Total Reward: {total_reward} | Steps: {steps}")
            time.sleep(0.1)  # Pause for a brief moment
    except Exception as e:
        print("An error occurred during training:")
        traceback.print_exc()
        input("Press Enter to exit...")

if __name__ == "__main__":
    train_agent() make the alphazero implementations implmenatiation better and add alot more logging for more things to easier track progress