#!/usr/bin/env python

import re, sys, argparse

import numpy as np
from ase.io import read
from scipy.sparse import block_diag

from vaspwfc import vaspwfc
from vasp_constant import *
from paw import pawpotcar, nonlq


'''
Dipole transition within the electric dipole approximation (EDA).  Please refer
to this post for more details.

  https://qijingzheng.github.io/posts/Light-Matter-Interaction-and-Dipole-Transition-Matrix/

The dipole transition matrix elements in the length gauge is given by:

        <psi_nk | e r | psi_mk>

where | psi_nk > is the pseudo-wavefunction.  In periodic systems, the position
operator "r" is not well-defined.  Therefore, we first evaluate the momentum
operator matrix in the velocity gauge, i.e.

        <psi_nk | p | psi_mk>

And then use simple "p-r" relation to apprimate the dipole transition matrix
element

                                  -i⋅h
    <psi_nk | r | psi_mk> =  -------------- ⋅ <psi_nk | p | psi_mk>
                               m⋅(En - Em)

Apparently, the above equaiton is not valid for the case Em == En. In this case,
we just set the dipole matrix element to be 0.

################################################################################
NOTE that, the simple "p-r" relation only applies to molecular or finite system,
and there might be problem in directly using it for periodic system. Please
refer to this paper for more details.

  "Relation between the interband dipole and momentum matrix elements in
  semiconductors"
  (https://journals.aps.org/prb/pdf/10.1103/PhysRevB.87.125301)

################################################################################

The momentum operator matrix in the velocity gauge

        <psi_nk | p | psi_mk> = hbar <u_nk | k - i nabla | u_mk>

In PAW, the matrix element can be divided into plane-wave parts and one-center
parts, i.e.

    <u_nk | k - i nabla | u_mk> = <tilde_u_nk | k - i nabla | tilde_u_mk>
                                 - \sum_ij <tilde_u_nk | p_i><p_j | tilde_u_mk>
                                   \times i [
                                     <phi_i | nabla | phi_j>
                                     -
                                     <tilde_phi_i | nabla | tilde_phi_j>
                                   ]

where | u_nk > and | tilde_u_nk > are cell-periodic part of the AE/PS
wavefunctions, | p_j > is the PAW projector function and | phi_j > and
| tilde_phi_j > are PAW AE/PS partial waves.

The nabla operator matrix elements between the pseudo-wavefuncitons

    <tilde_u_nk | k - i nabla | tilde_u_mk>

   = \sum_G C_nk(G).conj() * C_mk(G) * [k + G]

where C_nk(G) is the plane-wave coefficients for | u_nk >.
'''

def parse_band_idx(band_str_lst):
    '''
    Generate the band indices.
    '''
    ibands = []
    for bstr in band_str_lst:
        m = re.match(r'(\d+)(?::(\d+))?(?::(\d+))?$', bstr)
        if not m:
            raise ValueError(
                "\nThe band index should be assigned with the combination of the following ways:\n\n"
                "-> 10, a single band with index 20\n"
                "-> 20:30, a continuous range from 1 to 10, 10 included\n"
                "-> 30:50:2, a continuous range from 30 to 50, with step size 2\n"
                "\nFor example, '1 4:6 8:16:2' will convert to '1 4 5 6 8 10 12 14 16'."
            )
        ii = m.group(1)
        jj = m.group(2) or ii
        ss = m.group(3) or 1
        ibands += [x for x in range(int(ii), int(jj)+1, int(ss))]

    return np.array(ibands, dtype=int)


def parse_cml_args(cml):
    '''
    CML parser.
    '''
    arg = argparse.ArgumentParser(add_help=True)

    arg.add_argument('-w', '--wavecar', dest='wavecar', action='store', type=str,
                     default='WAVECAR',
                     help='Location of the WAVECAR')

    arg.add_argument('-p', '--poscar', dest='poscar', action='store', type=str,
                     default='POSCAR',
                     help='Location of the POSCAR')

    arg.add_argument('--potcar', dest='potcar', action='store', type=str,
                     default='POTCAR',
                     help='Location of the POTCAR')

    arg.add_argument('-s', dest='ispin', action='store', type=int,
                     default=1,
                     help='the spin index, starting from 1')

    arg.add_argument('-k', dest='ikpt', action='store', type=int,
                     default=1,
                     help='the k-point index, starting from 1')

    arg.add_argument('-i', '--ibands', dest='ibands', action='store', type=str,
                     nargs='+',
                     help='the initial band indices, e.g. 1:10 15 20:30:2')

    arg.add_argument("-j", '--jbands', dest='jbands', action='store', type=str,
                     nargs='+',
                     help='the final band indices, e.g. 1:10 15 20:30:2')

    arg.add_argument('--lgamma', dest='lgamma', action='store_true',
                     help='WAVECAR from GAMMA-ONLY VASP')

    arg.add_argument('--gamma_half', dest='gamma_half', action='store',
                     type=str, default='x', choices=['x', 'z'],
                     help='GAMMA half in x/z direction')

    arg.add_argument('--lpseudo', dest='lpseudo', action='store_true',
                     help='Dipole matrix from pseudo-wavefunctions')

    arg.add_argument('--lsorbit', dest='lsorbit', action='store_true',
                     help='WAVECAR from NONCOLLINEAR VASP')

    arg.add_argument('--sigma', dest='sigma',
                     action='store', type=float,
                     default=0.05,
                     help='smearing parameter, default 0.05')

    arg.add_argument('--nedos', dest='nedos',
                     action='store', type=int,
                     default=5000,
                     help='number of point in DOS plot')

    arg.add_argument('--figsize', dest='figsize',
                     action='store', type=float,
                     nargs=2, default=(4.8, 3.0),
                     help='number of point in DOS plot')

    return arg.parse_args(cml)


def get_dipole_mat(
        pswfc, ibands, jbands, ispin=1, ikpt=1,
        lae=True,
        projector=None,
    ):
    '''
    Get the dipole transition matrix using the "p-r" relation.
    '''

    if lae:
        assert (projector is not None) and isinstance(projector, nonlq)

    # initial states
    ibands = np.asarray(ibands, dtype=int)
    # final states
    jbands = np.asarray(jbands, dtype=int)
    assert jbands.min() > ibands.max()
    nx, ny = np.meshgrid(ibands, jbands)

    E_i = pswfc._bands[ispin-1, ikpt-1, ibands-1]
    E_j = pswfc._bands[ispin-1, ikpt-1, jbands-1]
    Ex, Ey = np.meshgrid(E_i, E_j)
    dE = Ey - Ex

    k0 = pswfc._kvecs[ikpt - 1]
    Gk = np.dot(pswfc.gvectors(ikpt=ikpt) + k0, pswfc._Bcell*TPI)
    Cg_i = np.array([
        pswfc.readBandCoeff(ispin=ispin, ikpt=ikpt, iband=ii)
        for ii in ibands
    ])
    Cg_j = np.array([
        pswfc.readBandCoeff(ispin=ispin, ikpt=ikpt, iband=jj)
        for jj in jbands
    ])

    nplw = Gk.shape[0]
    if pswfc._lsoc:
        dp_mat  = np.array([
            Cg_j[:,:nplw].conj() @ (Cg_i[:,:nplw] * Gk[:,ii]).T
            +
            Cg_j[:,nplw:].conj() @ (Cg_i[:,nplw:] * Gk[:,ii]).T
            for ii in range(3)
        ])
    elif pswfc._lgam:
        dp_mat = np.array([
            (
                Cg_j.conj() @ (Cg_i * Gk[:, ii]).T
                -
                Cg_j @ (Cg_i.conj() * Gk[:, ii]).T
            ) / 2.
            for ii in range(3)
        ])
    else:
        dp_mat = np.array([
            Cg_j.conj() @ (Cg_i * Gk[:, ii]).T
            for ii in range(3)
        ])

    if lae:
        nablaij = [
            block_diag([
                projector.pawpp[projector.element_idx[iatom]].get_nablaij()[ii]
                for iatom in range(projector.natoms)
            ])
            for ii in range(3)
        ]

        # one-center correction
        if pswfc._lsoc:
            # spin-up part of the spinor
            beta_i = np.array(
                [projector.proj(Cg) for Cg in Cg_i[:,:nplw]]
            )
            beta_j = np.array(
                [projector.proj(Cg) for Cg in Cg_j[:,:nplw]]
            )

            # spin-down part of the spinor
            beta_i2 = np.array(
                [projector.proj(Cg) for Cg in Cg_i[:,nplw:]]
            )
            beta_j2 = np.array(
                [projector.proj(Cg) for Cg in Cg_j[:,nplw:]]
            )
        else:
            beta_i = np.array(
                [projector.proj(Cg) for Cg in Cg_i]
            )
            beta_j = np.array(
                [projector.proj(Cg) for Cg in Cg_j]
            )

        for ii in range(3):
            dp_mat[ii] += -1j*(beta_j.conj() @ (nablaij[ii] @ beta_i.T))
            if pswfc._lsoc:
                dp_mat[ii] += -1j*(beta_j2.conj() @ (nablaij[ii] @ beta_i2.T))

    # if dE = 0, set dp_mat to zero
    dp_mat = np.where(
        np.isclose(dE, 0),
        0,
        -1j / (dE / (2*RYTOEV)) * dp_mat * AUTOA * AUTDEBYE
    )

    dump = [
        nx.flatten(), ny.flatten(),
        Ex.flatten(), Ey.flatten(), dE.flatten(),
        dp_mat[0].flatten(),
        dp_mat[1].flatten(),
        dp_mat[2].flatten()
    ]

    return dump


def showtdm(tdm, p):
    '''
    '''
    dE = tdm[0]
    emin = dE.min()
    emax = dE.max()
    eran = emax - emin
    e0 = np.linspace(emin - 0.10 * eran, emax + 0.10 * eran, p.nedos)
    tdm_abs = np.abs(tdm[-3:]).T

    SMEAR = 1. / (np.sqrt(2 * np.pi) * p.sigma) * \
        np.exp(
        -(dE[None, :] - e0[:, None])**2 / (2*p.sigma**2)
    )
    Tx = np.sum(tdm_abs[:, 0] * SMEAR, axis=1)
    Ty = np.sum(tdm_abs[:, 1] * SMEAR, axis=1)
    Tz = np.sum(tdm_abs[:, 2] * SMEAR, axis=1)
    Tt = Tx + Ty + Tz

    import matplotlib as mpl
    import matplotlib.pyplot as plt
    plt.style.use('ggplot')

    fig = plt.figure(
        figsize=p.figsize, dpi=300
    )
    ax = plt.subplot()

    ax.plot(e0, Tx, color='r', lw=0.8, ls=':', label='x', alpha=0.8)
    ax.plot(e0, Ty, color='g', lw=0.8, ls=':', label='y', alpha=0.8)
    ax.plot(e0, Tz, color='b', lw=0.8, ls=':', label='z', alpha=0.8)
    ax.plot(e0, Tt, color='k', lw=0.8, ls='-', label='tot', alpha=0.8)

    tdm_abs_tot = np.sum(tdm_abs, axis=1)
    fac = Tt.max() * 0.20 / tdm_abs_tot.max()
    ax.vlines(dE, ymin=0.0, ymax=tdm_abs_tot * fac, lw=1.0, color='k')

    ax.set_xlabel('Energy [eV]', labelpad=5)
    ax.set_ylabel('Transition Dipole Moment', labelpad=5)
    plt.legend(loc='best')

    plt.tight_layout()
    plt.savefig('tdm.png', dpi=300)
    plt.show()


def main(cml):
    p = parse_cml_args(cml)

    pswfc = vaspwfc(p.wavecar,
                  lsorbit=p.lsorbit, lgamma=p.lgamma,
                  gamma_half=p.gamma_half)
    ibands = parse_band_idx(p.ibands)
    jbands = parse_band_idx(p.jbands)

    # dipole matrix elements from AE or PS wavefucntions
    lae = not p.lpseudo

    # If all-electron dipole matrix elements are needed, must also provide the
    # projector function
    projector = None
    if lae:
        pawpp = [pawpotcar(potstr) for potstr in
                 open(p.potcar).read().split('End of Dataset')[:-1]]
        projector = nonlq(
            read(p.poscar),
            pswfc._encut,
            pawpp,
            k=pswfc._kvecs[p.ikpt-1],
            lgam=pswfc._lgam,
            gamma_half=pswfc._gam_half,
        )

    dp_mat = get_dipole_mat(
        pswfc, ibands, jbands,
        ikpt=p.ikpt, ispin=p.ispin,
        lae=lae, projector=projector
    )

    header = ('{:5s} {:5s}' + ' {:>12s}' * 6).format(
        "iband", "jband", "E_i", "E_j", "dE", "Tx", "Ty", "Tz")
    
    tmp = np.r_[dp_mat[:5], np.abs(dp_mat[-3:])]
    np.savetxt('TDM.dat', tmp.T,
               header=header,
               fmt='  %5d %5d' + ' %12.6f' * 6)
    showtdm(dp_mat[-4:], p)


if __name__ == "__main__":
    main(sys.argv[1:])
    pass
