"""
Optimization Algorithms with SciPy
Author: RSK World
Website: https://rskworld.in
Email: help@rskworld.in
Phone: +91 93305 39277
"""

import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import minimize, minimize_scalar, differential_evolution
from scipy.optimize import curve_fit

# Example 1: Univariate Optimization
def example_univariate_optimization():
    """
    Find the minimum of a univariate function
    """
    print("=" * 60)
    print("Example 1: Univariate Optimization")
    print("=" * 60)
    
    # Define a function: f(x) = x^4 - 5x^2 + 4
    def objective(x):
        return x**4 - 5*x**2 + 4
    
    # Find minimum using minimize_scalar
    result = minimize_scalar(objective, method='brent')
    
    print(f"Minimum value: {result.fun:.4f}")
    print(f"At x = {result.x:.4f}")
    
    # Visualize
    x = np.linspace(-3, 3, 100)
    y = objective(x)
    
    plt.figure(figsize=(10, 6))
    plt.plot(x, y, 'b-', linewidth=2, label='f(x) = x⁴ - 5x² + 4')
    plt.plot(result.x, result.fun, 'ro', markersize=10, label=f'Minimum at x={result.x:.4f}')
    plt.xlabel('x', fontsize=12)
    plt.ylabel('f(x)', fontsize=12)
    plt.title('Univariate Optimization', fontsize=14, fontweight='bold')
    plt.grid(True, alpha=0.3)
    plt.legend()
    plt.tight_layout()
    plt.savefig('optimization_univariate.png', dpi=300, bbox_inches='tight')
    print("\nPlot saved as 'optimization_univariate.png'")
    plt.close()


# Example 2: Multivariate Optimization
def example_multivariate_optimization():
    """
    Find the minimum of a multivariate function
    """
    print("\n" + "=" * 60)
    print("Example 2: Multivariate Optimization")
    print("=" * 60)
    
    # Rosenbrock function: f(x,y) = (1-x)² + 100(y-x²)²
    def rosenbrock(x):
        return (1 - x[0])**2 + 100 * (x[1] - x[0]**2)**2
    
    # Initial guess
    x0 = np.array([-1.2, 1.0])
    
    # Minimize using different methods
    methods = ['BFGS', 'CG', 'Nelder-Mead']
    results = {}
    
    for method in methods:
        result = minimize(rosenbrock, x0, method=method)
        results[method] = result
        print(f"\n{method} Method:")
        print(f"  Minimum value: {result.fun:.6f}")
        print(f"  At point: ({result.x[0]:.6f}, {result.x[1]:.6f})")
        print(f"  Iterations: {result.nit}")
    
    # Visualize Rosenbrock function
    x = np.linspace(-2, 2, 100)
    y = np.linspace(-1, 3, 100)
    X, Y = np.meshgrid(x, y)
    Z = (1 - X)**2 + 100 * (Y - X**2)**2
    
    plt.figure(figsize=(12, 5))
    
    # Contour plot
    plt.subplot(1, 2, 1)
    contour = plt.contour(X, Y, Z, levels=50, cmap='viridis')
    plt.colorbar(contour)
    plt.plot(1, 1, 'r*', markersize=15, label='Global Minimum (1,1)')
    for method, result in results.items():
        plt.plot(result.x[0], result.x[1], 'o', markersize=8, label=f'{method}')
    plt.xlabel('x', fontsize=12)
    plt.ylabel('y', fontsize=12)
    plt.title('Rosenbrock Function Contour', fontsize=12, fontweight='bold')
    plt.legend()
    plt.grid(True, alpha=0.3)
    
    # 3D surface plot
    ax = plt.subplot(1, 2, 2, projection='3d')
    surf = ax.plot_surface(X, Y, Z, cmap='viridis', alpha=0.8)
    ax.set_xlabel('x', fontsize=10)
    ax.set_ylabel('y', fontsize=10)
    ax.set_zlabel('f(x,y)', fontsize=10)
    ax.set_title('Rosenbrock Function 3D', fontsize=12, fontweight='bold')
    
    plt.tight_layout()
    plt.savefig('optimization_multivariate.png', dpi=300, bbox_inches='tight')
    print("\nPlot saved as 'optimization_multivariate.png'")
    plt.close()


# Example 3: Constrained Optimization
def example_constrained_optimization():
    """
    Optimization with constraints
    """
    print("\n" + "=" * 60)
    print("Example 3: Constrained Optimization")
    print("=" * 60)
    
    # Objective function: f(x,y) = x² + y²
    def objective(x):
        return x[0]**2 + x[1]**2
    
    # Constraint: x + y >= 1
    def constraint(x):
        return x[0] + x[1] - 1
    
    # Bounds: x >= 0, y >= 0
    bounds = [(0, None), (0, None)]
    
    # Constraint dictionary
    constraints = {'type': 'ineq', 'fun': constraint}
    
    # Initial guess
    x0 = [0.5, 0.5]
    
    # Minimize with constraints
    result = minimize(objective, x0, method='SLSQP', bounds=bounds, constraints=constraints)
    
    print(f"Minimum value: {result.fun:.4f}")
    print(f"At point: ({result.x[0]:.4f}, {result.x[1]:.4f})")
    print(f"Constraint value: {constraint(result.x):.4f} (should be >= 0)")
    
    # Visualize
    x = np.linspace(0, 2, 100)
    y = np.linspace(0, 2, 100)
    X, Y = np.meshgrid(x, y)
    Z = X**2 + Y**2
    
    plt.figure(figsize=(10, 8))
    contour = plt.contour(X, Y, Z, levels=20, cmap='viridis')
    plt.colorbar(contour, label='f(x,y) = x² + y²')
    
    # Plot constraint line: x + y = 1
    constraint_line = 1 - x
    plt.plot(x, constraint_line, 'r--', linewidth=2, label='Constraint: x + y = 1')
    plt.fill_between(x, constraint_line, 2, alpha=0.3, color='red', label='Feasible region')
    
    # Plot optimal point
    plt.plot(result.x[0], result.x[1], 'ro', markersize=12, label=f'Optimum ({result.x[0]:.3f}, {result.x[1]:.3f})')
    
    plt.xlabel('x', fontsize=12)
    plt.ylabel('y', fontsize=12)
    plt.title('Constrained Optimization', fontsize=14, fontweight='bold')
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.xlim(0, 2)
    plt.ylim(0, 2)
    plt.tight_layout()
    plt.savefig('optimization_constrained.png', dpi=300, bbox_inches='tight')
    print("\nPlot saved as 'optimization_constrained.png'")
    plt.close()


# Example 4: Global Optimization
def example_global_optimization():
    """
    Global optimization using differential evolution
    """
    print("\n" + "=" * 60)
    print("Example 4: Global Optimization")
    print("=" * 60)
    
    # Function with multiple local minima: f(x) = x² + 10*sin(x)
    def objective(x):
        return x[0]**2 + 10 * np.sin(x[0])
    
    # Bounds
    bounds = [(-5, 5)]
    
    # Global optimization
    result = differential_evolution(objective, bounds, seed=42)
    
    print(f"Global minimum value: {result.fun:.4f}")
    print(f"At x = {result.x[0]:.4f}")
    print(f"Function evaluations: {result.nfev}")
    
    # Visualize
    x = np.linspace(-5, 5, 200)
    y = x**2 + 10 * np.sin(x)
    
    plt.figure(figsize=(10, 6))
    plt.plot(x, y, 'b-', linewidth=2, label='f(x) = x² + 10sin(x)')
    plt.plot(result.x[0], result.fun, 'ro', markersize=12, label=f'Global minimum at x={result.x[0]:.4f}')
    plt.xlabel('x', fontsize=12)
    plt.ylabel('f(x)', fontsize=12)
    plt.title('Global Optimization with Differential Evolution', fontsize=14, fontweight='bold')
    plt.grid(True, alpha=0.3)
    plt.legend()
    plt.tight_layout()
    plt.savefig('optimization_global.png', dpi=300, bbox_inches='tight')
    print("\nPlot saved as 'optimization_global.png'")
    plt.close()


# Example 5: Curve Fitting
def example_curve_fitting():
    """
    Curve fitting using optimization
    """
    print("\n" + "=" * 60)
    print("Example 5: Curve Fitting")
    print("=" * 60)
    
    # Generate noisy data
    np.random.seed(42)
    x_data = np.linspace(0, 10, 50)
    y_true = 2 * np.sin(0.5 * x_data) + 0.5 * x_data
    y_data = y_true + 0.5 * np.random.randn(len(x_data))
    
    # Define model function: y = a * sin(b * x) + c * x
    def model(x, a, b, c):
        return a * np.sin(b * x) + c * x
    
    # Fit the curve
    popt, pcov = curve_fit(model, x_data, y_data, p0=[2, 0.5, 0.5])
    
    print(f"Fitted parameters:")
    print(f"  a (amplitude) = {popt[0]:.4f}")
    print(f"  b (frequency) = {popt[1]:.4f}")
    print(f"  c (slope) = {popt[2]:.4f}")
    
    # Generate fitted curve
    x_fit = np.linspace(0, 10, 200)
    y_fit = model(x_fit, *popt)
    
    # Visualize
    plt.figure(figsize=(10, 6))
    plt.scatter(x_data, y_data, alpha=0.6, color='blue', label='Noisy data')
    plt.plot(x_data, y_true, 'g--', linewidth=2, label='True function')
    plt.plot(x_fit, y_fit, 'r-', linewidth=2, label='Fitted curve')
    plt.xlabel('x', fontsize=12)
    plt.ylabel('y', fontsize=12)
    plt.title('Curve Fitting Example', fontsize=14, fontweight='bold')
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.savefig('optimization_curve_fitting.png', dpi=300, bbox_inches='tight')
    print("\nPlot saved as 'optimization_curve_fitting.png'")
    plt.close()


def main():
    """
    Main function to run all optimization examples
    """
    print("\n" + "=" * 60)
    print("SciPy Optimization Examples")
    print("Author: RSK World - https://rskworld.in")
    print("=" * 60 + "\n")
    
    example_univariate_optimization()
    example_multivariate_optimization()
    example_constrained_optimization()
    example_global_optimization()
    example_curve_fitting()
    
    print("\n" + "=" * 60)
    print("All optimization examples completed!")
    print("=" * 60)


if __name__ == "__main__":
    main()