##
## Authors: MN (FP adapted)
##
import numpy as np
import os
from yambopy.dbs.excitondb import YamboExcitonDB
from yambopy.bse.exciton_matrix_elements import exciton_X_matelem
from yambopy.bse.rotate_excitonwf import rotate_exc_wf
from tqdm import tqdm

def exciton_phonon_matelem(latdb,elphdb,wfdb,Qrange=[0,1],BSE_dir='bse',BSE_Lin_dir=None,
                           nexc_in=-1,nexc_out=-1,dmat_mode='run',save_files=True,exph_file='Ex-ph.npy',overwrite=False):
    """
    This function calculates the exciton-phonon matrix elements

    - Q is the exciton momentum
    - q is the phonon momentum
    - exc_in represent the "initial" exciton states in the scattering process (at mom. Q)
    - exc_out represents the "final" exciton states in the scattering process (at mom. Q+q)

    Parameters
    ----------
    latdb : YamboLatticeDB
        The YamboLatticeDB object which contains the lattice information.
    elphdb : LetzElphElectronPhononDB
        The LetzElphElectronPhononDB object which contains the electron-phonon matrix
        elements.
    wfdb : YamboWFDB
        The YamboWFDB object which contains the wavefunction information.
    BSE_dir : str, optional
        The name of the folder which contains the BSE calculation. Default is 'bse'.
    BSE_Lin_dir : str, optional
        The name of the folder which contains the BSE Q=0 calculation (for optical spectra). Default is BSE_dir.
    Qrange : int list, optional
        Exciton Qpoint index range [iQ_initial, iQ_final] (python counting). Default is [0,0] (Gamma point only).
        Note that the indexing is in full BZ and not in iBZ. See wfc.kBZ to see the kpoints in full BZ
    nexc_in : int, optional
        Number of excitonic states included in Lin calculation. Default is -1 (all).
    nexc_out : int, optional
        Number of excitonic states included in Lout calculation. Default is -1 (all).
    dmat_mode : str, optional
        If 'save', print dmats on .npy file for faster recalculation. If 'load', load from .npy file. Else, calculate Dmats at runtime.
    save_files : bool, optional
        If True, the matrix elements will be saved in .npy file `exph_file`. Default is True.
    overwrite : bool, optional
        If False and `exph_file` is found, the matrix elements will be loaded from file. Default is False.
    """

    # Check if we just need to load
    if os.path.exists(exph_file) and overwrite==False:
        print(f'Loading EXCPH matrices from {exph_file}...')
        exph_mat_loaded = np.load(exph_file)
        return exph_mat_loaded

    # Load exc dbs
    exdbs = []
    for ik in range(wfdb.nkpoints):
        filename = 'ndb.BS_diago_Q%d' % (ik+1)
        excdb = YamboExcitonDB.from_db_file(latdb,filename=filename,folder=BSE_dir,\
                                            Load_WF=True, neigs=nexc_out)
        exdbs.append(excdb)
        #
        # NM : Add a sanity check to avoid a disasterous consequence
        # if the user gives wrong bse band indices.
        min_bnd_bse = np.min(excdb.unique_vbands)
        max_bnd_bse = np.max(excdb.unique_cbands)+1
        assert (wfdb.min_bnd == min_bnd_bse) and ((wfdb.min_bnd + wfdb.nbands) == max_bnd_bse), \
            print("Error: BSE bands mismatch. Given bands range : [%d, %d]. " %(
                wfdb.min_bnd,wfdb.min_bnd + wfdb.nbands) +
                "Bse band range found (expected) : [%d %d]" %( min_bnd_bse,max_bnd_bse))

    # get D matrices
    Dmats = save_or_load_dmat(wfdb,mode=dmat_mode,dmat_file='Dmats.npy')

    # Calculation
    print('Calculating EXCPH matrix elements...')
    exph_mat = []
    for iQ in tqdm(range(Qrange[0],Qrange[1])):
        Q_in = wfdb.kBZ[iQ]
        exph_mat.append( exciton_phonon_matelem_iQ(elphdb,wfdb,exdbs,Dmats,\
                                                   BSE_Lin_dir=BSE_Lin_dir,Q_in=Q_in,neigs=nexc_in) )
    # IO
    if len(exph_mat)<2: exph_mat = exph_mat[0] # single Q-point calculation (suppress axis)
    else:               exph_mat = np.array(exph_mat) #[nQ,nq,nmodes,nexc_in (Qexc),nexc_out (Qexc+q)]
    
    if save_files: 
        if exph_file[-4:]!='.npy': exph_file = exph_file+'.npy'
        print(f'Excph coupling file saved to {exph_file}')
        np.save(exph_file,exph_mat)
    
    return exph_mat

def exciton_phonon_matelem_iQ(elphdb,wfdb,exdbs,Dmats,BSE_Lin_dir=None,
                              Q_in=np.zeros(3),neigs=-1,dmat_mode='run'): 
    """
    This function calculates the exciton-phonon matrix element per Q 

    - Q is the exciton momentum
    - q is the phonon momentum
    - exc_in represent the "initial" exciton states in the scattering process (at mom. Q)
    - exc_out represents the "final" exciton states in the scattering process (at mom. Q+q)

    Parameters
    ----------
    elphdb : LetzElphElectronPhononDB
        The LetzElphElectronPhononDB object which contains the electron-phonon matrix
        elements.
    wfdb : YamboWFDB
        The YamboWFDB object which contains the wavefunction information.
    exdbs : YamboExcitonDB list
        List of Q+q YamboExcitonDB objects containing the BSE calculation
    BSE_Lin_dir : str, optional
        The name of the folder which contains the BSE q=0 calculation (for optical spectra). Default is exdbs[Q].
    Q_in : np.ndarray, optional
        Excitonic momentum in reduced units. Default np.array([0.0,0.0,0.0]) 
    neigs : int, optional
        Number of excitonic states included in calculation. Default is -1 (all).
    """
    latdb = wfdb.ydb
    # Determine Lkind(in)
    Ak = rotate_Akcv_Q(wfdb, exdbs, Q_in, Dmats, neigs=neigs, folder=BSE_Lin_dir)
    # Compute ex-ph
    exph_mat = []
    bse_bnds_range = [wfdb.min_bnd,wfdb.min_bnd + wfdb.nbands]
    for iq in range(elphdb.nq):
        ph_eig, elph_mat = elphdb.read_iq(iq,bands_range=bse_bnds_range,convention='standard')
        elph_mat = elph_mat.transpose(1,0,2,4,3)
        #
        Akq = rotate_Akcv_Q(wfdb, exdbs, Q_in + elphdb.qpoints[iq], Dmats) # q+Q
        tmp_exph = exciton_X_matelem(Q_in, elphdb.qpoints[iq], \
                                     Akq, Ak, elph_mat, wfdb.kBZ, \
                                     contribution='b', diagonal_only=False, ktree=wfdb.ktree)
        exph_mat.append(tmp_exph)

    ## 0.5 for Ry to Ha
    exph_mat = 0.5 * np.array(exph_mat).transpose(0,1,3,2) #[nq,nmodes,nexc_in (Qexc),nexc_out (Qexc+q)]

    return exph_mat

def save_or_load_dmat(wfdb, mode='run', dmat_file='Dmats.npy'):
    """
    Save or load Dmats to/from .npy file `dmat_file` for faster recalculation.

     If mode=='save', print dmats on .npy file for faster recalculation. 
     If mode=='load', load from .npy file. 
     Else, calculate Dmats at runtime.
    """
    if dmat_file[-4:]!='.npy': dmat_file = dmat_file+'.npy'
    if mode=='save':
        print('Saving D matrices...')
        Dmats = wfdb.Dmat()
        np.save(dmat_file,Dmats)
        return Dmats
    elif mode=='load': 
        print('Loading D matrices...')
        if not os.path.exists(dmat_file):
            raise FileNotFoundError(f"Cannot load '{dmat_file}' - file does not exist.")
        Dmats_loaded = np.load(dmat_file)
        return Dmats_loaded
    else:
        return wfdb.Dmat()


def rotate_Akcv_Q(wfdb, exdbs, Qpt, Dmats, neigs=-1, folder=None):
    '''
    Qpt reduced coordinates in BZ or whatever

    wfdb : wavefunction object
    exdbs : previously read list of exciton objects for Lout (used if folder is None)
    Qpt: reduced coordinates in BZ or whatever
    Dmats: precalculated Dmat
    neigs: number of states (used if folder is not None, otherwise it's nexc_out)
    folder: where to load the Lin exciton states (if needed). 
            can be used also to reload the same in case we want nexc_in/=nexc_out for same L
    '''
    latdb = wfdb.ydb
    idx_BZQ = wfdb.kptBZidx(Qpt)
    iQ_isymm = latdb.symmetry_indexes[idx_BZQ]
    iQ_iBZ = latdb.kpoints_indexes[idx_BZQ]
    trev  = (iQ_isymm >= len(latdb.sym_car) / (1 + int(np.rint(latdb.time_rev))))
    symm_mat_red = latdb.lat@latdb.sym_car[iQ_isymm]@np.linalg.inv(latdb.lat)
    exe_iQIBZ = wfdb.kpts_iBZ[iQ_iBZ]
    #
    if folder is not None:
        if neigs==-1:
            neigs = len(exdbs[0].eigenvalues) # Set neigs equal to nexc_out
        filename = 'ndb.BS_diago_Q%d' % (iQ_iBZ+1)
        excdbin = YamboExcitonDB.from_db_file(latdb,filename=filename,folder=folder,\
                                              Load_WF=True, neigs=neigs)
        AQibz = excdbin.get_Akcv()
    else : AQibz = exdbs[iQ_iBZ].get_Akcv()
    #
    AQ_rot = rotate_exc_wf(AQibz,symm_mat_red,wfdb.kBZ,exe_iQIBZ,Dmats[iQ_isymm],trev,wfdb.ktree)
    
    return AQ_rot
