// Copyright (c) 2013-2018 Mathieu Taillefumier, 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 apply_hubbard_potential.hpp
 *
 *  \brief Contains implementation of sirius::Hubbard::apply_hubbard_potential() function.
 */

// this function computes the hubbard contribution to the hamiltonian
// and add it to ophi.

// the S matrix is already applied to phi_i

void Hubbard::apply_hubbard_potential(K_point&        kp__,
                                      const int       ispn__,
                                      const int       idx__,
                                      const int       n__,
                                      Wave_functions& phi,
                                      Wave_functions& ophi)
{
    auto& hub_wf = kp__.hubbard_wave_functions();

    dmatrix<double_complex> dm(this->number_of_hubbard_orbitals(), n__);
    dm.zero();

    #ifdef __GPU
    if (ctx_.processing_unit() == GPU) {
        dm.allocate(memory_t::device);
    }
    #endif

    // First calculate the local part of the projections
    // dm(i, n)  = <S phi_i | psi_{nk}>

    inner(ctx_.processing_unit(),
          ispn__,
          hub_wf,
          0,
          this->number_of_hubbard_orbitals(),
          phi,
          idx__,
          n__,
          dm,
          0,
          0);

// free memory on GPU
    #ifdef __GPU
    if (ctx_.processing_unit() == GPU) {
        dm.copy<memory_t::device, memory_t::host>();
    }
    #endif

    dmatrix<double_complex> Up;

    Up = dmatrix<double_complex>(this->number_of_hubbard_orbitals(), n__);
    Up.zero();

    #ifdef __GPU
    // the communicator is always of size 1.  I need to allocate memory
    // on the device manually

    if (ctx_.processing_unit() == GPU) {
        Up.allocate(memory_t::device);
    }
    #endif

    #pragma omp parallel for schedule(static)
    for (int ia = 0; ia < ctx_.unit_cell().num_atoms(); ++ia) {
        const auto& atom = ctx_.unit_cell().atom(ia);
        if (atom.type().hubbard_correction()) {
            const int lmax_at = 2 * atom.type().hubbard_orbital(0).hubbard_l() + 1;
            // we apply the hubbard correction. For now I have no papers
            // giving me the formula for the SO case so I rely on QE for it
            // but I do not like it at all
            if (ctx_.num_mag_dims() == 3) {
                for (int s1 = 0; s1 < ctx_.num_spins(); s1++) {
                    for (int s2 = 0; s2 < ctx_.num_spins(); s2++) {
                        const int ind = (s1 == s2) * s1 + (1 + 2 * s2 + s1) * (s1 != s2);

                        // !!! Replace this with matrix matrix multiplication

                        for (int nbd = 0; nbd < n__; nbd++) {
                            for (int m1 = 0; m1 < lmax_at; m1++) {
                                for (int m2 = 0; m2 < lmax_at; m2++) {
                                    Up(this->offset[ia] + s1 * lmax_at + m1, nbd) += this->hubbard_potential_(m2, m1, ind, ia, 0) *
                                        dm(this->offset[ia] + s2 * lmax_at + m2, nbd);
                                }
                            }
                        }
                    }
                }
            } else {
                // Conventional LDA or colinear magnetism
                for (int nbd = 0; nbd < n__; nbd++) {
                    for (int m1 = 0; m1 < lmax_at; m1++) {
                        for (int m2 = 0; m2 < lmax_at; m2++) {
                            Up(this->offset[ia] + m1, nbd) += this->hubbard_potential_(m2, m1, ispn__, ia, 0) *
                                dm(this->offset[ia] + m2, nbd);
                        }
                    }
                }
            }
        }
    }

    #ifdef __GPU
    if (ctx_.processing_unit() == GPU) {
        Up.copy<memory_t::host, memory_t::device>();
    }
    #endif

    transform<double_complex>(ctx_.processing_unit(),
                              ispn__,
                              1.0,
                              hub_wf,
                              0,
                              this->number_of_hubbard_orbitals(),
                              Up,
                              0,
                              0,
                              1.0,
                              ophi,
                              idx__,
                              n__);

    #ifdef __GPU
    if (ctx_.processing_unit() == GPU) {
        Up.deallocate(memory_t::device);
    }
    #endif
}
