import numpy as np
import tensorflow as tf
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_shape
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import check_ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import gen_linalg_ops
from tensorflow.python.ops import linalg_ops
from tensorflow.python.ops import map_fn
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import special_math_ops
from tensorflow.python.ops import stateless_random_ops
from tensorflow.python.util import dispatch
from tensorflow.python.util.tf_export import tf_export
import scipy.linalg
import scipy.sparse.linalg
from scipy.sparse.linalg import aslinearoperator


class LazyOperatorNormInfo:
    """
    Information about an operator is lazily computed.
    The information includes the exact 1-norm of the operator,
    in addition to estimates of 1-norms of powers of the operator.
    This uses the notation of Computing the Action (2011).
    This class is specialized enough to probably not be of general interest
    outside of this module.
    """
    def __init__(self, A, A_1_norm=None, ell=2, scale=1):
        """
        Provide the operator and some norm-related information.
        Parameters
        ----------
        A : linear operator
            The operator of interest.
        A_1_norm : float, optional
            The exact 1-norm of A.
        ell : int, optional
            A technical parameter controlling norm estimation quality.
        scale : int, optional
            If specified, return the norms of scale*A instead of A.
        """
        self._A = A
        self._A_1_norm = A_1_norm
        self._ell = ell
        self._d = np.zeros(35)
        self._scale = scale

    def set_scale(self,scale):
        """
        Set the scale parameter.
        """
        self._scale = scale

    def onenorm(self):
        """
        Compute the exact 1-norm.
        """
        if self._A_1_norm is None:
            self._A_1_norm = tf.norm(self._A, ord=1)
        return self._scale*self._A_1_norm

    def d(self, p):
        """
        Lazily estimate d_p(A) ~= || A^p ||^(1/p) where ||.|| is the 1-norm.
        """
        if self._d[p] == 0:
            self._d[p] = _onenormest_matrix_power(self._A, p, self._ell)
        return self._scale * self._d[p] ** (1.0 / float(p))

    def alpha(self, p):
        """
        Lazily compute max(d(p), d(p+1)).
        """
        d1 = self.d(p)
        d2 = self.d(p+1)
        return tf.where(d1 > d2, d1, d2)

def _onenormest_matrix_power(A, p,
        t=2, itmax=5, compute_v=False, compute_w=False):
    """
    Efficiently estimate the 1-norm of A^p.
    Parameters
    ----------
    A : ndarray
        Matrix whose 1-norm of a power is to be computed.
    p : int
        Non-negative integer power.
    t : int, optional
        A positive parameter controlling the tradeoff between
        accuracy versus time and memory usage.
        Larger values take longer and use more memory
        but give more accurate output.
    itmax : int, optional
        Use at most this many iterations.
    compute_v : bool, optional
        Request a norm-maximizing linear operator input vector if True.
    compute_w : bool, optional
        Request a norm-maximizing linear operator output vector if True.
    Returns
    -------
    est : float
        An underestimate of the 1-norm of the sparse matrix.
    v : ndarray, optional
        The vector such that ||Av||_1 == est*||v||_1.
        It can be thought of as an input to the linear operator
        that gives an output with particularly large norm.
    w : ndarray, optional
        The vector Av which has relatively large 1-norm.
        It can be thought of as an output of the linear operator
        that is relatively large in norm compared to the input.
    """
    #XXX Eventually turn this into an API function in the  _onenormest module,
    #XXX and remove its underscore,
    #XXX but wait until expm_multiply goes into scipy.
    A1 = A
    for i in range(p-1):
        A1 = A*A1
    return float(tf.norm(A1, ord=1))

# This table helps to compute bounds.
# They seem to have been difficult to calculate, involving symbolic
# manipulation of equations, followed by numerical root finding.
_theta = {
        # The first 30 values are from table A.3 of Computing Matrix Functions.
        1: 2.29e-16,
        2: 2.58e-8,
        3: 1.39e-5,
        4: 3.40e-4,
        5: 2.40e-3,
        6: 9.07e-3,
        7: 2.38e-2,
        8: 5.00e-2,
        9: 8.96e-2,
        10: 1.44e-1,
        # 11
        11: 2.14e-1,
        12: 3.00e-1,
        13: 4.00e-1,
        14: 5.14e-1,
        15: 6.41e-1,
        16: 7.81e-1,
        17: 9.31e-1,
        18: 1.09,
        19: 1.26,
        20: 1.44,
        # 21
        21: 1.62,
        22: 1.82,
        23: 2.01,
        24: 2.22,
        25: 2.43,
        26: 2.64,
        27: 2.86,
        28: 3.08,
        29: 3.31,
        30: 3.54,
        # The rest are from table 3.1 of
        # Computing the Action of the Matrix Exponential.
        35: 4.7,
        40: 6.0,
        45: 7.2,
        50: 8.5,
        55: 9.9,
        }


_theta_keys = tf.convert_to_tensor([
        # The first 30 values are from table A.3 of Computing Matrix Functions.
        1,
        2,
        3,
        4,
        5,
        6,
        7,
        8,
        9,
        10,
        # 11
        11,
        12,
        13,
        14,
        15,
        16,
        17,
        18,
        19,
        20,
        # 21
        21,
        22,
        23,
        24,
        25,
        26,
        27,
        28,
        29,
        30,
        # The rest are from table 3.1 of
        # Computing the Action of the Matrix Exponential.
        35,
        40,
        45,
        50,
        55,
        ])


_theta_values = tf.convert_to_tensor([
        # The first 30 values are from table A.3 of Computing Matrix Functions.
        2.29e-16,
        2.58e-8,
        1.39e-5,
        3.40e-4,
        2.40e-3,
        9.07e-3,
        2.38e-2,
        5.00e-2,
        8.96e-2,
        1.44e-1,
        # 11
        2.14e-1,
        3.00e-1,
        4.00e-1,
        5.14e-1,
        6.41e-1,
        7.81e-1,
        9.31e-1,
        1.09,
        1.26,
        1.44,
        # 21
        1.62,
        1.82,
        2.01,
        2.22,
        2.43,
        2.64,
        2.86,
        3.08,
        3.31,
        3.54,
        # The rest are from table 3.1 of
        # Computing the Action of the Matrix Exponential.
        4.7,
        6.0,
        7.2,
        8.5,
        9.9,
        ], dtype=tf.dtypes.float64)



def _onenormest_matrix_power(A, p,
        t=2, itmax=5, compute_v=False, compute_w=False):
    """
    Efficiently estimate the 1-norm of A^p.
    Parameters
    ----------
    A : ndarray
        Matrix whose 1-norm of a power is to be computed.
    p : int
        Non-negative integer power.
    t : int, optional
        A positive parameter controlling the tradeoff between
        accuracy versus time and memory usage.
        Larger values take longer and use more memory
        but give more accurate output.
    itmax : int, optional
        Use at most this many iterations.
    compute_v : bool, optional
        Request a norm-maximizing linear operator input vector if True.
    compute_w : bool, optional
        Request a norm-maximizing linear operator output vector if True.
    Returns
    -------
    est : float
        An underestimate of the 1-norm of the sparse matrix.
    v : ndarray, optional
        The vector such that ||Av||_1 == est*||v||_1.
        It can be thought of as an input to the linear operator
        that gives an output with particularly large norm.
    w : ndarray, optional
        The vector Av which has relatively large 1-norm.
        It can be thought of as an output of the linear operator
        that is relatively large in norm compared to the input.
    """
    #XXX Eventually turn this into an API function in the  _onenormest module,
    #XXX and remove its underscore,
    #XXX but wait until expm_multiply goes into scipy.
    A1 = A
    for i in range(p-1):
        A1 = A*A1
    return float(tf.norm(A1, ord=1))


def _expm_multiply_est(A, B, t=1.0, balance=False):
    if balance:
            raise NotImplementedError
    if len(A.shape) != 2 or A.shape[0] != A.shape[1]:
        raise ValueError('expected A to be like a square matrix')
    if A.shape[1] != B.shape[0]:
        raise ValueError('shapes of matrices A {} and B {} are incompatible'
                         .format(A.shape, B.shape))
    
    n = A.shape[0]
    ident = tf.eye(n, n, dtype=A.dtype)
    
    if len(B.shape) == 1:
        n0 = 1
    elif len(B.shape) == 2:
        n0 = B.shape[1]
    else:
        raise ValueError('expected B to be like a matrix or a vector')
    
    u_d = 2**-53
    tol = u_d
    mu = tf.linalg.trace(A) / tf.cast(n, dtype=tf.dtypes.complex128)
    A = A - mu * ident
    A_1_norm = tf.norm(A, ord=1)
    if t*A_1_norm == 0:
        m_star, s = 0, 1
    else:
        ell = 2
        norm_info = LazyOperatorNormInfo(t*A, A_1_norm=t*A_1_norm, ell=ell)
        m_star, s = _fragment_3_1(norm_info, n0, tol, ell=ell)
    return m_star, s


def _fragment_3_1(norm_info, n0, tol, m_max=55, ell=2):
    """
    A helper function for the _expm_multiply_* functions.
    Parameters
    ----------
    norm_info : LazyOperatorNormInfo
        Information about norms of certain linear operators of interest.
    n0 : int
        Number of columns in the _expm_multiply_* B matrix.
    tol : float
        Expected to be
        :math:`2^{-24}` for single precision or
        :math:`2^{-53}` for double precision.
    m_max : int
        A value related to a bound.
    ell : int
        The number of columns used in the 1-norm approximation.
        This is usually taken to be small, maybe between 1 and 5.
    Returns
    -------
    best_m : int
        Related to bounds for error control.
    best_s : int
        Amount of scaling.
    Notes
    -----
    This is code fragment (3.1) in Al-Mohy and Higham (2011).
    The discussion of default values for m_max and ell
    is given between the definitions of equation (3.11)
    and the definition of equation (3.12).
    """
    if ell < 1:
        raise ValueError('expected ell to be a positive integer')
    best_m = tf.constant(np.iinfo(np.int32).max)
    best_s = tf.constant(np.iinfo(np.int32).max)
    if _condition_3_13(norm_info.onenorm(), n0, m_max, ell):
        for m, theta in _theta.items():
            s = int(tf.math.ceil(norm_info.onenorm() / theta))
            if best_m == np.iinfo(np.int32).max or m * s < best_m * best_s:
                best_m = m
                best_s = s
    else:
        # Equation (3.11).
        for p in range(2, _compute_p_max(m_max) + 1):
            z = tf.zeros(tf.shape(_theta_keys), dtype=tf.dtypes.int32)
            ms = tf.where(p*(p-1) <= _theta_keys, _theta_keys, z)
            ms = tf.where( _theta_keys <= m_max+1, ms, z)

            z = tf.zeros(tf.shape(_theta_keys), dtype=tf.dtypes.float64)
            vs = tf.where(p*(p-1) <= _theta_keys, _theta_values, z)
            vs = tf.where( _theta_keys <= m_max+1, vs, z)
            for i in range(tf.shape(_theta_keys)[0]):
                if tf.math.equal(ms[i], 0):
                    continue
                else:
                    s = _compute_cost_div_m(vs[i], p, norm_info)
                    if best_m == np.iinfo(np.int32).max or ms[i] * s < tf.math.abs(best_m * best_s):
                        best_m = ms[i]
                        best_s = s
        best_s = tf.where(best_s > 1, best_s, 1)
    return best_m, best_s


def _condition_3_13(A_1_norm, n0, m_max, ell):
    """
    A helper function for the _expm_multiply_* functions.
    Parameters
    ----------
    A_1_norm : float
        The precomputed 1-norm of A.
    n0 : int
        Number of columns in the _expm_multiply_* B matrix.
    m_max : int
        A value related to a bound.
    ell : int
        The number of columns used in the 1-norm approximation.
        This is usually taken to be small, maybe between 1 and 5.
    Returns
    -------
    value : bool
        Indicates whether or not the condition has been met.
    Notes
    -----
    This is condition (3.13) in Al-Mohy and Higham (2011).
    """

    # This is the rhs of equation (3.12).
    p_max = _compute_p_max(m_max)
    a = 2 * ell * p_max * (p_max + 3)

    # Evaluate the condition (3.13).
    b = _theta[m_max] / float(n0 * m_max)
    return float(A_1_norm) <= float(a) * b

def _compute_cost_div_m(m, p, norm_info):
    """
    A helper function for computing bounds.
    This is equation (3.10).
    It measures cost in terms of the number of required matrix products.
    Parameters
    ----------
    m : int
        A valid key of _theta.
    p : int
        A matrix power.
    norm_info : LazyOperatorNormInfo
        Information about 1-norms of related operators.
    Returns
    -------
    cost_div_m : int
        Required number of matrix products divided by m.
    """

    return int(tf.math.ceil(norm_info.alpha(p) / m))

def _compute_p_max(m_max):
    """
    Compute the largest positive integer p such that p*(p-1) <= m_max + 1.
    Do this in a slightly dumb way, but safe and not too slow.
    Parameters
    ----------
    m_max : int
        A count related to bounds.
    """
    sqrt_m_max = np.math.sqrt(float(m_max))
    p_low = np.math.floor(sqrt_m_max)
    p_high = np.math.ceil(sqrt_m_max + 1)
    r = tf.range(p_low, p_high+1, 1, dtype=tf.dtypes.float64)
    r = tf.where(r*(r-1) <= m_max + 1, r, tf.zeros(tf.shape(r),dtype=tf.dtypes.float64))
    return int(tf.reduce_max(r))

@tf.function
def _expm_multiply_simple_core(A, B, t, m_star, s, tol=None, balance=False):
    """
    A helper function.
    """
    if balance:
        raise NotImplementedError
    if tol is None:
        u_d = 2 ** -53
        tol = u_d
        
    n = A.shape[0]
    F = B
    mu = tf.linalg.trace(A) / tf.cast(n, dtype=tf.dtypes.complex128)
    tc = tf.cast(t, dtype=tf.dtypes.complex128)
    sc = tf.cast(s, dtype=tf.dtypes.complex128)
    eta = tf.math.exp(tc*mu/sc)
    for i in range(s):
        for j in range(m_star):
            coeff = tf.cast(tf.math.divide(float(t) , float(s*(j+1))), dtype=tf.dtypes.complex128)
            B = tf.math.scalar_mul(coeff, tf.tensordot(A, B, 1))
            F = F + B
            # if c1 + c2 <= tol * tf.norm(F, ord=np.inf):
            #     break
            # c1 = c2
        F = eta * F
        B = F
    return F