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, 100]
    
    # Solve the ODE
    print('AVERAGING SOLUTION')
    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',
        atol = 1e-3
    )
    #### DIRECT INTEGRATION
        # Compute nondimensional parameters
    mu = Cg / Cs
    alpha_1 = np.sqrt(Lm * Cm) / (Cg * Rg)
    alpha_2 = np.sqrt(Lm * Cm) / (Cs * Rs)
    epsilon = np.sqrt(Cm / Cg)
    lambduh = Rm * Cg / np.sqrt(Lm * Cm)
    gamma = beta * abs(Vth) * np.sqrt(Lm * Cm) / Cs
    tspan_2 = [0, 1e6]

    def g(value):
        if value >= 0:
            output = gamma*value**2
        else:
            output = 0
        return output
            

    def colpitts(t, X):
        x_1, x_2, x_3, x_4 = X
        dx1_dt = x_2
        dx2_dt = -alpha_1*epsilon*mu*x_4 - alpha_1*epsilon*x_4 - \
                    alpha_2*epsilon*x_3 - epsilon**2*x_1*mu - epsilon**2*x_1 - \
                    epsilon**2*x_2*lambduh + epsilon*g(x_4-x_3+1) - x_1
        dx3_dt = -alpha_1*mu*x_4 - alpha_2*x_3 - epsilon*x_1*mu + g(x_4-x_3+1)
        dx4_dt = -alpha_1*mu*x_4 - alpha_1*x_4 - alpha_2*x_3 - epsilon*x_1*mu - epsilon*x_1 + g(x_4-x_3+1)
        print(f't: {t*epsilon**2}')
        return [dx1_dt, dx2_dt, dx3_dt, dx4_dt]

        
# Solve the ODE
    print('DIRECT SOLUTION')
    sol2 = solve_ivp(
        fun=colpitts,
        t_span=tspan_2,
        y0=[0,0,0,0],
        method='RK45',
        atol = 1e-13,
        rtol = 1e-6
    )

    # Plot the results
    plt.plot(sol.t / params['ep']**2, np.abs(sol.y[0]), sol2.t, np.abs(sol2.y[0]))
    plt.xlabel('Scaled Time (t / ep^2)')
    plt.ylabel('Amplitude')
    plt.title('Averaging solution compared to direct numerical solution')
    plt.legend(['A0','||f||'])
    plt.grid()
    plt.savefig('test.png')

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