import random
import numpy as np
from tensorflow import keras
from sklearn.preprocessing import MinMaxScaler
from tensorflow.keras.optimizers import Adam
from collections import deque

numbers_range = (1, 10)
large_numbers = [25, 50, 75, 100]
input_size = 7
output_size = 4  # We have four possible operations: addition(+), subtraction(-), multiplication(*), and division(/)
memory_buffer = deque(maxlen=1000)  # Experience replay buffer

# Define possible operations
operations = [np.add, np.subtract, np.multiply, np.divide]


def generate_game():
    target = random.randint(100, 999)
    numbers = random.sample(range(*numbers_range), k=5) + random.choices(large_numbers, k=2)
    return target, numbers


def store_experience(state, action, reward, next_state, done):
    memory_buffer.append((state, action, reward, next_state, done))


def sample_from_memory(batch_size=32):
    samples = random.sample(memory_buffer, min(batch_size, len(memory_buffer)))
    states, actions, rewards, next_states, dones = zip(*samples)
    return np.array(states), actions, rewards, np.array(next_states), dones


# Create the DQN model
model = keras.Sequential()
model.add(keras.layers.Dense(64, input_dim=input_size, activation='relu'))
model.add(keras.layers.Dense(32, activation='relu'))
model.add(keras.layers.Dense(output_size, activation='linear'))

model.compile(loss='mse', optimizer = Adam(learning_rate=0.001))

# Normalize input data by fitting a scaler only on the possible range of values
scaler = MinMaxScaler()
scaler.fit(np.array([[0] + [1]*6, [999] + large_numbers*2]))  # Min and max states


def play_and_learn(iters=1000, batch_size=32):
    for _ in range(iters):
        # Initialize game
        target, numbers = generate_game()
        done = False
        current_sum = 0

        while not done:
            # Prepare the state
            state = np.array([[target] + numbers + [current_sum]])

            # Predict the best operation
            state_scaled = scaler.transform(state)
            action_values = model.predict(state_scaled)
            action = np.argmax(action_values[0])

            # Choose two numbers to perform the operation on
            num1, num2 = random.sample(numbers, 2)
            if action == 3 and num2 == 0:  # Avoid division by zero
                continue

            # Apply the chosen operation
            next_sum = operations[action](num1, num2)
            next_state = np.array([[target] + numbers + [next_sum]])
            reward = 1.0 if abs(next_sum - target) < abs(current_sum - target) else 0.0

            # Check if the game is done
            if len(memory_buffer) >= memory_buffer.maxlen:
                done = True

            # Store the experience
            store_experience(state_scaled[0], action, reward, scaler.transform(next_state)[0], done)

            # Sample a minibatch from the memory
            if len(memory_buffer) >= batch_size:
                minibatch_states, minibatch_actions, minibatch_rewards, minibatch_next_states, minibatch_dones = sample_from_memory(
                    batch_size)
                # Compute Q values for minibatch
                Q_values_next = model.predict_on_batch(minibatch_next_states)
                Q_targets = minibatch_rewards + (np.amax(Q_values_next, axis=1) * (1 - minibatch_dones))
                # Update the Q values for the actions that were taken
                Q_values = model.predict_on_batch(minibatch_states)
                for i, action in enumerate(minibatch_actions):
                    Q_values[i, action] = Q_targets[i]
                # Train the model on the minibatch
                model.train_on_batch(minibatch_states, Q_values)

    print("Training Completed")

    play_and_learn()