#cython: language_level=3
from libc.math cimport fabs, fmin
from libc.float cimport DBL_MAX

cimport numpy as cnp
import numpy as np

cimport cython

from qutip.core.cy.qobjevo cimport QobjEvo
from qutip.core.cy.coefficient cimport Coefficient
from qutip.core.data cimport Data, Dense, idxint
import qutip.core.data as _data
from qutip.core import Qobj
from scipy.linalg cimport cython_blas as blas


__all__ = ['SpectraCoefficient']


cdef class SpectraCoefficient(Coefficient):
    """
    Change a Coefficient with `t` dependence to one with `w` dependence to use
    in Bloch-Redfield tensor to allow array based coefficients to be used as
    spectral functions.
    If 2 coefficients are passed, the first one is the frequence response and
    the second is the time response.
    """
    def __init__(self, Coefficient coeff_w, Coefficient coeff_t=None, double w=0):
        self.coeff_t = coeff_t
        self.coeff_w = coeff_w
        self.w = w

    cdef complex _call(self, double t) except *:
        if self.coeff_t is None:
            return self.coeff_w(self.w)
        return self.coeff_t(t) * self.coeff_w(self.w)

    cpdef Coefficient copy(self):
        """Return a copy of the :obj:`.Coefficient`."""
        return SpectraCoefficient(self.coeff_t, self.coeff_w, self.w)

    def replace_arguments(self, _args=None, *, w=None, **kwargs):
        if _args:
            kwargs.update(_args)
        if kwargs:
            return SpectraCoefficient(
                self.coeff_w.replace(**kwargs),
                self.coeff_t.replace(**kwargs) if self.coeff_t else None,
                kwargs.get('w', w or self.w)
              )
        if w is not None:
            return SpectraCoefficient(self.coeff_w, self.coeff_t, w)
        return self


cdef Data _apply_trans(Data original, int trans):
    """helper function for matmul_var_data, apply transform."""
    cdef Data out
    if trans == 0:
        out = original
    elif trans == 1:
        out = original.transpose()
    elif trans == 2:
        out = original.conj()
    elif trans == 3:
        out = original.adjoint()
    return out


cdef char _fetch_trans_code(int trans):
    """helper function for matmul_var_Dense, fetch the blas flag byte"""
    if trans == 0:
        return b'N'
    elif trans == 1:
        return b'T'
    elif trans == 3:
        return b'C'


cpdef Data matmul_var_data(Data left, Data right,
                           int transleft, int transright):
    """
    matmul which input matrices can be transposed or adjoint.
    out = transleft(left) @ transright(right)

    with trans[left, right]:
        0 : Normal
        1 : Transpose
        2 : Conjugate
        3 : Adjoint
    """
    # Should this be supported in data.matmul?
    # A.dag() @ A is quite common.
    cdef int size
    cdef double complex beta
    cdef char left_code, right_code
    if (
        type(left) is Dense
        and type(right) is Dense
        and left.shape[0] == left.shape[1]
        and left.shape[0] == right.shape[0]
        and left.shape[0] == right.shape[1]
    ):
        return matmul_var_Dense(left, right, transleft, transright)
    left = _apply_trans(left, transleft)
    right = _apply_trans(right, transright)
    return _data.matmul(left, right)


cdef Dense matmul_var_Dense(Dense left, Dense right,
                            int transleft, int transright):
    """
    matmul which input matrices can be transposed or adjoint.
    out = transleft(left) @ transright(right)

    with trans[left, right]:
       0 : Normal
       1 : Transpose
       2 : Conjugate
       3 : Adjoint
    """
    # blas support matmul for normal, transpose, adjoint for fortran ordered
    # matrices.
    if not (
        left.shape[0] == left.shape[1]
        and left.shape[0] == right.shape[0]
        and left.shape[0] == right.shape[1]
    ):
        raise ValueError("Only implemented for square operators")
    cdef Dense out, a, b
    cdef double complex alpha = 1., beta = 0.
    cdef int tleft, tright, size = left.shape[0]

    # Since fortran to C equivalent to transpose, fix the codes to fortran
    tleft = transleft ^ (not left.fortran)
    tright = transright ^ (not right.fortran)

    # blas.zgemm does not support adjoint.
    if tleft + tright == 5:
        # adjoint and conjugate, we can't transpose the output to use zgemm
        out = matmul_var_data(left, right, transleft-2, transright-2)
        return out.conj()
    if tleft == 2 or tright == 2:
        # Need a conjugate, we compute the transpose of the desired results.
        # A.conj @ B^op -> (B^T^op @ A.dag)^T
        out = _data.dense.empty(left.shape[0], right.shape[1], False)
        a, b = right, left
        tleft, tright = tright ^ 1, tleft ^ 1
    else:
        out = _data.dense.empty(left.shape[0], right.shape[1], True)
        a, b = left, right

    left_code = _fetch_trans_code(tleft)
    right_code = _fetch_trans_code(tright)
    blas.zgemm(&left_code, &right_code, &size, &size, &size,
               &alpha, (<Dense> a).data, &size, (<Dense> b).data, &size,
               &beta, (<Dense> out).data, &size)

    return out


class _eigen_qevo:
    """
    Callable function to represent the eigenvectors of a QobjEvo at a time
    ``t``.
    """
    def __init__(self, qevo):
        self.qevo = QobjEvo(qevo)  # Force a copy
        self.args = None
        # This is a base conversion operator, the eigen basis part of the dims
        # are flat.
        self.out_dims = [qevo.dims[0], [qevo.shape[1]]]

    def __call__(self, t, args):
        if args is not self.args:
            self.args = args
            self.qevo.arguments(self.args)
        _, data = _data.eigs(self.qevo._call(t), True, True)
        return Qobj(data, copy=False, dims=self.out_dims)


cdef class _EigenBasisTransform:
    """
    For an hermitian operator, compute the eigenvalues and eigenstates and do
    the base change to and from that eigenbasis.

    parameter
    ---------
    oper : QobjEvo
        Hermitian operator for which to compute the eigenbasis.

    sparse : bool [False]
        Whether to use sparse solver for eigen decomposition.
    """
    def __init__(self, QobjEvo oper, bint sparse=False):
        if oper.dims[0] != oper.dims[1]:
            raise ValueError
        if type(oper(0).data) in (_data.CSR, _data.Dia) and not sparse:
            oper = oper.to(Dense)
        self.oper = oper
        self.isconstant = oper.isconstant
        self.size = oper.shape[0]

        if oper.isconstant:
            self._eigvals, self._evecs = _data.eigs(self.oper._call(0),
                                                    True, True)
        else:
            self._evecs = None
            self._eigvals = None

        self._t = np.nan
        self._evecs_inv = None

    def as_Qobj(self):
        """Make an Qobj or QobjEvo of the eigenvectors."""
        if self.isconstant:
            return Qobj(self.evecs(0), dims=self.oper.dims)
        else:
            return QobjEvo(_eigen_qevo(self.oper))

    cdef void _compute_eigen(self, double t) except *:
        if self._t != t and not self.isconstant:
            self._t = t
            self._evecs_inv = None
            self._eigvals, self._evecs = _data.eigs(self.oper._call(t), True, True)

    cpdef object eigenvalues(self, double t):
        """
        Return the eigenvalues at ``t``.
        """
        self._compute_eigen(t)
        return self._eigvals

    cpdef Data evecs(self, double t):
        """
        Return the eigenstates at ``t``.
        """
        self._compute_eigen(t)
        return self._evecs

    cdef Data _inv(self, double t):
        if self._evecs_inv is None:
            self._evecs_inv = self.evecs(t).adjoint()
        return self._evecs_inv

    cdef Data _S_converter(self, double t):
        return _data.kron_transpose(self.evecs(t), self._inv(t))

    cdef Data _S_converter_inverse(self, double t):
        return _data.kron_transpose(self._inv(t), self.evecs(t))

    cpdef Data to_eigbasis(self, double t, Data fock):
        """
        Do the transformation of the :cls:`Qobj` ``fock`` to the basis where
        ``oper(t)`` is diagonalized.
        """
        # For Hermitian operator, the inverse of evecs is the adjoint matrix.
        # Blas include A.dag @ B in one operation. We use it if we can so we
        # don't make unneeded copy of evecs.
        cdef Data temp
        if fock.shape[0] == self.size and fock.shape[1] == 1:
            return matmul_var_data(self.evecs(t), fock, 3, 0)

        elif fock.shape[0] == self.size**2 and fock.shape[1] == 1:
            if type(fock) is Dense and (<Dense> fock).fortran:
                fock = _data.column_unstack_dense(fock, self.size, True)
                temp = _data.matmul(matmul_var_data(self.evecs(t), fock, 3, 0),
                                    self.evecs(t))
                fock = _data.column_stack_dense(fock, True)
            else:
                fock = _data.column_unstack(fock, self.size)
                temp = _data.matmul(matmul_var_data(self.evecs(t), fock, 3, 0),
                                    self.evecs(t))
            if type(temp) is Dense:
                return _data.column_stack_dense(temp, True)
            return _data.column_stack(temp)

        elif fock.shape[0] == self.size and fock.shape[0] == fock.shape[1]:
            return _data.matmul(matmul_var_data(self.evecs(t), fock, 3, 0),
                                self.evecs(t))

        elif fock.shape[0] == self.size**2 and fock.shape[0] == fock.shape[1]:
            temp = self._S_converter_inverse(t)
            return _data.matmul(matmul_var_data(temp, fock, 3, 0), temp)

        raise ValueError("Could not convert the Qobj's data to eigenbasis: "
                         "can't guess type from shape.")

    cpdef Data from_eigbasis(self, double t, Data eig):
        """
        Do the transformation of the :cls:`Qobj` ``eig`` in the basis where
        ``oper(t)`` is diagonalized to the outside basis.
        """
        cdef Data temp
        if eig.shape[0] == self.size and eig.shape[1] == 1:
            return _data.matmul(self.evecs(t), eig)

        elif eig.shape[0] == self.size**2 and eig.shape[1] == 1:
            if type(eig) is Dense and (<Dense> eig).fortran:
                eig = _data.column_unstack_dense(eig, self.size, True)
                temp = matmul_var_data(_data.matmul(self.evecs(t), eig),
                                  self.evecs(t), 0, 3)
                eig = _data.column_stack_dense(eig, True)
            else:
                eig = _data.column_unstack(eig, self.size)
                temp = matmul_var_data(_data.matmul(self.evecs(t), eig),
                                  self.evecs(t), 0, 3)
            if type(temp) is Dense:
                return _data.column_stack_dense(temp, True)
            return _data.column_stack(temp)

        elif eig.shape[0] == self.size and eig.shape[0] == eig.shape[1]:
            temp = self.evecs(t)
            return matmul_var_data(_data.matmul(temp, eig), temp, 0, 3)

        elif eig.shape[0] == self.size**2 and eig.shape[0] == eig.shape[1]:
            temp = self._S_converter_inverse(t)
            return _data.matmul(temp, matmul_var_data(eig, temp, 0, 3))

        raise ValueError("Could not convert the Qobj's data from eigenbasis: "
                         "can't guess type from shape.")
