# author: Eduardo Gutierrez Prieto, September 2025
# Based on Matlab algorithm provided in ME-411 Studio 03
#
# The goal of this code is to compute the equilibrium shape of an elastic
# sheet deformed by its own distributed weight.
#
# To compute this shape, we numerically solve the ODE:
# theta''(s) = Delta * (1 - s) * cos(theta(s))
# with the boundary conditions:
# theta(s=0) = pi/2   and   theta'(s=1) = 0
#
# This Boundary Value Problem (BVP) is solved using a "shooting method".
# It is converted into an Initial Value Problem (IVP) by guessing the
# value of theta'(s=0). The scipy.optimize.fsolve function (a root-finder)
# iteratively calls an ODE solver (scipy.integrate.solve_ivp) to adjust
# the initial guess until the solution satisfies the boundary condition
# at s=1 (i.e., theta'(s=1) = 0).
#
# Finally, the resulting shape is plotted.

# Library Requirements:
# - numpy: for numerical operations (pi, arrays, trigonometric functions)
# - scipy: for the ODE solver (solve_ivp) and the root-finder (fsolve)
# - matplotlib: for plotting the results

import numpy as np
from scipy.integrate import solve_ivp
from scipy.optimize import fsolve
import matplotlib.pyplot as plt

def ODE_complete(s, X, delta):
    """
    Defines the system of first-order ordinary differential equations.

    Args:
        s (float): The independent variable (arc length).
        X (list or np.array): The state vector [x, y, theta, theta_prime].
        delta (float): The system parameter Delta.

    Returns:
        list: The derivatives [dx/ds, dy/ds, dtheta/ds, dtheta'/ds].
    """
    theta = X[2]
    
    dxds = np.cos(theta)
    dyds = np.sin(theta)
    dthetads = X[3]
    ddthetads = delta * (1 - s) * np.cos(theta)
    
    return [dxds, dyds, dthetads, ddthetads]

def theta_prime_generator_complete(theta_prime0, x0, y0, theta0, Delta):
    """
    Integrates the ODE for a given guess of the initial condition theta_prime0
    and returns the value of theta_prime at the end of the interval (s=1).
    This function acts as the objective function for the root-finder.

    Args:
        theta_prime0 (array-like): Guessed initial value for dtheta/ds from fsolve.
        x0, y0, theta0 (float): Known initial conditions at s=0.
        Delta (float): The system parameter.

    Returns:
        float: The value of dtheta/ds at s=1, which we want to be zero.
    """
    # --- FIX ---
    # fsolve passes an array (e.g., [-14.0]), not a scalar. We extract the scalar value.
    theta_prime0_scalar = theta_prime0[0]
    
    # Define the integration interval
    s_span = [0, 1]
    
    # Define the complete set of initial conditions using the scalar guess
    initial_conditions = [x0, y0, theta0, theta_prime0_scalar]
    
    # Integrate the ODE system
    sol = solve_ivp(ODE_complete, s_span, initial_conditions, args=(Delta,))
    
    # Return the value of theta_prime at the final point.
    theta_prime1 = sol.y[3, -1]
    
    return theta_prime1

# --- Main script ---

# KNOWN BOUNDARY CONDITIONS AT s=0
x0 = 0.0
y0 = 0.0
theta0 = np.pi / 2

# SYSTEM PARAMETERS
L_L_c = 3.0
Delta = L_L_c**3

# Create a lambda function for fsolve
theta_prime_generator = lambda theta_prime0: theta_prime_generator_complete(
    theta_prime0, x0, y0, theta0, Delta
)

# Find the right theta_prime0 using fsolve
theta_prime0_guess = -14.0
theta_prime0_good = fsolve(theta_prime_generator, theta_prime0_guess)[0]

print(f"Initial guess for theta'(0): {theta_prime0_guess}")
print(f"Found value for theta'(0) that satisfies BC: {theta_prime0_good:.6f}")

# Now, we integrate the ODE one last time with the correct theta_prime0
s_span = [0, 1]
s_eval = np.linspace(s_span[0], s_span[1], 200)
initial_conditions_good = [x0, y0, theta0, theta_prime0_good]

sol = solve_ivp(ODE_complete, s_span, initial_conditions_good, args=(Delta,), t_eval=s_eval)

# Extract the results for plotting
x = sol.y[0]
y = sol.y[1]

# Plotting the results
plt.figure(figsize=(8, 8))
plt.plot(x * L_L_c, y * L_L_c, '-k', linewidth=1.5)

# Cosmetics of the figure
plt.axis('equal')
W_window = 8.0
plt.xlim([-W_window / 2, W_window / 2])
plt.ylim([-W_window / 2, W_window / 2])

the_title = f'$L/L_c = {L_L_c:.2f} \\, (\\Delta = {Delta:.2f})$'
plt.title(the_title, fontsize=14)
plt.gcf().set_facecolor('white')

plt.xlabel('$X/L_c$', fontsize=12)
plt.ylabel('$Y/L_c$', fontsize=12)
plt.grid(True, linestyle='--', alpha=0.6)
plt.show()