# 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.

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

from grogupy.core import (
    build_hh_ss,
    calc_Vu,
    commutator,
    onsite_projection,
    parallel_Gk,
    sequential_GK,
)


# Helper function to generate complex arrays
def complex_arrays(shape):
    return st.complex_numbers(
        min_magnitude=0.0, max_magnitude=1e6, allow_infinity=False, allow_nan=False
    ).map(lambda x: np.full(shape, x))


# Test commutator function
def test_commutator_hermitian():
    """Test that commutator of Hermitian matrices is anti-Hermitian"""
    a = np.array([[1, 1j], [-1j, 2]], dtype=np.complex128)
    b = np.array([[3, -2j], [2j, 4]], dtype=np.complex128)

    result = commutator(a, b)

    # For Hermitian matrices, commutator should be anti-Hermitian
    assert_array_almost_equal(result, -result.conj().T)


@given(st.integers(min_value=2, max_value=10))
def test_commutator_zero_identity(n):
    """Test that commutator of identity matrix with any matrix is zero"""
    # Create random complex matrix
    rng = np.random.default_rng(42)
    a = rng.random((n, n)) + 1j * rng.random((n, n))
    identity = np.eye(n)

    result = commutator(a, identity)

    assert_array_almost_equal(result, np.zeros_like(result))


# Test parallel_Gk and sequential_GK implementations
@pytest.mark.parametrize("size", [2, 4, 8])
def test_green_function_implementations_match(size):
    """Test that parallel and sequential Green's function calculations match"""
    rng = np.random.default_rng(42)

    # Generate test inputs
    HK = rng.random((size, size)) + 1j * rng.random((size, size))
    HK = (HK + HK.conj().T) / 2  # Make Hermitian
    SK = np.eye(size) + 0.1 * (rng.random((size, size)) + 1j * rng.random((size, size)))
    SK = (SK + SK.conj().T) / 2  # Make Hermitian
    eran = rng.random(5) + 1j * rng.random(5)
    eset = len(eran)

    # Calculate using both methods
    G_parallel = parallel_Gk(HK, SK, eran, eset)
    G_sequential = sequential_GK(HK, SK, eran, eset)

    assert_array_almost_equal(G_parallel, G_sequential)


# Test calc_Vu function
def test_calc_Vu_hermiticity():
    """Test that Vu1 and Vu2 maintain expected Hermiticity properties"""
    rng = np.random.default_rng(42)
    size = 4

    # Create Hermitian Hamiltonian
    H = rng.random((size, size)) + 1j * rng.random((size, size))
    H = (H + H.conj().T) / 2

    # Create unitary rotation matrix
    Tu = rng.random((size, size)) + 1j * rng.random((size, size))
    Tu = (Tu + Tu.conj().T) / 2

    Vu1, Vu2 = calc_Vu(H, Tu)

    # Vu1 should be anti-Hermitian (from commutator properties)
    assert_array_almost_equal(Vu1, -Vu1.conj().T)

    # Vu2 should be Hermitian
    assert_array_almost_equal(Vu2, Vu2.conj().T)


# Test onsite_projection function
def test_onsite_projection():
    """Test basic properties of onsite projection"""
    size = 4
    idx1 = np.array([0, 1])
    idx2 = np.array([2, 3])

    # Create test matrix
    matrix = np.arange(size * size).reshape(size, size)

    result = onsite_projection(matrix, idx1, idx2)

    # Check shape
    assert result.shape == (len(idx1), len(idx2))

    # Check values
    expected = matrix[np.ix_(idx1, idx2)]
    assert_array_almost_equal(result, expected)


@given(st.integers(min_value=2, max_value=5))
def test_build_hh_ss_hermiticity(size):
    """Test that built Hamiltonians maintain Hermiticity"""
    from unittest.mock import MagicMock

    # Create mock DFT Hamiltonian class
    class MockDH:
        def __init__(self, size):
            self.no = size
            self.n_s = 1
            self.M11r = np.eye(size)
            self.M11i = np.zeros((size, size))
            self.M22r = np.eye(size)
            self.M22i = np.zeros((size, size))
            self.M12r = np.zeros((size, size))
            self.M12i = np.zeros((size, size))
            self.M21r = np.zeros((size, size))
            self.M21i = np.zeros((size, size))
            self.S_idx = np.eye(size)

            self.lattice = MagicMock()
            self.lattice.nsc.prod.return_value = 1

        def tocsr(self, matrix):
            return matrix

    dh = MockDH(size)
    hh, ss = build_hh_ss(dh)

    # Check Hermiticity of Hamiltonian and overlap matrices
    for h in hh:
        assert_array_almost_equal(h, h.conj().T)

    for s in ss:
        assert_array_almost_equal(s, s.conj().T)


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