import numpy as np
import matplotlib.pyplot as plt
import sys

# Configuration for plotting
plt.style.use('seaborn-v0_8')
plt.rcParams['figure.figsize'] = (18, 6)

def newton_method(f, df, x0, tolerance=1e-12, max_iter=20):
    """
    Implements Newton's Method with overflow protection.
    Returns: list of iteration history [(iter, x, error)]
    """
    history = []
    x = x0
    
    for i in range(max_iter):
        try:
            # Check for divergence first
            if abs(x) > 1e100:
                break

            fx = f(x)
            dfx = df(x)
            
            if abs(dfx) < 1e-15: # Avoid division by zero
                break
                
            history.append(x)
            
            x_new = x - fx / dfx
            
            # Check convergence
            if abs(x_new - x) < tolerance:
                x = x_new
                history.append(x)
                break
                
            x = x_new
        except OverflowError:
            # Catch errors if numbers become too large
            break
        
    return history

def chord_method(f, df_x0, x0, tolerance=1e-12, max_iter=50):
    """
    Implements Chord Method using fixed slope m = f'(x0) with overflow protection.
    Returns: list of iteration history
    """
    history = []
    x = x0
    m = df_x0 # Fixed slope
    
    for i in range(max_iter):
        try:
            if abs(x) > 1e100:
                break
                
            fx = f(x)
            history.append(x)
            
            if abs(m) < 1e-15:
                break
                
            x_new = x - fx / m
            
            if abs(x_new - x) < tolerance:
                x = x_new
                history.append(x)
                break
            x = x_new
        except OverflowError:
            break
            
    return history

def secant_method(f, x0, x_minus_1, tolerance=1e-12, max_iter=20):
    """
    Implements Secant Method with overflow protection.
    Returns: list of iteration history
    """
    history = []
    x_curr = x0
    x_prev = x_minus_1
    
    for i in range(max_iter):
        try:
            if abs(x_curr) > 1e100:
                break

            history.append(x_curr)
            
            fx_curr = f(x_curr)
            fx_prev = f(x_prev)
            
            if abs(fx_curr - fx_prev) < 1e-15:
                break
                
            # Secant update
            x_new = x_curr - fx_curr * (x_curr - x_prev) / (fx_curr - fx_prev)
            
            if abs(x_new - x_curr) < tolerance:
                x_curr = x_new
                history.append(x_curr)
                break
                
            x_prev = x_curr
            x_curr = x_new
        except OverflowError:
            break
        
    return history

def solve_and_save():
    # Definition of problems
    problems = [
        {
            'id': 'a',
            'title': r'$f(x) = \cos(x) - x, x_0 = 0.5$',
            'func_name': 'cos(x) - x',
            'f': lambda x: np.cos(x) - x,
            'df': lambda x: -np.sin(x) - 1,
            'x0': 0.5,
            'true_root': 0.73908513321516064165531208767387340401341175890075746496568063577328465488354759 # From wolframalpha 
        },
        {
            'id': 'b',
            'title': r'$f(x) = \sin(x), x_0 = 3$',
            'func_name': 'sin(x)',
            'f': lambda x: np.sin(x),
            'df': lambda x: np.cos(x),
            'x0': 3.0,
            'true_root': np.pi
        },
        {
            'id': 'c',
            'title': r'$f(x) = x^2 + 1, x_0 = 0.5$',
            'func_name': 'x^2 + 1',
            'f': lambda x: x**2 + 1,
            'df': lambda x: 2*x,
            'x0': 0.5,
            'true_root': None # No real root
        }
    ]
    
    # Open file for writing text results
    output_filename = 'numerical_results.txt'
    image_filename = 'convergence_results.png'
    
    with open(output_filename, 'w', encoding='utf-8') as txt_file:
        
        fig, axes = plt.subplots(1, 3)
        
        for idx, prob in enumerate(problems):
            header_str = f"\n{'='*60}\nProblem ({prob['id']}): {prob['func_name']}, x0={prob['x0']}\n{'='*60}\n"
            print(header_str) # Print to console to show progress
            txt_file.write(header_str)
            
            x0 = prob['x0']
            f = prob['f']
            df = prob['df']
            
            # 1. Run Newton
            hist_newton = newton_method(f, df, x0)
            
            # 2. Run Chord (m = f'(x0))
            hist_chord = chord_method(f, df(x0), x0)
            
            # 3. Run Secant (x_-1 = 0.99 * x0)
            x_minus_1 = 0.99 * x0
            hist_secant = secant_method(f, x0, x_minus_1)
            
            methods_data = {
                "Newton": hist_newton,
                "Chord": hist_chord,
                "Secant": hist_secant
            }
            
            ax = axes[idx]
            
            for m_name, history in methods_data.items():
                iters = list(range(len(history)))
                errors = []
                
                # Calculate errors for plotting
                # Note: Handling overflow for plotting large errors in problem c
                for x in history:
                    try:
                        if prob['true_root'] is not None:
                            val = abs(x - prob['true_root'])
                        else:
                            val = abs(f(x)) # Residual for problem c
                        errors.append(val)
                    except OverflowError:
                        errors.append(float('inf'))

                # Write to TXT file
                txt_file.write(f"\nMethod: {m_name}\n")
                txt_file.write(f"{'Iter':<5} | {'Value (x_k)':<25} | {'Error/Resid':<25}\n")
                txt_file.write("-" * 60 + "\n")
                
                for i, val in enumerate(history):
                    # Only write full history if short, or head/tail if long
                    if len(history) < 20 or i < 5 or i > len(history) - 3:
                        if prob['true_root'] is not None:
                            err_val = abs(val - prob['true_root'])
                        else:
                            try:
                                err_val = abs(f(val))
                            except OverflowError:
                                err_val = float('inf')
                                
                        txt_file.write(f"{i:<5} | {val:<25.10g} | {err_val:<25.6e}\n")
                    elif i == 5:
                        txt_file.write("...\n")

                # Add to plot
                if prob['true_root'] is not None:
                    ylabel_text = "Error |x_k - x*|"
                    log_scale = True
                    ax.plot(iters, errors, marker='o', label=m_name)
                else:
                    ylabel_text = "Residual |f(x_k)|"
                    log_scale = False
                    # For problem C, data explodes, so we limit plotting range to avoid ruining the graph
                    valid_errors = [e for e in errors if e < 1e10]
                    valid_iters = iters[:len(valid_errors)]
                    ax.plot(valid_iters, valid_errors, marker='x', label=m_name)

            ax.set_title(f"Problem ({prob['id']})\n{prob['title']}")
            ax.set_xlabel("Iteration")
            ax.set_ylabel(ylabel_text)
            if log_scale:
                ax.set_yscale('log')
            ax.legend()
            if prob['id'] == 'c':
                ax.set_ylim(0, 20)
                ax.set_xticks(range(0, len(iters), 1))
            ax.grid(True)

    print(f"\nText results saved to: {output_filename}")
    
    plt.tight_layout()
    plt.savefig(image_filename)
    print(f"Chart saved to: {image_filename}")

if __name__ == "__main__":
    solve_and_save()