"""
DIIS!!!
"""


import numpy as np


def update_B(mdiis):
    """Update the DIIS B matrix, B = np.outer(rs, rs), and increment
    self.ndiis by one.
    """
    n = mdiis.ndiis - mdiis.diis_start
    if n == 0:
        mdiis.B[n,n] = mdiis.rs[n] @ mdiis.rs[n]
    elif 0 < n < mdiis.max_diis:
        b = mdiis.rs[:n] @ mdiis.rs[n]
        mdiis.B[n,:n] = b
        mdiis.B[:n,n] = b
        mdiis.B[n,n] = mdiis.rs[n] @ mdiis.rs[n]
    else:
        m = n % mdiis.max_diis
        b = mdiis.rs @ mdiis.rs[m]
        mdiis.B[m,:] = b
        mdiis.B[:,m] = b
        mdiis.B[m,m] = mdiis.rs[m] @ mdiis.rs[m]


def diis_update(mdiis):
    """Solve the DIIS equation, Ac = b for the linear combination
    coefficients that is then used to make the new Fock matrix.
    """
    n = mdiis.ndiis - mdiis.diis_start
    if n == 0:
        c = np.array([1.])
        mdiis.comment = "FP"
    else:
        if n >= mdiis.max_diis:
            n = mdiis.max_diis
        A = np.zeros([n+1,n+1])
        A[:n,:n] = mdiis.B[:n,:n]
        A[:,n] = A[n,:] = -np.ones(n+1)
        A[n,n] = 0.

        b = np.zeros(n+1)
        b[-1] = -1.

        try:
            c = np.linalg.solve(A, b)
        except np.linalg.LinAlgError:
            c = np.linalg.pinv(A, b)

        mdiis.comment = "DIIS"

    mdiis.c = c


class DIIS:

    def __init__(self, mf, max_diis=20, diis_start=0):
        """Constructor for DIIS optimizer

        Inp:
            mf (scf.RHF/UHF/USSCF etc.):
                A scf object
            max_diis (int, optional):
                Maximal number of Fock and DIIS error matrices stored.
        """
        self.mf = mf
        self.max_diis = max_diis
        self.diis_start = diis_start
        self.tol = 10**-self.mf.conv

        self.ndiis = 0
        fac = 2 if self.mf.unrestricted else 1
        self.Fs = np.zeros([self.max_diis, self.mf.nmo**2*fac])
        self.rs = np.zeros([self.max_diis, self.mf.nmo**2*fac])
        self.B = np.zeros([self.max_diis, self.max_diis])
        self.c = None

        self.err = float("inf")
        self.comment = ""

    def next_step(self):
        """Given current Fock matrix and DIIS error matrix, as well as a list
        of up to self.max_diis previous such matrices, compute the new Fock
        matrix as a linear combination that aims at minimizing the DIIS error.
        """
        self.update_diis_vecs()
        if self.err < self.tol and self.ndiis > 0:
            # if self.ndiis == 0:
            #     self.mf.Roothaan_step()
            #     self.mf.update()
            return True

        self.update_B()
        self.diis_update()

        self.ndiis += 1

        self.update_fock()
        self.mf.Roothaan_step()
        self.mf.update()

        return False

    def update_diis_vecs(self):
        """Push current Fock matrix and DIIS error matrix to the list. Note that
        earlest matrices will be replaced if self.ndiis exceeds self.max_diis.
        """
        rnew = self.mf.get_diis_errvec()
        Fnew = self.mf.get_diis_fockvec()

        n = self.ndiis % self.max_diis
        self.rs[n] = rnew
        self.Fs[n] = Fnew

        self.err = np.mean(rnew**2.)**0.5

    update_B = update_B
    diis_update = diis_update

    def update_fock(self):
        n = self.ndiis if self.ndiis < self.max_diis else self.max_diis
        Fnew = self.c[:n] @ self.Fs[:n]

        nmo = self.mf.nmo
        if self.mf.unrestricted:
            Fnew = Fnew.reshape(2*nmo,nmo)
            self.mf.fock = [Fnew[:nmo,:], Fnew[nmo:,:]]
        else:
            self.mf.fock = Fnew.reshape(nmo, nmo)


class DIIS_gen:

    def __init__(self, prob, max_diis=20, diis_start=0):
        """
            get_diis_errvec
            get_diis_fockvec
            update_fock
        """
        self.prob = prob
        self.max_diis = max_diis
        self.diis_start = diis_start

        self.ndiis = 0
        self.Fs = None
        self.rs = None
        self.B = np.zeros([self.max_diis, self.max_diis])
        self.c = None

        self.comment = ""

    def next_step(self):
        if self.ndiis >= self.diis_start:
            self.update_diis_vecs()
            self.update_B()
            self.diis_update()
            self.update_fock()

        self.ndiis += 1

        self.prob.next_step()

        return self.prob.check_conv()

    def update_diis_vecs(self):
        rnew = self.prob.get_diis_errvec()
        Fnew = self.prob.get_diis_fockvec()
        self.err = np.mean(rnew**2.)**0.5

        if self.rs is None: self.rs = np.zeros((self.max_diis, *rnew.shape))
        if self.Fs is None: self.Fs = np.zeros((self.max_diis, *Fnew.shape))

        n = (self.ndiis-self.diis_start) % self.max_diis
        self.rs[n] = rnew
        self.Fs[n] = Fnew

    update_B = update_B
    diis_update = diis_update

    def update_fock(self):
        n = self.ndiis if self.ndiis < self.max_diis else self.max_diis
        Fnew = self.c[:n] @ self.Fs[:n]

        if self.prob.verbose > 1:
            print("  DIIS update with c: " +
                " ".join(["{: .4f}".format(cc) for cc in self.c]), flush=True)
            print("", flush=True)

        self.prob.update_fock(Fnew)
