# DiscreteMethod.py

import numpy as np

from Bar import Bar

class DiscreteMethod():
    '''Provides the class for solving via FDM or FEM.
    The bar is provided at initiation, as well as the desired order.
    The user can change the forcing frequency and num_sections as desired'''
    def __init__(self, bar: 'Bar', desired_order:str="2"):
        self.bar: 'Bar' = bar
        self.desired_order:str = desired_order

        # Finite method solving
        self.num_sections: int = None
        self.dx: float = None
        self.alpha: float = None
        self.stiffness_matrix: np.ndarray = None
        self.constant_vector: np.ndarray = None

        # Result vectors
        self.x_vals: np.ndarray = None
        self.forcing_freq: float = None
        self.disp_amp: np.ndarray = None
        self.strain_amp: np.ndarray = None 

    def set_boundary_conditions(self, U_0:float=0, U_L:float=100):
        '''Set the left and right amplitude boundary conditions'''
        self.U_0 = U_0
        self.U_L = U_L 

    def nodal_strain(self, node_idx: int):
        '''Calculates the strain (U') at a given node using the 4th-order scheme.'''
        dx = self.dx
        alpha = self.alpha
        if node_idx == 0:  # Left boundary
            # Use forward difference at the boundary
            return (
                (1 - dx**2 * alpha**2 / 2 + dx**4 * alpha**4 / 24) / (-dx + (dx**3*alpha**2/6)) * self.disp_amp[node_idx]
                - (1 / (-dx + dx**3 * alpha**2 / 6)) * self.disp_amp[node_idx + 1]
            )
        elif node_idx == len(self.disp_amp) - 1:  # Right boundary, use backward difference
            return (
                (1 - dx**2 * alpha**2 / 2 + dx**4 * alpha**4 / 24) / (dx - (dx**3*alpha**2/6)) * self.disp_amp[node_idx]
                - (1 / (dx - (dx**3*alpha**2/6))) * self.disp_amp[node_idx - 1]
            )
        else:   # Return the average of the backward and forward derivative
            return (
                ((1 - dx**2 * alpha**2 / 2 + dx**4 * alpha**4 / 24) / (dx - (dx**3*alpha**2/6)) * self.disp_amp[node_idx]
                - (1 / (dx - (dx**3*alpha**2/6))) * self.disp_amp[node_idx - 1]
                +
                (1 - dx**2 * alpha**2 / 2 + dx**4 * alpha**4 / 24) / (-dx + (dx**3*alpha**2/6)) * self.disp_amp[node_idx]
                - (1 / (-dx + dx**3 * alpha**2 / 6)) * self.disp_amp[node_idx + 1])
                / 2.0
            )
    
    def calc_all_nodal_strain(self):
        '''Calculates all of the nodal strains after running the method'''
        nodal_strains = np.array([])
        for i in range(self.num_sections+1):
            nodal_strains = np.append(nodal_strains, self.nodal_strain(i))
        self.strain_amp = nodal_strains

    def get_interpolated_disp_amp(self, x_val: float):
        """Returns the interpolated displacement amplitude for a given x_val."""
        # Find the closest node to x_val
        idx = np.argmin(np.abs(self.x_vals - x_val))  # Closest node index
        alpha = self.alpha

        # Calculate h (distance from the nearest node)
        h = x_val - self.x_vals[idx]

        U_i = self.disp_amp[idx]

        if h == 0:
            return U_i

        # Calculate U' at the closest node
        U_prime = self.strain_amp[idx]

        # Apply the fourth-order Taylor series
        
        interpolated_value = (
            (1 - h**2 * alpha**2 / 2 + h**4 * alpha**4 / 24) * U_i
            + (h - h**3 * alpha**2 / 6) * U_prime
        )

        return interpolated_value
    
    def get_interpolated_strain_amp(self, x_val: float):
        '''Returns the interpolated strain (U') amplitude for a given x_val.'''
        # Find the closest node to x_val
        idx = np.argmin(np.abs(self.x_vals - x_val))  # Closest node index
        dx = 1 / (len(self.x_vals) - 1)  # Spacing between nodes
        alpha = self.forcing_freq * (self.bar.density / self.bar.E)**0.5

        # Calculate h (distance from the nearest node)
        h = x_val - self.x_vals[idx]

        # Get U at the closest node
        U_at_node = self.disp_amp[idx]

        # Calculate U' at the closest node
        U_prime_at_node = self.strain_amp[idx]

        if h == 0:
            return U_prime_at_node

        # Apply the fourth-order Taylor series
        # Taylor series expansion for strain
        U_prime_interpolated = (
            (-h*alpha**2 + h**3 * alpha**4 / 6) * U_at_node
            + (1 - h**2 * alpha**2 / 2 + h**4 * alpha**4 / 24) * U_prime_at_node
        )

        return U_prime_interpolated
    
    def build_stiffness_matrix(self):
        pass

    def build_constant_vector(self):
        pass

    def run(self, forcing_freq:float, num_sections:int=4):
        '''Runs the method 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

    def get_disp_amp(self, x):
        '''Returns the displacement amplitude at x'''
        return self.get_interpolated_disp_amp(x)
    
    def get_strain_amp(self, x):
        '''Returns the strain amplitude at x'''
        return self.get_interpolated_strain_amp(x)
    
    def get_force_amp(self, x):
        '''Returns the force amplitude at x.
        Units are in MN'''
        return self.bar.area * self.bar.E* self.get_interpolated_strain_amp(x) / (1e9)