""" Defines the Polynomial class """
from __future__ import division, print_function, absolute_import, unicode_literals
#***************************************************************************************************
# Copyright 2015, 2019 National Technology & Engineering Solutions of Sandia, LLC (NTESS).
# Under the terms of Contract DE-NA0003525 with NTESS, the U.S. Government retains certain rights
# in this software.
# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
# in compliance with the License.  You may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0 or in the LICENSE file in the root pyGSTi directory.
#***************************************************************************************************

import numpy as _np
import platform as _platform
import collections as _collections
from . import replib

assert(_platform.architecture()[0].endswith("bit"))  # e.g. "64bit"
PLATFORM_BITS = int(_platform.architecture()[0].strip("bit"))


class Polynomial(dict):
    """
    Encapsulates a polynomial as a subclass of the standard Python dict.

    Variables are represented by integer indices, e.g. "2" means "x_2".
    Keys are tuples of variable indices and values are numerical
    coefficients (floating point or complex numbers).  To specify a variable
    to some power, its index is repeated in the key-tuple.

    E.g. x_0^2 + 3*x_1 + 4 is stored as {(0,0): 1.0, (1,): 3.0, (): 4.0}
    """

    @classmethod
    def get_vindices_per_int(cls, max_num_vars):
        """
        Returns the number of variable indices that can be compactly fit
        into a single int when there are at most `max_num_vars` variables.

        This quantity is needed to directly construct Polynomial representations
        and is thus useful internally for forward simulators.
        """
        return int(_np.floor(PLATFORM_BITS / _np.log2(max_num_vars + 1)))

    @classmethod
    def fromrep(cls, rep):
        """
        Creates a Polynomial from a "representation" (essentially a
        lite-version) of a Polynomial.

        Note: usually we only need to convert from full-featured Python objects
        to the lighter-weight "representation" objects.  Polynomials are an
        exception, since as the results of probability computations they need
        to be converted back from "representation-form" to "full-form".

        Parameters
        ----------
        rep : PolyRep
            A polynomial representation.

        Returns
        -------
        Polynomial
        """
        max_num_vars = rep.max_num_vars  # one of the few/only cases where a rep
        # max_order = rep.max_order        # needs to expose some python properties

        def int_to_vinds(indx_tup):
            ret = []
            for indx in indx_tup:
                while indx != 0:
                    nxt = indx // (max_num_vars + 1)
                    i = indx - nxt * (max_num_vars + 1)
                    ret.append(i - 1)
                    indx = nxt
            #assert(len(ret) <= max_order) #TODO: is this needed anymore?
            return tuple(sorted(ret))

        tup_coeff_dict = {int_to_vinds(k): val for k, val in rep.coeffs.items()}
        return cls(tup_coeff_dict)

    def __init__(self, coeffs=None):
        """
        Initializes a new Polynomial object (a subclass of dict).

        Internally (as a dict) a Polynomial represents variables by integer
        indices, e.g. "2" means "x_2".  Keys are tuples of variable indices and
        values are numerical coefficients (floating point or complex numbers).
        A variable to a power > 1 has its index repeated in the key-tuple.

        E.g. x_0^2 + 3*x_1 + 4 is stored as `{(0,0): 1.0, (1,): 3.0, (): 4.0}`

        Parameters
        ----------
        coeffs : dict
            A dictionary of coefficients.  Keys are tuples of integers that
            specify the polynomial term the coefficient value multiplies
            (see above). If None, the zero polynomial (no terms) is created.
        """
        super(Polynomial, self).__init__()
        if coeffs is not None:
            self.update(coeffs)

    def deriv(self, wrtParam):
        """
        Take the derivative of this Polynomial with respect to a single
        variable/parameter.  The result is another Polynomial.

        E.g. deriv(x_2^3 + 3*x_1, wrtParam=2) = 3x^2

        Parameters
        ----------
        wrtParam : int
            The variable index to differentiate with respect to.
            E.g. "4" means "differentiate w.r.t. x_4".

        Returns
        -------
        Polynomial
        """
        dcoeffs = {}
        for ivar, coeff in self.items():
            cnt = float(ivar.count(wrtParam))
            if cnt > 0:
                l = list(ivar)
                del l[l.index(wrtParam)]
                dcoeffs[tuple(l)] = cnt * coeff

        return Polynomial(dcoeffs)

    def get_degree(self):
        """
        Return the largest sum-of-exponents for any term (monomial) within this
        polynomial. E.g. for x_2^3 + x_1^2*x_0^2 has degree 4.
        """
        return 0 if len(self) == 0 else max([len(k) for k in self.keys()])

    def evaluate(self, variable_values):
        """
        Evaluate this polynomial for a given set of variable values.

        Parameters
        ----------
        variable_values : array-like
            An object that can be indexed so that `variable_values[i]` gives the
            numerical value for i-th variable (x_i).

        Returns
        -------
        float or complex
            Depending on the types of the coefficients and `variable_values`.
        """
        #FUTURE: make this function smarter (Russian peasant)
        ret = 0
        for ivar, coeff in self.items():
            ret += coeff * _np.product([variable_values[i] for i in ivar])
        return ret

    def compact(self, force_complex=False):
        """
        Generate a compact form of this polynomial designed for fast evaluation.

        The resulting "tapes" can be evaluated using
        :function:`bulk_eval_compact_polys`.

        Parameters
        ----------
        force_complex : bool, optional
            Whether the `ctape` returned array is forced to be of complex type,
            even if all of the polynomial coefficients are real.

        Returns
        -------
        vtape, ctape : numpy.ndarray
            These two 1D arrays specify an efficient means for evaluating this
            polynomial.
        """
        if force_complex:
            iscomplex = True
        else:
            iscomplex = any([abs(_np.imag(x)) > 1e-12 for x in self.values()])

        nTerms = len(self)
        nVarIndices = sum(map(len, self.keys()))
        vtape = _np.empty(1 + nTerms + nVarIndices, _np.int64)  # "variable" tape
        ctape = _np.empty(nTerms, complex if iscomplex else 'd')  # "coefficient tape"

        i = 0
        vtape[i] = nTerms; i += 1
        for iTerm, k in enumerate(sorted(self.keys())):
            l = len(k)
            ctape[iTerm] = self[k] if iscomplex else _np.real(self[k])
            vtape[i] = l; i += 1
            vtape[i:i + l] = k; i += l
        assert(i == len(vtape)), "Logic Error!"
        return vtape, ctape

    def copy(self):
        """
        Returns a copy of this polynomial.
        """
        return Polynomial(self)

    def map_indices(self, mapfn):
        """
        Performs a bulk find & replace on this polynomial's variable indices.

        This is useful when the variable indices have external significance
        (like being the indices of a gate's parameters) and one want to convert
        to another set of indices (like a parent model's parameters).

        Parameters
        ----------
        mapfn : function
            A function that takes as input an "old" variable-index-tuple
            (a key of this Polynomial) and returns the updated "new"
            variable-index-tuple.

        Returns
        -------
        Polynomial
        """
        return Polynomial({mapfn(k): v for k, v in self.items()})

    def map_indices_inplace(self, mapfn):
        """
        Performs an in-place find & replace on this polynomial's variable indices.

        This is useful when the variable indices have external significance
        (like being the indices of a gate's parameters) and one want to convert
        to another set of indices (like a parent model's parameters).

        Parameters
        ----------
        mapfn : function
            A function that takes as input an "old" variable-index-tuple
            (a key of this Polynomial) and returns the updated "new"
            variable-index-tuple.

        Returns
        -------
        None
        """
        new_items = {mapfn(k): v for k, v in self.items()}
        self.clear()
        self.update(new_items)

    def mult(self, x):
        """
        Multiplies this polynomial by another polynomial `x`.

        Parameters
        ----------
        x : Polynomial
            The polynomial to multiply by.

        Returns
        -------
        Polynomial
             The polynomial representing self * x.
        """
        newpoly = Polynomial()
        for k1, v1 in self.items():
            for k2, v2 in x.items():
                k = tuple(sorted(k1 + k2))
                if k in newpoly: newpoly[k] += v1 * v2
                else: newpoly[k] = v1 * v2
        return newpoly

    def scale(self, x):
        """
        Scale this polynomial by `x` (multiply all coefficients by `x`).

        Parameters
        ----------
        x : float or complex
            The value to scale by.

        Returns
        -------
        None
        """
        # assume a scalar that can multiply values
        for k in tuple(self.keys()):  # I think the tuple() might speed things up (why?)
            self[k] *= x

    def scalar_mult(self, x):
        """
        Multiplies this polynomial by a scalar `x`.

        Parameters
        ----------
        x : float or complex
            The value to multiply by.

        Returns
        -------
        Polynomial
        """
        newpoly = self.copy()
        newpoly.scale(x)
        return newpoly

    def __str__(self):
        def fmt(x):
            if abs(_np.imag(x)) > 1e-6:
                if abs(_np.real(x)) > 1e-6: return "(%.3f+%.3fj)" % (x.real, x.imag)
                else: return "(%.3fj)" % x.imag
            else: return "%.3f" % x.real

        termstrs = []
        sorted_keys = sorted(list(self.keys()))
        for k in sorted_keys:
            varstr = ""; last_i = None; n = 1
            for i in sorted(k):
                if i == last_i: n += 1
                elif last_i is not None:
                    varstr += "x%d%s" % (last_i, ("^%d" % n) if n > 1 else "")
                    n = 1
                last_i = i
            if last_i is not None:
                varstr += "x%d%s" % (last_i, ("^%d" % n) if n > 1 else "")
            #print("DB: k = ",k, " varstr = ",varstr)
            if abs(self[k]) > 1e-4:
                termstrs.append("%s%s" % (fmt(self[k]), varstr))
        if len(termstrs) > 0:
            return " + ".join(termstrs)
        else: return "0"

    def __repr__(self):
        return "Poly[ " + str(self) + " ]"

    def __add__(self, x):
        newpoly = self.copy()
        if isinstance(x, Polynomial):
            for k, v in x.items():
                if k in newpoly: newpoly[k] += v
                else: newpoly[k] = v
        else:  # assume a scalar that can be added to values
            for k in newpoly:
                newpoly[k] += x
        return newpoly

    def __iadd__(self, x):
        """ Does self += x more efficiently """
        if isinstance(x, Polynomial):
            for k, v in x.items():
                try:
                    self[k] += v
                except KeyError:
                    self[k] = v
        else:  # assume a scalar that can be added to values
            for k in self:
                self[k] += x
        return self

    def __mul__(self, x):
        #if isinstance(x, Polynomial):
        #    newpoly = Polynomial()
        #    for k1,v1 in self.items():
        #        for k2,v2 in x.items():
        #            k = tuple(sorted(k1+k2))
        #            if k in newpoly: newpoly[k] += v1*v2
        #            else: newpoly[k] = v1*v2
        #else:
        #    # assume a scalar that can multiply values
        #    newpoly = self.copy()
        #    for k in newpoly:
        #        newpoly[k] *= x
        #return newpoly
        if isinstance(x, Polynomial):
            return self.mult(x)
        else:  # assume a scalar that can multiply values
            return self.scalar_mult(x)

    def __rmul__(self, x):
        return self.__mul__(x)

    def __imul__(self, x):
        if isinstance(x, Polynomial):
            newcoeffs = {}
            for k1, v1 in self.items():
                for k2, v2 in x.items():
                    k = tuple(sorted(k1 + k2))
                    if k in newcoeffs: newcoeffs[k] += v1 * v2
                    else: newcoeffs[k] = v1 * v2
            self.clear()
            self.update(newcoeffs)
        else:
            self.scale(x)
        return self

    def __pow__(self, n):
        ret = Polynomial({(): 1.0})  # max_order updated by mults below
        cur = self
        for i in range(int(_np.floor(_np.log2(n))) + 1):
            rem = n % 2  # gets least significant bit (i-th) of n
            if rem == 1: ret *= cur  # add current power of x (2^i) if needed
            cur = cur * cur  # current power *= 2
            n //= 2  # shift bits of n right
        return ret

    def __copy__(self):
        return self.copy()

    def torep(self, max_order=None, max_num_vars=None):
        """
        Construct a representation of this polynomial.

        "Representations" are lightweight versions of objects used to improve
        the efficiency of intensely computational tasks.  Note that Polynomial
        representations must have the same `max_order` and `max_num_vars` in
        order to interact with each other (add, multiply, etc.).

        Parameters
        ----------
        max_order : int, optional
            The maximum order (degree) terms are allowed to have.  If None,
            then it is taken as the current degree of this polynomial.

        max_num_vars : int, optional
            The maximum number of variables the represenatation is allowed to
            have (x_0 to x_(`max_num_vars-1`)).  This sets the maximum allowed
            variable index within the representation.

        Returns
        -------
        PolyRep
        """
        # Set max_order (determines based on coeffs if necessary)
        default_max_order = self.get_degree()
        if max_order is None:
            max_order = default_max_order
        else:
            assert(default_max_order <= max_order)

        # Set max_num_vars (determines based on coeffs if necessary)
        default_max_vars = 0 if len(self) == 0 else \
            max([(max(k) + 1 if k else 0) for k in self.keys()])
        if max_num_vars is None:
            max_num_vars = default_max_vars
        else:
            assert(default_max_vars <= max_num_vars)

        #new.max_order = max_order
        #new.max_num_vars = max_num_vars
        vindices_per_int = Polynomial.get_vindices_per_int(max_num_vars)

        def vinds_to_int(vinds):
            """ Convert tuple index of ints to single int given max_order,max_numvars """
            ints_in_key = int(_np.ceil(len(vinds) / vindices_per_int))
            ret_tup = []
            for k in range(ints_in_key):
                ret = 0; m = 1
                for i in vinds[k * vindices_per_int:(k + 1) * vindices_per_int]:  # last tuple index is most significant
                    assert(i < max_num_vars), "Variable index exceed maximum!"
                    ret += (i + 1) * m
                    m *= max_num_vars + 1
                assert(ret >= 0), "vinds = %s -> %d!!" % (str(vinds), ret)
                ret_tup.append(ret)
            return tuple(ret_tup)

        int_coeffs = {vinds_to_int(k): v for k, v in self.items()}

        # (max_num_vars+1) ** vindices_per_int <= 2**PLATFORM_BITS, so:
        # vindices_per_int * log2(max_num_vars+1) <= PLATFORM_BITS
        vindices_per_int = int(_np.floor(PLATFORM_BITS / _np.log2(max_num_vars + 1)))

        return replib.PolyRep(int_coeffs, max_order, max_num_vars, vindices_per_int)


def bulk_eval_compact_polys(vtape, ctape, paramvec, dest_shape):
    """
    Evaluate many compact polynomial forms at a given set of variable values.

    Parameters
    ----------
    vtape, ctape : numpy.ndarray
        Specifies "variable" and "coefficient" 1D numpy arrays to evaluate.
        These "tapes" can be generated by concatenating the tapes of individual
        complact-polynomial tuples returned by :method:`Polynomial.compact`.

    paramvec : array-like
        An object that can be indexed so that `paramvec[i]` gives the
        numerical value to substitute for i-th polynomial variable (x_i).

    dest_shape : tuple
        The shape of the final array of evaluated polynomials.  The resulting
        1D array of evaluated polynomials is reshaped accordingly.

    Returns
    -------
    numpy.ndarray
        An array of the same type as the coefficient tape, with shape given
        by `dest_shape`.
    """
    result = _np.empty(dest_shape, ctape.dtype)  # auto-determine type?
    res = result.flat  # for 1D access

    c = 0; i = 0; r = 0
    while i < vtape.size:
        poly_val = 0
        nTerms = vtape[i]; i += 1
        #print("POLY w/%d terms (i=%d)" % (nTerms,i))
        for m in range(nTerms):
            nVars = vtape[i]; i += 1  # number of variable indices in this term
            a = ctape[c]; c += 1
            #print("  TERM%d: %d vars, coeff=%s" % (m,nVars,str(a)))
            for k in range(nVars):
                a *= paramvec[vtape[i]]; i += 1
            poly_val += a
            #print("  -> added %s to poly_val = %s" % (str(a),str(poly_val))," i=%d, vsize=%d" % (i,vtape.size))
        res[r] = poly_val; r += 1
    assert(c == ctape.size), "Coeff Tape length error: %d != %d !" % (c, ctape.size)
    assert(r == result.size), "Result/Tape size mismatch: only %d result entries filled!" % r
    return result


def bulk_load_compact_polys(vtape, ctape, keep_compact=False):
    """
    Create a list of Polynomial objects from a "tape" of their compact versions.

    Parameters
    ----------
    vtape, ctape : numpy.ndarray
        Specifies "variable" and "coefficient" 1D numpy arrays to load.
        These "tapes" can be generated by concatenating the tapes of individual
        complact-polynomial tuples returned by :method:`Polynomial.compact`.

    keep_compact : bool, optional
        If True the returned list has elements which are (vtape,ctape) tuples
        for each individual polynomial.  If False, then the elements are
        :class:`Polynomial` objects.

    Returns
    -------
    list
    """
    result = []
    c = 0; i = 0

    if keep_compact:
        while i < vtape.size:
            i2 = i  # increment i2 instead of i for this poly
            nTerms = vtape[i2]; i2 += 1
            for m in range(nTerms):
                nVars = vtape[i2]  # number of variable indices in this term
                i2 += nVars + 1
            result.append((vtape[i:i2], ctape[c:c + nTerms]))
            i = i2; c += nTerms
    else:
        while i < vtape.size:
            poly_coeffs = {}
            nTerms = vtape[i]; i += 1
            #print("POLY w/%d terms (i=%d)" % (nTerms,i))
            for m in range(nTerms):
                nVars = vtape[i]; i += 1  # number of variable indices in this term
                a = ctape[c]; c += 1
                #print("  TERM%d: %d vars, coeff=%s" % (m,nVars,str(a)))
                poly_coeffs[tuple(vtape[i:i + nVars])] = a; i += nVars
            result.append(Polynomial(poly_coeffs))
    return result


def compact_deriv(vtape, ctape, wrtParams):
    """
    Take the derivative of one or more compact Polynomials with respect
    to one or more variables/parameters.

    Parameters
    ----------
    vtape, ctape : numpy.ndarray
        Specifies "variable" and "coefficient" 1D numpy arrays to differentiate.
        These "tapes" can be generated by concatenating the tapes of individual
        complact-polynomial tuples returned by :method:`Polynomial.compact`.

    wrtParams : list
        The variable indices to differentiate with respect to.  They
        must be sorted in ascending order. E.g. "[0,3]" means separatey
        differentiate w.r.t x_0 and x_3 (concatenated first by wrtParam
        then by poly).

    Returns
    -------
    vtape, ctape : numpy.ndarray
    """
    result_vtape = []
    result_ctape = []
    wrt = sorted(wrtParams)
    assert(wrt == list(wrtParams)), "`wrtParams` (%s) must be in ascending order!" % wrtParams
    #print("TAPE SIZE = ",vtape.size)

    c = 0; i = 0
    while i < vtape.size:
        j = i  # increment j instead of i for this poly
        nTerms = vtape[j]; j += 1
        dctapes = [list() for x in range(len(wrt))]
        dvtapes = [list() for x in range(len(wrt))]
        dnterms = [0] * len(wrt)
        #print("POLY w/%d terms (i=%d)" % (nTerms,i))
        for m in range(nTerms):
            coeff = ctape[c]; c += 1
            nVars = vtape[j]; j += 1  # number of variable indices in this term

            #print("  TERM%d: %d vars, coeff=%s" % (m,nVars,str(coeff)))
            cur_iWrt = 0
            j0 = j  # the vtape index where the current term starts

            #Loop to get counts of each variable index that is also in `wrt`.
            # Once we've passed an element of `wrt` process it, since there can't
            # see it any more (the var indices are sorted).
            while j < j0 + nVars:  # loop over variable indices for this term
                # can't be while True above in case nVars == 0 (then vtape[j] isn't valid)

                #find an iVar that is also in wrt.
                # - increment the cur_iWrt or j as needed
                while cur_iWrt < len(wrt) and vtape[j] > wrt[cur_iWrt]:  # condition to increment cur_iWrt
                    cur_iWrt += 1  # so wrt[cur_iWrt] >= vtape[j]
                if cur_iWrt == len(wrt): break  # no more possible iVars we're interested in;
                # we're done with all wrt elements
                # - at this point we know wrt[cur_iWrt] is valid and wrt[cur_iWrt] >= tape[j]
                while j < j0 + nVars and vtape[j] < wrt[cur_iWrt]:
                    j += 1  # so vtape[j] >= wrt[cur_iWrt]
                if j == j0 + nVars: break  # no more iVars - we're done

                #print(" check j=%d, val=%d, wrt=%d, cur_iWrt=%d" % (j,vtape[j],wrt[cur_iWrt],cur_iWrt))
                if vtape[j] == wrt[cur_iWrt]:
                    #Yay! a value we're looking for is present in the vtape.
                    # Figure out how many there are (easy since vtape is sorted
                    # and we'll always stop on the first one)
                    cnt = 0
                    while j < j0 + nVars and vtape[j] == wrt[cur_iWrt]:
                        cnt += 1; j += 1
                    #Process cur_iWrt: add a term to tape for cur_iWrt
                    dvars = list(vtape[j0:j - 1]) + list(vtape[j:j0 + nVars])  # removes last wrt[cur_iWrt] var
                    dctapes[cur_iWrt].append(coeff * cnt)
                    dvtapes[cur_iWrt].extend([nVars - 1] + dvars)
                    dnterms[cur_iWrt] += 1
                    # print(" wrt=%d found cnt=%d: adding deriv term coeff=%f vars=%s" \
                    #       % (wrt[cur_iWrt], cnt, coeff*cnt, [nVars-1] + dvars))

                    cur_iWrt += 1  # processed this wrt param - move to next one

            #Now term has been processed, adding derivative terms to the dctapes and dvtapes "tape-lists"
            # We continue processing terms, adding to these tape lists, until all the terms of the
            # current poly are processed.  Then we can concatenate the tapes for each wrtParams element.
            # Move to next term; j may not have been incremented if we exited b/c of cur_iWrt reaching end
            j = j0 + nVars

        #Now all terms are processed - concatenate tapes for wrtParams and add to resulting tape.
        for nTerms, dvtape, dctape in zip(dnterms, dvtapes, dctapes):
            result_vtape.extend([nTerms] + dvtape)
            result_ctape.extend(dctape)
        i = j  # update location in vtape after processing poly - actually could just use i instead of j it seems??

    return _np.array(result_vtape, _np.int64), _np.array(result_ctape, complex)
