"""
Interpolation and Fitting 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.interpolate import interp1d, interp2d, griddata, UnivariateSpline, RBFInterpolator
from scipy.optimize import curve_fit

# Example 1: 1D Interpolation
def example_1d_interpolation():
    """
    One-dimensional interpolation methods
    """
    print("=" * 60)
    print("Example 1: 1D Interpolation")
    print("=" * 60)
    
    # Generate sample data
    np.random.seed(42)
    x_known = np.linspace(0, 10, 10)
    y_known = np.sin(x_known) + 0.1 * np.random.randn(len(x_known))
    
    # Create interpolation functions
    linear = interp1d(x_known, y_known, kind='linear')
    cubic = interp1d(x_known, y_known, kind='cubic')
    spline = UnivariateSpline(x_known, y_known, s=0)
    
    # Evaluate at new points
    x_new = np.linspace(0, 10, 100)
    y_linear = linear(x_new)
    y_cubic = cubic(x_new)
    y_spline = spline(x_new)
    y_true = np.sin(x_new)
    
    print(f"Interpolation methods: linear, cubic, spline")
    print(f"Original points: {len(x_known)}")
    print(f"Interpolated points: {len(x_new)}")
    
    # Visualize
    plt.figure(figsize=(12, 6))
    plt.scatter(x_known, y_known, color='red', s=100, zorder=5, label='Known data points')
    plt.plot(x_new, y_true, 'k--', linewidth=2, alpha=0.5, label='True function: sin(x)')
    plt.plot(x_new, y_linear, 'b-', linewidth=1.5, label='Linear interpolation')
    plt.plot(x_new, y_cubic, 'g-', linewidth=1.5, label='Cubic interpolation')
    plt.plot(x_new, y_spline, 'm-', linewidth=1.5, label='Spline interpolation')
    plt.xlabel('x', fontsize=12)
    plt.ylabel('y', fontsize=12)
    plt.title('1D Interpolation Methods', fontsize=14, fontweight='bold')
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.savefig('interpolation_1d.png', dpi=300, bbox_inches='tight')
    print("\nPlot saved as 'interpolation_1d.png'")
    plt.close()


# Example 2: 2D Interpolation
def example_2d_interpolation():
    """
    Two-dimensional interpolation
    """
    print("\n" + "=" * 60)
    print("Example 2: 2D Interpolation")
    print("=" * 60)
    
    # Generate sample 2D data
    np.random.seed(42)
    x_known = np.random.rand(20) * 10
    y_known = np.random.rand(20) * 10
    z_known = np.sin(x_known) * np.cos(y_known) + 0.1 * np.random.randn(len(x_known))
    
    # Create regular grid for interpolation
    x_grid = np.linspace(0, 10, 50)
    y_grid = np.linspace(0, 10, 50)
    X_grid, Y_grid = np.meshgrid(x_grid, y_grid)
    
    # Interpolate using different methods
    z_linear = griddata((x_known, y_known), z_known, (X_grid, Y_grid), method='linear')
    z_cubic = griddata((x_known, y_known), z_known, (X_grid, Y_grid), method='cubic')
    
    print(f"2D Interpolation methods: linear, cubic")
    print(f"Known points: {len(x_known)}")
    print(f"Grid size: {len(x_grid)} × {len(y_grid)}")
    
    # Visualize
    fig = plt.figure(figsize=(14, 5))
    
    # Original data
    ax1 = plt.subplot(1, 3, 1)
    scatter = ax1.scatter(x_known, y_known, c=z_known, s=100, cmap='viridis', edgecolors='black')
    plt.colorbar(scatter, label='z')
    ax1.set_xlabel('x', fontsize=12)
    ax1.set_ylabel('y', fontsize=12)
    ax1.set_title('Original Data Points', fontsize=12, fontweight='bold')
    ax1.grid(True, alpha=0.3)
    
    # Linear interpolation
    ax2 = plt.subplot(1, 3, 2)
    contour = ax2.contourf(X_grid, Y_grid, z_linear, levels=20, cmap='viridis')
    plt.colorbar(contour, label='z')
    ax2.set_xlabel('x', fontsize=12)
    ax2.set_ylabel('y', fontsize=12)
    ax2.set_title('Linear Interpolation', fontsize=12, fontweight='bold')
    ax2.grid(True, alpha=0.3)
    
    # Cubic interpolation
    ax3 = plt.subplot(1, 3, 3)
    contour = ax3.contourf(X_grid, Y_grid, z_cubic, levels=20, cmap='viridis')
    plt.colorbar(contour, label='z')
    ax3.set_xlabel('x', fontsize=12)
    ax3.set_ylabel('y', fontsize=12)
    ax3.set_title('Cubic Interpolation', fontsize=12, fontweight='bold')
    ax3.grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig('interpolation_2d.png', dpi=300, bbox_inches='tight')
    print("\nPlot saved as 'interpolation_2d.png'")
    plt.close()


# Example 3: Curve Fitting
def example_curve_fitting():
    """
    Curve fitting with different models
    """
    print("\n" + "=" * 60)
    print("Example 3: Curve Fitting")
    print("=" * 60)
    
    # Generate noisy data
    np.random.seed(42)
    x_data = np.linspace(0, 10, 30)
    y_true = 2 * np.exp(-0.5 * x_data) * np.sin(2 * x_data)
    y_data = y_true + 0.2 * np.random.randn(len(x_data))
    
    # Define model functions
    def exponential_decay(x, a, b, c, d):
        return a * np.exp(-b * x) * np.sin(c * x + d)
    
    def polynomial(x, a, b, c, d, e):
        return a * x**4 + b * x**3 + c * x**2 + d * x + e
    
    # Fit exponential model
    popt_exp, pcov_exp = curve_fit(exponential_decay, x_data, y_data, 
                                   p0=[2, 0.5, 2, 0])
    
    # Fit polynomial model
    popt_poly, pcov_poly = curve_fit(polynomial, x_data, y_data)
    
    # Generate fitted curves
    x_fit = np.linspace(0, 10, 200)
    y_fit_exp = exponential_decay(x_fit, *popt_exp)
    y_fit_poly = polynomial(x_fit, *popt_poly)
    
    print(f"Exponential model parameters:")
    print(f"  a = {popt_exp[0]:.4f}, b = {popt_exp[1]:.4f}")
    print(f"  c = {popt_exp[2]:.4f}, d = {popt_exp[3]:.4f}")
    
    # Calculate R²
    ss_res_exp = np.sum((y_data - exponential_decay(x_data, *popt_exp))**2)
    ss_tot = np.sum((y_data - np.mean(y_data))**2)
    r2_exp = 1 - (ss_res_exp / ss_tot)
    
    print(f"R² (exponential): {r2_exp:.4f}")
    
    # Visualize
    plt.figure(figsize=(12, 6))
    plt.scatter(x_data, y_data, alpha=0.7, color='blue', s=50, label='Noisy data')
    plt.plot(x_data, y_true, 'k--', linewidth=2, alpha=0.7, label='True function')
    plt.plot(x_fit, y_fit_exp, 'r-', linewidth=2, label=f'Exponential fit (R²={r2_exp:.3f})')
    plt.plot(x_fit, y_fit_poly, 'g-', linewidth=2, label='Polynomial fit')
    plt.xlabel('x', fontsize=12)
    plt.ylabel('y', fontsize=12)
    plt.title('Curve Fitting Comparison', fontsize=14, fontweight='bold')
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.savefig('interpolation_curve_fitting.png', dpi=300, bbox_inches='tight')
    print("\nPlot saved as 'interpolation_curve_fitting.png'")
    plt.close()


# Example 4: Spline Interpolation
def example_spline_interpolation():
    """
    Spline interpolation with different smoothing
    """
    print("\n" + "=" * 60)
    print("Example 4: Spline Interpolation")
    print("=" * 60)
    
    # Generate noisy data
    np.random.seed(42)
    x_known = np.linspace(0, 10, 15)
    y_known = np.sin(x_known) + 0.2 * np.random.randn(len(x_known))
    
    # Create splines with different smoothing factors
    spline_smooth = UnivariateSpline(x_known, y_known, s=len(x_known))
    spline_interp = UnivariateSpline(x_known, y_known, s=0)
    
    # Evaluate
    x_new = np.linspace(0, 10, 200)
    y_smooth = spline_smooth(x_new)
    y_interp = spline_interp(x_new)
    y_true = np.sin(x_new)
    
    print(f"Smoothing spline (s={len(x_known)}): smoother, less accurate")
    print(f"Interpolating spline (s=0): passes through all points")
    
    # Visualize
    plt.figure(figsize=(12, 6))
    plt.scatter(x_known, y_known, color='red', s=100, zorder=5, label='Known data points')
    plt.plot(x_new, y_true, 'k--', linewidth=2, alpha=0.5, label='True function: sin(x)')
    plt.plot(x_new, y_smooth, 'b-', linewidth=2, label='Smoothing spline')
    plt.plot(x_new, y_interp, 'g-', linewidth=2, label='Interpolating spline')
    plt.xlabel('x', fontsize=12)
    plt.ylabel('y', fontsize=12)
    plt.title('Spline Interpolation: Smoothing vs Interpolation', fontsize=14, fontweight='bold')
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.savefig('interpolation_spline.png', dpi=300, bbox_inches='tight')
    print("\nPlot saved as 'interpolation_spline.png'")
    plt.close()


# Example 5: Radial Basis Function Interpolation
def example_rbf_interpolation():
    """
    Radial Basis Function (RBF) interpolation
    """
    print("\n" + "=" * 60)
    print("Example 5: RBF Interpolation")
    print("=" * 60)
    
    # Generate sample data
    np.random.seed(42)
    x_known = np.random.rand(30) * 10
    y_known = np.random.rand(30) * 10
    z_known = np.sin(x_known) * np.cos(y_known) + 0.1 * np.random.randn(len(x_known))
    
    # Create RBF interpolator
    points = np.column_stack((x_known, y_known))
    rbf = RBFInterpolator(points, z_known, kernel='thin_plate_spline')
    
    # Create grid
    x_grid = np.linspace(0, 10, 50)
    y_grid = np.linspace(0, 10, 50)
    X_grid, Y_grid = np.meshgrid(x_grid, y_grid)
    grid_points = np.column_stack((X_grid.ravel(), Y_grid.ravel()))
    
    # Interpolate
    Z_rbf = rbf(grid_points).reshape(X_grid.shape)
    
    print(f"RBF Interpolation with thin plate spline kernel")
    print(f"Known points: {len(x_known)}")
    print(f"Grid size: {len(x_grid)} × {len(y_grid)}")
    
    # Visualize
    fig = plt.figure(figsize=(12, 5))
    
    # Original data
    ax1 = plt.subplot(1, 2, 1)
    scatter = ax1.scatter(x_known, y_known, c=z_known, s=100, cmap='viridis', 
                         edgecolors='black', vmin=Z_rbf.min(), vmax=Z_rbf.max())
    plt.colorbar(scatter, label='z')
    ax1.set_xlabel('x', fontsize=12)
    ax1.set_ylabel('y', fontsize=12)
    ax1.set_title('Original Data Points', fontsize=12, fontweight='bold')
    ax1.grid(True, alpha=0.3)
    
    # RBF interpolation
    ax2 = plt.subplot(1, 2, 2)
    contour = ax2.contourf(X_grid, Y_grid, Z_rbf, levels=20, cmap='viridis')
    plt.colorbar(contour, label='z')
    ax2.scatter(x_known, y_known, c='red', s=30, edgecolors='black', marker='x')
    ax2.set_xlabel('x', fontsize=12)
    ax2.set_ylabel('y', fontsize=12)
    ax2.set_title('RBF Interpolation', fontsize=12, fontweight='bold')
    ax2.grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig('interpolation_rbf.png', dpi=300, bbox_inches='tight')
    print("\nPlot saved as 'interpolation_rbf.png'")
    plt.close()


# Example 6: Extrapolation
def example_extrapolation():
    """
    Extrapolation beyond data range
    """
    print("\n" + "=" * 60)
    print("Example 6: Extrapolation")
    print("=" * 60)
    
    # Generate data
    x_known = np.linspace(0, 5, 10)
    y_known = np.exp(-x_known) * np.sin(x_known)
    
    # Create interpolation functions
    linear = interp1d(x_known, y_known, kind='linear', 
                      fill_value='extrapolate', bounds_error=False)
    cubic = interp1d(x_known, y_known, kind='cubic',
                     fill_value='extrapolate', bounds_error=False)
    
    # Evaluate including extrapolation region
    x_new = np.linspace(-1, 7, 100)
    y_linear = linear(x_new)
    y_cubic = cubic(x_new)
    y_true = np.exp(-x_new) * np.sin(x_new)
    
    # Identify regions
    mask_interp = (x_new >= x_known.min()) & (x_new <= x_known.max())
    mask_extrap = ~mask_interp
    
    print(f"Interpolation region: [{x_known.min():.1f}, {x_known.max():.1f}]")
    print(f"Extrapolation regions: [-∞, {x_known.min():.1f}] and [{x_known.max():.1f}, +∞]")
    
    # Visualize
    plt.figure(figsize=(12, 6))
    plt.scatter(x_known, y_known, color='red', s=100, zorder=5, label='Known data points')
    plt.plot(x_new, y_true, 'k--', linewidth=2, alpha=0.5, label='True function')
    plt.plot(x_new[mask_interp], y_linear[mask_interp], 'b-', linewidth=2, label='Linear (interpolation)')
    plt.plot(x_new[mask_extrap], y_linear[mask_extrap], 'b--', linewidth=1.5, alpha=0.7, label='Linear (extrapolation)')
    plt.plot(x_new[mask_interp], y_cubic[mask_interp], 'g-', linewidth=2, label='Cubic (interpolation)')
    plt.plot(x_new[mask_extrap], y_cubic[mask_extrap], 'g--', linewidth=1.5, alpha=0.7, label='Cubic (extrapolation)')
    plt.axvline(x_known.min(), color='gray', linestyle=':', alpha=0.5)
    plt.axvline(x_known.max(), color='gray', linestyle=':', alpha=0.5)
    plt.xlabel('x', fontsize=12)
    plt.ylabel('y', fontsize=12)
    plt.title('Interpolation vs Extrapolation', fontsize=14, fontweight='bold')
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.savefig('interpolation_extrapolation.png', dpi=300, bbox_inches='tight')
    print("\nPlot saved as 'interpolation_extrapolation.png'")
    plt.close()


def main():
    """
    Main function to run all interpolation examples
    """
    print("\n" + "=" * 60)
    print("SciPy Interpolation and Fitting Examples")
    print("Author: RSK World - https://rskworld.in")
    print("=" * 60 + "\n")
    
    example_1d_interpolation()
    example_2d_interpolation()
    example_curve_fitting()
    example_spline_interpolation()
    example_rbf_interpolation()
    example_extrapolation()
    
    print("\n" + "=" * 60)
    print("All interpolation examples completed!")
    print("=" * 60)


if __name__ == "__main__":
    main()