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 = 72  # Changed from 36 to 72
memory_buffer = deque(maxlen=1000)  # Adjust the buffer size as needed


def generate_game():
    target = random.randint(100, 999)
    numbers = random.choices(range(1, 10), k=5) + random.choices(large_numbers, k=1)  # Use random.choices to select the numbers
    return target, numbers


def generate_training_data(target, numbers):
    input_data = np.append(target, numbers)
    output_data = np.zeros(output_size)
    for i in range(6):
        for j in range(6):
            if i != j:
                # Generate output data for all 72 actions, not just 36
                output_data[i * 12 + j] = 1  # Addition
                output_data[i * 12 + j + 6] = 1  # Subtraction
                output_data[i * 12 + j + 36] = 1  # Multiplication
                if numbers[j] != 0:  # Avoid division by zero
                    output_data[i * 12 + j + 42] = 1  # Division

    return np.array([input_data]), output_data


def store_experience(target, numbers, output_data):
    input_data = np.append(target, numbers)
    memory_buffer.append((input_data, output_data))


def sample_from_memory(batch_size=32):
    samples = random.sample(memory_buffer, min(batch_size, len(memory_buffer)))
    x, y = zip(*samples)
    return np.array(x), np.array(y)


def get_solution(model, target, numbers):
    pred = model.predict(np.array([np.append(target, numbers)]))
    indices = np.argsort(pred.flatten())[-2:]  # Get the indices of the two highest predictions
    i, j = indices // 12, indices % 12  # Decode the indices of the actions from the output layer

    # Ensure each number is used only once in the solution
    i = list(set(i) & set(range(6)))  # Limit to the range of original numbers
    j = list(set(j) & set(range(6)))  # Limit to the range of original numbers

    operators = ['+', '+']  # Initialize with '+' for both operators
    for idx, val in enumerate(i):
        if val >= 5:
            operators[idx] = str(large_numbers[val - 5])
    for idx, val in enumerate(j):
        if val >= 5:
            operators[idx] = str(large_numbers[val - 5])

    # Take into account the new operators
    if j[0] >= 6:
        operators[0] = '-' if j[0] < 12 else '*' if j[0] < 18 else '/'
    if j[1] >= 6:
        operators[1] = '-' if j[1] < 12 else '*' if j[1] < 18 else '/'

    return i[0], operators[0], j[0] % 6, operators[1]


def calculate_reward(target, predicted):
    # Define a sliding scale for rewards based on the proximity to the target
    absolute_difference = abs(predicted - target)
    if absolute_difference == 0:
        return 1.0  # Full reward for an exact match
    elif absolute_difference <= 0.01 * target:  # Within 1% of the target
        return 0.8
    elif absolute_difference <= 0.02 * target:  # Within 2% of the target
        return 0.6
    # ... Add more conditions as needed for different levels of proximity
    else:
        return 0.0  # No reward if not within the defined thresholds


def play_and_train(model):
    target, numbers = generate_game()
    print("Target:", target)
    print("Numbers:", numbers)

    # Store current game experience in memory
    x, y = generate_training_data(target, numbers)
    store_experience(target, numbers, y)

    # Sample from memory for training
    x_memory, y_memory = sample_from_memory(batch_size=32)

    # Input Data Normalization
    scaler = MinMaxScaler()
    x_scaled = scaler.fit_transform(x_memory)

    # Train the model on both current and past experiences
    model.fit(x_scaled, y_memory, epochs=5, batch_size=32, verbose=0)

    # Predict the result
    i, operator_i, j, operator_j = get_solution(model, target, numbers)
    # Handle the cases where the predicted solution involves multiplication or division
    if operator_i == '*':
        result_i = numbers[i] * numbers[j % 6]
    elif operator_i == '/':
        result_i = numbers[i] / numbers[j % 6]
    else:
        result_i = numbers[i] + numbers[j % 6] if operator_i == '+' else numbers[i] - numbers[j % 6]
    print(f"Predicted Solution: ({numbers[i]} {operator_i} {numbers[j % 6]}) = {result_i}")

    # Calculate the reward based on how close the prediction is to the target
    reward = calculate_reward(target, result_i)
    print("Reward:", reward)

    return reward


# Model Architecture
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='sigmoid'))

# Model Compilation
model.compile(loss='binary_crossentropy', optimizer=Adam(learning_rate=0.001), metrics=['accuracy'])

solutions_count = 0
while solutions_count < 100:
    reward = play_and_train(model)
    if reward == 1.0:
        solutions_count += 1