#!/usr/bin/env python3
"""Read coefficients files and generate numpy file.
"""

# Import statements
import numpy as np
import aux
from const import (TCELS, PATM, TTP, PTP, DBAR2PA, SAL0, SRED, TRED, PRED,
                   TXTFILE, NPZFILE)

# Constants
PREFIX = 'coefs'


# Helper functions
def readdict(suffix):
    """Read the coefficients from a file and save as dictionaries.

    Read a text file of coefficients and format the data into a dictionary. In
    the text file, each new list of coefficients is preceded by a line of the
    form
        '{name} i j ... c'
    which gives the name of the list and its number of indices. Each line below
    this is an entry in the list, with the form
        '  {i} {j} ... {c}'
    which gives the integer indices followed by the coefficient. This function
    saves these lists into a dictionary of dictionaries, such that
        coefdict[name][(i,j,...)] = c.

    Arguments:
        suffix (str): Suffix for the coefficients file. The full filename is
            'coefs{suffix}.txt'.

    Returns:
        coefdict (dict of dict of floats): Dictionary with the entries
                coefdict[name][(i,j,...)] = c
            where 'name' is the name of the coefficients list; (i,j,...) are
            the indices denoting the powers of each term; and c is the float
            value of the coefficient.
    """
    coefdict = dict()
    name = None
    with open(f'{PREFIX}{suffix}.txt', mode='rt') as f:
        for line in f:
            words = line.split()
            # Is this a header line?
            if not line.startswith('  '):
                name = words[0]
                coefdict[name] = dict()
                continue

            # This is a data line; add to the dictionary
            inds = tuple(int(i) for i in words[:-1])
            coef = float(words[-1])
            if name is not None:
                coefdict[name][inds] = coef
    return coefdict


def fmt1darray(cdict):
    """Convert a 1d coefficient dictionary to an array.

    Convert a dictionary with the entries
        cdict[(j,)] = c
    to an array carr[j] = c.

    Arguments:
        cdict (dict of float): Dictionary of coefficients, expected to be
            generated by `readdict`.

    Returns:
        carr (array of float): Array with the same coefficients.
    """
    # Make sure this represents a 1d array
    if any(len(inds) != 1 for inds in cdict):
        raise ValueError(
            'The keys of the input dictionary must be length-1 tuples')

    jmax = max((inds[-1] for inds in cdict), default=0)
    carr = np.zeros(jmax+1)
    for ((j,), c) in cdict.items():
        carr[j] = c
    return carr


def fmt2darray(cdict):
    """Convert a 2d dictionary to an array.

    Convert a dictionary with entries
        cdict[(j,k)] = c
    to a flat (1d) array. The array is structured by having the last index (k)
    varying slowest and the first index (j) varying fastest. The dictionary can
    be reconstructed by using the nested loop:
        ind = 0
        for k in range(kmax+1):
            for j in range(jmax+1):
                cdict[(j,k)] = carr[ind]
                ind += 1
    Along with the flattened array, the maximum value of j for each value of k
    is returned.

    Arguments:
        cdict (dict of float): Dictionary of coefficients, expected to be
            generated by `readdict`.

    Returns:
        jmaxs (array of int): The maximum value of the first index for each
            value of the second index.
        carr (array of float): Flattened array of coefficients.
    """
    # Make sure this represents a 2d array
    if any(len(inds) != 2 for inds in cdict):
        raise ValueError(
            'The keys of the input dictionary must be length-2 tuples')

    # Find ranges of each index
    kmax = max((inds[-1] for inds in cdict), default=0)
    jmaxs = np.array(
        [max((j for (j, k1) in cdict if k1 == k), default=0)
         for k in range(kmax+1)],
        dtype=int)

    # Construct flattened array
    ncoefs = sum((jmax+1) for jmax in jmaxs)
    carr = np.zeros(ncoefs)
    ind = 0
    for k in range(kmax+1):
        for j in range(jmaxs[k]+1):
            carr[ind] = cdict.get((j, k), 0.)
            ind += 1
    return (jmaxs, carr)


def fmt3darray(cdict):
    """Convert a 3d dictionary to an array.

    Convert a dictionary with entries
        cdict[(i,j,k)] = c
    to a flat (1d) array. The array is structured by having the last index (k)
    varying slowest. For each value of k, the sum (i+j) has a maximum value,
    given by ijmaxs[k]. In the array, i varies fastest, up to the point where
    (i+j) = ijmax; then j increments. The dictionary can be reconstructed using
    the nested loop:
        ind = 0
        for k in range(kmax+1):
            ijmax = ijmaxs[k]
            for j in range(ijmax+1):
                for i in range(ijmax+1-j):
                    cdict[(i,j,k)] = carr[ind]
                    ind += 1
    Along with the flattened array, the array of ijmaxs is returned.

    Arguments:
        cdict (dict of float): Dictionary of coefficients, expected to be
            generated by `readdict`.

    Returns:
        ijmaxs (array of int): The maximum value of (i+j) for each value of k.
        carr (array of float): Flattened array of coefficients.
    """
    # Make sure this represents a 3d array
    if any(len(inds) != 3 for inds in cdict):
        raise ValueError(
            'The keys of the input dictionary must be length-3 tuples')

    # Find ranges of each index
    kmax = max((inds[-1] for inds in cdict), default=0)
    ijmaxs = np.array(
        [max(((i+j) for (i, j, k1) in cdict if k1 == k), default=0)
         for k in range(kmax+1)],
        dtype=int)

    # Construct flattened array
    ncoefs = sum(((ijmax+1)*(ijmax+2)//2) for ijmax in ijmaxs)
    carr = np.zeros(ncoefs)
    ind = 0
    for k in range(kmax+1):
        ijmax = ijmaxs[k]
        for j in range(ijmax+1):
            for i in range(ijmax+1-j):
                carr[ind] = cdict.get((i, j, k), 0.)
                ind += 1
    return (ijmaxs, carr)


# Main functions
def readtxt():
    """Read coefs.txt and format its contents.

    Read the plain-text file of coefficients 'coefs.txt'. Format the lists of
    coefficients, first as a plain dictionary and then as flat (1d) arrays that
    can be used by the functions in the `poly` module.

    Arguments: None.

    Returns:
        coefdict (dict of dict of float): A dictionary with the entries
                coefdict[name] = {(i,j,...): c}
            where 'name' is the name of the coefficients list; (i,j,...) are
            the indices denoting the powers of each term; and c is the float
            value of the coefficient. This is the same dictionary returned by
            `readdict`.
        jmaxdict (dict of array of int): A dictionary with the entries
                jmaxdict[name] = [jmax[k]]
            where 'name' is the name of the coefficients list; k is the power
            of the term used to flatten the array; and jmax[k] is the maximum
            power of j for each value of k.
        ijmaxdict (dict of array of int): A dictionary with the entries
                ijmaxdict[name] = [ijmax[k]]
            where 'name' is the name of the coefficients list; k is the power
            of the term used to flatten the array; and ijmax[k] is the maximum
            power of the other terms for that value of k.
        arraydict (dict of array of float): A dictionary with the flattened
            arrays of coefficients. The original dictionaries can be
            reconstructed by looping through indices in the manner specified
            above.
    """
    coefdict = dict()
    name = None
    with open(TXTFILE, mode='rt') as f:
        for line in f:
            res = line.split()
            if not res:
                # Skip empty lines
                name = None
                continue
            # Is this a header line?
            if not line.startswith('  '):
                name = res[0]
                coefdict[name] = dict()
                continue
            # Data line; add to current
            if name is None:
                # Shouldn't have gotten here!
                raise Exception('No name found for the associated line')
            inds = tuple(int(ind) for ind in res[:-1])
            val = float(res[-1])
            coefdict[name][inds] = val

    # Reformat GSAL dict to have i varying last
    cdict = coefdict.pop('GSAL')
    coefdict['GSAL'] = {(j, k, i): c for ((i, j, k), c) in cdict.items()}

    # Format each entry of the dictionary
    arraydict, jmaxdict, ijmaxdict = dict(), dict(), dict()
    for (name, cdict) in coefdict.items():
        # Figure out number of dimensions
        nd = len(list(cdict.keys())[0])
        if nd == 1:
            arraydict[name] = fmt1darray(cdict)
        elif nd == 2:
            jmaxdict[name], arraydict[name] = fmt2darray(cdict)
        else:
            ijmaxdict[name], arraydict[name] = fmt3darray(cdict)

    # Return all results
    return (coefdict, jmaxdict, ijmaxdict, arraydict)


def adjustcoefs(coefdict, jmaxdict, ijmaxdict, arraydict):
    """Adjust the coefficients for the GSW functions.

    Adjust the coefficients for the GWAT and GSAL lists to match the reference
    conditions. For pure water, the entropy and internal energy are 0 at the
    triple point of water. For salt, the entropy and enthalpy of seawater are 0
    at standard salinity, 0 Celsius, and atmospheric pressure.

    Arguments:
        coefdict, jmaxdict, ijmaxdict, arraydict: The dictionaries returned by
            `readtxt`. All dictionaries except coefdict are modified in-place.

    Returns:
        coefsw0adj (length-2 array): The adjusted values for the (0,0) and
            (1,0) terms in GWAT(j,k).
        coefss1adj (length-2 array): The adjusted values for the (0,0,2) and
            (1,0,2) terms in GSAL(j,k,i).
    """
    # Get the old coefficients for pure water
    name = 'GWAT'
    jmaxs = jmaxdict.pop(name)
    cw_all = arraydict.pop(name)
    cw0_orig = cw_all[:2]
    arraydict[f'{name}0_ORIG'] = cw0_orig
    cw1 = cw_all.copy()
    cw1[:2] = 0.
    arraydict[f'{name}1'] = cw1
    jmaxdict[f'{name}1'] = jmaxs
    gw00_orig, gw10_orig = cw0_orig

    # Calculate the new coefficients from the triple point
    ytp = (TTP - TCELS) / TRED
    ztp = (PTP - PATM)/DBAR2PA / PRED
    gwtp1, gwtp1t, gwtp1p = aux.poly2d_1der(
        ytp, ztp, cw1, jmaxs, yscale=TRED, zscale=PRED)
    gw10_adj = -gwtp1t*TRED
    gw00_adj = (TTP - TCELS)*gwtp1t + PTP/DBAR2PA*gwtp1p - gwtp1

    # Add the new coefficients back into the array
    cw0_adj = np.array([gw00_adj, gw10_adj])
    arraydict[f'{name}0_ADJ'] = cw0_adj

    # Get the old coefficients for salinity
    name = 'GSAL'
    jkmaxs = ijmaxdict.pop(name)
    cs_all = arraydict.pop(name)
    cs0 = cs_all[1:3]
    gs000, gs010 = cs0[:]
    arraydict[f'{name}0'] = cs0
    cs1_orig = cs_all[4:6]
    gs100_orig, gs110_orig = cs1_orig[:]
    arraydict[f'{name}1_ORIG'] = cs1_orig
    cs2 = cs_all[4:].copy()
    cs2[:2] = 0.
    ijmaxdict[f'{name}2'] = jkmaxs[2:]
    arraydict[f'{name}2'] = cs2

    # Calculate the new coefficients from standard conditions
    sso = SAL0/SRED
    xso = sso**.5
    i1max = len(jkmaxs)-4
    # cx2 = np.array([cdict.get((0, 0, i1+3), 0.) for i1 in range(i1max+1)])
    # cx2t = np.array([cdict.get((1, 0, i1+3), 0.) for i1 in range(i1max+1)])
    cx2, cx2t = np.zeros(i1max+1), np.zeros(i1max+1)
    ind = (jkmaxs[2]+1)*(jkmaxs[2]+2)//2
    for i1 in range(i1max+1):
        jkmax = jkmaxs[3+i1]
        cx2[i1] = cs2[ind]
        if jkmax > 0:
            cx2t[i1] = cs2[ind+1]
        ind += (jkmax+1)*(jkmax+2)//2

    gs2_os = aux.poly1d(xso, cx2)
    gs2t_ttos = aux.poly1d(xso, cx2t)
    lnx = np.log(xso)
    gs110_adj = -gs2t_ttos*xso - gs010*lnx - gw10_adj/sso
    gs100_adj = -gs2_os*xso - gs000*lnx - gw00_adj/sso
    cs1_adj = np.array([gs100_adj, gs110_adj])
    arraydict[f'{name}1_ADJ'] = cs1_adj
    return (cw0_adj, cs1_adj)


def savecoefs():
    """Save the coefficients in a .npz file.

    Read all of the individual coefficient files. Save the resulting arrays in
    a numpy .npz file for easier access later.

    Arguments: None
    Returns: None
    """
    # Read the text file of coefficients
    coefdict, jmaxdict, ijmaxdict, arraydict = readtxt()

    # Calculate the adjusted coefficients
    adjustcoefs(coefdict, jmaxdict, ijmaxdict, arraydict)

    # Collect all arrays into a single set of keyword arguments
    kwargs = dict()
    kwargs.update(arraydict)
    for (name, jmaxs) in jmaxdict.items():
        kwargs[f'jmaxs_{name}'] = jmaxs
    for (name, ijmaxs) in ijmaxdict.items():
        if name == 'GSAL2':
            mname = f'jkmaxs_{name}'
        else:
            mname = f'ijmaxs_{name}'
        kwargs[mname] = ijmaxs

    # Save results to file
    np.savez(NPZFILE, **kwargs)
    return None


# If called as a script
if __name__ == '__main__':
    # Save the coefficients
    savecoefs()
