"""Tests of cluster orbit search."""

from pathlib import Path

import numpy as np

from pyclupan.core.cell_utils import supercell_general
from pyclupan.features.orbit import get_orbit_supercell, get_orbit_unitcell

cwd = Path(__file__).parent


def test_fcc(fcc_binary_clusters):
    """Test cluster orbit search in fcc."""
    lattice_unitcell, clusters, _, _ = fcc_binary_clusters
    unitcell = lattice_unitcell.cell

    orbit_fracs_unitcell = get_orbit_unitcell(clusters, unitcell)
    n_orbits = [len(orbit[0]) for orbit in orbit_fracs_unitcell]

    assert n_orbits == [
        1,
        12,
        12,
        24,
        6,
        18,
        72,
        24,
        144,
        36,
        72,
        72,
        24,
        36,
        36,
        24,
        72,
        192,
        96,
        48,
        48,
        48,
        96,
        96,
        32,
        96,
        96,
        192,
        192,
        192,
        192,
        8,
        96,
        32,
        48,
        48,
        96,
        192,
        96,
        192,
        96,
        96,
        12,
        96,
        24,
        192,
        48,
        8,
        32,
        96,
        12,
        24,
    ]

    orbit_positions = orbit_fracs_unitcell[31]
    positions = orbit_positions[0]
    assert list(positions[0][0]) == [-1, 0, 0, 0]
    assert list(positions[0][1]) == [1, 0, 1, 1]
    assert list(positions[0][2]) == [0, 0, -1, 0]

    assert list(positions[1][0]) == [0, 1, 1, 1]
    assert list(positions[1][1]) == [0, -1, 0, 0]
    assert list(positions[1][2]) == [0, 0, -1, 0]

    assert list(positions[2][0]) == [-1, 0, 0, 0]
    assert list(positions[2][1]) == [0, -1, 0, 0]
    assert list(positions[2][2]) == [0, 0, -1, 0]

    assert list(positions[3][0]) == [0, 0, 0, 1]
    assert list(positions[3][1]) == [-1, -1, 0, -1]
    assert list(positions[3][2]) == [0, 1, 0, 0]

    assert list(positions[4][0]) == [-1, -1, -1, 0]
    assert list(positions[4][1]) == [0, 0, 1, 0]
    assert list(positions[4][2]) == [0, 1, 0, 0]

    assert list(positions[5][0]) == [-1, 0, 0, 0]
    assert list(positions[5][1]) == [0, -1, 0, 0]
    assert list(positions[5][2]) == [1, 1, 0, 1]

    assert list(positions[6][0]) == [0, 0, 0, 1]
    assert list(positions[6][1]) == [0, 0, 1, 0]
    assert list(positions[6][2]) == [-1, 0, -1, -1]

    assert list(positions[7][0]) == [0, 0, 0, 1]
    assert list(positions[7][1]) == [0, 0, 1, 0]
    assert list(positions[7][2]) == [0, 1, 0, 0]


def test_perovskite_orbit_unitcell(perovskite_binary_clusters):
    """Test orbits in perovskite."""

    lattice_unitcell, clusters, _, _ = perovskite_binary_clusters
    unitcell = lattice_unitcell.cell
    orbit_fracs_unitcell = get_orbit_unitcell(clusters, unitcell)
    n_orbits = [len(orbit[2]) for orbit in orbit_fracs_unitcell]
    n_orbits_true = [
        1,
        4,
        4,
        2,
        8,
        8,
        16,
        12,
        24,
        12,
        24,
        48,
        24,
        48,
        12,
        24,
        12,
        24,
        48,
        8,
        24,
        24,
        8,
        4,
        32,
        64,
        64,
        8,
        32,
        64,
        64,
        64,
        32,
        4,
        32,
        64,
        64,
        8,
        64,
        64,
        16,
        32,
        32,
        32,
        32,
        8,
        64,
        64,
        16,
        32,
        32,
        32,
        32,
        4,
        32,
        16,
        32,
        8,
        16,
        64,
        64,
        16,
    ]
    assert n_orbits == n_orbits_true


def test_perovskite_orbit_supercell(perovskite_binary_clusters):
    """Test orbits in perovskite."""
    lattice_unitcell, clusters, _, _ = perovskite_binary_clusters
    unitcell = lattice_unitcell.cell

    hnf = np.array([[1, 0, 0], [0, 1, 0], [1, 0, 2]])
    supercell = supercell_general(unitcell, hnf)
    lattice_supercell = lattice_unitcell.lattice_supercell(supercell)
    orbit_fracs_unitcell = get_orbit_unitcell(clusters, unitcell)
    orbit = get_orbit_supercell(
        lattice_unitcell,
        lattice_supercell,
        orbit_fracs_unitcell,
        return_array=False,
    )
    orbit_test = orbit[20]

    orbit_true = np.array(
        [
            [4, 7, 8],
            [4, 7, 9],
            [4, 7, 9],
            [4, 7, 9],
            [4, 6, 9],
            [4, 7, 8],
            [4, 6, 8],
            [4, 7, 8],
            [4, 7, 8],
            [4, 7, 8],
            [4, 7, 9],
            [4, 7, 9],
            [4, 6, 8],
            [4, 6, 8],
            [4, 6, 9],
            [4, 6, 8],
            [4, 7, 8],
            [4, 6, 9],
            [4, 6, 9],
            [4, 6, 8],
            [4, 6, 8],
            [4, 6, 9],
            [4, 7, 9],
            [4, 6, 9],
        ]
    )
    assert list(orbit_test.keys()) == [0, 1, 2, 3, 4, 5]
    np.testing.assert_equal(np.array(orbit_test[0]), orbit_true - 4)
