# Autogenerated with SMOP 0.29
import numpy as np
from typing import Union
from geomaglib import legendre
#

def design_SHA_Sq_e_Re_v2(rho: Union[float, np.ndarray], theta: Union[float, np.ndarray], phi: Union[float, np.ndarray], nmax: int, mmax:int) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    # [A_r, A_theta, A_phi] = design_SHA_Sq_e_Re_v2(rho,theta,phi,t,t_ut,...
    #    nmax,mmax,p_vec,s_vec);

    # Calculate design matrices A_i that connects the vector of REAL,
    # Schmidt-normalized, spherical harmonic expansion coefficients,
    #        x = q_{nsp}^{m} and s_{nsp}^{m}
    # and the magnetic component B_r, B_theta, B_phi:
    #        B_r     = x'*A_r
    #        B_theta = x'*A_theta
    #        B_phi   = x'*A_phi
    # for the region below (primary) Sq currents.

    # Inputs:   rho(:)                  radius [units of reference radius]
    #           theta(:), phi(:)        co-latitude, longitude [deg]
    #           nmax, mmax              maximum degree and order

    # (Optimized version)

    # A. Chulliat, 2016-09-22
    # (from an earlier version dated 2011-04-21, with inputs from N. Olsen)
    rad = np.pi / 180
    N_data = np.size(theta, 0)
    if (np.size(phi, 0) != N_data):
        raise ValueError(
            "design_SHA_Sq_e: theta and phi have different dimensions"
        )
    if np.isscalar(theta) is not True and theta.ndim > 1:
        dim2 = np.size(theta, 1)
    else:
        dim2 = 1
    if dim2 > 1:
        theta = theta.transpose()
        phi = phi.transpose()

    cos_theta = np.cos(theta * rad)
    sin_theta = np.sin(theta * rad)
    # convert to row vector if input parameter is scalar
    if np.isscalar(rho) is True:
        rho = rho * np.ones(N_data)

    # number of parameters
    N_nm = mmax * (mmax + 2) + (nmax - mmax) * (2 * mmax + 1)
    # calculate sub-matrices for s=0, p=0
    A_r_0 = np.zeros((N_nm, N_data))
    A_theta_0 = np.zeros((N_nm, N_data))
    A_phi_0 = np.zeros((N_nm, N_data))
    # Cycles through all the coefficients for degree n, order m model.
    k = 0


    Leg = legendre.Flattened_Chaos_Legendre1(nmax, theta)
    Pnm, dPnm = Leg[0], Leg[1]

    # Cycle through all of n and m
    for n in range(1, nmax+1):
        rn1 = rho**(n-1)
        # The maximum m is less than the maximum n in the default DIFI
        Mmax = min(n, mmax)
        for m in range(Mmax+1):
            index = legendre.get_index(n, m, nmax)
            if m == 0:
                # no h terms for g10, g20, g30 etc...
                A_r_0[k] = (-n)*rn1*Pnm[index]
                A_theta_0[k] = -rn1*dPnm[index]
                A_phi_0[k] = -rn1*0
                k += 1
            else:
                alpha = m*phi*rad
                cos_phi = np.cos(alpha)
                sin_phi = np.sin(alpha)
                # g terms, as in g11, g21, g22 etc...
                A_r_0[k] = (-n)*rn1*Pnm[index]*cos_phi
                A_theta_0[k] = -rn1*dPnm[index]*cos_phi
                A_phi_0[k] = -rn1*Pnm[index]/sin_theta*(-m*sin_phi)
                k += 1
                # h terms as in h11, h21, h22 etc...
                A_r_0[k] = (-n)*rn1*Pnm[index]*sin_phi
                A_theta_0[k] = -rn1*dPnm[index]*sin_phi
                A_phi_0[k] = -rn1*Pnm[index]/sin_theta*m*cos_phi
                k += 1

    return A_r_0, A_theta_0, A_phi_0

