/* This file is part of SIRIUS electronic structure library.
 *
 * Copyright (c), ETH Zurich.  All rights reserved.
 *
 * Please, refer to the LICENSE file in the root directory.
 * SPDX-License-Identifier: BSD-3-Clause
 */

/** \file atom_symmetry_class.cpp
 *
 *  \brief Partial implementation of sirius::Atom_symmetry_class class.
 */

#include "atom_symmetry_class.hpp"
#include "core/la/eigensolver.hpp"

namespace sirius {

Atom_symmetry_class::Atom_symmetry_class(int id__, Atom_type const& atom_type__)
    : id_(id__)
    , atom_type_(atom_type__)
{
    if (!atom_type_.initialized()) {
        RTE_THROW("atom type is not initialized");
    }

    int nl        = atom_type_.indexr().lmax() + 1;
    int max_order = atom_type_.indexr().max_order();
    int nrf       = atom_type_.mt_radial_basis_size();

    surface_derivatives_ = mdarray<double, 2>({3, atom_type_.mt_radial_basis_size() - atom_type_.num_lo_descriptors()});

    radial_functions_ = mdarray<double, 3>({atom_type_.num_mt_points(), nrf, 2});

    h_spherical_integrals_ = mdarray<double, 2>({nrf, nrf});
    h_spherical_integrals_.zero();

    o_radial_integrals_ = mdarray<double, 3>({nl, max_order, max_order});
    o_radial_integrals_.zero();

    so_radial_integrals_ = mdarray<double, 3>({nl, max_order, max_order});
    so_radial_integrals_.zero();

    if (atom_type_.parameters().valence_relativity() == relativity_t::iora) {
        o1_radial_integrals_ = mdarray<double, 2>({nrf, nrf});
        o1_radial_integrals_.zero();
    }

    /* copy descriptors because enu is different between atom classes */
    for (int i = 0; i < atom_type_.num_aw_descriptors(); i++) {
        aw_descriptors_.push_back(atom_type_.aw_descriptor(i));
    }

    for (int i = 0; i < atom_type_.num_lo_descriptors(); i++) {
        lo_descriptors_.push_back(atom_type_.lo_descriptor(i));
    }
}

int
Atom_symmetry_class::generate_aw_radial_functions(relativity_t rel__, mdarray<double, 3>& rf__,
                                                  mdarray<double, 2>& sd__) const
{
    int nmtp = atom_type_.num_mt_points();

    Radial_solver solver(atom_type_.zn(), spherical_potential_, atom_type_.radial_grid());

    auto compute_all_orders = [&](int l, double enu_shift) -> bool {
        Spline<double> s(atom_type_.radial_grid());

        for (int order = 0; order < (int)aw_descriptor(l).size(); order++) {
            auto rsd = aw_descriptor(l)[order];

            auto idxrf = atom_type_.indexr().index_of(angular_momentum(l), order);

            /* integrate radial equation forward and find radial solution */
            auto result = solver.solve(rel__, rsd.dme, rsd.l, rsd.enu + enu_shift);

            for (int ir = 0; ir < nmtp; ir++) {
                rf__(ir, idxrf, 0) = result.p[ir];
                rf__(ir, idxrf, 1) = result.rdudr[ir];
            }
            for (int i : {0, 1, 2}) {
                sd__(i, idxrf) = result.uderiv[i];
            }

            /* orthogonalize to previous radial functions */
            for (int order1 = 0; order1 < order; order1++) {
                auto idxrf1 = atom_type_.indexr().index_of(angular_momentum(l), order1);

                for (int ir = 0; ir < nmtp; ir++) {
                    s(ir) = rf__(ir, idxrf, 0) * rf__(ir, idxrf1, 0);
                }

                /* <u_{\nu'}|u_{\nu}> */
                double ovlp = s.interpolate().integrate(0);

                for (int ir = 0; ir < nmtp; ir++) {
                    rf__(ir, idxrf, 0) -= rf__(ir, idxrf1, 0) * ovlp;
                    rf__(ir, idxrf, 1) -= rf__(ir, idxrf1, 1) * ovlp;
                }
                for (int i : {0, 1, 2}) {
                    sd__(i, idxrf) -= sd__(i, idxrf1) * ovlp;
                }
            }

            /* normalize orthogonal function */
            for (int ir = 0; ir < nmtp; ir++) {
                s(ir) = std::pow(rf__(ir, idxrf, 0), 2);
            }
            auto norm = s.interpolate().integrate(0);

            /* in case of linear dependency return failure */
            if (std::abs(norm) < 1e-8) {
                return false;
            }

            norm = 1.0 / std::sqrt(norm);

            for (int ir = 0; ir < nmtp; ir++) {
                rf__(ir, idxrf, 0) *= norm;
                rf__(ir, idxrf, 1) *= norm;
            }
            /* aw radial function can't be zero at MT boundary */
            if (std::abs(rf__(nmtp - 1, idxrf, 0)) < 1e-2) {
                return false;
            }
            for (int i : {0, 1, 2}) {
                sd__(i, idxrf) *= norm;
            }
        } // order
        return true;
    };

    bool found{true};

    #pragma omp parallel for schedule(dynamic, 1) default(shared)
    for (int l = 0; l < num_aw_descriptors(); l++) {
        /* try to increase linearisation energy several times in order to find linearly independent
         * radial functions */
        bool success{false};
        for (int k = 0; k < 100; k++) {
            if (l <= 3) {
                /* for low l numbers Enu finder will find the top of the band;
                 * in this case we need to go down in energy to remove any degeneracy of radial functions */
                success = compute_all_orders(l, -k * 0.05);
            } else {
                /* for high l values, Enu is typically set in the species files and is not searched;
                 * in case of trouble with them we need to increase linearisation energies */
                success = compute_all_orders(l, k * 0.25);
            }
            if (success) {
                break;
            }
        }
        if (success) {
            /* divide by r */
            for (int order = 0; order < (int)aw_descriptor(l).size(); order++) {
                auto idxrf = atom_type_.indexr().index_of(angular_momentum(l), order);
                for (int ir = 0; ir < nmtp; ir++) {
                    rf__(ir, idxrf, 0) *= atom_type_.radial_grid().x_inv(ir);
                }
            }
        } else {
            found = false;
        }
    } // l
    return found ? 0 : 1;
}

int
Atom_symmetry_class::generate_lo_radial_functions(relativity_t rel__, mdarray<double, 3>& rf__) const
{
    int nmtp = atom_type_.num_mt_points();

    Radial_solver solver(atom_type_.zn(), spherical_potential_, atom_type_.radial_grid());

    bool found{true};

    #pragma omp parallel for schedule(dynamic, 1)
    for (int idxlo = 0; idxlo < num_lo_descriptors(); idxlo++) {
        Spline<double> s(atom_type_.radial_grid());
        double a[3][3];
        double rderiv[3][3];

        /* number of radial functions */
        int num_rf = static_cast<int>(lo_descriptor(idxlo).rsd_set.size());
        RTE_ASSERT(num_rf <= 3);

        std::vector<std::vector<double>> u(num_rf);
        std::vector<std::vector<double>> rdudr(num_rf);

        for (int irf = 0; irf < num_rf; irf++) {
            auto rsd = lo_descriptor(idxlo).rsd_set[irf];

            auto result = solver.solve(rel__, rsd.dme, rsd.l, rsd.enu);

            u[irf]     = result.p;
            rdudr[irf] = result.rdudr;

            /* divide by r */
            for (int ir = 0; ir < nmtp; ir++) {
                /* store u(r) = p(r)/r */
                u[irf][ir] *= atom_type_.radial_grid().x_inv(ir);
            }

            for (int i = 0; i < num_rf; i++) {
                /* matrix of derivatives */
                a[irf][i] = rderiv[irf][i] = result.uderiv[i];
            }
        }

        double b[]    = {0, 0, 0};
        b[num_rf - 1] = 1.0;

        int info = la::wrap(la::lib_t::lapack).gesv(num_rf, 1, &a[0][0], 3, b, 3);

        if (info) {
            std::stringstream s;
            s << "a[i][j] = ";
            for (int i = 0; i < num_rf; i++) {
                for (int j = 0; j < num_rf; j++) {
                    s << rderiv[i][j] << " ";
                }
            }
            s << std::endl;
            s << "atom: " << atom_type_.label() << std::endl
              << "zn: " << atom_type_.zn() << std::endl
              << "l: " << lo_descriptor(idxlo).am.l() << std::endl;
            s << "gesv returned " << info;
            RTE_THROW(s);
        }

        /* index of local orbital radial function */
        auto idxrf = atom_type_.indexr().index_of(rf_lo_index(idxlo));
        /* take linear combination of radial solutions */
        for (int order = 0; order < num_rf; order++) {
            for (int ir = 0; ir < nmtp; ir++) {
                /* u(r) function */
                rf__(ir, idxrf, 0) += b[order] * u[order][ir];
                /* r(du/dr) function */
                rf__(ir, idxrf, 1) += b[order] * rdudr[order][ir];
            }
        }

        /* find norm of constructed local orbital */
        for (int ir = 0; ir < nmtp; ir++) {
            s(ir) = std::pow(rf__(ir, idxrf, 0), 2);
        }
        double norm = 1.0 / std::sqrt(s.interpolate().integrate(2));

        /* normalize */
        for (int ir = 0; ir < nmtp; ir++) {
            rf__(ir, idxrf, 0) *= norm;
            rf__(ir, idxrf, 1) *= norm;
        }

        if (std::abs(rf__(nmtp - 1, idxrf, 0)) > 1e-10) {
            std::stringstream s;
            s << "local orbital " << idxlo << " is not zero at MT boundary" << std::endl
              << "  atom symmetry class id : " << id() << " (" << atom_type().symbol() << ")" << std::endl
              << "  value : " << rf__(nmtp - 1, idxrf, 0) << std::endl
              << "  number of MT points: " << nmtp << std::endl
              << "  MT radius: " << atom_type_.radial_grid().last() << std::endl
              << "  matrix of derivatives:" << std::endl;
            for (int i = 0; i < num_rf; i++) {
                for (int j = 0; j < num_rf; j++) {
                    s << rderiv[i][j] << " ";
                }
                s << std::endl;
            }
            s << "  b_coeffs: ";
            for (int j = 0; j < num_rf; j++) {
                s << b[j] << " ";
            }
            s << std::endl;
            s << "  norm: " << norm << std::endl;
            double d{0};
            for (int i = 0; i < num_rf; i++) {
                d += b[i] * rderiv[i][0];
            }
            s << "  expected value at MT boundary from the linear equations: " << d << std::endl;
            for (int i = 0; i < num_rf; i++) {
                s << " rderiv, u: " << rderiv[i][0] << " " << u[i][nmtp - 1] << std::endl;
            }
            RTE_WARNING(s);
        }
    }

    if (found && atom_type_.parameters().cfg().control().verification() > 0 && num_lo_descriptors() > 0) {
        check_lo_linear_independence(0.0001);
    }

    return found ? 0 : 1;
}

void
Atom_symmetry_class::orthogonalize_radial_functions()
{
    int nmtp = atom_type().num_mt_points();
    Spline<double> s(atom_type().radial_grid());
    /* orthogonalize local orbitals */
    for (int l = 0; l <= atom_type().indexr().lmax_lo(); l++) {
        for (int j = 0; j < atom_type().indexr().num_lo(l); j++) {
            int idxrf = atom_type().indexr().index_of(angular_momentum(l), j + atom_type_.aw_order(l));
            for (int j1 = 0; j1 < j; j1++) {
                int idxrf1 = atom_type().indexr().index_of(angular_momentum(l), j1 + atom_type().aw_order(l));

                for (int ir = 0; ir < nmtp; ir++) {
                    s(ir) = radial_functions_(ir, idxrf, 0) * radial_functions_(ir, idxrf1, 0);
                }
                /* <u_{j}|u_{j'}> */
                double ovlp = s.interpolate().integrate(2);

                for (int ir = 0; ir < nmtp; ir++) {
                    radial_functions_(ir, idxrf, 0) -= radial_functions_(ir, idxrf1, 0) * ovlp;
                    radial_functions_(ir, idxrf, 1) -= radial_functions_(ir, idxrf1, 1) * ovlp;
                }
            }
            /* normalize again */
            for (int ir = 0; ir < nmtp; ir++) {
                s(ir) = std::pow(radial_functions_(ir, idxrf, 0), 2);
            }
            auto norm = s.interpolate().integrate(2);

            if (std::abs(norm) < 1e-10) {
                std::stringstream s;
                s << "LO radial function for atom " << atom_type_.label() << " is linearly dependent" << std::endl
                  << "  l : " << l << std::endl
                  << "  index for a given l : " << j;
            }

            norm = 1.0 / std::sqrt(norm);

            for (int ir = 0; ir < nmtp; ir++) {
                radial_functions_(ir, idxrf, 0) *= norm;
                radial_functions_(ir, idxrf, 1) *= norm;
            }
        }
    }
}

std::vector<int>
Atom_symmetry_class::check_lo_linear_independence(double tol__) const
{
    int nmtp = atom_type_.num_mt_points();

    Spline<double> s(atom_type_.radial_grid());
    la::dmatrix<double> loprod(num_lo_descriptors(), num_lo_descriptors());
    loprod.zero();
    for (int idxlo1 = 0; idxlo1 < num_lo_descriptors(); idxlo1++) {

        int idxrf1 = atom_type_.indexr().index_of(rf_lo_index(idxlo1));

        for (int idxlo2 = 0; idxlo2 < num_lo_descriptors(); idxlo2++) {

            int idxrf2 = atom_type_.indexr().index_of(rf_lo_index(idxlo2));

            if (lo_descriptor(idxlo1).am == lo_descriptor(idxlo2).am) {

                for (int ir = 0; ir < nmtp; ir++) {
                    s(ir) = radial_functions_(ir, idxrf1, 0) * radial_functions_(ir, idxrf2, 0);
                }
                loprod(idxlo1, idxlo2) = s.interpolate().integrate(2);
            }
        }
    }

    mdarray<double, 2> ovlp({num_lo_descriptors(), num_lo_descriptors()});
    copy(loprod, ovlp);

    auto stdevp = la::Eigensolver_factory("lapack");

    std::vector<double> loprod_eval(num_lo_descriptors());
    la::dmatrix<double> loprod_evec(num_lo_descriptors(), num_lo_descriptors());

    stdevp->solve(num_lo_descriptors(), loprod, &loprod_eval[0], loprod_evec);

    if (loprod_eval[0] < tol__) {
        std::printf("\n");
        std::printf("local orbitals for atom symmetry class %i are almost linearly dependent\n", id_);
        std::printf("local orbitals overlap matrix:\n");
        for (int i = 0; i < num_lo_descriptors(); i++) {
            for (int j = 0; j < num_lo_descriptors(); j++) {
                std::printf("%12.6f", ovlp(i, j));
            }
            std::printf("\n");
        }
        std::printf("overlap matrix eigen-values:\n");
        for (int i = 0; i < num_lo_descriptors(); i++) {
            std::printf("%12.6f", loprod_eval[i]);
        }
        std::printf("\n");
        std::printf("smallest eigenvalue: %20.16f\n", loprod_eval[0]);
    }

    std::vector<int> inc(num_lo_descriptors(), 0);

    /* try all local orbitals */
    for (int i = 0; i < num_lo_descriptors(); i++) {
        inc[i] = 1;

        std::vector<int> ilo;
        for (int j = 0; j < num_lo_descriptors(); j++) {
            if (inc[j] == 1) {
                ilo.push_back(j);
            }
        }

        std::vector<double> eval(ilo.size());
        la::dmatrix<double> evec(static_cast<int>(ilo.size()), static_cast<int>(ilo.size()));
        la::dmatrix<double> tmp(static_cast<int>(ilo.size()), static_cast<int>(ilo.size()));
        for (int j1 = 0; j1 < (int)ilo.size(); j1++) {
            for (int j2 = 0; j2 < (int)ilo.size(); j2++) {
                tmp(j1, j2) = ovlp(ilo[j1], ilo[j2]);
            }
        }

        stdevp->solve(static_cast<int>(ilo.size()), tmp, &eval[0], evec);

        if (eval[0] < tol__) {
            std::printf("local orbital %i can be removed\n", i);
            inc[i] = 0;
        }
    }
    return inc;
}

void
Atom_symmetry_class::dump_lo()
{
    std::stringstream s;
    s << "local_orbitals_" << id_ << ".dat";
    FILE* fout = fopen(s.str().c_str(), "w");

    for (int ir = 0; ir < atom_type_.num_mt_points(); ir++) {
        fprintf(fout, "%f ", atom_type_.radial_grid(ir));
        for (int idxlo = 0; idxlo < num_lo_descriptors(); idxlo++) {
            int idxrf = atom_type_.indexr().index_of(rf_lo_index(idxlo));
            fprintf(fout, "%f ", radial_functions_(ir, idxrf, 0));
        }
        fprintf(fout, "\n");
    }
    fclose(fout);

    s.str("");
    s << "local_orbitals_deriv_" << id_ << ".dat";
    fout = fopen(s.str().c_str(), "w");

    for (int ir = 0; ir < atom_type_.num_mt_points(); ir++) {
        fprintf(fout, "%f ", atom_type_.radial_grid(ir));
        for (int idxlo = 0; idxlo < num_lo_descriptors(); idxlo++) {
            int idxrf = atom_type_.indexr().index_of(rf_lo_index(idxlo));
            fprintf(fout, "%f ", radial_functions_(ir, idxrf, 1));
        }
        fprintf(fout, "\n");
    }
    fclose(fout);
}

void
Atom_symmetry_class::set_spherical_potential(std::vector<double> const& vs__)
{
    if (atom_type_.num_mt_points() != (int)vs__.size()) {
        RTE_THROW("wrong size of effective potential array");
    }

    spherical_potential_ = vs__;
}

int
Atom_symmetry_class::find_enu(relativity_t rel__)
{
    PROFILE("sirius::Atom_symmetry_class::find_enu");

    /* {n,l} -> enu map */
    std::map<std::pair<int, int>, double> nl_enu;

    /* find which aw functions need auto enu */
    for (int l = 0; l < num_aw_descriptors(); l++) {
        for (auto const& d : aw_descriptor(l)) {
            if (d.auto_enu) {
                nl_enu[{d.n, d.l}] = d.enu;
            }
        }
    }

    /* find which lo functions need auto enu */
    for (int idxlo = 0; idxlo < num_lo_descriptors(); idxlo++) {
        for (auto const& d : lo_descriptor(idxlo).rsd_set) {
            if (d.auto_enu) {
                nl_enu[{d.n, d.l}] = d.enu;
            }
        }
    }

    /* unroll {n,l} -> enu map to enable omp for loop */
    std::vector<std::pair<std::pair<int, int>, double>> nl_enu_vec(nl_enu.begin(), nl_enu.end());

    int ierr{0};
    #pragma omp parallel for reduction(+:ierr)
    for (size_t i = 0; i < nl_enu_vec.size(); i++) {
        try {
            int n                = nl_enu_vec[i].first.first;
            int l                = nl_enu_vec[i].first.second;
            nl_enu_vec[i].second = Enu_finder(rel__, atom_type_.zn(), n, l, atom_type_.radial_grid(),
                                              spherical_potential_, nl_enu_vec[i].second, 1)
                                           .enu();
        } catch (std::exception const& e) {
            std::cout << e.what() << std::endl;
            ierr++;
        }
    }
    /* update the {n,l} -> enu map */
    for (auto& e : nl_enu_vec) {
        nl_enu[e.first] = e.second;
    }

    for (int l = 0; l < num_aw_descriptors(); l++) {
        for (auto& d : aw_descriptor(l)) {
            if (d.auto_enu) {
                d.enu = nl_enu[{d.n, d.l}];
            }
        }
    }
    for (int idxlo = 0; idxlo < num_lo_descriptors(); idxlo++) {
        for (auto& d : lo_descriptor(idxlo).rsd_set) {
            if (d.auto_enu) {
                d.enu = nl_enu[{d.n, d.l}];
            }
        }
    }

    return ierr;
}

void
Atom_symmetry_class::generate_radial_functions(relativity_t rel__)
{
    PROFILE("sirius::Atom_symmetry_class::generate_radial_functions");

    mdarray<double, 2> sd({3, atom_type_.mt_radial_basis_size() - atom_type_.num_lo_descriptors()});
    mdarray<double, 3> rf({atom_type_.num_mt_points(), atom_type_.mt_radial_basis_size(), 2});
    sd.zero();
    rf.zero();

    auto ierr_enu = find_enu(rel__);
    if (ierr_enu) {
        std::stringstream s;
        s << "find_enu() failed for atom class " << id_;
        RTE_WARNING(s);
    }

    auto ierr_aw = generate_aw_radial_functions(rel__, rf, sd);
    if (ierr_aw) {
        std::stringstream s;
        s << "generate_aw_radial_functions() failed for atom class " << id_;
        RTE_WARNING(s);
    }

    auto ierr_lo = generate_lo_radial_functions(rel__, rf);
    if (ierr_lo) {
        std::stringstream s;
        s << "generate_lo_radial_functions() failed for atom class " << id_;
        RTE_WARNING(s);
    }

    if (ierr_aw + ierr_lo == 0) {
        copy(rf, radial_functions_);
        copy(sd, surface_derivatives_);
        if (atom_type().parameters().cfg().control().ortho_rf()) {
            orthogonalize_radial_functions();
        }
    }

    if (atom_type().parameters().cfg().control().save_rf()) {
        static int count{0};

        if (true) {
            std::stringstream s;
            s << "radial_functions_class_" << id_ << "_step_" << count << ".dat";
            std::ofstream ofs(s.str(), std::ofstream::out | std::ofstream::trunc);

            for (int ir = 0; ir < atom_type_.num_mt_points(); ir++) {
                ofs << std::setprecision(12) << atom_type_.radial_grid(ir) << " ";
                for (int idxrf = 0; idxrf < atom_type_.indexr().size(); idxrf++) {
                    ofs << std::setprecision(12) << radial_functions_(ir, idxrf, 0) << " "
                        << radial_functions_(ir, idxrf, 1) << " ";
                }
                ofs << std::endl;
            }
        }

        if (true) {
            std::stringstream s;
            s << "radial_potential_class_" << id_ << "_step_" << count << ".dat";
            std::ofstream ofs(s.str(), std::ofstream::out | std::ofstream::trunc);

            for (int ir = 0; ir < atom_type_.num_mt_points(); ir++) {
                ofs << std::setprecision(12) << atom_type_.radial_grid(ir) << " " << spherical_potential_[ir] << " "
                    << spherical_potential_[ir] + atom_type_.zn() / atom_type_.radial_grid(ir) << std::endl;
            }
        }

        count++;
    }
}

void
Atom_symmetry_class::sync_radial_functions(mpi::Communicator const& comm__, int const rank__)
{
    /* don't broadcast Hamiltonian radial functions, because they are used locally */
    int size = (int)(radial_functions_.size(0) * radial_functions_.size(1));
    comm__.bcast(radial_functions_.at(memory_t::host), size, rank__);
    comm__.bcast(surface_derivatives_.at(memory_t::host), (int)surface_derivatives_.size(), rank__);
}

void
Atom_symmetry_class::sync_radial_integrals(mpi::Communicator const& comm__, int const rank__)
{
    comm__.bcast(h_spherical_integrals_.at(memory_t::host), (int)h_spherical_integrals_.size(), rank__);
    comm__.bcast(o_radial_integrals_.at(memory_t::host), (int)o_radial_integrals_.size(), rank__);
    comm__.bcast(so_radial_integrals_.at(memory_t::host), (int)so_radial_integrals_.size(), rank__);
    if (atom_type_.parameters().valence_relativity() == relativity_t::iora) {
        comm__.bcast(o1_radial_integrals_.at(memory_t::host), (int)o1_radial_integrals_.size(), rank__);
    }
}

void
Atom_symmetry_class::generate_radial_integrals(relativity_t rel__)
{
    PROFILE("sirius::Atom_symmetry_class::generate_radial_integrals");

    int nmtp = atom_type_.num_mt_points();

    double sq_alpha_half = 0.5 * std::pow(speed_of_light, -2);
    if (rel__ == relativity_t::none) {
        sq_alpha_half = 0;
    }

    h_spherical_integrals_.zero();
    #pragma omp parallel default(shared)
    {
        Spline<double> s(atom_type_.radial_grid());
        #pragma omp for
        for (int i1 = 0; i1 < atom_type_.mt_radial_basis_size(); i1++) {
            for (int i2 = 0; i2 < atom_type_.mt_radial_basis_size(); i2++) {
                /* for spherical part of potential integrals are diagonal in l */
                if (atom_type_.indexr(i1).am.l() == atom_type_.indexr(i2).am.l()) {
                    int ll = atom_type_.indexr(i1).am.l() * (atom_type_.indexr(i1).am.l() + 1);
                    for (int ir = 0; ir < nmtp; ir++) {
                        double Minv = 1.0 / (1 - spherical_potential_[ir] * sq_alpha_half);
                        /* u_1(r) * u_2(r) */
                        double t0 = radial_functions_(ir, i1, 0) * radial_functions_(ir, i2, 0);
                        /* r*u'_1(r) * r*u'_2(r) */
                        double t1 = radial_functions_(ir, i1, 1) * radial_functions_(ir, i2, 1);
                        s(ir)     = 0.5 * t1 * Minv +
                                t0 * (0.5 * ll * Minv +
                                      spherical_potential_[ir] * std::pow(atom_type_.radial_grid(ir), 2));
                    }
                    h_spherical_integrals_(i1, i2) = s.interpolate().integrate(0) / y00;
                }
            }
        }
    }

    o_radial_integrals_.zero();
    #pragma omp parallel default(shared)
    {
        Spline<double> s(atom_type_.radial_grid());
        #pragma omp for
        for (int l = 0; l <= atom_type_.indexr().lmax(); l++) {
            int nrf = atom_type_.indexr().max_order(l);

            for (int order1 = 0; order1 < nrf; order1++) {
                auto idxrf1 = atom_type_.indexr().index_of(angular_momentum(l), order1);
                for (int order2 = 0; order2 < nrf; order2++) {
                    auto idxrf2 = atom_type_.indexr().index_of(angular_momentum(l), order2);
                    if (order1 == order2) {
                        o_radial_integrals_(l, order1, order2) = 1.0;
                    } else {
                        if (atom_type_.parameters().cfg().settings().simple_lapw_ri()) {
                            for (int ir = 0; ir < nmtp; ir++) {
                                s(ir) = radial_functions_(ir, idxrf1, 0) * radial_functions_(ir, idxrf2, 0) *
                                        std::pow(atom_type_.radial_grid(ir), 2);
                            }
                            o_radial_integrals_(l, order1, order2) = s.interpolate().integrate(0);
                        } else {
                            for (int ir = 0; ir < nmtp; ir++) {
                                s(ir) = radial_functions_(ir, idxrf1, 0) * radial_functions_(ir, idxrf2, 0);
                            }
                            o_radial_integrals_(l, order1, order2) = s.interpolate().integrate(2);
                        }
                    }
                }
            }
        }
    }
    if (atom_type_.parameters().valence_relativity() == relativity_t::iora) {
        o1_radial_integrals_.zero();
        #pragma omp parallel for
        for (int i1 = 0; i1 < atom_type_.mt_radial_basis_size(); i1++) {
            Spline<double> s(atom_type_.radial_grid());
            for (int i2 = 0; i2 < atom_type_.mt_radial_basis_size(); i2++) {
                /* for spherical part of potential integrals are diagonal in l */
                if (atom_type_.indexr(i1).am.l() == atom_type_.indexr(i2).am.l()) {
                    int ll = atom_type_.indexr(i1).am.l() * (atom_type_.indexr(i1).am.l() + 1);
                    for (int ir = 0; ir < nmtp; ir++) {
                        double Minv2 = std::pow(1 - spherical_potential_[ir] * sq_alpha_half, -2);
                        /* u_1(r) * u_2(r) */
                        double t0 = radial_functions_(ir, i1, 0) * radial_functions_(ir, i2, 0);
                        /* r*u'_1(r) * r*u'_2(r) */
                        double t1 = radial_functions_(ir, i1, 1) * radial_functions_(ir, i2, 1);
                        s(ir)     = sq_alpha_half * 0.5 * Minv2 * (t1 + t0 * ll);
                    }
                    o1_radial_integrals_(i1, i2) = s.interpolate().integrate(0);
                }
            }
        }
    }

    if (atom_type_.parameters().so_correction()) {
        double soc = std::pow(2 * speed_of_light, -2);

        Spline<double> s(atom_type_.radial_grid());
        Spline<double> s1(atom_type_.radial_grid());
        Spline<double> ve(atom_type_.radial_grid());

        for (int i = 0; i < nmtp; i++) {
            ve(i) = spherical_potential_[i] + atom_type_.zn() / atom_type_.radial_grid(i);
        }
        ve.interpolate();

        so_radial_integrals_.zero();
        for (int l = 0; l <= atom_type_.indexr().lmax(); l++) {
            int nrf = atom_type_.indexr().max_order(l);

            for (int order1 = 0; order1 < nrf; order1++) {
                auto idxrf1 = atom_type_.indexr().index_of(angular_momentum(l), order1);
                for (int order2 = 0; order2 < nrf; order2++) {
                    auto idxrf2 = atom_type_.indexr().index_of(angular_momentum(l), order2);

                    for (int ir = 0; ir < nmtp; ir++) {
                        double M = 1.0 - 2 * soc * spherical_potential_[ir];
                        /* first part <f| dVe / dr |f'> */
                        s(ir) = radial_functions_(ir, idxrf1, 0) * radial_functions_(ir, idxrf2, 0) * soc *
                                ve.deriv(1, ir) / pow(M, 2);

                        /* second part <f| d(z/r) / dr |f'> */
                        s1(ir) = radial_functions_(ir, idxrf1, 0) * radial_functions_(ir, idxrf2, 0) * soc *
                                 atom_type_.zn() / pow(M, 2);
                    }
                    s.interpolate();
                    s1.interpolate();

                    so_radial_integrals_(l, order1, order2) = s.integrate(1) + s1.integrate(-1);
                }
            }
        }
    }
}

void
Atom_symmetry_class::write_enu(mpi::pstdout& pout) const
{
    pout << "Atom : " << atom_type_.symbol() << ", class id : " << id_ << std::endl;
    pout << "augmented waves" << std::endl;
    for (int l = 0; l < num_aw_descriptors(); l++) {
        for (size_t order = 0; order < aw_descriptor(l).size(); order++) {
            auto& rsd = aw_descriptor(l)[order];
            if (rsd.auto_enu) {
                pout << rsd << std::endl;
            }
        }
    }

    pout << "local orbitals" << std::endl;
    for (int idxlo = 0; idxlo < num_lo_descriptors(); idxlo++) {
        for (size_t order = 0; order < lo_descriptor(idxlo).rsd_set.size(); order++) {
            auto& rsd = lo_descriptor(idxlo).rsd_set[order];
            if (rsd.auto_enu) {
                pout << rsd << std::endl;
            }
        }
    }
    pout << std::endl;
}

} // namespace sirius
