# main.py

import numpy as np
import matplotlib.pyplot as plt

from Bar import Bar
from Analytical import Analytical
from DiscreteMethod import DiscreteMethod
from FDM import FDM
from FEM import FEM
from Convergence import Convergence
from common import freq_from_alpha

bar = Bar()
alpha_values = np.array([0.25, 0.5, 1.0, 2.0, 4.0, 10.0, 20.0])
alpha_sweep_values = np.linspace(0.1, 50, 1000)

def analytical():
    analytical = Analytical(bar)

    x_vals = np.linspace(0, 1, 1000)
    x_vals_table = np.array([0.0, 0.125, 1/(2*np.pi), 0.25, 0.375, 0.5, 0.625, 0.75, 0.875, 1.0])
    displacement_data = []
    strain_data = []

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

    plt.title("Analytical Displacement Amplitude")
    plt.xlabel("x")
    plt.ylabel("Displacement Amplitude")
    plt.legend()
    plt.grid(True)
    plt.savefig("Images/1_Displacement_vs_Position.png")
    plt.close()

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

    plt.title("Analytical Strain Amplitude")
    plt.xlabel("x")
    plt.ylabel("Strain Amplitude")
    plt.legend()
    plt.grid(True)
    plt.savefig("Images/1_Strain_vs_Position.png")
    plt.close()

    # Display Displacement Table
    print("\nAnalytical Results")
    print("Analytical Displacement")
    header = ["x"] + [f"α = {alpha:.2f}" for alpha in alpha_values]
    print("{:<10}".format(header[0]) + "".join(f"{val:<15}" for val in header[1:]))
    print("-" * (10 + 15 * len(alpha_values)))

    for i, x in enumerate(x_vals_table):
        row = [f"{x:.3f}"] + [f"{disp[i]:.4f}" for disp in displacement_data]
        print("{:<10}".format(row[0]) + "".join(f"{val:<15}" for val in row[1:]))

    # Display Strain Table
    print("\nAnalytical Strain")
    print("{:<10}".format(header[0]) + "".join(f"{val:<15}" for val in header[1:]))
    print("-" * (10 + 15 * len(alpha_values)))

    for i, x in enumerate(x_vals_table):
        row = [f"{x:.3f}"] + [f"{strain[i]:.4f}" for strain in strain_data]
        print("{:<10}".format(row[0]) + "".join(f"{val:<15}" for val in row[1:]))

    # Frequency sweep of strain at x = 1
    plt.figure(figsize=(12, 6))
    disp_amplitudes = np.array([])
    for alpha in alpha_sweep_values:
        freq = freq_from_alpha(bar, alpha)
        disp_amplitude = analytical.get_disp_amp(freq, x_vals=np.array([1/(2*np.pi)]))
        disp_amplitudes = np.append(disp_amplitudes, disp_amplitude)

    plt.plot(alpha_sweep_values, disp_amplitudes)
    plt.ylim((-10000, 10000))
    plt.title("Analytical Frequency Sweep for Displacement at x = 1/(2pi)")
    plt.xlabel("α")
    plt.ylabel("Displacement Amplitude")
    plt.legend()
    plt.grid(True)
    #plt.show()
    plt.savefig("Images/1_Strain_Frequency_Sweep.png")
    plt.close()


def convergence():
    def convergence_metric(method: 'DiscreteMethod', forcing_freq: float = None):
        """Extract U(1/(2pi)) from the method."""
        x = 1 / (2 * np.pi)
        if forcing_freq is None:  # Finite method
            return method.get_disp_amp(x)
        else:
            return method.get_disp_amp(forcing_freq, x)  # Analytical

    print("\nConvergence Results")
    print("\n2nd Order FDM and P=1 FEM Convergence of U(x=1/(2pi))")

    # Initialize Bar, FDM, FEM
    bar = Bar()
    analy = Analytical(bar)
    fdm = FDM(bar, desired_order="2")
    fem = FEM(bar, desired_order="2")

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

    # Initialize Convergence
    convergence = Convergence(analy, fdm, fem)

    num_sections_range = [2**n for n in range(1, 12)]
    alpha_values = [0.25, 0.5, 1.0, 2.0, 4.0]

    # Create a figure with 2 columns of subplots
    fig, axs = plt.subplots((len(alpha_values) + 1) // 2, 2, figsize=(14, 12), sharex=True, sharey=True)
    axs = axs.flatten()

    for i, alpha in enumerate(alpha_values):
        forcing_freq = freq_from_alpha(bar, alpha)

        # Run convergence
        results = convergence.run_convergence(forcing_freq, num_sections_range, convergence_metric)

        # Plot results in the subplot
        ax = axs[i]
        ax.loglog(results['fdm']['num_sections'], results['fdm']['errors'], '-o', label='FDM Analytical Error')
        ax.loglog(results['fem']['num_sections'], results['fem']['errors'], '--s', label='FEM Analytical Error')
        ax.loglog(results['fdm']['num_sections'], results['fdm']['extrapolated_errors'], '-x', label='FDM Extrapolated Error')
        ax.loglog(results['fem']['num_sections'], results['fem']['extrapolated_errors'], '--d', label='FEM Extrapolated Error')
        ax.set_title(f"Convergence for α = {alpha}")
        ax.set_ylabel("Error (Relative)")
        ax.grid(True, which="both", linestyle="--")
        if i >= len(alpha_values) - 2:
            ax.set_xlabel("Number of Sections")
        ax.legend()

        # Output table for this alpha
        print(f"\nConvergence Table for α = {alpha}")
        print("{:<15} {:<20} {:<20} {:<20} {:<15} {:<15} {:<15} {:<15}".format(
            "dx", "Analytical Value", "Method Value", "Extrap Value",
            "% Error", "β", "% Extrap Error", "β Extrap"
        ))
        print("-" * 135)
        for j in range(len(results['fdm']['num_sections'])):
            num_sections = results['fdm']['num_sections'][j]
            dx = 1 / num_sections
            analytical_value = results['fdm']['analytical_values'][j]
            fdm_value = results['fdm']['extrapolated_values'][j]
            fdm_error = results['fdm']['errors'][j]
            fdm_beta = results['fdm']['betas'][j] if j >= 1 else float('nan')  # Beta starts at index 1
            fdm_extrap_err = results['fdm']['extrapolated_errors'][j]
            fdm_extrap_beta = results['fdm']['extrapolated_betas'][j] if j >= 2 else float('nan')  # Extrap Beta starts at index 2

            print("{:<15} {:<15.6f} {:<15.6f} {:<15.6f} {:<15.6f} {:<15.6f} {:<15.6f} {:<15.6f}".format(
                dx, analytical_value, fdm_value, fdm_value, fdm_error, fdm_beta, fdm_extrap_err, fdm_extrap_beta
            ))

    # Hide any unused subplots
    for j in range(len(alpha_values), len(axs)):
        fig.delaxes(axs[j])

    # Save and show the plot
    plt.tight_layout()
    plt.savefig("Images/2_Second_Order_Convergence.png")
    plt.show()



    print("\n4th Order FDM and P=2 FEM Convergence of U(x=1/(2pi))")
    # Initialize Bar, FDM, FEM for 4th order convergence
    bar = Bar()
    analy = Analytical(bar)
    fdm = FDM(bar, desired_order="4")
    fem = FEM(bar, desired_order="4")

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

    # Initialize Convergence
    convergence = Convergence(analy, fdm, fem)

    num_sections_range = [2**n for n in range(1, 12)]
    alpha_values = [0.25, 0.5, 1.0, 2.0, 4.0]

    # Create a figure with 2 columns of subplots
    fig, axs = plt.subplots((len(alpha_values) + 1) // 2, 2, figsize=(14, 12), sharex=True, sharey=True)
    axs = axs.flatten()

    for i, alpha in enumerate(alpha_values):
        forcing_freq = freq_from_alpha(bar, alpha)

        # Run convergence
        results = convergence.run_convergence(forcing_freq, num_sections_range, convergence_metric)

        # Plot results in the subplot
        ax = axs[i]
        ax.loglog(results['fdm']['num_sections'], results['fdm']['errors'], '-o', label='FDM Analytical Error')
        ax.loglog(results['fem']['num_sections'], results['fem']['errors'], '--s', label='FEM Analytical Error')
        ax.loglog(results['fdm']['num_sections'], results['fdm']['extrapolated_errors'], '-x', label='FDM Extrapolated Error')
        ax.loglog(results['fem']['num_sections'], results['fem']['extrapolated_errors'], '--d', label='FEM Extrapolated Error')
        ax.set_title(f"Convergence for α = {alpha}")
        ax.set_ylabel("Error (Relative)")
        ax.grid(True, which="both", linestyle="--")
        if i >= len(alpha_values) - 2:
            ax.set_xlabel("Number of Sections")
        ax.legend()

        # Output table for this alpha
        print(f"\nConvergence Table for α = {alpha}")
        print("{:<15} {:<20} {:<20} {:<20} {:<15} {:<15} {:<15} {:<15}".format(
            "dx", "Analytical Value", "Method Value", "Extrap Value",
            "% Error", "β", "% Extrap Error", "β Extrap"
        ))
        print("-" * 135)
        for j in range(len(results['fdm']['num_sections'])):
            num_sections = results['fdm']['num_sections'][j]
            dx = 1 / num_sections
            analytical_value = results['fdm']['analytical_values'][j]
            fdm_value = results['fdm']['extrapolated_values'][j]
            fdm_error = results['fdm']['errors'][j]
            fdm_beta = results['fdm']['betas'][j] if j >= 1 else float('nan')  # Beta starts at index 1
            fdm_extrap_err = results['fdm']['extrapolated_errors'][j]
            fdm_extrap_beta = results['fdm']['extrapolated_betas'][j] if j >= 2 else float('nan')  # Extrap Beta starts at index 2

            print("{:<15} {:<15.6f} {:<15.6f} {:<15.6f} {:<15.6f} {:<15.6f} {:<15.6f} {:<15.6f}".format(
                dx, analytical_value, fdm_value, fdm_value, fdm_error, fdm_beta, fdm_extrap_err, fdm_extrap_beta
            ))

    # Hide any unused subplots
    for j in range(len(alpha_values), len(axs)):
        fig.delaxes(axs[j])

    # Save and show the plot
    plt.tight_layout()
    plt.savefig("Images/2_Fourth_Order_Convergence.png")
    plt.show()


import numpy as np
import matplotlib.pyplot as plt

def frequency_sweep():
    # Alpha sweep values
    alpha_sweep_values = np.linspace(0.1, 50, 1000)

    # Section sweep values
    num_sections_range = [2**n for n in range(1, 10)]

    # Initialize Bar, Analytical, FDM, FEM
    bar = Bar()
    analytical = Analytical(bar)

    # Plot 1: Analytical, FDM 2nd Order, FEM P=1
    fdm_2nd = FDM(bar, desired_order="2")
    fem_p1 = FEM(bar, desired_order="2")

    fdm_2nd.set_boundary_conditions(U_0=0, U_L=100)
    fem_p1.set_boundary_conditions(U_0=0, U_L=100)

    fig, axs = plt.subplots(len(num_sections_range), 1, figsize=(16, 20))  # 10 rows, 1 column
    axs = axs.flatten()

    for idx, num_sections in enumerate(num_sections_range):
        dx = 1 / num_sections
        disp_amplitudes_analytical = []
        disp_amplitudes_fdm_2nd = []
        disp_amplitudes_fem_p1 = []

        for alpha in alpha_sweep_values:
            freq = freq_from_alpha(bar, alpha)

            # Analytical
            disp_amplitudes_analytical.append(analytical.get_disp_amp(freq, x_vals=np.array([1 / (2 * np.pi)])))

            # FDM 2nd Order
            fdm_2nd.run(freq, num_sections)
            disp_amplitudes_fdm_2nd.append(fdm_2nd.get_disp_amp(x=np.array([1 / (2 * np.pi)])))

            # FEM P=1
            fem_p1.run(freq, num_sections)
            disp_amplitudes_fem_p1.append(fem_p1.get_disp_amp(x=np.array([1 / (2 * np.pi)])))

        # Subplot for this number of sections
        ax = axs[idx]
        ax.plot(alpha_sweep_values, disp_amplitudes_analytical, label="Analytical", linestyle="-")
        ax.plot(alpha_sweep_values, disp_amplitudes_fdm_2nd, label="FDM 2nd Order", linestyle="--")
        ax.plot(alpha_sweep_values, disp_amplitudes_fem_p1, label="FEM P=1", linestyle=":")
        ax.set_ylim((-10000, 10000))  # Set uniform y-limits for all subplots
        ax.set_title(f"dx = {dx:.5f} (Num Sections = {num_sections})")
        ax.set_xlabel("α")
        ax.set_ylabel("Displacement Amplitude")
        ax.legend()
        ax.grid(True)

    plt.tight_layout()
    plt.suptitle("Frequency Sweep (Displacement Amplitude, FDM 2nd Order, FEM P=1)", y=1.02)
    plt.savefig("Images/3_Frequency_Sweep_FDM2_FEM1.png")
    plt.close()

    # Plot 2: Analytical, FDM 4th Order, FEM P=2
    fdm_4th = FDM(bar, desired_order="4")
    fem_p2 = FEM(bar, desired_order="4")

    fdm_4th.set_boundary_conditions(U_0=0, U_L=100)
    fem_p2.set_boundary_conditions(U_0=0, U_L=100)

    fig, axs = plt.subplots(len(num_sections_range), 1, figsize=(16, 20))  # 10 rows, 1 column
    axs = axs.flatten()

    for idx, num_sections in enumerate(num_sections_range):
        dx = 1 / num_sections
        disp_amplitudes_analytical = []
        disp_amplitudes_fdm_4th = []
        disp_amplitudes_fem_p2 = []

        for alpha in alpha_sweep_values:
            freq = freq_from_alpha(bar, alpha)

            # Analytical
            disp_amplitudes_analytical.append(analytical.get_disp_amp(freq, x_vals=np.array([1 / (2 * np.pi)])))

            # FDM 4th Order
            fdm_4th.run(freq, num_sections)
            disp_amplitudes_fdm_4th.append(fdm_4th.get_disp_amp(x=np.array([1 / (2 * np.pi)])))

            # FEM P=2
            fem_p2.run(freq, num_sections)
            disp_amplitudes_fem_p2.append(fem_p2.get_disp_amp(x=np.array([1 / (2 * np.pi)])))

        # Subplot for this number of sections
        ax = axs[idx]
        ax.plot(alpha_sweep_values, disp_amplitudes_analytical, label="Analytical", linestyle="-")
        ax.plot(alpha_sweep_values, disp_amplitudes_fdm_4th, label="FDM 4th Order", linestyle="--")
        ax.plot(alpha_sweep_values, disp_amplitudes_fem_p2, label="FEM P=2", linestyle=":")
        ax.set_ylim((-10000, 10000))  # Set uniform y-limits for all subplots
        ax.set_title(f"dx = {dx:.5f} (Num Sections = {num_sections})")
        ax.set_xlabel("α")
        ax.set_ylabel("Displacement Amplitude")
        ax.legend()
        ax.grid(True)

    plt.tight_layout()
    plt.suptitle("Frequency Sweep (Displacement Amplitude, FDM 4th Order, FEM P=2)", y=1.02)
    plt.savefig("Images/3_Frequency_Sweep_FDM4_FEM2.png")
    plt.close()





if __name__ == "__main__":
    import os
    os.system('cls' if os.name == 'nt' else 'clear')
    analytical()

    convergence()

    frequency_sweep()

# if __name__ == "__main__":
#     import os
#     import sys
#     os.system('cls' if os.name == 'nt' else 'clear')
#     # Redirect stdout and stderr to a file
#     with open("output.txt", "w", encoding="utf-8") as output_file:
#         # Redirect stdout
#         sys.stdout = output_file
#         # Redirect stderr
#         sys.stderr = output_file

#         # Run the main functions
#         try:
#             analytical()
#             convergence()
#             # frequency_sweep()
#         except Exception as e:
#             print(f"An error occurred: {e}")

#     # Restore stdout and stderr
#     sys.stdout = sys.__stdout__
#     sys.stderr = sys.__stderr__

#     print("Execution completed. Output saved to output.txt.")