# Copyright (c) [2024] []
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

from unittest.mock import Mock

import numpy as np
import pytest
from hypothesis import given
from hypothesis import strategies as st
from numpy.testing import assert_allclose, assert_array_almost_equal

from grogupy.magnetism import (
    blow_up_orbindx,
    calculate_anisotropy_tensor,
    calculate_exchange_tensor,
    parse_magnetic_entity,
    spin_tracer,
)


# Utility to create mock sisl atoms with orbitals
def create_mock_atom(orbitals):
    atom = Mock()
    atom.orbitals = []
    for l in orbitals:
        orbital = Mock()
        orbital.l = l
        atom.orbitals.append(orbital)
    return atom


# Test blow_up_orbindx function
def test_blow_up_orbindx_basic():
    """Test basic orbital index expansion"""
    orb_indices = np.array([0, 1, 2])
    result = blow_up_orbindx(orb_indices)
    expected = np.array([0, 1, 2, 3, 4, 5])  # Each index expands to two indices
    assert_array_almost_equal(result, expected)


@given(st.lists(st.integers(min_value=0, max_value=100), min_size=1, max_size=10))
def test_blow_up_orbindx_properties(indices):
    """Property-based test for orbital index expansion"""
    orb_indices = np.array(indices)
    result = blow_up_orbindx(orb_indices)

    # Result should be twice as long as input
    assert len(result) == 2 * len(orb_indices)

    # Each pair should be consecutive numbers
    for i in range(0, len(result), 2):
        assert result[i + 1] == result[i] + 1


# Test spin_tracer function
def test_spin_tracer_basic():
    """Test basic spin tracing functionality"""
    # Create a test matrix in SPIN-BOX representation
    size = 2
    M = np.zeros((2 * size, 2 * size), dtype=complex)
    M[0::2, 0::2] = 1.0  # M11
    M[1::2, 1::2] = -1.0  # M22
    M[0::2, 1::2] = 0.5j  # M12
    M[1::2, 0::2] = -0.5j  # M21

    result = spin_tracer(M)

    # Check components
    assert_array_almost_equal(result["z"], np.eye(size) * 2.0)  # M11 - M22
    assert_array_almost_equal(result["x"], np.eye(size))  # M12 + M21
    assert_array_almost_equal(result["y"], np.eye(size) * 1.0j)  # i(M12 - M21)
    assert_array_almost_equal(result["c"], np.zeros((size, size)))  # M11 + M22


def test_spin_tracer_hermiticity():
    """Test that spin tracer preserves Hermiticity"""
    size = 2
    M = np.random.random((2 * size, 2 * size)) + 1j * np.random.random(
        (2 * size, 2 * size)
    )
    M = M + M.conj().T  # Make Hermitian

    result = spin_tracer(M)

    # Check Hermiticity of components
    for component in ["x", "y", "z", "c"]:
        assert_array_almost_equal(result[component], result[component].conj().T)


# Test parse_magnetic_entity function
def test_parse_magnetic_entity_single_atom():
    """Test parsing of single atom magnetic entity"""
    # Create mock geometry
    geometry = Mock()
    atom = create_mock_atom([0, 2, 2, 2])  # s and d orbitals
    geometry.atoms = [atom]
    geometry.a2o.return_value = np.array([0, 1, 2, 3])

    dh = Mock()
    dh.geometry = geometry

    # Test with d orbitals only
    entity = {"atom": 0, "l": 2}
    result = parse_magnetic_entity(dh, **entity)
    assert len(result) == 3  # Should only include d orbitals

    # Test with all orbitals
    entity = {"atom": 0, "l": None}
    result = parse_magnetic_entity(dh, **entity)
    assert len(result) == 4  # Should include all orbitals


def test_parse_magnetic_entity_multiple_atoms():
    """Test parsing of multiple atom magnetic entity"""
    # Create mock geometry with two atoms
    geometry = Mock()
    atom1 = create_mock_atom([0, 2])
    atom2 = create_mock_atom([0, 2])
    geometry.atoms = [atom1, atom2]
    geometry.a2o.side_effect = [np.array([0, 1]), np.array([2, 3])]

    dh = Mock()
    dh.geometry = geometry

    entity = {"atom": [0, 1], "l": 2}
    result = parse_magnetic_entity(dh, **entity)
    assert len(result) == 2  # Should include d orbitals from both atoms


# Test calculate_anisotropy_tensor function
def test_calculate_anisotropy_tensor():
    """Test calculation of anisotropy tensor"""
    # Create mock magnetic entity with energies
    mag_ent = {
        "energies": np.array(
            [
                [0.1, 0.2],  # First rotation axis
                [0.3, 0.4],  # Second rotation axis
                [0.5, 0.6],  # Third rotation axis
            ]
        )
    }

    Kxx, Kyy, Kzz, consistency = calculate_anisotropy_tensor(mag_ent)

    # Check basic properties
    assert isinstance(Kxx, float)
    assert isinstance(Kyy, float)
    assert isinstance(Kzz, float)
    assert isinstance(consistency, float)

    # Verify calculations
    assert_allclose(Kxx, mag_ent["energies"][1, 1] - mag_ent["energies"][1, 0])
    assert_allclose(Kyy, mag_ent["energies"][0, 1] - mag_ent["energies"][0, 0])
    assert_allclose(Kzz, 0)


# Test calculate_exchange_tensor function
def test_calculate_exchange_tensor():
    """Test calculation of exchange tensor"""
    # Create mock pair with energies
    pair = {
        "energies": np.zeros((3, 4))  # 3 quantization axes, 4 rotation combinations
    }
    # Set some test values
    pair["energies"][0] = [0.1, 0.2, 0.3, 0.4]  # First quantization axis
    pair["energies"][1] = [0.2, 0.3, 0.4, 0.5]  # Second quantization axis
    pair["energies"][2] = [0.3, 0.4, 0.5, 0.6]  # Third quantization axis

    J_iso, J_S, D, J = calculate_exchange_tensor(pair)

    # Check dimensions
    assert isinstance(J_iso, float)
    assert J_S.shape == (5,)  # Symmetric exchange has 5 independent components
    assert D.shape == (3,)  # DM vector has 3 components
    assert J.shape == (3, 3)  # Full exchange tensor is 3x3

    # Check properties
    assert_allclose(np.trace(J) / 3, J_iso)  # Isotropic exchange
    assert_array_almost_equal(J, J.T)  # Symmetric part

    # Verify DM vector components
    assert_allclose(D[2], (pair["energies"][2, 1] - pair["energies"][2, 2]) / 2)


def test_exchange_tensor_symmetries():
    """Test symmetry properties of exchange tensor calculations"""
    pair = {"energies": np.zeros((3, 4))}
    # Set symmetric energies
    pair["energies"][0] = [0.1, 0.1, 0.1, 0.1]
    pair["energies"][1] = [0.2, 0.2, 0.2, 0.2]
    pair["energies"][2] = [0.3, 0.3, 0.3, 0.3]

    J_iso, J_S, D, J = calculate_exchange_tensor(pair)

    # For symmetric energies, DM vector should be zero
    assert_array_almost_equal(D, np.zeros(3))

    # Exchange tensor should be symmetric
    assert_array_almost_equal(J, J.T)


@pytest.mark.parametrize(
    "test_energies,expected_D",
    [
        (
            np.array(
                [[0.1, 0.2, -0.2, 0.1], [0.1, 0.2, -0.2, 0.1], [0.1, 0.2, -0.2, 0.1]]
            ),
            np.array([0.2, 0.2, 0.2]),
        ),
        (np.zeros((3, 4)), np.zeros(3)),
    ],
)
def test_exchange_tensor_dm_vector(test_energies, expected_D):
    """Test DM vector calculation with different energy configurations"""
    pair = {"energies": test_energies}
    _, _, D, _ = calculate_exchange_tensor(pair)
    assert_array_almost_equal(D, expected_D)


if __name__ == "__main__":
    pytest.main([__file__])