"""Solve Schrodinger's equation.
"""

import numpy as np
from pydft.crystal import *

def diagouter(A,B):
    """Perform diag(AB*) more efficiently.

    Args:
        A (numpy.ndarray)
        B (numpy.ndarray)

    Return:
        diag(AB*) where * is the Hermitian conjugate.
    """

    return np.sum(np.dot(A,B.conj().T), axis=0)

def getE(crystal, basis, potential, potargs):
    """Calculate the energies that are the solution to Schrodinger's equation.
    
    Args:
        crystal (Crystal): a Crystal object that has had the wave function
            coefficients added in a given basis.
        basis (string): the basis for the expansion of the wave function,
            charge density and potential.
        potential (strinc): a potential energy function such as a simple
            harmonic oscillator.
        potargs: various parameters required for the potential.
    
    Return:
        (float): the energy obtained by solving Schrodinger's equation 
            for a crystal.
    """

    if basis == "planewave":
        from pydft.bases.planewave import g, W, I, O, L
        if potential == "sho":
            from pydft.potentials import sho_dual
            g(crystal)
            W(crystal)

            U = np.dot(crystal._W.conj().T,
                       O(crystal, crystal._W))            
            n = diagouter(I(crystal, np.dot(crystal._W, np.linalg.inv(U))),
                          I(crystal, crystal._W))
            V = sho_dual(crystal, basis, crystal._r, *potargs)

            return -1./2*np.trace(
                crystal._W.conj().T.dot(
                    L(crystal, crystal._W.dot(np.linalg.inv(U))))) + (
                        np.dot(np.conj(V),n))

        else:
            ValueError("The SHO potential is the only supported basis.")
    else:
        ValueError("The planewave basis is the only supported basis.")
