"""Gradient descent
"""


import numpy as np

from frankenstein.tools.perf_utils import TIMER

from pyscf.lib import logger


""" Helper functions
"""
def get_gHp_fd(get_grad, p, order=1, eps=1.E-4):
    """ Compute gradient-Hessian product using finite difference

    Inps:
        get_grad (callable):
            grad(p) --> gradient given a direction p
        p (np.ndarray):
            initial gradient
        order (int, default=1):
            order 1 --> foward FD (err ~ O(eps))
            order 2 --> central FD (err ~ O(eps^2))
        eps (float, default=1.E-4):
            strength of perturbation
    """
    p_f = get_grad(eps*p)
    if order == 1:
        return 2. * (p_f-p) / eps
    elif order == 2:
        p_b = get_grad(-eps*p)
        return (p_f-p_b) / eps
    else:
        raise ValueError("Invalid order (must be 1 or 2)!")

        
# Newton-raphson (for debug)
class NR:
    def __init__(self, mf, eps=1.E-3, fd=2):
        self.verbose = mf.verbose
        self.stdout = mf.stdout
        self.comment = ""
        self.eps = eps
        self.fd = fd

        try:
            stdout = mf.stdout
        except:
            stdout = None
        self.timer = TIMER(4, stdout=stdout)

        self.iteration = 0

    def next_step(self, mf):
        f = mf.get_value_gdm()
        g = mf.get_grad_gdm()

        # build fd hessian
        def dphi(i, eps):
            mf.back_to_origin()
            mf.ov = np.zeros([mf.ov_size])
            mf.ov[i] = eps
            mf.update_all()
            mf.ov[i] = 0.

            return mf.get_grad_gdm()

        self.timer.start(0)
        mf.save_new_origin()
        H = np.zeros([mf.ov_size]*2)
        for i in range(mf.ov_size):
            if self.fd == 1:
                H[i] = (dphi(i,self.eps) - g) / self.eps
            elif self.fd == 2:
                H[i] = (dphi(i,self.eps) - dphi(i,-self.eps)) / (2.*self.eps)
            else:
                raise ValueError("fd must be 1 or 2.")
        mf.back_to_origin()
        self.timer.stop(0)

        # get raw NR step
        self.timer.start(1)
        lbd = 1.E-5
        du = -np.linalg.solve(H+lbd*np.eye(H.shape[1]), g)
        self.timer.stop(1)

        # line search
        fc = [0]
        def phi(alp):
            fc[0] += 1
            mf.back_to_origin()
            mf.ov = alp * mf.regularize_step_gdm(du)
            mf.update_all(skip_grad=True)

            return mf.get_value_gdm()

        self.timer.start(2)
        mf.save_new_origin()
        fold = f
        dphi0 = g @ du
        alp, fnew = scopt_linsrc.scalar_search_armijo(
            phi, fold, dphi0, c1=1.E-4, alpha0=1.)
        self.timer.stop(2)

        fc = fc[0]

        if alp is None:
            raise RuntimeError("Line search failed.")

        if fc == 1:
            self.comment = "NR"
        else:
            self.comment = "LnSr (%d,%.2f)"%(fc,alp)

        self.timer.start(3)
        mf.update_gdm()
        self.timer.stop(3)

        self.iteration += 1

    def report_timing(self):
        self.timer.report(tnames=["hess", "linsolve", "linsrch", "grad"])


# Direct minimization (for debug)
class DM:
    def __init__(self, mf, bounds=[-1,0], method="bf", plot=False):

        if method == "bf":
            self.alps = np.arange(*bounds, 0.05)
        elif method == "interpolate":
            self.amin = min(bounds)
            self.amax = max(bounds)
            self.ninter = 5
            self.neval = 100
        else:
            raise ValueError("Unknown method '%s'." % method)

        self.method = method
        self.plot = plot

        self.verbose = mf.verbose
        self.stdout = mf.stdout
        self.comment = ""

        try:
            stdout = mf.stdout
        except:
            stdout = None
        self.timer = TIMER(2, stdout=stdout)

        self.iteration = 0

    def next_step(self, mf):
        from scipy import interpolate as itplt
        from matplotlib import pyplot as plt

        g = mf.get_grad_gdm()

        def phi(alp):
            


        mf.save_new_origin()
        E0 = mf.get_value_gdm()

        self.timer.start(0)
        if self.method == "bf":
            alps = self.alps
            Es = np.asarray([phi(alp) for alp in alps]) - E0
        elif self.method == "interpolate":
            amin = self.amin
            amax = self.amax
            err_g = np.mean(g**2)**0.5
            if err_g > 1.E-3:
                xs = np.linspace(amin, amax, self.ninter)
                ys = np.asarray([phi(x) for x in xs])
                xyrep = itplt.splrep(xs, ys)
                fp = lambda x: itplt.splev(x, xyrep)
            else:
                xs = np.linspace(amin, amax, 3)
                ys = np.asarray([phi(x) for x in xs])
                p = np.polyfit(xs, ys, 2)
                fp = np.poly1d(p)
            alps = np.linspace(amin, amax, self.neval)
            Es = fp(alps)
        idmin = np.argmin(Es)
        alp = alps[idmin]
        E = Es[idmin]
        self.timer.stop(0)

        if self.plot:
            plt.plot(alps, Es, "-")
            if self.method == "interpolate": plt.plot(xs, ys, "o")
            plt.plot(alp, E, "rx")
            plt.show()

        self.comment = "alp = % .2f" % alp

        self.timer.start(1)
        mf.back_to_origin()
        mf.ov = alp * g
        mf.update_all()
        self.timer.stop(1)

        self.iteration += 1

    def report_timing(self):
        self.timer.report(["lnsrch", "update me"])


# Direct inversion of iterative subspace (DIIS)
from pyscf.lib.diis import DIIS as pyDIIS
class DIIS:
    def __init__(self, mf, ndiis=50, diis_start=1):

        self.adiis = pyDIIS()
        self.adiis.space = ndiis
        self.adiis.min_space = diis_start
        self.iteration = 0
        self.comment = ""

        try:
            stdout = mf.stdout
        except:
            stdout = None
        self.timer = TIMER(4, stdout=stdout)

    def next_step(self, mf):
        self.iteration += 1

        self.timer.start(0)
        f = mf.get_fock_diis()
        ferr = mf.get_err_diis()
        self.timer.stop(0)

        self.timer.start(1)
        f = self.adiis.update(f, ferr)
        self.timer.stop(1)

        self.timer.start(2)
        if hasattr(mf, "mom_start"):
            mom = self.iteration >= mf.mom_start
        else:
            mom = False
        comment = mf.update_diis(f, mom=mom)
        self.timer.stop(2)

        self.timer.start(3)
        mf.update_all()
        self.timer.stop(3)

        self.comment = "DIIS" if self.iteration > 0 else "Roothaan"
        self.comment += " %s" % comment

    def report_timing(self):
        self.timer.report(tnames=["diis prep", "diis extrap", "roothaan",
            "fock build"])

mf.back_to_origin()
            mf.ov = alp * g
            mf.update_all(skip_grad=True)
            mf.ov = np.zeros(mf.ov_size)

            return mf.get_value_gdm()