import numpy as np
import matplotlib.pyplot as plt
from iterative_solvers import *

def mat_vec_prod_p2(x_flat: np.ndarray, n: int, h_sq: float) -> np.ndarray:
    """
    Computes the matrix-vector product y = Ax for the 2D PDE problem.
    """
    x_2d = x_flat.reshape((n, n))
    x_padded = np.zeros((n + 2, n + 2))
    x_padded[1:-1, 1:-1] = x_2d
    
    y_2d = np.zeros((n, n))
    
    # Apply the finite difference stencil: (4+h^2)u_ij - u_{i-1,j} - u_{i+1,j} - u_{i,j-1} - u_{i,j+1}
    for i in range(n):
        for j in range(n):
            y_2d[i, j] = (4.0 + h_sq) * x_padded[i + 1, j + 1] \
                       - x_padded[i, j + 1] \
                       - x_padded[i + 2, j + 1] \
                       - x_padded[i + 1, j] \
                       - x_padded[i + 1, j + 2]
                       
    return y_2d.flatten()

def get_diag_inv_p2(i: int, n: int, h_sq: float) -> float:
    """Returns 1/A_ii for Problem 2. Diagonal is constant: 4 + h^2."""
    return 1.0 / (4.0 + h_sq)

def get_off_diag_sum_p2(x_flat: np.ndarray, i_flat: int, n: int, **kwargs) -> float:
    """Computes the off-diagonal sum for a given row `i_flat` in Problem 2."""
    row = i_flat // n
    col = i_flat % n
    sum_val = 0.0
    
    if row > 0: sum_val += -x_flat[i_flat - n]
    if row < n - 1: sum_val += -x_flat[i_flat + n]
    if col > 0: sum_val += -x_flat[i_flat - 1]
    if col < n - 1: sum_val += -x_flat[i_flat + 1]
        
    return sum_val

def solve_and_report_p2(n: int):
    """Solves the 2D PDE linear system for a given n and reports results."""
    print(f"\n--- Solving for n = {n} (System size: {n*n} x {n*n}) ---")
    h = 1.0 / (n + 1)
    h_sq = h**2
    N = n * n
    
    b = np.full(N, h_sq)
    omega = 2.0 / (1.0 + np.sin(np.pi * h))
    
    iteration_counts = []
    all_residuals = {}
    all_solutions = {}

    methods_to_run = {
        'Jacobi': ('jacobi', {}),
        'Gauss-Seidel': ('gauss_seidel', {}),
        'SOR': ('sor', {'omega': omega}),
        'Steepest Descent': (steepest_descent, {}),
        'Conjugate Gradient': (conjugate_gradient, {})
    }
    
    print("\nStarting iterative solvers...")
    for name, (method_func, params) in methods_to_run.items():
        print(f"Running {name}...")
        
        solver_kwargs = {'n': n, 'h_sq': h_sq}
        
        if isinstance(method_func, str):
            solver_args = {'b': b, 'method': method_func, 'get_diag_inv': get_diag_inv_p2, 
                           'get_off_diag_sum': get_off_diag_sum_p2, 'mat_vec_prod': mat_vec_prod_p2,
                           **solver_kwargs, **params}
            solution, iters, res_hist = stationary_method(**solver_args)
        else:
            solver_args = {'b': b, 'mat_vec_prod': mat_vec_prod_p2, **solver_kwargs}
            solution, iters, res_hist = method_func(**solver_args)
        
        iteration_counts.append((name, iters))
        all_residuals[name] = res_hist
        all_solutions[name] = solution if iters < MAX_ITER else None

    # Report iteration counts
    print("\n--- Iteration Counts ---")
    for name, iters in iteration_counts:
        status = "converged" if iters < MAX_ITER else "did NOT converge"
        print(f"{name:<20}: {iters} iterations ({status})")

    # Plot residue history
    plt.figure(figsize=(12, 8))
    for name, res_hist in all_residuals.items():
        res_log = np.log10(np.array(res_hist) + 1e-20)
        plt.plot(res_log, label=name, linestyle='-')
    
    plt.title(f'Convergence History for PDE Problem, n = {n}')
    plt.xlabel('Iteration Step (m)')
    plt.ylabel('log10(||r^(m)||_2)')
    plt.legend()
    plt.grid(True)
    
    current_ylim = plt.ylim()
    plt.ylim(bottom=max(current_ylim[0], -13), top=current_ylim[1])
    
    plt.savefig(f'problem2_convergence_n{n}.png')
    print(f"\nConvergence plot saved to 'problem2_convergence_n{n}.png'")

    if all_solutions:
        print("Generating 3D surface plots of the solutions...")
        
        # Create the X, Y grid for the surface plot, including boundaries
        plot_grid = np.linspace(0, 1, n + 2)
        X, Y = np.meshgrid(plot_grid, plot_grid)
        
        # Create a figure for the subplots
        fig = plt.figure(figsize=(18, 10))
        fig.suptitle(f'3D Surface Plots of Numerical Solutions for n = {n}', fontsize=16)
        # Loop through the stored solutions and create a subplot for each
        for i, (name, solution_flat) in enumerate(all_solutions.items()):
            # Reshape the 1D solution vector into a 2D grid
            solution_2d = solution_flat.reshape((n, n))
            
            # Pad the solution with zero boundaries for plotting
            u_padded = np.zeros((n + 2, n + 2))
            u_padded[1:-1, 1:-1] = solution_2d
            
            # Add a new 3D subplot
            ax = fig.add_subplot(2, 3, i + 1, projection='3d')
            
            # Plot the surface
            ax.plot_surface(X, Y, u_padded, cmap='viridis', edgecolor='none')
            
            # Set titles and labels for each subplot
            ax.set_title(name)
            ax.set_xlabel('x')
            ax.set_ylabel('y')
            ax.set_zlabel('u(x,y)')
            ax.view_init(elev=30, azim=-60) # Set a nice viewing angle
        plt.tight_layout(rect=[0, 0.03, 1, 0.95]) # Adjust layout to make room for suptitle
        plt.savefig(f'problem2_3d_solutions_n{n}.png')
        print(f"3D surface comparison plot saved to 'problem2_3d_solutions_n{n}.png'")

if __name__ == "__main__":
    # Solve Problem 2 for n=16 and n=32
    solve_and_report_p2(n=16)
    solve_and_report_p2(n=32)