"""
dvr.py
"""

from .legDVR import _legDVR, _legDVRwfs
from .fourDVR import _fourDVR, _fourDVRwfs
from .hoDVR import _hoDVR,_hoDVRwfs
from .sincDVR import _sincDVR,_sincDVRwfs

import numpy as np 

__all__ = ['GenericDVR','SimpleDVR','Contracted']


class GenericDVR:
    """
    A super-class for generic 1D DVRs.
    
    Attributes
    ----------
    num : int
        The number of DVR grid points
    grid : ndarray
        The DVR grid points
    D : ndarray
        First derivative operator.
    D2 : ndarray
        Second derivative operator. 
    
    """
    
    def __init__(self, grid, D, D2):
        
        if grid.ndim != 1:
            raise ValueError('grid must be 1-dimensional')
        
        num = len(grid) 
        
        if D.ndim != 2 or D.shape[0] != num or D.shape[1] != num:
            raise ValueError('D must be (num,num) array')
            
        if D2.ndim != 2 or D2.shape[0] != num or D2.shape[1] != num:
            raise ValueError('D2 must be (num,num) array')
        
        self.num = len(grid)
        self.grid = grid 
        self.D = D
        self.D2 = D2
        
        return 
    
    def wfs(self, q):
        """
        Evaluate DVR basis wavefunctions.

        Parameters
        ----------
        q : ndarray
            A 1-D array of coordinate values.

        Returns
        -------
        wfs : ndarray
            An array of shape (`q`.size, :attr:`num`) containing
            the values of the DVR wavefunctions evaluated at 
            coordinate values `q`.

        """
        return self._wfs(q) 
        
    def _wfs(self, q):
        
        raise NotImplementedError("")
    
    def matchfun(self, f):
        """
        Calculate DVR coefficients to match a function at 
        DVR grid points.

        Parameters
        ----------
        f : function
            The function to be matched.

        Returns
        -------
        coeff : ndarray
            A (`num`,1) array with the DVR basis function coefficients.

        """
        
        # Calculate the values of the DVR basis functions
        # at their respective grid-points
        wgts = np.diag(self.wfs(self.grid))
        
        coeffs = (f(self.grid) / wgts).reshape((self.num,1))
        
        return coeffs 
    
    def contract(self, u):
        return Contracted(u, self)


class SimpleDVR(GenericDVR):
    """
    Simple, standard 1D DVRs.
    
    Attributes
    ----------
    start : float
        The DVR grid starting value.
    stop : float
        The DVR grid stopping value.
    basis : {'sinc','ho','fourier', 'legendre'}
        The DVR basis type.
    
    """
    
    def __init__(self, start = 0, stop = 1, num = 10, basis = 'sinc'):
        """
        Create a DVR object.

        Parameters
        ----------
        start : float, optional
            The DVR grid starting value. The default is 0.
        stop : float, optional
            The DVR grid stopping value. The default is 1.
        num : int, optional
            The number of DVR grid points. The default is 10.
        basis : {'sinc','ho','fourier','legendre'}, optional
            The DVR basis type. The default is 'sinc'.

        """
        
        if stop <= start:
            raise ValueError("stop value must be larger than start value")
        if num < 2 :
            raise ValueError("num must be >= 2")
        

        basis = basis.lower()
        
        if basis == 'sinc' : 
            grid, D, D2 = _sincDVR(start, stop, num)
            _wfs_fun = lambda q : _sincDVRwfs(q, start, stop, num)
        elif basis == 'ho' : 
            grid, D, D2 = _hoDVR(start, stop, num)
            _wfs_fun = lambda q : _hoDVRwfs(q, start, stop, num)
        elif basis == 'fourier' : 
            grid, D, D2 = _fourDVR(start, stop, num)
            _wfs_fun = lambda q : _fourDVRwfs(q, start, stop, num)
        elif basis == 'legendre' : 
            grid, D, D2 = _legDVR(start, stop, num)
            _wfs_fun = lambda q : _legDVRwfs(q, start, stop, num)
        else:
            raise ValueError("basis type '{}' not recognized".format(basis))
            
        super().__init__(grid, D, D2)

        self.basis = basis
        self.start = start 
        self.stop = stop 
        
        self._wfs_fun = _wfs_fun
        
    def _wfs(self, q):
        
        return self._wfs_fun(q)
        
    
class Contracted(GenericDVR):
    """
    A contracted basis DVR class (usually for PO-DVRs).
    
    Attributes
    ----------
    prim_dvr : GenericDVR
        The primitive basis
    W : ndarray
        The transformation matrix from the contracted DVR
        to the primitive DVR 
    
    """
    
    def __init__(self, U, prim_dvr):
        """
        

        Parameters
        ----------
        U : (n,m) ndarray
            A unitary operator defining the contracted basis.
            `n` is the size of the primitive basis. 
            `m` is the size of the contracted basis. 
        prim_dvr : GenericDVR
            The primitive DVR


        """
        
        #
        # Construct the coordinate operator in the 
        # contracted basis 
        Q = (U.conj().T * prim_dvr.grid) @ U 
        q,u = np.linalg.eigh(Q)   # u takes a vector from the new DVR 
                                  # to the contracted "fbr"
        
        q = np.reshape(q, (-1,))  # The new DVR grid 
        
        W = U @ u # w takes a vector from the new DVR to the old DVR 
        
        d = W.conj().T @ prim_dvr.D @ W
        d2 = W.conj().T @ prim_dvr.D2 @ W
        
        super().__init__(q, d, d2) 
        
        self.prim_dvr = prim_dvr 
        self.W = W 
    
    def _wfs(self, q):
        
        f = self.prim_dvr.wfs(q) # f is (q.size, primitive num)
    
        return f @ self.W 
    
        