"""Simple magnons calculation in a 1D hydrogen chain."""

# General modules
import pytest
import numpy as np

# Script modules
from ase import Atoms
from ase.dft.kpoints import monkhorst_pack

from gpaw import PW, GPAW
from gpaw.mpi import world
from gpaw.test import findpeak
from gpaw.response.tms import TransverseMagneticSusceptibility
from gpaw.response.susceptibility import read_macroscopic_component


@pytest.mark.kspair
@pytest.mark.response
def test_response_afm_hchain_gssALDA(in_tmp_dir):
    # ---------- Inputs ---------- #

    # Part 1: Ground state calculation
    # Define atomic structure
    a = 2.5
    vfactor = 0.8
    mm = 1.
    # Ground state calculator options
    xc = 'LDA'
    kpts = 12
    nbands = 2 * (1 + 0)  # 1s + 0 empty shell bands
    ebands = 2 * 1  # Include also 2s bands for numerical consistency
    pw = 250
    conv = {'bands': nbands}

    # # Part 2: Magnetic response calculation
    q_qc = [[0., 0., 0.],
            [1. / 6., 0., 0.],
            [1. / 3., 0., 0.]]
    fxc = 'ALDA'
    fxc_scaling = [True, None, 'afm']
    rshelmax = -1
    rshewmin = 1e-8
    ecut = 120
    frq_w = np.linspace(-0.6, 0.6, 41)
    eta = 0.24
    if world.size > 1:
        nblocks = 2
    else:
        nblocks = 1

    # ---------- Script ---------- #

    # Part 1: Ground state calculation

    Hatom = Atoms('H',
                  cell=[a, 0, 0],
                  pbc=[1, 0, 0])
    Hatom.center(vacuum=vfactor * a, axis=(1, 2))
    Hchain = Hatom.repeat((2, 1, 1))
    Hchain.set_initial_magnetic_moments([mm, -mm])

    calc = GPAW(xc=xc,
                mode=PW(pw),
                kpts=monkhorst_pack((kpts, 1, 1)),
                nbands=nbands + ebands,
                convergence=conv,
                symmetry={'point_group': True},
                parallel={'domain': 1})

    Hchain.calc = calc
    Hchain.get_potential_energy()

    # Part 2: Magnetic response calculation
    fxckwargs = {'rshelmax': rshelmax,
                 'rshewmin': rshewmin,
                 'fxc_scaling': fxc_scaling}
    tms = TransverseMagneticSusceptibility(calc,
                                           nbands=nbands,
                                           fxc=fxc,
                                           eta=eta,
                                           ecut=ecut,
                                           fxckwargs=fxckwargs,
                                           gammacentered=True,
                                           nblocks=nblocks)
    for q, q_c in enumerate(q_qc):
        tms.get_macroscopic_component('+-', q_c, frq_w,
                                      filename=('h-chain_macro_tms'
                                                + '_q%d.csv' % q))

    tms.write_timer()
    world.barrier()

    # Part 3: Identify magnon peak in finite q scattering function
    w0_w, _, chi0_w = read_macroscopic_component('h-chain_macro_tms_q0.csv')
    w1_w, _, chi1_w = read_macroscopic_component('h-chain_macro_tms_q1.csv')
    w2_w, _, chi2_w = read_macroscopic_component('h-chain_macro_tms_q2.csv')

    wpeak1, Ipeak1 = findpeak(w1_w, -chi1_w.imag / np.pi)
    wpeak2, Ipeak2 = findpeak(w2_w, -chi2_w.imag / np.pi)
    mw1 = calculate_afm_mw(wpeak1, eta) * 1000  # to meV
    mw2 = calculate_afm_mw(wpeak2, eta) * 1000  # to meV

    # Part 4: compare new results to test values
    test_fxcs = 1.04744
    test_mw1 = 285.  # meV
    test_mw2 = 494.  # meV  # Remark that mw2 < 2 * mw1 (linear dispersion)
    test_Ipeak1 = 0.0131  # a.u.
    test_Ipeak2 = 0.0290

    # Test fxc_scaling:
    assert abs(fxc_scaling[1] - test_fxcs) < 0.005

    # Magnon peak at q=1/3 q_X:
    assert abs(mw1 - test_mw1) < 10.

    # Magnon peak at q=2/3 q_X:
    assert abs(mw2 - test_mw2) < 10.

    # Scattering function intensity:
    assert abs(Ipeak1 - test_Ipeak1) < 0.005
    assert abs(Ipeak2 - test_Ipeak2) < 0.005

    # Check that spectrum at q=0 vanishes
    chitol = 1e-3 * np.abs(chi1_w.imag).max()
    assert np.abs(chi0_w.imag).max() < chitol

    # Check that the spectrum is antisymmetric around q=0
    assert np.allclose(w0_w[19::-1] + w0_w[21:], 0.)
    assert np.allclose(chi0_w.imag[19::-1] + chi0_w.imag[21:], 0.,
                       atol=0.01 * chitol)
    assert np.allclose(chi1_w.imag[19::-1] + chi1_w.imag[21:], 0.,
                       atol=chitol)
    assert np.allclose(chi2_w.imag[19::-1] + chi2_w.imag[21:], 0.,
                       atol=chitol)


# ---------- Script functionality ---------- #


def calculate_afm_mw(peak_frequency, eta):
    """Assume lorentzian lineshape to compute the magnon frequency
    from the afm magnon lineshape peak position."""
    return np.sqrt(2 * np.sqrt(peak_frequency**4 + eta**2 * peak_frequency**2)
                   - peak_frequency**2 - eta**2)
