"""Illustration of Simulated Annealing

  Code framework generated with chatGPT:
  Q: 'Do you have an example in two dimensions with a graphical 
  illustration of the path to convergence?'

  Parameter adjustments ( pertubation_scale, num_iterations) 
  and graphics improvements by G.Q. 
"""

import math
import random
import matplotlib.pyplot as plt
import numpy as np

def rosenbrock(x, y):
    # Rosenbrock function
    return (1 - x)**2 + 100 * (y - x**2)**2

def perturb(x, y, perturbation_scale=1.):
    # Perturb the current solution by adding random values scaled by perturbation_scale
    perturbed_x = x + perturbation_scale * random.uniform(-1, 1)
    perturbed_y = y + perturbation_scale * random.uniform(-1, 1)
    return perturbed_x, perturbed_y

def acceptance_probability(delta_E, temperature):
    # Metropolis acceptance probability formula
    if delta_E < 0:
        return 1.0
    return math.exp(-delta_E / temperature)

def simulated_annealing_2d():
    current_solution = initial_solution
    current_temperature = initial_temperature
    scale = initial_scale
    
    path_x = [current_solution[0]]
    path_y = [current_solution[1]]

    for iteration in range(num_iterations):
        # Perturb the current solution
        new_solution = perturb(*current_solution, perturbation_scale=scale)

        # Calculate the change in objective function
        delta_E = rosenbrock(*new_solution) - rosenbrock(*current_solution)

        # Accept or reject the new solution based on acceptance probability
        if random.random() < acceptance_probability(delta_E, current_temperature):
            current_solution = new_solution

        # Decrease the temperature
        current_temperature *= cooling_rate
        
        # Record the path for visualization
        path_x.append(current_solution[0])
        path_y.append(current_solution[1])

    return current_solution, rosenbrock(*current_solution), path_x, path_y

# Set initial parameters
initial_solution = (-1., 3.5)
initial_temperature = 1.0
initial_scale = 1.
cooling_rate = 0.97
num_iterations = 500

# Run simulated annealing in 2D
best_solution, best_value, path_x, path_y = simulated_annealing_2d()

# Visualize the path to convergence
x_vals = np.linspace(-2.5, 2.5, 150)
y_vals = np.linspace(-2, 5.5, 150)
X, Y = np.meshgrid(x_vals, y_vals)
Z = rosenbrock(X, Y)

plt.contour(X, Y, Z, levels=50, cmap='viridis')
plt.plot(path_x, path_y, '--o', color='red', marker='x', label='Simulated Annealing Path'
         , alpha=0.1)
plt.plot([best_solution[0]], [best_solution[1]], color='blue', marker='o',
         markersize=4, label='Best Solution')
plt.plot([1.], [1.], color='green', marker='x', markersize=10., label='true Solution')
plt.title('Simulated Annealing Path to Convergence (Rosenbrock Function)')
plt.xlabel('X')
plt.ylabel('Y')
plt.legend()
plt.show()