import numpy as np

from kliff.dataset import Configuration
from kliff.descriptors import Bispectrum

zeta_ref = [
    [
        9.67233961e00,
        7.11696281e00,
        3.64267853e00,
        9.67186730e00,
        5.16158022e00,
        1.87921544e00,
        6.30815842e00,
        8.92398717e00,
        4.20835447e00,
        2.21451314e00,
        2.50615741e00,
        1.33881116e01,
        4.60013516e00,
        5.06161573e00,
    ],
    [
        4.12842991e01,
        9.07588360e00,
        1.53655662e00,
        6.20291097e00,
        1.87679232e00,
        1.07391872e00,
        1.98717119e00,
        3.29696884e01,
        7.70422907e00,
        -3.31491113e-02,
        1.03289442e01,
        4.90925660e01,
        2.82436942e00,
        1.24081835e01,
    ],
    [
        1.73741554e01,
        5.52458737e00,
        1.96946128e-01,
        1.04292104e01,
        3.55805716e00,
        1.18314825e00,
        3.91265100e00,
        1.58723840e01,
        4.48329643e00,
        2.81638023e00,
        5.91001876e00,
        2.22956933e01,
        6.61593854e00,
        8.78960892e00,
    ],
    [
        7.54202120e00,
        3.88998085e00,
        1.84728448e00,
        7.92505231e00,
        3.17562384e00,
        2.57786432e00,
        5.79129219e00,
        8.38915899e00,
        4.19169728e00,
        2.52893629e00,
        4.54391666e00,
        1.59250759e01,
        7.75529446e00,
        8.71117842e00,
    ],
    [
        1.81428646e01,
        5.58779679e00,
        2.42492943e-01,
        1.06442756e01,
        3.60183088e00,
        5.19421191e-01,
        4.45926283e00,
        1.55305607e01,
        4.06832692e00,
        2.64394591e00,
        5.65285302e00,
        2.25947275e01,
        6.72632618e00,
        8.82676665e00,
    ],
    [
        7.20942183e00,
        3.35387725e00,
        1.29834905e00,
        7.40333237e00,
        3.04806432e00,
        2.62397206e00,
        5.47358955e00,
        8.82578443e00,
        4.07363243e00,
        3.14645616e00,
        5.32317864e00,
        1.48942847e01,
        7.29826349e00,
        8.05789635e00,
    ],
    [
        1.55876864e01,
        4.96629002e00,
        7.20501227e-02,
        9.24460991e00,
        3.26832361e00,
        1.05017170e00,
        4.04959834e00,
        1.59503590e01,
        4.56659280e00,
        3.47242538e00,
        7.21475752e00,
        2.12985011e01,
        6.19157000e00,
        9.13757773e00,
    ],
    [
        6.98474121e00,
        3.25489619e00,
        1.38505477e00,
        7.53831385e00,
        2.97284093e00,
        2.72033762e00,
        5.87839298e00,
        8.39930075e00,
        4.03605570e00,
        3.02566585e00,
        5.18825294e00,
        1.53404945e01,
        7.81374089e00,
        8.65233779e00,
    ],
]

dzeta_dr_001 = [
    [
        [-6.39268316e00, -6.86589539e00, -6.62024394e00],
        [2.16503211e00, 2.36645084e00, 2.38444841e00],
        [-3.21580741e-02, 2.11132242e00, 1.95472362e00],
        [0.00000000e00, 0.00000000e00, 0.00000000e00],
        [2.23695231e00, 1.20777528e-01, 2.34747873e00],
        [0.00000000e00, 0.00000000e00, 0.00000000e00],
        [2.02285681e00, 2.26734460e00, -6.64068138e-02],
        [0.00000000e00, 0.00000000e00, 0.00000000e00],
    ],
    [
        [-5.22011564e00, -5.33684406e00, -5.49690457e00],
        [3.14644477e00, 3.51638646e00, 3.53197204e00],
        [-1.12782961e-01, 8.68000925e-01, 7.94524758e-01],
        [0.00000000e00, 0.00000000e00, 0.00000000e00],
        [1.26962538e00, -8.51094111e-02, 1.33159255e00],
        [0.00000000e00, 0.00000000e00, 0.00000000e00],
        [9.16828446e-01, 1.03756608e00, -1.61184779e-01],
        [0.00000000e00, 0.00000000e00, 0.00000000e00],
    ],
]

dzeta_dr_minus_121 = [
    [
        [0.00000000e00, 0.00000000e00, 0.00000000e00],
        [3.12475437e-01, 3.24356814e-01, -6.68023945e-02],
        [-8.63269237e-01, -2.76358711e-01, 2.44682793e-01],
        [6.96102118e-01, 7.11951913e-02, -6.90939013e-01],
        [-3.35455831e-01, -8.98079831e-01, 2.69389345e-01],
        [6.06252875e-02, 1.32942947e00, -1.31600607e00],
        [-8.74938767e-01, -7.52900651e-01, -7.02958799e-01],
        [1.00446099e00, 2.02357721e-01, 2.26263414e00],
    ],
    [
        [0.00000000e00, 0.00000000e00, 0.00000000e00],
        [5.55828779e-01, 5.59311995e-01, 2.84957153e-02],
        [-5.05085367e-01, -1.54057451e-01, 1.50878576e-01],
        [1.51893023e00, 7.81720721e-02, -1.52518695e00],
        [-1.96182375e-01, -5.48614995e-01, 1.71857203e-01],
        [2.70303840e-02, 1.74694325e00, -1.75640560e00],
        [-1.78858747e00, -1.53682143e00, -1.72681136e00],
        [3.88065824e-01, -1.44933440e-01, 4.65717241e00],
    ],
]


def test_desc(test_data_dir):
    conf = Configuration.from_file(test_data_dir / "configs" / "Si.xyz")

    cut_name = "cos"
    cut_dists = {"Si-Si": 5.0}
    hyperparams = {
        "jmax": 2,
        "rfac0": 1,
        "diagonalstyle": 3,
        "rmin0": 0,
        "switch_flag": 1,
        "bzero_flag": 0,
        "weight": {"Si": 1.0},
    }

    desc = Bispectrum(cut_dists, cut_name, hyperparams)
    zeta, dzeta_dr = desc.transform(conf, grad=True)
    assert desc.get_size() == 14
    assert np.allclose(zeta, zeta_ref)
    assert np.allclose(dzeta_dr[0][:2], dzeta_dr_001)
    assert np.allclose(dzeta_dr[-1][-2:], dzeta_dr_minus_121)

    zeta, dzeta_dr = desc.transform(conf, grad=False)
    assert np.allclose(zeta, zeta_ref)
    assert dzeta_dr is None
