"""
PySpinW Test Suite

Based on SpinW MATLAB tutorials and examples.
"""

import numpy as np
import sys
import os
# Add parent directory to path
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))

from pyspinw import SpinW, sw_model
from pyspinw.core.lattice import Lattice, SpaceGroup
from pyspinw.core.magnetic import (
    MagneticStructure, MagneticAtom, 
    generate_ferromagnetic, generate_antiferromagnetic, generate_120_structure
)
from pyspinw.core.coupling import Coupling, ExchangeMatrix, Bond
from pyspinw.backends import get_backend, PYTORCH_AVAILABLE


class TestLattice:
    """Test lattice and crystallographic operations."""
    
    def test_cubic_lattice(self):
        lat = Lattice(lat_const=(3.0, 3.0, 3.0), angles=(90, 90, 90))
        assert lat.a == 3.0
        assert np.isclose(lat.volume, 27.0)
    
    def test_hexagonal_lattice(self):
        lat = Lattice(lat_const=(3.0, 3.0, 6.0), angles=(90, 90, 120))
        expected_volume = 3.0 * 3.0 * 6.0 * np.sin(np.radians(120))
        assert np.isclose(lat.volume, expected_volume, rtol=1e-5)
    
    def test_reciprocal_lattice(self):
        lat = Lattice(lat_const=(2.0, 3.0, 4.0), angles=(90, 90, 90))
        rl = lat.reciprocal_vectors
        bv = lat.basis_vectors
        product = rl.T @ bv
        expected = 2 * np.pi * np.eye(3)
        assert np.allclose(product, expected, rtol=1e-5)
    
    def test_coordinate_transforms(self):
        lat = Lattice(lat_const=(3.0, 4.0, 5.0), angles=(90, 90, 90))
        frac = np.array([0.5, 0.5, 0.5])
        cart = lat.frac_to_cart(frac)
        frac_back = lat.cart_to_frac(cart)
        assert np.allclose(frac, frac_back)


class TestMagneticStructure:
    """Test magnetic structure generation."""
    
    def test_ferromagnetic(self):
        magstr = generate_ferromagnetic(3, direction=[0, 0, 1])
        assert magstr.S.shape == (3, 3)
        assert magstr.is_collinear
    
    def test_120_structure(self):
        magstr = generate_120_structure(3)
        assert magstr.S.shape == (3, 3)
        for i in range(3):
            for j in range(i+1, 3):
                dot = np.dot(magstr.S[i], magstr.S[j])
                assert np.isclose(dot, -0.5, atol=1e-5)


class TestExchangeMatrix:
    """Test exchange matrix operations."""
    
    def test_heisenberg(self):
        J = ExchangeMatrix.heisenberg(1.0)
        assert np.allclose(J.matrix, np.eye(3))
        assert np.isclose(J.isotropic_part, 1.0)
    
    def test_dzyaloshinskii_moriya(self):
        D = np.array([0, 0, 0.1])
        dm = ExchangeMatrix.dzyaloshinskii_moriya(D)
        assert np.allclose(dm.dm_vector, D)


class TestSpinWBasic:
    """Test basic SpinW functionality."""
    
    def test_creation(self):
        sw = SpinW()
        assert sw.n_atoms == 0
    
    def test_genlattice(self):
        sw = SpinW()
        sw.genlattice(lat_const=[3, 3, 6], angled=[90, 90, 120])
        assert sw.lattice.a == 3
        assert sw.lattice.gamma == 120
    
    def test_addatom(self):
        sw = SpinW()
        sw.genlattice(lat_const=[3, 3, 6])
        sw.addatom(r=[0, 0, 0], S=1, label='Cu1')
        sw.addatom(r=[0.5, 0.5, 0], S=0, label='O')
        assert sw.n_atoms == 2
        assert sw.n_mag_atoms == 1
    
    def test_gencoupling(self):
        sw = SpinW()
        sw.genlattice(lat_const=[3, 3, 6], angled=[90, 90, 90])
        sw.addatom(r=[0, 0, 0], S=1, label='Cu1')
        sw.gencoupling(max_distance=5)
        assert len(sw._coupling.bonds) > 0


class TestSpinWModels:
    """Test predefined spin wave models."""
    
    def test_chain_model(self):
        sw = sw_model('chain', 1.0)
        assert sw.n_mag_atoms == 1
    
    def test_square_model(self):
        sw = sw_model('squareAF', 1.0)
        assert sw.n_mag_atoms == 1
    
    def test_triangular_model(self):
        sw = sw_model('triAF', 1.0)
        assert sw.n_mag_atoms == 1
        assert not np.allclose(sw.magstr.k, 0)


class TestSpinWaveCalculation:
    """Test spin wave calculations."""
    
    def test_afm_chain_dispersion(self):
        sw = sw_model('chain', 1.0)
        spec = sw.spinwave([[0, 0, 0], [1, 0, 0]], n_pts=100)
        omega = np.real(spec['omega'])
        assert omega.shape[0] == 100
        assert omega.shape[1] >= 1
    
    def test_triangular_afm(self):
        sw = sw_model('triAF', 1.0)
        spec = sw.spinwave([[0, 0, 0], [1, 1, 0]], n_pts=50)
        omega = np.real(spec['omega'])
        assert np.all(omega >= -1e-6)
    
    def test_spinwave_output_format(self):
        sw = sw_model('chain', 1.0)
        spec = sw.spinwave([[0, 0, 0], [1, 0, 0]], n_pts=10)
        assert 'omega' in spec
        assert 'hkl' in spec
        assert spec['hkl'].shape == (10, 3)


class TestEnergyGrid:
    """Test energy grid convolution."""
    
    def test_sw_egrid_basic(self):
        from pyspinw.utils import sw_egrid
        sw = sw_model('chain', 1.0)
        spec = sw.spinwave([[0, 0, 0], [1, 0, 0]], n_pts=50)
        spec = sw_egrid(spec, dE=0.1)
        assert 'swConv' in spec
        assert 'Evect' in spec


class TestBackends:
    """Test computational backends."""
    
    def test_pytorch_backend(self):
        if not PYTORCH_AVAILABLE:
            print("(skipped - PyTorch not installed)")
            return
        backend = get_backend('pytorch')
        a = backend.to_array([[1, 2], [3, 4]])
        b = backend.to_array([[5, 6], [7, 8]])
        c = backend.matmul(a, b)
        c_np = backend.to_numpy(c)
        expected = np.array([[19, 22], [43, 50]])
        assert np.allclose(c_np, expected)
    
    def test_numpy_backend(self):
        backend = get_backend('numpy')
        a = backend.to_array([[1, 2], [3, 4]])
        b = backend.to_array([[5, 6], [7, 8]])
        c = backend.matmul(a, b)
        expected = np.array([[19, 22], [43, 50]])
        assert np.allclose(c, expected)


def run_all_tests():
    """Run all tests and print summary."""
    import traceback
    
    test_classes = [
        TestLattice,
        TestMagneticStructure,
        TestExchangeMatrix,
        TestSpinWBasic,
        TestSpinWModels,
        TestSpinWaveCalculation,
        TestEnergyGrid,
        TestBackends,
    ]
    
    total = 0
    passed = 0
    failed = 0
    
    print("=" * 60)
    print("PySpinW Test Suite")
    print("=" * 60)
    
    for test_class in test_classes:
        print(f"\n{test_class.__name__}:")
        
        instance = test_class()
        methods = [m for m in dir(instance) if m.startswith('test_')]
        
        for method_name in methods:
            total += 1
            try:
                method = getattr(instance, method_name)
                method()
                print(f"  ✓ {method_name}")
                passed += 1
            except Exception as e:
                print(f"  ✗ {method_name}: {str(e)[:50]}")
                failed += 1
    
    print("\n" + "=" * 60)
    print(f"Results: {passed}/{total} passed, {failed} failed")
    print("=" * 60)
    
    return failed == 0


if __name__ == '__main__':
    success = run_all_tests()
    exit(0 if success else 1)
