import numpy as np
from scipy.optimize import fsolve

# Parameters
n = 3.0
R0 = 5.0
r_max = 1.0
num_points = 100

# Discretization
r = np.linspace(0, r_max, num_points)
dr = r[1] - r[0]

# Initial guess for f(r)
f_initial = R0**(2/(n-1)) * (1 - r**2)**2

# Boundary conditions
f_initial[0] = f_initial[1]  # Regularity condition at r=0
f_initial[-1] = 0.0  # Dirichlet condition at r=1

# Finite difference method
def lane_emden_eq(f):
    dfdr = np.zeros_like(f)
    d2fdr2 = np.zeros_like(f)
    
    # Central difference for the second derivative
    for i in range(1, num_points - 1):
        dfdr[i] = (f[i+1] - f[i-1]) / (2 * dr)
        d2fdr2[i] = (f[i+1] - 2*f[i] + f[i-1]) / (dr**2)
    
    # Regularity condition at r=0
    d2fdr2[0] = 2 * (f[1] - f[0]) / (dr**2)
    
    # Lane-Emden equation
    return d2fdr2 + f**n

# Solve the nonlinear system
f_solution = fsolve(lane_emden_eq, f_initial)

# Save the solution
np.save('/home/weichao/Downloads/Code_Generation_Benchmark/PDE_Benchmark/results/prediction/gpt-4o/prompts/f_solution_Lane_Emden_Equation.npy', f_solution)