import numpy as np

def read_gfc(file_name):
    """
    Read a set of potential coefficients from a GFC file.

    Parameters
    ----------
    file_name : str
        name of GFC file

    Returns
    -------
    anm : ndarray(max_degree + 1, max_degree + 1)
        Dimensionless potential coefficients. Cosine coefficients of degree n and order m are stored in the lower triangle as anm[n, m],
        sine coefficients of degree n and order m > 0 are stored in the upper triangle as anm[m - 1, n].
    """
    with open(file_name) as f:
        for line in f:
            if line.startswith('gfc'):
                sline = line.split()
                n, m, cnm, snm = int(sline[1]), int(sline[2]), float(sline[3]), float(sline[4])
                anm[n, m] = cnm
                if m > 0:
                    anm[m - 1, n] = snm
            elif line.startswith('max_degree'):
                array_size = int(line.split()[-1]) + 1
                anm = np.zeros((array_size, array_size))
    return anm


def doodson_arguments(mjd):
    """
    Compute the six Doodson fundamental arguments (tau, s, h, p, N0, ps), for a given time.

    Parameters
    ----------
    mjd : float
        Time in Modified Julian Date (MJD).

    Returns
    -------
    doodson : ndarray(6)
        Doodson fundamental arguments (tau, s, h, p, N0, ps).
    """
    t    = (mjd - 51544.5) / 365250.0
    era  = 2*np.pi * (np.mod(mjd-51544.5, 1) + 0.7790572732640 + 0.00273781191135448 * (mjd-51544.5));
    gmst = era + (0.014506 + (4612.156534 + (1.3915817 + (-0.00000044 + (-0.000029956 + -0.0000000368*t)*t)*t)*t)*t) * np.pi/(180*3600);

    dood = np.zeros(6)
    dood[1] = (218.31664562999 + (4812678.81195750 + (-0.14663889 + ( 0.00185140 + -0.00015355*t)*t)*t)*t) * np.pi / 180
    dood[2] = (280.46645016002 + ( 360007.69748806 + ( 0.03032222 + ( 0.00002000 + -0.00006532*t)*t)*t)*t) * np.pi / 180
    dood[3] = ( 83.35324311998 + (  40690.13635250 + (-1.03217222 + (-0.01249168 +  0.00052655*t)*t)*t)*t) * np.pi / 180
    dood[4] = (234.95544499000 + (  19341.36261972 + (-0.20756111 + (-0.00213942 +  0.00016501*t)*t)*t)*t) * np.pi / 180
    dood[5] = (282.93734098001 + (     17.19457667 + ( 0.04568889 + (-0.00001776 + -0.00003323*t)*t)*t)*t) * np.pi / 180
    dood[0] = gmst + np.pi - dood[1]
    return dood


if __name__ == "__main__":

    path                = '../../models/EOT20_OCN/';
    basename            = 'EOT20_OCN';
    doodson_multipliers = np.loadtxt(path+basename+'_002doodson.txt')
    admittance_matrix   = np.loadtxt(path+basename+'_003admittance.txt')
    file_list           = np.loadtxt(path+basename+'_001fileList.txt', dtype=str)
    anm_cos             = [read_gfc(path+file_names[0]) for file_names in file_list]
    anm_sin             = [read_gfc(path+file_names[1]) for file_names in file_list]

    mjd = 51544.5 # date for test computation
    thetaf     = doodson_multipliers @ doodson_arguments(mjd)
    factor_cos = admittance_matrix @ np.cos(thetaf)
    factor_sin = admittance_matrix @ np.sin(thetaf)

    anm = np.zeros(anm_cos[0].shape)
    for k in range(factor_sin.size):
        anm += factor_cos[k] * anm_cos[k] + factor_sin[k] * anm_sin[k]

    # set degree one zero: Center of mass (CM)
    anm[0:2, 0:2] = 0.0;

    # anm contain now the stokes coeffcients of the ocean tides at time mjd
    # Cnm coefficients of degree n and order m are stored in the lower triangle of anm,
    # Snm coefficients of degree n and order m>0 are stored in the upper triangle of anm.