#include "cal_ldos.h"

#include "cal_dos.h"
#include "cube_io.h"
#include "source_estate/module_dm/cal_dm_psi.h"
#include "source_lcao/module_gint/temp_gint/gint_interface.h"

#include 

namespace ModuleIO
{

#ifdef __LCAO
template 
void Cal_ldos::cal_ldos_lcao(const elecstate::ElecStateLCAO* pelec,
                                const psi::Psi& psi,
                                const Parallel_Grid& pgrid,
                                const UnitCell& ucell)
{
    for (int ie = 0; ie < PARAM.inp.stm_bias[2]; ie++)
    {
        // energy range for ldos (efermi as reference)
        const double en = PARAM.inp.stm_bias[0] + ie * PARAM.inp.stm_bias[1];
        const double emin = en < 0 ? en : 0;
        const double emax = en > 0 ? en : 0;

        // calculate weight (for bands not in the range, weight is zero)
        ModuleBase::matrix weight(pelec->ekb.nr, pelec->ekb.nc);
        for (int ik = 0; ik < pelec->ekb.nr; ++ik)
        {
            const double efermi = pelec->eferm.get_efval(pelec->klist->isk[ik]);

            for (int ib = 0; ib < pelec->ekb.nc; ib++)
            {
                const double eigenval = (pelec->ekb(ik, ib) - efermi) * ModuleBase::Ry_to_eV;
                if (eigenval >= emin && eigenval  0 ? pelec->klist->wk[ik] - pelec->wg(ik, ib) : pelec->wg(ik, ib);
                }
            }
        }

        // calculate dm-like for ldos
        const int nspin_dm = PARAM.inp.nspin == 2 ? 2 : 1;
        elecstate::DensityMatrix dm_ldos(pelec->DM->get_paraV_pointer(),
                                                    nspin_dm,
                                                    pelec->klist->kvec_d,
                                                    pelec->klist->get_nks() / nspin_dm);

        elecstate::cal_dm_psi(pelec->DM->get_paraV_pointer(), weight, psi, dm_ldos);
        dm_ldos.init_DMR(*(pelec->DM->get_DMR_pointer(1)));
        dm_ldos.cal_DMR();

        // allocate ldos space
        std::vector ldos_space(PARAM.inp.nspin * pelec->charge->nrxx);
        double** ldos = new double*[PARAM.inp.nspin];
        for (int is = 0; is < PARAM.inp.nspin; ++is)
        {
            ldos[is] = &ldos_space[is * pelec->charge->nrxx];
        }

    // calculate ldos
#ifdef __OLD_GINT
        ModuleBase::WARNING_QUIT("Cal_ldos::dm2ldos",
                                 "do not support old grid integral, please recompile with __NEW_GINT");
#else
        ModuleGint::cal_gint_rho(dm_ldos.get_DMR_vector(), PARAM.inp.nspin, ldos);
#endif

        // I'm not sure whether ldos should be output for each spin or not
        // ldos[0] += ldos[1] for nspin_dm == 2
        if (nspin_dm == 2)
        {
            BlasConnector::axpy(pelec->charge->nrxx, 1.0, ldos[1], 1, ldos[0], 1);
        }

        // write ldos to cube file
        std::stringstream fn;
        fn nrxx);

        for (int ik = 0; ik < pelec->klist->get_nks(); ++ik)
        {
            psi.fix_k(ik);
            const double efermi = pelec->eferm.get_efval(pelec->klist->isk[ik]);
            const int nbands = psi.get_nbands();

            for (int ib = 0; ib < nbands; ib++)
            {
                pelec->basis->recip2real(&psi(ib, 0), wfcr.data(), ik);

                const double eigenval = (pelec->ekb(ik, ib) - efermi) * ModuleBase::Ry_to_eV;
                double weight = en > 0 ? pelec->klist->wk[ik] - pelec->wg(ik, ib) : pelec->wg(ik, ib);
                weight /= ucell.omega;

                if (eigenval >= emin && eigenval basis->nrxx; ir++)
                    {
                        ldos[ir] += weight * norm(wfcr[ir]);
                    }
                }
            }
        }

        std::stringstream fn;
        fn charge->nrxx);
    std::vector wfcr(pelec->basis->nrxx);
    for (int ik = 0; ik < pelec->klist->get_nks(); ++ik)
    {
        psi.fix_k(ik);
        const double efermi = pelec->eferm.get_efval(pelec->klist->isk[ik]);
        const int nbands = psi.get_nbands();

        for (int ib = 0; ib < nbands; ib++)
        {
            pelec->basis->recip2real(&psi(ib, 0), wfcr.data(), ik);
            const double weight = pelec->klist->wk[ik] / ucell.omega;

            for (int ir = 0; ir < pelec->basis->nrxx; ir++)
            {
                tmp[ir] += weight * norm(wfcr[ir]);
            }

            std::vector results(npoints, 0);
            trilinear_interpolate(points, shifts, pgrid, tmp, results);

            const double eigenval = pelec->ekb(ik, ib) * ModuleBase::Ry_to_eV;

            for (int ie = 0; ie < ndata; ++ie)
            {
                const double en = emin + ie * PARAM.inp.dos_edelta_ev;
                const double de = en - eigenval;
                const double de2 = de * de;
                const double gauss = exp(-de2 / sigma2) / sigma_PI;
                for (int ip = 0; ip < npoints; ++ip)
                {
                    ldos[ip][ie] += results[ip] * gauss;
                }
            }
        }
    }

    std::ofstream ofs_ldos;
    std::stringstream fn;
    fn