// Copyright (c) 2013-2017 Anton Kozhevnikov, Thomas Schulthess
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without modification, are permitted provided that
// the following conditions are met:
//
// 1. Redistributions of source code must retain the above copyright notice, this list of conditions and the
//    following disclaimer.
// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions
//    and the following disclaimer in the documentation and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED
// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
// PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
// ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
// OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

/** \file poisson.hpp
 *
 *  \brief Implementation of the Poisson solver for the full-potential muffin-tin case.
 */

inline void Potential::poisson_add_pseudo_pw(mdarray<double_complex, 2>& qmt__,
                                             mdarray<double_complex, 2>& qit__,
                                             double_complex* rho_pw__)
{
    PROFILE("sirius::Potential::poisson_add_pseudo_pw");

    int lmmax = ctx_.lmmax_rho();
    int ngv = ctx_.gvec().count();

    /* The following term is added to the plane-wave coefficients of the charge density:
     * Integrate[SphericalBesselJ[l,a*x]*p[x,R]*x^2,{x,0,R},Assumptions->{l>=0,n>=0,R>0,a>0}] / 
     *  Integrate[p[x,R]*x^(2+l),{x,0,R},Assumptions->{h>=0,n>=0,R>0}]
     * i.e. contributon from pseudodensity to l-th channel of plane wave expansion multiplied by 
     * the difference bethween true and interstitial-in-the-mt multipole moments and divided by the 
     * moment of the pseudodensity.
     */
    for (int iat = 0; iat < unit_cell_.num_atom_types(); iat++) {
        double R = unit_cell_.atom_type(iat).mt_radius();
        int na = unit_cell_.atom_type(iat).num_atoms();

        mdarray<double_complex, 2> pf;
        mdarray<double_complex, 2> qa;
        mdarray<double_complex, 2> qapf;

        switch (ctx_.processing_unit()) {
            case CPU: {
                pf   = mdarray<double_complex, 2>(ctx_.mem_pool().allocate<double_complex, memory_t::host>(ngv * na),    ngv,   na);
                qa   = mdarray<double_complex, 2>(ctx_.mem_pool().allocate<double_complex, memory_t::host>(lmmax * na),  lmmax, na);
                qapf = mdarray<double_complex, 2>(ctx_.mem_pool().allocate<double_complex, memory_t::host>(lmmax * ngv), lmmax, ngv);
                break;
            }
            case GPU: {
                /* allocate on GPU */
                pf   = mdarray<double_complex, 2>(nullptr, ctx_.mem_pool().allocate<double_complex, memory_t::device>(ngv * na), ngv, na);
                /* allocate on CPU & GPU */
                qa   = mdarray<double_complex, 2>(ctx_.mem_pool().allocate<double_complex, memory_t::host>(lmmax * na),
                                                  ctx_.mem_pool().allocate<double_complex, memory_t::device>(lmmax * na), lmmax, na);
                /* allocate on CPU & GPU */
                qapf = mdarray<double_complex, 2>(ctx_.mem_pool().allocate<double_complex, memory_t::host>(ngv * lmmax), 
                                                  ctx_.mem_pool().allocate<double_complex, memory_t::device>(ngv * lmmax), lmmax, ngv);
                break;
            }
        }

        ctx_.generate_phase_factors(iat, pf);

        for (int i = 0; i < unit_cell_.atom_type(iat).num_atoms(); i++) {
            int ia = unit_cell_.atom_type(iat).atom_id(i);
            for (int lm = 0; lm < ctx_.lmmax_rho(); lm++) {
                qa(lm, i) = qmt__(lm, ia) - qit__(lm, ia);
            }
        }

        switch (ctx_.processing_unit()) {
            case CPU: {
                linalg<CPU>::gemm(0, 2, ctx_.lmmax_rho(), ctx_.gvec().count(), unit_cell_.atom_type(iat).num_atoms(),
                                  qa, pf, qapf);
                break;
            }
            case GPU: {
#if defined(__GPU)
                qa.copy<memory_t::host, memory_t::device>();
                linalg<GPU>::gemm(0, 2, ctx_.lmmax_rho(), ctx_.gvec().count(), unit_cell_.atom_type(iat).num_atoms(),
                                  qa.at<GPU>(), qa.ld(),
                                  pf.at<GPU>(), pf.ld(),
                                  qapf.at<GPU>(), qapf.ld());
                qapf.copy<memory_t::device, memory_t::host>();
#endif
                break;
            }
        }

        /* add pseudo_density to interstitial charge density so that rho(G) has the correct 
         * multipole moments in the muffin-tins */
        #pragma omp parallel for schedule(static)
        for (int igloc = 0; igloc < ctx_.gvec().count(); igloc++) {
            int ig = ctx_.gvec().offset() + igloc;

            double gR = ctx_.gvec().gvec_len(ig) * R;
            double gRn = std::pow(2.0 / gR, pseudo_density_order_ + 1);

            double_complex rho_G(0, 0);
            if (ig) { // G!=0
                /* loop over atoms of the same type */
                for (int l = 0, lm = 0; l <= ctx_.lmax_rho(); l++) {
                    double_complex zt1(0, 0);
                    for (int m = -l; m <= l; m++, lm++) {
                        zt1 += gvec_ylm_(lm, igloc) * qapf(lm, igloc);
                    }
                    rho_G += (fourpi / unit_cell_.omega()) * std::conj(zil_[l]) * zt1 * gamma_factors_R_(l, iat) * 
                             sbessel_mt_(l + pseudo_density_order_ + 1, igloc, iat) * gRn;
                } // l
            } else { // G=0
                for (int i = 0; i < unit_cell_.atom_type(iat).num_atoms(); i++) {
                    int ia = unit_cell_.atom_type(iat).atom_id(i);
                    rho_G += (fourpi / unit_cell_.omega()) * y00 * (qmt__(0, ia) - qit__(0, ia));
                }
            }
            rho_pw__[igloc] += rho_G;
        }
        ctx_.mem_pool().reset<memory_t::host>();
        if (ctx_.processing_unit() == GPU) {
            ctx_.mem_pool().reset<memory_t::device>();
        }
    }
}

inline void Potential::poisson(Periodic_function<double> const& rho)
{
    PROFILE("sirius::Potential::poisson");

    /* in case of full potential we need to do pseudo-charge multipoles */
    if (ctx_.full_potential()) {

        /* true multipole moments */
        mdarray<double_complex, 2> qmt(ctx_.lmmax_rho(), unit_cell_.num_atoms());
        poisson_vmt(rho, qmt);
        
        //== for (int ia = 0; ia < unit_cell_.num_atoms(); ia++) {
        //==     for (int lm = 0; lm < ctx_.lmmax_rho(); lm++) {
        //==         printf("qmt(%2i, %2i) = %18.12f %18.12f\n", lm, ia, qmt(lm, ia).real(), qmt(lm, ia).imag());
        //==     }
        //==     printf("\n");
        //== }

        if (ctx_.control().print_checksum_) {
            if (ctx_.comm().rank() == 0) {
                utils::print_checksum("qmt", qmt.checksum());
            }
        }

        /* compute multipoles of interstitial density in MT region */
        //mdarray<double_complex, 2> qit(ctx_.lmmax_rho(), unit_cell_.num_atoms());
        //poisson_sum_G(ctx_.lmmax_rho(), &rho.f_pw_local(0), sbessel_mom_, qit);
        auto qit = ctx_.sum_fg_fl_yg(ctx_.lmax_rho(), &rho.f_pw_local(0), sbessel_mom_, gvec_ylm_);

        //== for (int ia = 0; ia < unit_cell_.num_atoms(); ia++) {
        //==     for (int lm = 0; lm < ctx_.lmmax_rho(); lm++) {
        //==         printf("qi(%2i, %2i) = %18.12f %18.12f\n", lm, ia, qit(lm, ia).real(), qit(lm, ia).imag());
        //==     }
        //==     printf("\n");
        //== }

        if (ctx_.control().print_checksum_) {
            if (ctx_.comm().rank() == 0) {
                utils::print_checksum("qit", qit.checksum());
            }
        }

        /* add contribution from the pseudo-charge */
        poisson_add_pseudo_pw(qmt, qit, const_cast<double_complex*>(&rho.f_pw_local(0)));

        if (ctx_.control().verification_ >= 2) {
            auto qit = ctx_.sum_fg_fl_yg(ctx_.lmax_rho(), &rho.f_pw_local(0), sbessel_mom_, gvec_ylm_);

            double d = 0.0;
            for (int ia = 0; ia < unit_cell_.num_atoms(); ia++) {
                for (int lm = 0; lm < ctx_.lmmax_rho(); lm++) {
                    d += abs(qmt(lm, ia) - qit(lm, ia));
                }
            }
            printf("pseudocharge error: %18.10f\n", d);
        }
    }

    /* compute pw coefficients of Hartree potential */
    int ig0{0};
    if (ctx_.gvec().comm().rank() == 0) {
        hartree_potential_->f_pw_local(0) = 0.0;
        ig0 = 1;
    }
    if (!ctx_.molecule()) {
        #pragma omp parallel for schedule(static)
        for (int igloc = ig0; igloc < ctx_.gvec().count(); igloc++) {
            int ig = ctx_.gvec().offset() + igloc;
            hartree_potential_->f_pw_local(igloc) = (fourpi * rho.f_pw_local(igloc) / std::pow(ctx_.gvec().gvec_len(ig), 2));
        }
    } else {
        double R_cut = 0.5 * std::pow(unit_cell_.omega(), 1.0 / 3);
        #pragma omp parallel for schedule(static)
        for (int igloc = ig0; igloc < ctx_.gvec().count(); igloc++) {
            int ig = ctx_.gvec().offset() + igloc;
            hartree_potential_->f_pw_local(igloc) = (fourpi * rho.f_pw_local(igloc) / std::pow(ctx_.gvec().gvec_len(ig), 2)) *
                                                    (1.0 - std::cos(ctx_.gvec().gvec_len(ig) * R_cut));
        }
    }

    //if (ctx_.control().print_checksum_) {
    //    auto z4 = mdarray<double_complex, 1>(&vh->f_pw(0), ctx_.gvec().num_gvec()).checksum();
    //    if (ctx_.comm().rank() == 0) {
    //        DUMP("checksum(vh_pw): %20.14f %20.14f", z4.real(), z4.imag());
    //    }
    //}

    /* boundary condition for muffin-tins */
    if (ctx_.full_potential()) {
        /* compute V_lm at the MT boundary */
        auto vmtlm = ctx_.sum_fg_fl_yg(ctx_.lmax_pot(), &hartree_potential_->f_pw_local(0), sbessel_mt_, gvec_ylm_);

        /* add boundary condition and convert to Rlm */
        utils::timer t1("sirius::Potential::poisson|bc");
        mdarray<double, 3> rRl(unit_cell_.max_num_mt_points(), ctx_.lmax_pot() + 1, unit_cell_.num_atom_types());
        for (int iat = 0; iat < unit_cell_.num_atom_types(); iat++) {
            int nmtp = unit_cell_.atom_type(iat).num_mt_points();
            double R = unit_cell_.atom_type(iat).mt_radius();

            #pragma omp parallel for default(shared)
            for (int l = 0; l <= ctx_.lmax_pot(); l++) {
                for (int ir = 0; ir < nmtp; ir++) {
                    rRl(ir, l, iat) = std::pow(unit_cell_.atom_type(iat).radial_grid(ir) / R, l);
                }
            }
        }

        for (int ialoc = 0; ialoc < unit_cell_.spl_num_atoms().local_size(); ialoc++) {
            int ia = unit_cell_.spl_num_atoms(ialoc);
            int nmtp = unit_cell_.atom(ia).num_mt_points();

            std::vector<double> vlm(ctx_.lmmax_pot());
            SHT::convert(ctx_.lmax_pot(), &vmtlm(0, ia), &vlm[0]);

            #pragma omp parallel for default(shared)
            for (int lm = 0; lm < ctx_.lmmax_pot(); lm++) {
                int l = l_by_lm_[lm];

                for (int ir = 0; ir < nmtp; ir++) {
                    hartree_potential_->f_mt<index_domain_t::local>(lm, ir, ialoc) += 
                        vlm[lm] * rRl(ir, l, unit_cell_.atom(ia).type_id());
                }
            }
            /* save electronic part of the potential at the point of origin */
#ifdef __VHA_AUX
            vh_el_(ia) = y00 * hartree_potential_->f_mt<index_domain_t::local>(0, 0, ialoc) + unit_cell_.atom(ia).zn() / unit_cell_.atom(ia).radial_grid(0);
#else
            vh_el_(ia) = y00 * hartree_potential_->f_mt<index_domain_t::local>(0, 0, ialoc);
#endif
        }
        ctx_.comm().allgather(vh_el_.at<CPU>(), unit_cell_.spl_num_atoms().global_offset(),
                              unit_cell_.spl_num_atoms().local_size());
    }

    /* transform Hartree potential to real space */
    hartree_potential_->fft_transform(1);

    if (ctx_.control().print_checksum_) {
        auto cs = hartree_potential_->checksum_rg();
        if (ctx_.comm().rank() == 0) {
            utils::print_checksum("vha_rg", cs);
        }
    }

    /* compute contribution from the smooth part of Hartree potential */
    energy_vha_ = rho.inner(hartree_potential());

#ifndef __VHA_AUX
    /* add nucleus potential and contribution to Hartree energy */
    if (ctx_.full_potential()) {
        double evha_nuc{0};
        for (int ialoc = 0; ialoc < unit_cell_.spl_num_atoms().local_size(); ialoc++) {
            int ia = unit_cell_.spl_num_atoms(ialoc);
            auto& atom = unit_cell_.atom(ia);
            Spline<double> srho(atom.radial_grid());
            for (int ir = 0; ir < atom.num_mt_points(); ir++) {
                double r = atom.radial_grid(ir);
                hartree_potential_->f_mt<index_domain_t::local>(0, ir, ialoc) -= atom.zn() / r / y00;
                srho(ir) = rho.f_mt<index_domain_t::local>(0, ir, ialoc);
            }
            evha_nuc -= atom.zn() * srho.interpolate().integrate(1) / y00;
        }
        ctx_.comm().allreduce(&evha_nuc, 1);
        energy_vha_ += evha_nuc;
    }
#endif
}
