# coding=utf-8
# Copyright 2019 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Non-interacting Solver
######################

**Summary** 
    Solve non-interacting 1d system numerically on a grid
    (position-space basis). Each eigenstate will be occupied by one electron.

.. moduleauthor::
    `Ryan Pederson <pedersor@uci.edu>`_ ORCID: `0000-0002-7228-9478 <https://orcid.org/0000-0002-7228-9478>`_,
    `Chris (Jielun) Chen`,
    `Li Li`

.. note::
    Both solvers (EigenSolver, SparseEigenSolver) here are based on directly
    diagonalizing the Hamiltonian matrix.

.. todo::

    * Figure out what to do about joint copyright holders (Google + other)
    * Clean out unused example rst content in here.

"""

import numpy as np
from scipy import sparse
from scipy.sparse import linalg
from utils import get_dx, quadratic


class SolverBase:
    """Base Solver for non-interacting 1d system. Subclasses should define
    `solve_ground_state` method.
    """

    def __init__(self,
                 grids,
                 potential_fn=None,
                 num_electrons=1,
                 k_point=None,
                 boundary_condition='open',
                 n_point_stencil=3,
                 perturbation=None):
        """Initialize the solver with potential function and grid.

        Args:
            grids: numpy array of grid points for evaluating 1d potential.
                Shape: (num_grids,)
            potential_fn: potential function taking grids as argument.
            num_electrons: integer, the number of electrons in the system. Must be
                greater or equal to 1.
            k_point: the k-point in reciprocal space used to evaluate Hamiltonian
                for the case of a periodic potential. K should be chosen to be
                within the first Brillouin zone.
            boundary_condition:
                'closed': forward/backward finite difference methods will be used
                near the boundaries to ensure the wavefunction is zero at
                boundaries. This should only be used when the grid interval is
                purposefully small.

                'open': all ghost points outside of the grid are set to zero. This
                should be used whenever the grid interval is sufficiently large.
                Setting to false also results in a faster computational time due
                to matrix symmetry.

                'periodic': periodic (wrap-around) boundary conditions. Requires
                a k_point to be specified.
            n_point_stencil: total number of points used in the central finite
                difference method. The default 5-points results in a convergence
                rate of 4 for most systems. Suggested: use 3-point stencil for
                potentials with cusps as n_point_stencil > 3 will not improve
                convergence rates.
            perturbation: Provide a perturbation operator H' (a matrix) to append
                to the non-interacting Hamiltonian, H + H'.

        Raises:
            ValueError: If num_electrons is less than 1; or num_electrons is not
                an integer.
        """
        self.k = k_point
        self.boundary_condition = boundary_condition
        self.n_point_stencil = n_point_stencil
        self.grids = grids
        self.dx = get_dx(grids)
        self.num_grids = len(grids)
        self.potential_fn = potential_fn
        if self.potential_fn != None:
            self.vp = np.nan_to_num(potential_fn(grids))
        self.perturbation = perturbation

        if not isinstance(num_electrons, int):
            raise ValueError('num_electrons is not an integer.')
        elif num_electrons < 1:
            raise ValueError(
                'num_electrons must be greater or equal to 1, but got %d'
                % num_electrons)
        else:
            self.num_electrons = num_electrons
        # Solver is unsolved by default.
        self._solved = False

    def is_solved(self):
        """Returns whether this solver has been solved."""
        return self._solved

    def solve_ground_state(self):
        """Solve ground state. Need to be implemented in subclasses.

        Compute attributes:
        total_energy, kinetic_energy, potential_energy, density, wave_function.
        """
        raise NotImplementedError('Must be implemented by subclass.')


class EigenSolver(SolverBase):
    """Represents the Hamiltonian as a matrix and diagonalizes it directly.

    This is the most stable and accurate eigensolver. Use SparseEigenSolver
    for a faster iterative eigensolver.
    """

    def __init__(self,
                 grids,
                 potential_fn=None,
                 num_electrons=1,
                 k_point=None,
                 boundary_condition='open',
                 n_point_stencil=3,
                 perturbation=None):
        """Initialize the solver with potential function and grid.

        Args:
          See SolverBase for args. discriptions.
        """
        super(EigenSolver, self).__init__(grids, potential_fn, num_electrons,
                                          k_point, boundary_condition,
                                          n_point_stencil, perturbation)
        self._set_matrices()
        self.quadratic_function = quadratic

    def _diagonal_matrix(self, form):
        """Creates diagonal matrix.

        Attributes:
          form: string, creates identity matrix if form == 'identity'
                        creates potential matrix if form == 'potential'
        """
        if form == 'identity':
            return np.eye(self.num_grids, dtype=complex)
        elif form == 'potential':
            return np.diag(self.vp)

    def _set_matrices(self):
        """Sets matrix attributes.

        Attributes:
          _t_mat: numpy matrix, kinetic matrix in Hamiltonian.
          _v_mat: numpy matrix, potential matrix in Hamiltonian.
          _h: numpy matrix, Hamiltonian matrix.
        """

        # Kinetic matrix
        self._t_mat = self.get_kinetic_matrix()

        if self.potential_fn != None:
            # Potential matrix
            self._v_mat = self.get_potential_matrix()
            # Hamiltonian matrix
            self._h = self._t_mat + self._v_mat

        # Perturbation matrix
        if self.perturbation is not None:
            self._h += self.perturbation

    def update_potential(self, potential_fn):
        """Replace the current potential grids with a new potential grids.
        Delete all attributes created by solving eigenvalues, set _solved to
        False.
        
        Args:
          potential_fn: potential function taking grids as argument.
        """

        self.potential_fn = potential_fn
        self.vp = np.nan_to_num(potential_fn(self.grids))
        # Potential matrix
        self._v_mat = self.get_potential_matrix()
        # Hamiltonian matrix
        self._h = self._t_mat + self._v_mat
        # Perturbation matrix
        if self.perturbation is not None:
            self._h += self.perturbation

        if self._solved:
            del self.total_energy
            del self.wave_function
            del self.density
            del self.kinetic_energy
            del self.potential_energy
            del self.eigenvalues
            self._solved = False

    def get_kinetic_matrix(self):
        """Kinetic matrix. Here the finite difference method is used to
        generate a kinetic energy operator in discrete space while satisfying
        desired boundary conditions.

        Returns:
            `ndarray`: Kinetic matrix with shape (num_grids, num_grids)
        """

        # n-point centered difference formula coefficients
        # these coefficients are for centered 2-order derivatives
        if self.n_point_stencil == 5:
            A_central = [-5 / 2, 4 / 3, -1 / 12]
        elif self.n_point_stencil == 3:
            A_central = [-2., 1.]
        else:
            raise ValueError(
                'n_point_stencil = %d is not supported' % self.n_point_stencil)

        mat = self._diagonal_matrix('identity')

        # get pentadiagonal KE matrix
        idx = np.arange(self.num_grids)
        for i, A_n in enumerate(A_central):
            mat[idx[i:], idx[i:] - i] = A_n
            mat[idx[:-i], idx[:-i] + i] = A_n

        # open-boundary
        if self.boundary_condition == 'open':
            mat = -0.5 * mat / (self.dx * self.dx)
            return np.real(mat)

        # append end-point forward/backward difference formulas
        elif self.boundary_condition == 'closed':

            if self.n_point_stencil == 5:
                # 0 means the first row, 1 means the second row
                # 0/1 are for forward/backward formulas at the end of the grid
                A_end_0 = [15 / 4, -77 / 6, 107 / 6, -13., 61 / 12, -5 / 6]
                A_end_1 = [5 / 6, -5 / 4, -1 / 3, 7 / 6, -1 / 2, 1 / 12]
            elif self.n_point_stencil == 3:
                # 0 means the same with 5 point
                A_end_0 = [2., -5., 4., -1.]

            # replace two ends of the matrix with forward/backward formulas
            for i, A_n_0 in enumerate(A_end_0):
                mat[0, i] = A_n_0
                mat[-1, -1 - i] = A_n_0
            if self.n_point_stencil == 5:
                for i, A_n_1 in enumerate(A_end_1):
                    mat[1, i] = A_n_1
                    mat[-2, -1 - i] = A_n_1

            # also change two end points as 0
            mat[0, 0] = 0
            mat[-1, -1] = 0

            mat = -0.5 * mat / (self.dx * self.dx)
            return np.real(mat)

        # periodic (no end point formulas needed)
        elif self.boundary_condition == 'periodic' and self.k is not None:
            k = self.k

            # assign central FDM formula (without center point)
            # also change 2nd-order end points
            if self.n_point_stencil == 3:
                D1_central = [1 / 2]

                mat[0, -1] = 1
                mat[-1, 0] = 1

            elif self.n_point_stencil == 5:
                D1_central = [2 / 3, -1 / 12]

                mat[0, -1] = 4 / 3
                mat[0, -2] = -1 / 12
                mat[1, -1] = -1 / 12

                mat[-1, 0] = 4 / 3
                mat[-1, 1] = -1 / 12
                mat[-2, 0] = -1 / 12

            # scale 2nd-order derivative matrix
            mat = -0.5 * mat / (self.dx * self.dx)

            # create identity matrix
            mat1 = 0.5 * (k ** 2) * self._diagonal_matrix('identity')

            # add 1st-order derivative matrix to identity
            idy = np.arange(self.num_grids)
            for i, D1_n in enumerate(D1_central):
                j = i + 1
                mat1[idy[0:], (idy[0:] - j) % self.num_grids] = \
                    complex(0., D1_n * k / self.dx)
                mat1[idy[0:], (idy[0:] + j) % self.num_grids] = \
                    complex(0., -D1_n * k / self.dx)

            # add all to second order matrix
            mat = mat + mat1

            return mat

        else:
            raise ValueError('boundary_condition = %s is not supported' %
                             self.boundary_condition)

    def get_potential_matrix(self):
        """Potential matrix. A diagonal matrix corresponding to the one-body
        potential input.

        Returns:
            `ndarray`: Potential matrix with shape (num_grids, num_grids)
        """

        if self.potential_fn == None:
            raise ValueError(
                'potential_fn is None, unable to get potential matrix.')

        return self._diagonal_matrix('potential')

    def get_kinetic_energy(self, wave_function):
      kinetic_energy = 0
      for phi in self.wave_function:
          kinetic_energy += self.quadratic_function(
              self._t_mat, phi) * self.dx

      return kinetic_energy

    def _update_ground_state(self, eigenvalues, eigenvectors, occupation_per_state):
        """Helper function to solve_ground_state() method.

        Updates the attributes total_energy, wave_function, density,
        kinetic_energy, potential_enenrgy and _solved from the eigensolver's
        output (w, v).

        Args:
          eigenvalues: Numpy array with shape [num_eigenstates,], the
              eigenvalues in ascending order.
          eigenvectors: Numpy array with shape [num_grids, num_eigenstates],
              each column eigenvectors[:, i] is the normalized eigenvector
              corresponding to the eigenvalue eigenvalues[i].

        Returns:
          self
        """
        self.wave_function = np.repeat(eigenvectors.T,
                                       repeats=occupation_per_state,
                                       axis=0)[:self.num_electrons]
        self.wave_function /= np.sqrt(self.dx)
        self.eigenvalues = np.repeat(
            eigenvalues, repeats=occupation_per_state)[:self.num_electrons]
        self.density = np.sum(self.wave_function ** 2, axis=0)
        self.total_energy = np.sum(self.eigenvalues)
        self.potential_energy = np.dot(self.density, self.vp)*self.dx
        self.kinetic_energy = self.get_kinetic_energy(self.wave_function)

        self._solved = True
        return self

    def solve_ground_state(self, occupation_per_state=1):
        """Solve ground state by diagonalize the Hamiltonian matrix directly.

        Compute attributes:
        total_energy, kinetic_energy, potential_energy, density, wave_function.

        Returns:
          self
        """

        if self.potential_fn == None:
            raise ValueError(
                'potential_fn is None, unable to solve for ground state.')

        if (self.boundary_condition == 'open'
                or self.boundary_condition == 'periodic'):
            eigenvalues, eigenvectors = np.linalg.eigh(self._h)
        else:
            eigenvalues, eigenvectors = np.linalg.eig(self._h)
            idx = eigenvalues.argsort()
            eigenvalues = eigenvalues[idx]
            eigenvectors = eigenvectors[:, idx]

        return self._update_ground_state(
            eigenvalues,eigenvectors, occupation_per_state)


class SparseEigenSolver(EigenSolver):
    """Represents the Hamiltonian as a matrix and solve with sparse eigensolver.

    This eigensolver is iterative and approximate. It is in general much
    faster than EigenSolver but less robust.
    """

    def __init__(self,
                 grids,
                 potential_fn=None,
                 num_electrons=1,
                 k_point=None,
                 boundary_condition='open',
                 n_point_stencil=3,
                 tol=10 ** -6):
        """Initialize the solver with potential function and grid.

        Args:
          tol: Relative accuracy for eigenvalues (stopping criterion).

          See SolverBase for additional args. discriptions.
        """
        super(SparseEigenSolver, self).__init__(grids, potential_fn,
                                                num_electrons, k_point,
                                                boundary_condition,
                                                n_point_stencil)
        self._tol = tol
        self.quadratic_function = self._sparse_quadratic

    def _diagonal_matrix(self, form):
        """Creates diagonal matrix.

        Attributes:
          form: string, creates identity matrix if form == 'identity'
                        creates potential matrix if form == 'potential'
        """

        if form == 'identity':
            return sparse.eye(self.num_grids, dtype=complex, format="lil")
        elif form == 'potential':
            return sparse.diags(self.vp, offsets=0, format='lil')

    def _sparse_quadratic(self, sparse_matrix, vector):
        """Compute quadratic of a sparse matrix and a dense vector.

        As of Numpy 1.7, np.dot is not aware of sparse matrices, scipy suggests
        to use the matrix dot method: sparse_matrix.dot(vector).

        Args:
          sparse_matrix: Scipy sparse matrix with shape [dim, dim].
          vector: Numpy array with shape [dim].

        Returns:
          Float, quadratic form of the input matrix and vector.
        """
        return np.dot(vector, sparse_matrix.dot(vector))

    def solve_ground_state(self):
        """Solve ground state by sparse eigensolver.

        Compute attributes:
        total_energy, kinetic_energy, potential_energy, density, wave_function.

        Returns:
          self
        """
        if (self.boundary_condition == 'open'
                or self.boundary_condition == 'periodic'):
            eigenvalues, eigenvectors = linalg.eigsh(
                self._h, k=self.num_electrons,
                which='SA', tol=self._tol)
        else:
            eigenvalues, eigenvectors = linalg.eigs(
                self._h, k=self.num_electrons,
                which='SR', tol=self._tol)
            idx = eigenvalues.argsort()
            eigenvalues = eigenvalues[idx]
            eigenvectors = eigenvectors[:, idx]

        return self._update_ground_state(
            eigenvalues, eigenvectors)
