# FDM.py

import numpy as np

from Bar import Bar
from DiscreteMethod import DiscreteMethod

class FDM(DiscreteMethod):
    def __init__(self, bar:'Bar', desired_order="2"):
        super().__init__(bar, desired_order)

    def get_kappa(self):
        '''Returns the 2nd or 4th order kappa value'''
        dx = self.dx
        alpha = self.alpha

        kappa = 2 - alpha**2 * dx**2

        if self.desired_order == "4": # Add the 4th order term
            kappa += alpha**4 * dx**4 / 12
        
        return kappa

    def build_stiffness_matrix(self):
        '''Given a number of sections, build the stiffness_matrix'''
        # The stiffness matrix follows the form [1, -kappa, 1] in a diagonal
        kappa = self.get_kappa()

        size = self.num_sections + 1  # Include boundary points
        self.stiffness_matrix = np.zeros((size, size))

        for i in range(1, size - 1):  # Internal points
            self.stiffness_matrix[i, i - 1] = 1  # Left neighbor
            self.stiffness_matrix[i, i] = -kappa  # Center
            self.stiffness_matrix[i, i + 1] = 1  # Right neighbor

        # Enforce boundary conditions in the stiffness matrix
        self.stiffness_matrix[0, 0] = 1  # Left boundary
        self.stiffness_matrix[-1, -1] = 1  # Right boundary
    
    def build_constant_vector(self):
        '''Given a number of sections, build the constant_vector where the 
        first and last elements come from boundary conditions'''
        self.constant_vector = np.zeros(self.num_sections + 1)  # Include endpoints
        self.constant_vector[0] = self.U_0  # Left boundary
        self.constant_vector[-1] = self.U_L  # Right boundary

    def run(self, forcing_freq:float, num_sections:int=4):
        '''Runs FDM for the given forcing frequency and num_sections.
        Stores displacment ampltitude data in self.disp_amp
        Stores strain amplitude data in self.strain_amp'''
        # Save the appropriate values
        self.num_sections = num_sections
        self.dx = 1 / num_sections
        self.alpha = forcing_freq*(self.bar.density / self.bar.E)**0.5
        self.forcing_freq = forcing_freq

        self.build_stiffness_matrix()
        self.build_constant_vector()

        # Solve the system
        self.disp_amp = np.linalg.solve(self.stiffness_matrix, self.constant_vector)

        # Calculate all of the nodal strains
        self.calc_all_nodal_strain()

        # Save the x values for the nodal values
        self.x_vals = np.linspace(0, 1, num_sections+1)  # sections + endpoints
        

if __name__ == "__main__":
    import matplotlib.pyplot as plt
    from common import freq_from_alpha
    bar = Bar()

    # Initialize FDM
    fdm_2 = FDM(bar, desired_order="2")

    # Define alpha values for testing
    alpha_values = np.linspace(0, 20, 6)

    # Set boundary conditions
    fdm_2.set_boundary_conditions(U_0=0, U_L=100)

    # Plot displacement amplitudes for varying alpha values
    plt.figure(figsize=(12, 6))
    for alpha in [100]:
        forcing_freq = freq_from_alpha(bar, alpha)
        fdm_2.run(forcing_freq, num_sections=4)
        plt.plot(fdm_2.x_vals, fdm_2.disp_amp, label=f'α = {alpha:.2f}')

    plt.title("FDM Displacement Amplitude vs x for Various α Values")
    plt.xlabel("x (Normalized Position)")
    plt.ylabel("Displacement Amplitude")
    plt.legend()
    plt.grid(True)
    plt.show()

    # Set boundary conditions and run the solver
    fdm_4 = FDM(bar, desired_order="4")
    fdm_4.set_boundary_conditions(U_0=0, U_L=100)
    freq = freq_from_alpha(bar, 4)  
    num_sections = 10
    fdm_4.run(freq, num_sections)
    fdm_2.run(freq, num_sections)

    # Interpolate at a specific x_val
    x_val = 0.31
    interpolated_strain = fdm_4.get_interpolated_strain_amp(x_val)
    print(f"Interpolated Strain Amplitude at x = {x_val}: {interpolated_strain}")

    interpolated_strain = fdm_2.get_interpolated_strain_amp(x_val)
    print(f"Interpolated Strain Amplitude at x = {x_val}: {interpolated_strain}")