import numpy as np
import matplotlib.pyplot as plt
from sympy import symbols, expand
import time
from tqdm import tqdm
from mpl_toolkits.mplot3d import Axes3D

# Function to generate a random 3SAT formula
def generate_3sat_formula(n, m):
    np.random.seed(0)
    variables = symbols(f'x1:{n+1}')
    clauses = []
    for _ in range(m):
        clause = []
        for _ in range(3):
            var = np.random.choice(variables)
            if np.random.rand() > 0.5:
                clause.append(var)
            else:
                clause.append(1 - var)
        clauses.append(clause)
    return variables, clauses

# Function to convert a clause to polynomial
def clause_to_polynomial(clause):
    x_j, x_k, x_l = clause
    polynomial = expand(x_j + (1 - x_k) + x_l - 2*x_j*(1 - x_k) - 2*x_j*x_l - 2*(1 - x_k)*x_l + 4*x_j*(1 - x_k)*x_l)
    return polynomial

# Controlled simplification of polynomials
def efficient_simplify(polynomials):
    simplified_polynomials = []
    for poly in tqdm(polynomials, desc='Simplification'):
        simplified_polynomials.append(expand(poly))
    return simplified_polynomials

# Measure time for each step with large input
def measure_large_input_time(n, m):
    variables, clauses = generate_3sat_formula(n, m)

    # Measure time for Input Transformation
    start_time = time.time()
    polynomials = [clause_to_polynomial(clause) for clause in tqdm(clauses, desc='Input Transformation')]
    transformation_time = time.time() - start_time

    # Measure time for Simplification
    start_time = time.time()
    simplified_polynomials = efficient_simplify(polynomials)
    simplification_time = time.time() - start_time

    return transformation_time, simplification_time

# Set input sizes for testing
n = 500  # initial number of variables
m = 1000  # initial number of clauses

# Prepare lists to store results
variables_list = []
clauses_list = []
transformation_times = []
simplification_times = []

# Measure times for increasing sizes
for i in range(1, 6):
    current_n = n * i
    current_m = m * i
    transformation_time, simplification_time = measure_large_input_time(current_n, current_m)
    variables_list.append(current_n)
    clauses_list.append(current_m)
    transformation_times.append(transformation_time)
    simplification_times.append(simplification_time)

# Plot the results
fig = plt.figure(figsize=(14, 8))

# Input Transformation times 3D plot
ax = fig.add_subplot(121, projection='3d')
ax.scatter(variables_list, clauses_list, transformation_times, c='blue', marker='o')
ax.set_xlabel('Number of Variables')
ax.set_ylabel('Number of Clauses')
ax.set_zlabel('Transformation Time (seconds)')
ax.set_title('3D Plot of Input Transformation Time')

# Simplification times 3D plot
ax = fig.add_subplot(122, projection='3d')
ax.scatter(variables_list, clauses_list, simplification_times, c='orange', marker='o')
ax.set_xlabel('Number of Variables')
ax.set_ylabel('Number of Clauses')
ax.set_zlabel('Simplification Time (seconds)')
ax.set_title('3D Plot of Simplification Time')

plt.tight_layout()
plt.show()