"""
Top-level script for calling simulate_maglev_control
Ported from topSimulate.m to Python
"""

import numpy as np
import matplotlib.pyplot as plt
from parameters import QuadParams, Constants
from utils import euler2dcm, fmag2
from simulate import simulate_maglev_control
from visualize import visualize_quad
import os


def main():
    """Main simulation script"""
    
    # Create output directory if it doesn't exist
    output_dir = 'sim_results'
    os.makedirs(output_dir, exist_ok=True)
    
    # Total simulation time, in seconds
    Tsim = 2
    
    # Update interval, in seconds. This value should be small relative to the
    # shortest time constant of your system.
    delt = 0.005  # sampling interval
    
    # Maglev parameters and constants
    quad_params = QuadParams()
    constants = Constants()
    
    m = quad_params.m
    g = constants.g
    J = quad_params.Jq
    
    # Time vector, in seconds
    N = int(np.floor(Tsim / delt))
    tVec = np.arange(N) * delt
    
    # Matrix of disturbance forces acting on the body, in Newtons, expressed in I
    distMat = np.random.normal(0, 0, (N-1, 3))
    
    # Oversampling factor
    oversampFact = 10
    
    # Check nominal gap
    print(f"Force check: {4*fmag2(0, 10.830e-3) - m*g}")
    
    # SET REFERENCE HERE
    ref_gap = 10.830e-3  # from python code
    z0 = ref_gap - 2e-3
    
    # Create reference trajectories
    rIstar = np.zeros((N, 3))
    vIstar = np.zeros((N, 3))
    aIstar = np.zeros((N, 3))
    xIstar = np.zeros((N, 3))
    
    for k in range(N):
        rIstar[k, :] = [0, 0, -ref_gap]
        vIstar[k, :] = [0, 0, 0]
        aIstar[k, :] = [0, 0, 0]
        xIstar[k, :] = [0, 1, 0]
    
    # Setup reference structure
    R = {
        'tVec': tVec,
        'rIstar': rIstar,
        'vIstar': vIstar,
        'aIstar': aIstar,
        'xIstar': xIstar
    }
    
    # Initial state
    state0 = {
        'r': np.array([0, 0, -(z0 + quad_params.yh)]),
        'v': np.array([0, 0, 0]),
        'e': np.array([0.01, 0.01, np.pi/2]),  # xyz euler angles
        'omegaB': np.array([0.00, 0.00, 0])
    }
    
    # Setup simulation structure
    S = {
        'tVec': tVec,
        'distMat': distMat,
        'oversampFact': oversampFact,
        'state0': state0
    }
    
    # Setup parameters structure
    P = {
        'quadParams': quad_params,
        'constants': constants
    }
    
    # Run simulation
    print("Running simulation...")
    P0 = simulate_maglev_control(R, S, P)
    print("Simulation complete!")
    
    # Extract results
    tVec_out = P0['tVec']
    state = P0['state']
    rMat = state['rMat']
    eMat = state['eMat']
    vMat = state['vMat']
    gaps = state['gaps']
    currents = state['currents']
    omegaBMat = state['omegaBMat']
    
    # Calculate forces
    Fm = fmag2(currents[:, 0], gaps[:, 0])
    
    # Calculate ex and ey for visualization
    N_out = len(eMat)
    ex = np.zeros(N_out)
    ey = np.zeros(N_out)
    
    for k in range(N_out):
        Rk = euler2dcm(eMat[k, :])
        x = Rk @ np.array([1, 0, 0])
        ex[k] = x[0]
        ey[k] = x[1]
    
    # Visualize the quad motion
    print("Generating 3D visualization...")
    S2 = {
        'tVec': tVec_out,
        'rMat': rMat,
        'eMat': eMat,
        'plotFrequency': 20,
        'makeGifFlag': True,
        'gifFileName': 'sim_results/testGif.gif',
        'bounds': [-1, 1, -1, 1, -300e-3, 0.000]
    }
    visualize_quad(S2)
    
    # Create plots
    fig = plt.figure(figsize=(12, 8))
    
    # Plot 1: Gaps
    ax1 = plt.subplot(3, 1, 1)
    plt.plot(tVec_out, gaps * 1e3)
    plt.axhline(y=ref_gap * 1e3, color='k', linestyle='-', linewidth=1)
    plt.ylabel('Vertical (mm)')
    plt.title('Gaps')
    plt.grid(True)
    plt.xticks([])
    
    # Plot 2: Currents
    ax2 = plt.subplot(3, 1, 2)
    plt.plot(tVec_out, currents)
    plt.ylabel('Current (Amps)')
    plt.title('Power')
    plt.grid(True)
    plt.xticks([])
    
    # Plot 3: Forces
    ax3 = plt.subplot(3, 1, 3)
    plt.plot(tVec_out, Fm)
    plt.xlabel('Time (sec)')
    plt.ylabel('Fm (N)')
    plt.title('Forcing')
    plt.grid(True)
    
    plt.tight_layout()
    plt.savefig('sim_results/simulation_results.png', dpi=150)
    print("Saved plot to simulation_results.png")
    
    # Commented out horizontal position plot (as in MATLAB)
    # fig_horizontal = plt.figure()
    # plt.plot(rMat[:, 0], rMat[:, 1])
    # spacing = 300
    # plt.quiver(rMat[::spacing, 0], rMat[::spacing, 1], 
    #           ex[::spacing], -ey[::spacing])
    # plt.axis('equal')
    # plt.grid(True)
    # plt.xlabel('X (m)')
    # plt.ylabel('Y (m)')
    # plt.title('Horizontal position of CM')
    
    # Frequency analysis (mech freq)
    Fs = 1/delt * oversampFact  # Sampling frequency
    T = 1/Fs  # Sampling period
    L = len(tVec_out)  # Length of signal
    
    Y = np.fft.fft(Fm)
    frequencies = Fs / L * np.arange(L)
    
    fig2 = plt.figure(figsize=(10, 6))
    plt.semilogx(frequencies, np.abs(Y), linewidth=3)
    plt.title("Complex Magnitude of fft Spectrum")
    plt.xlabel("f (Hz)")
    plt.ylabel("|fft(X)|")
    # Exclude DC component (first element) from scaling to see AC components better
    plt.ylim([0, np.max(np.abs(Y[1:])) * 1.05])  # Dynamic y-axis excluding DC
    plt.grid(True)
    plt.savefig('sim_results/fft_spectrum.png', dpi=150)
    print("Saved FFT plot to fft_spectrum.png")
    
    # Show all plots
    plt.show()
    
    return P0


if __name__ == '__main__':
    results = main()
    print("\nSimulation completed successfully!")
    print(f"Final time: {results['tVec'][-1]:.3f} seconds")
    print(f"Number of time points: {len(results['tVec'])}")