"""
pdeco.py

Functions associated with PDE-Constrained Optimization.
"""

import numpy as np
from opt_einsum import contract
from scipy.optimize import minimize

class PDECO():
    """
    Performs Optimization as in: 10.1063/1.1535422 - Qin Wu + Weitao Yang

    Attributes:
    -----------
    lambda_rgl: {None, float}. If float, lambda-regularization is added with lambda=lambda_rgl.
    """

    regul_norm = None  # Regularization norm: ||v||^2
    lambda_reg = None  # Regularization constant

    def pdeco(self, opt_max_iter, reg=None, gtol=1e-3,
              opt_method='L-BFGS-B', opt=None):
        """
        Calls scipy minimizer to minimize lagrangian. 
        """
        self.lambda_reg = reg

        self.lambda_reg = reg
        if opt is None:
            opt = {"disp": False}
        opt['maxiter'] = opt_max_iter
        opt['gtol'] = gtol

        # Initialization for D and C
        self._diagonalize_with_potential_pbs(self.v_pbs)

        if self.S4 is None:
            self.S4 = self.eng.get_S4()

        if opt_method.lower() == 'bfgs' or opt_method.lower() == 'l-bfgs-b':
            opt_results = minimize( fun = self.lagrangian_pbeco,
                                    x0  = self.v_pbs, 
                                    jac = self.gradient_pbeco,
                                    method  = opt_method,
                                    options = opt
                                    )
        else:
            raise ValueError(F'{opt_method} is not supported. Only BFGS '
                             F'and L-BFGS is supported.')

        if opt_results.success == False:
            self.v_pbs = opt_results.x
            self.opt_info = opt_results
            raise ValueError("Optimization was unsucessful (|grad|=%.2e) within %i iterations, "
                             "try a different initial guess. %s"% (np.linalg.norm(opt_results.jac), opt_results.nit, opt_results.message)
                             )
        else:
            print(f"Optimization Successful within {opt_results.nit} iterations! |grad|={np.linalg.norm(opt_results.jac):.2e}." )
            self.v_pbs = opt_results.x
            self.opt_info = opt_results

    def lagrangian_pbeco(self, v):
        """
        Lagrangian to be minimized wrt external potential
        Equation (5) of main reference
        """

        # If v is not updated, will not re-calculate.
        if not np.allclose(v, self.v_pbs, atol=1e-15):
            self._diagonalize_with_potential_pbs(v)

        # self._diagonalize_with_potential_pbs(v)

        if self.ref == 1:
            L = 4 * contract("ijkl,ij,kl", self.S4, self.Da - self.Dt[0], self.Da- self.Dt[0])
        else:
            L = contract("ijkl,ij,kl", self.S4, self.Da+self.Db-self.Dt[0]-self.Dt[1], self.Da+self.Db-self.Dt[0]-self.Dt[1])
        # Add lambda-regularization
        if self.lambda_reg is not None:
            T = self.T_pbs
            if self.ref == 1:
                norm = 2 * (v[:self.npbs] @ T @ v[:self.npbs])
            else:
                norm = (v[self.npbs:] @ T @ v[self.npbs:]) + (v[:self.npbs] @ T @ v[:self.npbs])

            L += norm * self.lambda_reg
            self.regul_norm = norm
        return L

    def gradient_pbeco(self, v):
        """
        Calculates gradient wrt target density
        Equation (11) of main reference
        """
        # If v is not updated, will not re-calculate.
        if not np.allclose(v, self.v_pbs, atol=1e-15):
            self._diagonalize_with_potential_pbs(v)

        if self.ref == 1:
            grad_temp = np.zeros((self.nbf, self.nbf))
            g = 8 * contract("ijkl,ij,km->lm", self.S4, 2 * (self.Dt[0] - self.Da), self.Coca)  # shape (ao, mo)
            u = 0.5 * contract("lm,lm->m", self.Coca, g)  # shape (mo, )
            g -= 2 * contract('m,ij,jm->im', u, self.S2, self.Coca) # shape (ao, mo)
            for idx in range(self.nalpha):
                LHS = self.Fock - self.S2 * self.eigvecs_a[idx]
                p_i = np.linalg.solve(LHS, g[:, idx])
                # Gram–Schmidt rotation
                p_i -= np.sum(p_i * np.dot(self.S2, self.Coca[:,idx])) * self.Coca[:,idx]
                assert np.allclose([np.sum(p_i * (self.S2 @ self.Coca[:,idx])), np.linalg.norm(np.dot(LHS,p_i)-g[:, idx]), np.sum(g[:, idx]*self.Coca[:,idx])], 0, atol=1e-4)
                grad_temp += p_i[:, np.newaxis] * self.Coca[:,idx]

            self.grad = contract("ij,ijk->k", grad_temp, self.S3)
        else:
            grad_temp_a = np.zeros((self.nbf, self.nbf))
            g_a = 4 * contract("ijkl,ij,km->lm", self.S4, (self.Dt[0] - self.Da) + (self.Dt[1] - self.Db), self.Coca)  # shape (ao, mo)
            u_a = 0.5 * contract("lm,lm->m", self.Coca, g_a)  # shape (mo, )
            g_a -= 2 * contract('m,ij,jm->im', u_a, self.S2, self.Coca) # shape (ao, mo)
            for idx in range(self.nalpha):
                LHS = self.Fock[0] - self.S2 * self.eigvecs_a[idx]
                p_i = np.linalg.solve(LHS, g_a[:, idx])
                # Gram–Schmidt rotation
                p_i -= np.sum(p_i * np.dot(self.S2, self.Coca[:,idx])) * self.Coca[:,idx]
                assert np.allclose([np.sum(p_i * (self.S2 @ self.Coca[:,idx])), np.linalg.norm(np.dot(LHS,p_i)-g_a[:, idx]), np.sum(g_a[:, idx]*self.Coca[:,idx])], 0, atol=1e-4)
                grad_temp_a += p_i[:, np.newaxis] * self.Coca[:,idx]

            grad_temp_b = np.zeros((self.nbf, self.nbf))
            g_b = 4 * contract("ijkl,ij,km->lm", self.S4, (self.Dt[0] - self.Da) + (self.Dt[1] - self.Db), self.Cocb)  # shape (ao, mo)
            u_b = 0.5 * contract("lm,lm->m", self.Cocb, g_b)  # shape (mo, )
            g_b -= 2 * contract('m,ij,jm->im', u_b, self.S2, self.Cocb) # shape (ao, mo)
            for idx in range(self.nbeta):
                LHS = self.Fock[1] - self.S2 * self.eigvecs_b[idx]
                p_i = np.linalg.solve(LHS, g_b[:, idx])
                # Gram–Schmidt rotation
                p_i -= np.sum(p_i * np.dot(self.S2, self.Cocb[:,idx])) * self.Cocb[:,idx]
                assert np.allclose([np.sum(p_i * (self.S2 @ self.Cocb[:,idx])), np.linalg.norm(np.dot(LHS,p_i)-g_b[:, idx]), np.sum(g_b[:, idx]*self.Cocb[:,idx])], 0, atol=1e-4)
                grad_temp_b += p_i[:, np.newaxis] * self.Cocb[:,idx]

            self.grad = np.concatenate((contract("ij,ijk->k", grad_temp_a, self.S3), contract("ij,ijk->k", grad_temp_b, self.S3)))

        if self.lambda_reg is not None:
            T = self.T_pbs
            if self.ref == 1:
                rgl_vector = 4 * self.lambda_reg*np.dot(T, v[:self.npbs])
                self.grad += rgl_vector
            else:
                self.grad[:self.npbs] += 2 * self.lambda_reg*np.dot(T, v[:self.npbs])
                self.grad[self.npbs:] += 2 * self.lambda_reg*np.dot(T, v[self.npbs:])
        return self.grad

    def find_regularization_constant_pdeco(self, opt_max_iter, opt_method="L-BFGS-B", gtol=1e-3,
                                     opt=None, lambda_list=None):
        """
        Finding regularization constant lambda.

        Note: it is recommend to set a specific convergence criteria by opt or tol,
                in order to control the same convergence
                for different lambda value.

        After the calculation is done, one can plot the returns to select a good lambda.

        Parameters:
        -----------
        guide_potential_components: a list of string
            the components for guide potential v_pbs.
            see Inverter.generate_components() for details.

        opt_method: string default: "trust-krylov"
            opt_methods available in scipy.optimize.minimize

        tol: float
            Tolerance for termination. See scipy.optimize.minimize for details.

        opt: dictionary, optional
            if given:
                scipy.optimize.minimize(method=opt_method, options=opt)

        lambda_list: np.ndarray, optional
            A array of lambda to search; otherwise, it will be 10 ** np.linspace(-1, -7, 7).

        Returns:
        --------
        lambda_list: np.ndarray
            A array of lambda searched.

        P_list: np.ndarray
            The value defined by [Bulat, Heaton-Burgess, Cohen, Yang 2007] eqn (21).
            Corresponding to lambda in lambda_list.

        error_list: np.ndarray
            The Ts value for each lambda.


        """

        error_list = []
        v_norm_list = []


        if lambda_list is None:
            lambda_list = 10 ** np.linspace(-3, -9, 7)

        if opt is None:
            opt = {"disp"    : False}
        opt['maxiter'] = opt_max_iter
        opt['gtol'] = gtol

        self.lambda_reg = None
        # Initial calculation with no regularization
        # Initialization for D and C
        self._diagonalize_with_potential_pbs(self.v_pbs)
        if opt_method.lower() == 'bfgs' or opt_method.lower() == 'l-bfgs-b':
            initial_result = minimize(fun=self.lagrangian_pbeco,
                                   x0=self.v_pbs,
                                   jac=self.gradient_pbeco,
                                   method=opt_method,
                                   options=opt
                                   )
        else:
            raise ValueError(F'{opt_method} is not supported. Only BFGS '
                             F'and L-BFGS is supported.')

        if initial_result.success == False:
            raise ValueError("Optimization was unsucessful (|grad|=%.2e) within %i iterations, "
                             "try a different intitial guess"% (np.linalg.norm(initial_result.jac), initial_result.nit)
                             + initial_result.message)
        else:
            error0 = -initial_result.fun
            initial_v_pbs = initial_result.x  # This is used as the initial guess for with regularization calculation.

        for reg in lambda_list:
            self.lambda_reg = reg

            if opt_method.lower() == 'bfgs' or opt_method.lower() == 'l-bfgs-b':
                opt_results = minimize(fun=self.lagrangian_pbeco,
                                       x0=initial_v_pbs,
                                       jac=self.gradient_pbeco,
                                       method=opt_method,
                                       options=opt
                                       )

            else:
                raise ValueError(F'{opt_method} is not supported. Only BFGS '
                                 F'and L-BFGS is supported.')


            v_norm_list.append(self.regul_norm)
            error_list.append(opt_results.fun - self.lambda_reg * self.regul_norm)

        P_list = lambda_list * np.array(v_norm_list) / (np.array(error_list) - error0)

        return lambda_list, P_list, np.array(error_list)
