import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation, PillowWriter

# --- 1. Define Energy Function ---
# We design an asymmetric double well:
# A deep pit (global optimum), a shallow pit (local optimum)
def energy_function(x):
    # (x^2 - 1)^2 is standard W-shaped double well
    # + 0.3*x tilts it, making the left pit deeper than the right
    return (x**2 - 1)**2 + 0.3 * x

# --- 2. Prepare Data ---
x = np.linspace(-2.5, 2.5, 500)
E = energy_function(x)

# Find true global min for plotting
min_idx = np.argmin(E)
global_min_x = x[min_idx]
global_min_y = E[min_idx]

# --- 3. Setup Plot ---
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(8, 10), constrained_layout=True)

# Top: Energy Landscape E(x) - Static
ax1.plot(x, E, 'k-', linewidth=3, label='Energy $E(x)$')
ax1.scatter(global_min_x, global_min_y, color='gold', s=150, zorder=5, edgecolors='k', label='Global Min')
ax1.set_title("The Problem: Energy Landscape $E(x)$", fontsize=14)
ax1.set_ylabel("Energy")
ax1.grid(True, alpha=0.3)
ax1.legend()
ax1.set_xlim(-2.5, 2.5)

# Bottom: Probability Distribution f(x) - Changes with T
line, = ax2.plot([], [], 'r-', linewidth=3, alpha=0.8)
fill_poly = ax2.fill_between(x, np.zeros_like(x), np.zeros_like(x), color='red', alpha=0.3)
ax2.set_xlim(-2.5, 2.5)
ax2.set_ylabel("Probability Density $f(x)$")
ax2.set_xlabel("x")
ax2.grid(True, alpha=0.3)

# Dynamic Text: Lambda and Temperature
text_info = ax2.text(0.05, 0.9, '', transform=ax2.transAxes, fontsize=12, 
                     bbox=dict(facecolor='white', alpha=0.8))

# --- 4. Animation Logic ---
# Lambda from 0.1 (High T) to 15.0 (Low T)
# Log scale to show high T phase slowly, low T phase quickly
lambdas = np.logspace(np.log10(0.1), np.log10(15.0), 100)

def init():
    line.set_data([], [])
    return line,

def update(frame_lambda):
    global fill_poly
    
    # === Core Physics ===
    # 1. Boltzmann Factor (Unnormalized Prob)
    # e^(-lambda * E)
    unnormalized_prob = np.exp(-frame_lambda * E)
    
    # 2. Normalization Constant A (Partition Function Z)
    # Numerical integration using trapezoidal rule
    integral_Z = np.trapezoid(unnormalized_prob, x)
    A = 1.0 / integral_Z
    
    # 3. Final PDF f(x)
    # f(x) = A * e^(-lambda * E)
    pdf = A * unnormalized_prob
    # =================
    
    # Update curve
    line.set_data(x, pdf)
    
    # Update fill (remove old, add new)
    fill_poly.remove()
    fill_poly = ax2.fill_between(x, 0, pdf, color='red', alpha=0.3)
    
    # Dynamically adjust Y limit (peaks get higher)
    ax2.set_ylim(0, np.max(pdf) * 1.2)
    
    # Update title and text
    T = 1 / frame_lambda
    ax2.set_title(f"The Stochastic Solution: Probability Distribution", fontsize=14)
    text_info.set_text(f"$\lambda$ = {frame_lambda:.2f} (Inverse Temp)\n$T$ = {T:.2f} (Temperature)")
    
    return line, fill_poly, text_info

# --- 5. Generate and Save ---
print("Generating animation...")
ani = FuncAnimation(fig, update, frames=lambdas, init_func=init, blit=False)

# Save as GIF
ani.save('simulated_annealing.gif', writer=PillowWriter(fps=15))
print("✅ Saved as 'simulated_annealing.gif'")

import numpy as np
import matplotlib.pyplot as plt

# --- 1. Objective Function ---
def E(x):
    return x**2 - np.cos(np.pi * x)

# --- 2. Sampling Core: Metropolis Criterion ---
# This is a "sampler", generates sample based on current T
def sample_one_step(x_curr, T):
    # a. Propose: Random step left or right
    x_next = x_curr + np.random.uniform(-0.5, 0.5)
    
    # b. Delta E
    dE = E(x_next) - E(x_curr)
    
    # c. Accept/Reject
    # Core Logic: If new energy is lower, go; if higher, go with prob (depends on T)
    if dE < 0 or np.random.rand() < np.exp(-dE / T):
        return x_next # Accept
    else:
        return x_curr # Reject

# --- Main Flow ---
def run_stochastic_optimization():
    x = -2.5       # Start (Deliberately far local opt)
    T = 10.0       # T_max
    T_min = 0.01   # T_min
    alpha = 0.99   # Cooling rate
    
    path = []      
    temps = []     
    
    print(f"{'Step':<6} | {'Temp':<8} | {'Current x':<10} | {'Action'}")
    print("-" * 45)
    
    step = 0
    # Flow 4: Till very low T
    while T > T_min:
        
        # Flow 2 & 3: Sample & Cool
        x = sample_one_step(x, T)
        
        path.append(x)
        temps.append(T)
        
        # Print intermediate
        if step % 200 == 0:
            status = "Explore 🎲" if T > 1.0 else "Exploit 🎯"
            print(f"{step:<6} | {T:<8.4f} | {x:<10.4f} | {status}")
            
        T = T * alpha 
        step += 1
        
    # Flow 5: Avg Samples
    final_samples = path[-100:]
    estimated_min = np.mean(final_samples)
    
    print("-" * 45)
    print(f"✅ Final Estimate: x = {estimated_min:.4f}")
    print(f"✅ True Min: x = 0.0000")
    
    return path, temps

# --- Run & Visualize ---
path, temps = run_stochastic_optimization()

plt.figure(figsize=(10, 6))
plt.plot(path, alpha=0.6, label='Particle Path')
plt.xlabel('Iterations (Time)')
plt.ylabel('Position x')
plt.title('Algorithm Flow: From Exploration (High T) to Exploitation (Low T)')
plt.axhline(0, color='r', linestyle='--', label='Global Min (x=0)')
plt.grid(True, alpha=0.3)
plt.legend()
plt.show()

import numpy as np
import matplotlib.pyplot as plt

# --- 1. N-D Objective (Rastrigin) ---
def rastrigin(x):
    # x is vector
    A = 10
    n = len(x)
    return A * n + np.sum(x**2 - A * np.cos(2 * np.pi * x))

# --- 2. N-D Proposal ---
def get_neighbor(x_curr, step_size=0.5):
    # Key: Random walk in N-D space
    # size=len(x_curr) ensures dimension match
    perturbation = np.random.uniform(-step_size, step_size, size=len(x_curr))
    return x_curr + perturbation

# --- 3. Simulated Annealing Main ---
def simulated_annealing_nd(n_dim=2, n_iter=2000):
    # Init: Random start in range [-5.12, 5.12]
    current_x = np.random.uniform(-5.12, 5.12, size=n_dim)
    current_E = rastrigin(current_x)
    
    # Track Best So Far
    best_x = current_x.copy()
    best_E = current_E
    
    # Temps
    T = 100.0
    T_min = 1e-4
    alpha = 0.99 
    
    path = [current_x] 
    energy_history = [current_E]

    print(f"Start {n_dim}-D Optimization...")
    print(f"Start: {np.round(current_x, 2)}, Energy: {current_E:.2f}")

    iter_count = 0
    while T > T_min and iter_count < n_iter:
        # 1. Propose
        new_x = get_neighbor(current_x)
        new_x = np.clip(new_x, -5.12, 5.12) # Clip to domain
        
        new_E = rastrigin(new_x)
        
        # 2. Delta E
        dE = new_E - current_E
        
        # 3. Metropolis Criterion
        if dE < 0 or np.random.rand() < np.exp(-dE / T):
            current_x = new_x
            current_E = new_E
            
            if current_E < best_E:
                best_x = current_x.copy()
                best_E = current_E
        
        path.append(current_x)
        energy_history.append(current_E)
        
        T *= alpha
        iter_count += 1
        
    print(f"End. Final Loc: {np.round(best_x, 4)}")
    print(f"Final Energy: {best_E:.6f} (Theoretical 0.0)")
    
    return np.array(path), energy_history, best_x

# --- Run ---
DIMENSION = 2
path, energies, final_sol = simulated_annealing_nd(n_dim=DIMENSION, n_iter=3000)

# --- 4. Visualize (2D only) ---
if DIMENSION == 2:
    plt.figure(figsize=(12, 5))
    
    # Subplot 1: Terrain & Path
    plt.subplot(1, 2, 1)
    x_grid = np.linspace(-5.12, 5.12, 100)
    y_grid = np.linspace(-5.12, 5.12, 100)
    X, Y = np.meshgrid(x_grid, y_grid)
    Z = np.zeros_like(X)
    for i in range(X.shape[0]):
        for j in range(X.shape[1]):
            Z[i,j] = rastrigin(np.array([X[i,j], Y[i,j]]))
            
    plt.contourf(X, Y, Z, levels=50, cmap='viridis', alpha=0.7)
    plt.colorbar(label='Energy')
    
    # Path: White start, Red end
    plt.plot(path[:, 0], path[:, 1], 'w-', linewidth=0.5, alpha=0.6)
    plt.scatter(path[0, 0], path[0, 1], c='white', s=50, label='Start')
    plt.scatter(final_sol[0], final_sol[1], c='red', marker='*', s=200, label='End')
    plt.legend()
    plt.title(f"2D Rastrigin Optimization Path\n(Escaping many local minima)")
    
    # Subplot 2: Energy Drop
    plt.subplot(1, 2, 2)
    plt.plot(energies)
    plt.yscale('log') 
    plt.xlabel('Iteration')
    plt.ylabel('Energy (Log Scale)')
    plt.title('Energy Minimization Process')
    plt.grid(True, which="both", ls="--")
    
    plt.tight_layout()
    plt.show()

import numpy as np
import matplotlib.pyplot as plt

# ==========================================
# 1. Define Target Function
# ==========================================
# A constructed multi-modal function with multiple pits
def target_func(x):
    # Parameters (corresponding to a, b, c in formula)
    a = np.array([-40, -20, -40])
    b = np.array([150, 100, 300])
    c = np.array([-10, -40, 40])
    
    y = 0
    for i in range(3):
        y += a[i] / (b[i] + (x + c[i])**2)
    return y

# Plot range
xmin, xmax = -100, 100
x_vec = np.linspace(xmin, xmax, 1000)
y_vec = target_func(x_vec)

plt.figure(figsize=(10, 4))
plt.plot(x_vec, y_vec, linewidth=2)
plt.title('Target Function (Optimization Landscape)')
plt.grid(True)
plt.show()

# ==========================================
# 2. Define Temperature Schedule
# ==========================================
# We want:
# - At Tin: 80% prob to accept max energy jump (High Exploration)
# - At Tend: 5% prob to accept tiny jump (1e-3) (High Locking)

y_max, y_min = np.max(y_vec), np.min(y_vec)
delta_E_max = y_max - y_min

# Reverse calculate T
# P = exp(-dE / T)  =>  ln(P) = -dE / T  =>  T = -dE / ln(P)
T_in = -delta_E_max / np.log(0.8)
T_end = -1e-3 / np.log(0.05)

n_iter = 10000

# Geometric Cooling: T(k) = T_in * exp(-tau * k)
# At k = n_iter-1, T = T_end
tau = -np.log(T_end / T_in) / (n_iter - 1)

iterations = np.arange(n_iter)
T_schedule = T_in * np.exp(-tau * iterations)

plt.figure(figsize=(6, 4))
plt.plot(iterations, T_schedule)
plt.title('Temperature Law (Geometric Decay)')
plt.xlabel('Iteration')
plt.ylabel('Temperature')
plt.grid(True)
plt.show()

# ==========================================
# 3. Simulated Annealing Main Loop
# ==========================================
print("Starting Simulated Annealing...")
sx = np.zeros(n_iter)
alpha_hist = np.zeros(n_iter)

# Initial Point (Random)
current_x = np.random.randn() * 10
sx[0] = current_x
proposal_std = 50 # Standard deviation for proposal

for i in range(1, n_iter):
    # 1. Propose: Random Walk
    # Truncated to constraints
    while True:
        candidate = current_x + np.random.randn() * proposal_std
        if xmin <= candidate <= xmax:
            break
            
    # 2. Calculate Energy Diff
    E_curr = target_func(current_x)
    E_cand = target_func(candidate)
    dE = E_cand - E_curr # Minimize
    
    # 3. Acceptance Ratio
    T_curr = T_schedule[i]
    
    # alpha = exp(-dE / T)
    alpha = min(1, np.exp(-dE / T_curr))
    
    alpha_hist[i] = alpha
    
    # 4. Accept/Reject
    if np.random.rand() < alpha:
        current_x = candidate
    # else: keep current_x
    
    sx[i] = current_x

# Result (Mean of last 20% samples)
t_burn = int(0.8 * n_iter)
x_sol_stochastic = np.mean(sx[t_burn:])
y_sol_stochastic = target_func(x_sol_stochastic)

print(f"Stochastic Solution (SA): x = {x_sol_stochastic:.4f}, y = {y_sol_stochastic:.4f}")
# Plot on the image
plt.figure(figsize=(10, 4))
plt.plot(x_vec, y_vec, linewidth=2)
plt.plot(x_sol_stochastic, y_sol_stochastic, 'ro')
plt.title('Target Function With Solution')
plt.grid(True)
plt.show()

# ==========================================
# 4. Deterministic Contrast (Newton's Method)
# ==========================================
# Newton: x_new = x - f'(x) / f''(x)
# Deliberately start far from global opt, near a local opt
x_newton = -10.0 
newton_path = [x_newton]

# Central Difference for Derivatives
def get_derivatives(f, x, h=1e-5):
    # First derivative f'
    df = (f(x + h) - f(x - h)) / (2 * h)
    # Second derivative f''
    ddf = (f(x + h) - 2*f(x) + f(x - h)) / (h**2)
    return df, ddf

for _ in range(100):
    df, ddf = get_derivatives(target_func, x_newton)
    if abs(ddf) < 1e-6: break # Prevent div by zero
    
    step = df / ddf
    x_newton = x_newton - step
    newton_path.append(x_newton)
    
    if abs(step) < 1e-6: break # Converged

y_sol_deterministic = target_func(x_newton)
print(f"Deterministic Solution (Newton): x = {x_newton:.4f}, y = {y_sol_deterministic:.4f}")

# Plot on the image
plt.figure(figsize=(10, 4))
plt.plot(x_vec, y_vec, linewidth=2)
plt.plot(x_newton, y_sol_deterministic, 'ro')
plt.title('Target Function With Solution)')
plt.grid(True)
plt.show()

# ==========================================
# 5. Final Visualization
# ==========================================
plt.figure(figsize=(12, 6))
plt.plot(x_vec, y_vec, label='Target Function', color='blue', alpha=0.5)

# Plot SA samples (color by time)
plt.scatter(sx[::10], target_func(sx[::10]), c=np.arange(0, n_iter, 10), 
            cmap='Wistia', alpha=0.5, s=10, label='SA Trajectory')

# Mark Solutions
plt.plot(x_sol_stochastic, y_sol_stochastic, 'r*', markersize=20, label='SA Solution (Global)')
plt.plot(x_newton, y_sol_deterministic, 'kx', markersize=15, markeredgewidth=3, label='Newton Solution (Local)')

plt.colorbar(label='Iteration Time')
plt.legend()
plt.title('Global Optimization: Simulated Annealing vs Newton Method')
plt.xlim(xmin, xmax)
plt.show()

import numpy as np
import matplotlib.pyplot as plt

# Plot Prep
x_vec = np.linspace(-100, 100, 1000)
y_vec = target_func(x_vec)

# ==========================================
# Phase 1: Simulated Annealing (Global Search)
# ==========================================
print("🚀 Phase 1: Simulated Annealing (Global Search)...")

# 1. Auto Temp
y_max, y_min = np.max(y_vec), np.min(y_vec)
delta_E_max = y_max - y_min

# Init: 80% accept max jump; End: 5% accept tiny jump
T_in = -delta_E_max / np.log(0.8)
T_end = -1e-3 / np.log(0.05)

n_iter = 5000 
tau = -np.log(T_end / T_in) / (n_iter - 1)

# Init
current_x = np.random.uniform(-90, 90) # Random Start
sx = np.zeros(n_iter)
proposal_std = 20 

# SA Loop
for i in range(n_iter):
    T_curr = T_in * np.exp(-tau * i)
    
    while True:
        candidate = current_x + np.random.randn() * proposal_std
        if -100 <= candidate <= 100: break
    
    dE = target_func(candidate) - target_func(current_x)
    
    if dE < 0 or np.random.rand() < np.exp(-dE / T_curr):
        current_x = candidate
        
    sx[i] = current_x

# SA Result (Mean of last 10%)
t_burn = int(0.9 * n_iter)
x_sa_approx = np.mean(sx[t_burn:])
y_sa_approx = target_func(x_sa_approx)

print(f"✅ SA Done. Approx Opt: x ≈ {x_sa_approx:.4f}")


# ==========================================
# Phase 2: Newton Method (Fine Tuning)
# ==========================================
print("\n🎯 Phase 2: Hybrid Newton Method (Fine Tuning)...")

# --- Key: Use SA result as Newton Start ---
x_hybrid = x_sa_approx 
hybrid_path = [x_hybrid]

for _ in range(50): 
    df, ddf = get_derivatives(target_func, x_hybrid)
    
    if abs(ddf) < 1e-9: break 
    
    step = df / ddf
    x_hybrid = x_hybrid - step
    hybrid_path.append(x_hybrid)
    
    if abs(step) < 1e-10: # High precision stop
        print("   -> Newton Converged.")
        break

y_hybrid = target_func(x_hybrid)
print(f"✅ Hybrid Final Solution: x = {x_hybrid:.8f} (High Precision)")


# ==========================================
# Contrast: Newton Direct (Bad Start)
# ==========================================
print("\n⚠️ Contrast: Direct Newton (No SA)...")
x_bad_newton = -10.0 
for _ in range(50):
    df, ddf = get_derivatives(target_func, x_bad_newton)
    if abs(ddf) < 1e-9: break
    x_bad_newton = x_bad_newton - df / ddf

print(f"❌ Failed. Trapped in Local Opt: x = {x_bad_newton:.4f}")


# ==========================================
# Visualization
# ==========================================
plt.figure(figsize=(12, 6))
plt.plot(x_vec, y_vec, 'k-', alpha=0.3, label='Target Function')

# 1. SA Path
plt.scatter(sx[::10], target_func(sx[::10]), c='orange', s=10, alpha=0.3, label='Phase 1: SA Exploration')

# 2. Hybrid Result
plt.plot(x_hybrid, y_hybrid, 'g*', markersize=25, label='Phase 2: Hybrid Result (Global Opt)')

# 3. Bad Newton
plt.plot(x_bad_newton, target_func(x_bad_newton), 'rx', markersize=20, markeredgewidth=3, label='Newton Only (Local Trap)')

plt.title(f"Hybrid Optimization Strategy\nSA (Approx: {x_sa_approx:.2f}) + Newton (Final: {x_hybrid:.5f})")
plt.legend()
plt.grid(True)
plt.show()