import uncertainties.unumpy as unp
import numpy as np
import pandas as pd


# Helpers useful to the user
# --------------------------
def data_out(cbdat, path=None, include_constants=False):
    """
    Save output from cbsyst.

    Parameters
    ----------
    cbdat : dict / Bunch
        The output from Csys, Bsys, ABsys or CBsys.
    path : str
        The file name (and path) where you want to
        save the data. If not provided, data are not
        saved to a file.

        The extension of the file determines the output
        format. Can be 'csv', 'xls', 'html, 'tex', or 'pkl'.
    include_constants : bool
        If True, include pK and alpha constants in output.

    Returns
    -------
    * pandas.DataFrame of output
    * Saves file (if specified)

    """

    cols = [
        "pH",
        "DIC",
        "fCO2",
        "pCO2",
        "CO2",
        "HCO3",
        "CO3",
        "TA",
        "BT",
        "BO3",
        "BO4",
        "dBT",
        "dBO3",
        "dBO4",
        "ABT",
        "ABO3",
        "ABO4",
        "T",
        "S",
        "P",
        "Ca",
        "Mg",
    ]

    consts = ["K0", "K1", "K2", "KB", "KW", "KSO4", "KspA", "KspC"]

    size = cbdat.pH.size
    out = pd.DataFrame(index=range(size))

    for c in cols:
        if c in cbdat and cbdat[c] is not None:
            if (np.ndim(cbdat[c]) == 1) & (cbdat[c].size == 1):
                cbdat[c] = cbdat[c][0]
            if c in cbdat:
                out.loc[:, c] = cbdat[c]

    if include_constants:
        for c in consts:
            if c in cbdat.Ks and cbdat.Ks[c] is not None:
                if (np.ndim(cbdat.Ks[c]) == 1) & (cbdat.Ks[c].size == 1):
                    cbdat.Ks[c] = cbdat.Ks[c][0]
                out.loc[:, "p" + c] = -np.log10(cbdat.Ks[c])
        if "alphaB" in cbdat and cbdat.alphaB is not None:
            if (np.ndim(cbdat.alphaB) == 1) & (cbdat.alphaB.size == 1):
                cbdat.alphaB = cbdat.alphaB[0]
            out.loc[:, "alphaB"] = cbdat.alphaB

    if path is not None:
        fmt = path.split(".")[-1]
        fdict = {
            "csv": "to_csv",
            "html": "to_html",
            "xls": "to_excel",
            "pkl": "to_pickle",
            "tex": "to_latex",
        }

        if fmt not in fdict:
            raise ValueError(
                (
                    "File extension does not match available output\n"
                    + "options. Should be one of 'csv', 'html', 'xls',\n"
                    + "'pkl' (pickle) or 'tex' (LaTeX)."
                )
            )
        try:
            _ = getattr(out, fdict[fmt])(path, index=None)
        except TypeError:
            _ = getattr(out, fdict[fmt])(path)

    return out


# Programmatic helpers for code elsewhere
# ---------------------------------------

# Bunch modifies dict to allow item access using dot (.) operator
class Bunch(dict):
    def __init__(self, *args, **kwds):
        super(Bunch, self).__init__(*args, **kwds)
        self.__dict__ = self


def noms(*it):
    """
    Return nominal_values for provided objects.

    Parameters
    ----------
    *it : n objects
    """
    return [unp.nominal_values(i) for i in it]


def maxL(*it):
    """
    Calculate maximum length of provided items.

    Parameters
    ----------
    *it : objects
        Items of various lengths. Only lengths
        of iterables are returned.

    Returns
    -------
    Length of longest object (int).
    """
    m = set()
    for i in it:
        try:
            m.add(len(i))
        except TypeError:
            pass
    if len(m) > 0:
        return max(m)
    else:
        return 1


def cast_array(*it):
    """
    Recasts inputs into array of shape (len(it), maxL(*it))
    """
    new = np.zeros((len(it), maxL(*it)))
    for i, t in enumerate(it):
        new[i, :] = t
    return new


def NnotNone(*it):
    """
    Returns the number of elements of it tha are not None.

    Parameters
    ----------
    it : iterable
        iterable of elements that are either None, or not None

    Returns
    -------
    int
    """
    return sum([i is not None for i in it])


# pK <--> K converters
def ch(pK):
    """
    Convert pK to K
    """
    return np.power(10.0, np.multiply(pK, -1.0))


def cp(K):
    """
    Convert K to pK
    """
    return -np.log10(K)


# Helpers for aspects of seawater chemistry
# -----------------------------------------
def prescorr(P, Tc, a0, a1, a2, b0, b1):
    """
    Calculate pressore correction factor for thermodynamic Ks.

    From Millero et al (2007, doi:10.1021/cr0503557)
    Eqns 38-40

    Usage:
    K_corr / K_orig = [output]
    Kcorr = [output] * K_orig
    """
    dV = a0 + a1 * Tc + a2 * Tc ** 2
    dk = (b0 + b1 * Tc) / 1000
    # factor of 1000 not mentioned in Millero,
    # but present in Zeebe book, and used in CO2SYS
    RT = 83.1451 * (Tc + 273.15)
    return np.exp((-dV + 0.5 * dk * P) * P / RT)


def swdens(TempC, Sal):
    """
    Seawater Density (kg / L) from Temp (C) and Sal (PSU)

    Chapter 5, Section 4.2 of Dickson, Sabine and Christian
    (2007, http://cdiac.ornl.gov/oceans/Handbook_2007.html)

    Parameters
    ----------
    TempC : array-like
        Temperature in celcius.
    Sal : array-like
        Salinity in PSU

    Returns
    -------
    Density in kg / L
    """
    # convert temperature to IPTS-68
    T68 = (TempC + 0.0002) / 0.99975
    pSMOW = (
        999.842594
        + 6.793952e-2 * T68
        + -9.095290e-3 * T68 ** 2
        + 1.001685e-4 * T68 ** 3
        + -1.120083e-6 * T68 ** 4
        + 6.536332e-9 * T68 ** 5
    )
    A = (
        8.24493e-1
        + -4.0899e-3 * T68
        + 7.6438e-5 * T68 ** 2
        + -8.2467e-7 * T68 ** 3
        + 5.3875e-9 * T68 ** 4
    )
    B = -5.72466e-3 + 1.0227e-4 * T68 + -1.6546e-6 * T68 ** 2
    C = 4.8314e-4
    return (pSMOW + A * Sal + B * Sal ** 1.5 + C * Sal ** 2) / 1000


def calc_TS(Sal):
    """
    Calculate total Sulphur

    Morris, A. W., and Riley, J. P., Deep-Sea Research 13:699-705, 1966:
    this is .02824.*Sali./35. = .0008067.*Sali
    """
    a, b, c = (0.14, 96.062, 1.80655)
    return (a / b) * (Sal / c)  # mol/kg-SW


def calc_TF(Sal):
    """
    Calculate total Fluorine

    Riley, J. P., Deep-Sea Research 12:219-220, 1965:
    this is .000068.*Sali./35. = .00000195.*Sali
    """
    a, b, c = (0.000067, 18.998, 1.80655)
    return (a / b) * (Sal / c)  # mol/kg-SW


# def calc_TB(Sal):
#     """
#     Calculate total Boron

#     Lee, Kim, Byrne, Millero, Feely, Yong-Ming Liu. 2010.
#     Geochimica Et Cosmochimica Acta 74 (6): 1801-1811
#     """
#     a, b = (0.0004326, 35.)
#     return a * Sal / b


def calc_TB(Sal):
    """
    Calculate total Boron

    Directly from CO2SYS:
    Uppstrom, L., Deep-Sea Research 21:161-162, 1974:
    this is 0.000416 * Sal/35. = 0.0000119 * Sal
    TB(FF) = (0.000232 / 10.811) * (Sal / 1.80655) in mol/kg-SW
    """
    a, b = (0.0004157, 35.0)
    return a * Sal / b


def calc_fH(TempK, Sal):
    # Same as CO2SYS
    # Takahashi et al, Chapter 3 in GEOSECS Pacific Expedition,
    # v. 3, 1982 (p. 80)

    a, b, c, d = (1.2948, -2.036e-3, 4.607e-4, -1.475e-6)
    return a + b * TempK + (c + d * TempK) * Sal ** 2


# Convert between pH scales
def calc_pH_scales(pHtot, pHfree, pHsws, pHNBS, TS, TF, TempK, Sal, Ks):
    """
    Calculate pH on all scales, given one.
    """

    # check if any pH scale is given.
    npH = NnotNone(pHfree, pHsws, pHtot, pHNBS)

    if npH == 1:
        # pH scale conversions
        FREEtoTOT = -np.log10((1 + TS / Ks.KSO4))
        SWStoTOT = -np.log10((1 + TS / Ks.KSO4) / (1 + TS / Ks.KSO4 + TF / Ks.KF))
        fH = calc_fH(TempK, Sal)

        if pHtot is not None:
            return {
                "pHfree": pHtot - FREEtoTOT,
                "pHsws": pHtot - SWStoTOT,
                "pHNBS": pHtot - SWStoTOT - np.log10(fH),
            }
        elif pHsws is not None:
            return {
                "pHfree": pHsws + SWStoTOT - FREEtoTOT,
                "pHtot": pHsws + SWStoTOT,
                "pHNBS": pHsws - np.log10(fH),
            }
        elif pHfree is not None:
            return {
                "pHsws": pHfree + FREEtoTOT - SWStoTOT,
                "pHtot": pHfree + FREEtoTOT,
                "pHNBS": pHfree + FREEtoTOT - SWStoTOT - np.log10(fH),
            }
        elif pHNBS is not None:
            return {
                "pHsws": pHNBS + np.log10(fH),
                "pHtot": pHNBS + np.log10(fH) + SWStoTOT,
                "pHfree": pHNBS + np.log10(fH) + SWStoTOT - FREEtoTOT,
            }
    else:
        return {}
