import numpy as np
from scipy.optimize import fsolve
from scipy.integrate import quad
import matplotlib.pyplot as plt

# Persistent storage for fsolve initial conditions
nsolve_ics_storage = {'nsolve_ics': None}

# Helper functions for g1, g1timescos, and g1timessin
def g1(phi0, x, params):
    """
    Purpose of the function:
    Computes the transistor's nondimensionalized function `g1`, which represents 
    the piecewise relationship between input and output voltages.

    Inputs:
    1. phi0 (float or array-like): Phase angle(s) at which the function is evaluated.
    2. x (array-like): Fourier coefficients [U, G, H, X, Y].
    3. params (dict): Dictionary containing nondimensional parameters:
        - ep, gamma, u0.

    Outputs:
    1. g1out (array-like): The value of the `g1` function at the input phase angles.
    
    Notes:
    - Implements piecewise behavior for the `g1` function based on the conditional 
      `ep * v1mu1 >= u0 - 1`.
    - This is based on equations 28 and 34 in the referenced model.
    """

    U, V, G, H, X, Y = x[0], 0, x[1], x[2], x[3], x[4]
    ep = params['ep']
    gamma = params['gamma']
    u0 = params['u0']
    
    # Equation 34
    v1mu1 = (V - U) + (G - X) * np.cos(phi0) + (H - Y) * np.sin(phi0)
    
    # Equation 28
    cond = ep * v1mu1 >= u0 - 1
    g1out = np.where(
        cond,
        gamma * v1mu1 * (2 * (1 - u0) + ep * v1mu1),
        -gamma / ep * (1 - u0)**2
    )
    return g1out

def g1timescos(phi0, x, params):
    """Compute g1(phi0) * cos(phi0)."""
    return g1(phi0, x, params) * np.cos(phi0)

def g1timessin(phi0, x, params):
    """Compute g1(phi0) * sin(phi0)."""
    return g1(phi0, x, params) * np.sin(phi0)

# Harmonic balance system (hbsystem_case1)
def hbsystem_case1(x, params):
    """
    Purpose of the function:
    Computes the residuals for the harmonic balance system (Eq. 33a-f). 
    These residuals are minimized by `fsolve` to find the Fourier coefficients.

    Inputs:
    1. x (array-like): A vector of Fourier coefficients [U, G, H, X, Y].
    2. params (dict): Dictionary containing nondimensional parameters:
        - mu, a1, a2, ep, gamma, a0 (current state), and u0 (steady-state solution).

    Outputs:
    1. F (array-like): A vector of residuals corresponding to equations 33a-f.

    Notes:
    - The residuals involve integrals over `phi0` (0 to 2π), which are computed using 
      numerical quadrature (`quad`).
    - The function incorporates `g1`, `g1timescos`, and `g1timessin`.
    """ 
    mu = params['mu']
    a1 = params['a1']
    a2 = params['a2']
    ep = params['ep']
    gamma = params['gamma']
    a0 = params['a0cur']
    u0 = params['u0']
    
    # Fourier coefficients
    U, V, G, H, X, Y = x[0], 0, x[1], x[2], x[3], x[4]
    
    # Equations 33a-f
    F = np.zeros(5)
    F[0] = -a2 * U + (1 / (2 * np.pi)) * quad(lambda phi: g1(phi, x, params), 0, 2 * np.pi)[0]
    F[1] = H - Y + a1 * G
    F[2] = -mu * H + (mu + 1) * Y + a2 * X - (1 / np.pi) * quad(lambda phi: g1timescos(phi, x, params), 0, 2 * np.pi)[0]
    F[3] = -G + X + a1 * H + a0
    F[4] = mu * G - (mu + 1) * X + a2 * Y - (1 / np.pi) * quad(lambda phi: g1timessin(phi, x, params), 0, 2 * np.pi)[0]
    #print(f'F: {F}')
    return F

# ODE function for A0
def odefun_A0(t, A0, params):
    """
    Purpose of the function:
    Defines the ordinary differential equation (ODE) for `A0` as part of the 
    larger Colpitts oscillator system. Computes the time derivative of `A0`.

    Inputs:
    1. t (float): Current time (unused in this specific ODE but required by solver).
    2. A0 (float): Current value of `A0`.
    3. params (dict): Dictionary containing nondimensional parameters:
        - lambda, a0cur (current A0 state), and others used in `hbsystem_case1`.

    Outputs:
    1. dA0 (float): The time derivative of `A0`.

    Notes:
    - Uses `fsolve` to solve the harmonic balance equations (via `hbsystem_case1`).
    - Returns `dA0` using the computed `H` value (harmonic balance coefficient).
    """
    global nsolve_ics_storage
    
    # Initialize fsolve initial conditions if not set
    if nsolve_ics_storage['nsolve_ics'] is None:
        nsolve_ics_storage['nsolve_ics'] = np.array([100, 100, 100, 100, 100])
    
    # Update current A0 in parameters
    params['a0cur'] = A0
    lambda_ = params['lambda']
    
    # Solve harmonic balance system
    nsolve_ics = nsolve_ics_storage['nsolve_ics']
    x = fsolve(lambda x: hbsystem_case1(x, params), nsolve_ics, xtol=1e-4)
    
    # Update fsolve initial conditions for next iteration
    nsolve_ics_storage['nsolve_ics'] = x
    
    # Extract H from the solution
    H = x[2]
    
    # Compute dA0/d_eta_2
    dA0 = (-1 / 2) * (lambda_ * A0 - H)
    #print(f'dA0 {dA0}')
    print(f't {t}')
    return dA0

# Main function to solve the system and plot results
def main():
    # Dimensional parameters
    Cm = 2.201e-15
    Cs = 1.62e-11
    Cg = 2.32e-11
    Lm = 4.496e-2
    Rm = 5.62e1
    Rs = 7.444e3
    Rg = 5.168e4
    Vth = -0.95
    beta = 0.12
    B = 0.8
    
    # Compute nondimensional parameters
    mu = Cg / Cs
    a1 = np.sqrt(Lm * Cm) / (Cg * Rg)
    a2 = np.sqrt(Lm * Cm) / (Cs * Rs)
    ep = np.sqrt(Cm / Cg)
    lmbda = Rm * Cg / np.sqrt(Lm * Cm)
    gamma = beta * abs(Vth) * np.sqrt(Lm * Cm) / Cs
    u0 = (1 + a2 / (2 * gamma)) - np.sqrt((1 + a2 / (2 * gamma))**2 - 1)
    
    # Parameters
    params = {
        'mu': mu,
        'a1': a1,
        'a2': a2,
        'ep': ep,
        'lambda': lmbda,
        'gamma': gamma,
        'u0': u0,
        'a0cur': 0, 
        'v0': 0,
    }
    
    # Initial conditions and time span
    a0ic = 0
    tspan = [0, 200]
    
    # Solve the ODE
    from scipy.integrate import solve_ivp
    sol = solve_ivp(
        fun=lambda t, A0: odefun_A0(t, A0, params),
        t_span=tspan,
        y0=[a0ic],
        method='RK45',
        max_step=1e-0
    )
    
    # Plot the results
    plt.plot(sol.t / params['ep']**2, np.abs(sol.y[0]))
    plt.xlabel('Scaled Time (t / ep^2)')
    plt.ylabel('A0 (Amplitude)')
    plt.title('Dynamics of A0 Over Time')
    plt.grid()
    plt.savefig('test.png')

# Run the main function
if __name__ == "__main__":
    main()
