# analytical.py

import numpy as np
from Bar import Bar


class Analytical():
    def __init__(self, bar:'Bar'):
        self.bar:'Bar' = bar

    def get_disp_amp(self, freq, x_vals:np.ndarray=None):
        '''Gets the analytical displacement amplitude values for a given forcing frequency.
        Provides output for a given x_vals. If x_vals is not provided, returns 1000 points'''
        if x_vals is None:
            x_vals = np.linspace(0, 1, 1000)
        
        alpha = freq*(self.bar.density / self.bar.E)**0.5
        return 100 * np.sin(alpha*x_vals) / np.sin(alpha)

    def get_strain_amp(self, freq, x_vals:np.ndarray=None):
        '''Gets the analytical strain amplitude values for a given forcing frequency.
        Provides output for a given x_vals. If x_vals is not provided, returns 1000 points'''
        if x_vals is None:
            x_vals = np.linspace(0, 1, 1000)
        
        alpha = freq*(self.bar.density / self.bar.E)**0.5
        return 100 * alpha * np.cos(alpha*x_vals) / np.sin(alpha)
    
    def get_force_amp(self, freq, x_vals:np.ndarray=None):
        '''Gets the analytical force amplitude values for a given forcing frequency.
        Provides output for a given x_vals. If x_vals is not provided, returns 1000 points.
        Units are in MN'''
        return self.bar.area * self.bar.E * self.get_strain_amp(freq, x_vals) / (1e9)


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

    x_vals = np.linspace(0, 1, 1000)

    # Specific alpha values
    alpha_values = np.linspace(0, 20, 6)

    freq = freq_from_alpha(bar, 4)
    strain_amplitude = analytical.get_strain_amp(freq, x_vals=np.array([0.31]))
    print(f"Analytical strain at x = {0.31}: {strain_amplitude}")

    # Plotting displacement amplitude
    plt.figure(figsize=(12, 6))
    for alpha in [20]:
        freq = freq_from_alpha(bar, alpha)
        disp_amplitude = analytical.get_disp_amp(freq, x_vals)
        plt.plot(x_vals, disp_amplitude, label=f'α = {alpha:.2f}')

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

    # # Plotting strain amplitude
    # plt.figure(figsize=(12, 6))
    # for alpha in alpha_values:
    #     freq = freq_from_alpha(bar, alpha)
    #     strain_amplitude = analytical.strain_amp(freq, x_vals)
    #     plt.plot(x_vals, strain_amplitude, label=f'α = {alpha:.2f}')

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