import pytest
import numpy as np
from quantumwalks import DTQW1D, DTQW2D, DTQW3D

class TestDTQW1D:
    def test_initialization(self):
        qw = DTQW1D(N=10, coin_type='hadamard')
        assert qw.N == 10
        assert len(qw.positions) == 21
        assert qw.coin.shape == (2, 2)
    
    def test_state_initialization(self):
        qw = DTQW1D(N=10, coin_type='hadamard')
        qw.initialize_state(init_pos=0)
        probs = qw.get_probabilities()
        assert np.isclose(np.sum(probs), 1.0)
    
    def test_evolution(self):
        qw = DTQW1D(N=10, coin_type='hadamard')
        qw.initialize_state(init_pos=0)
        initial_probs = qw.get_probabilities().copy()
        qw.evolve(5)
        final_probs = qw.get_probabilities()
        assert not np.array_equal(initial_probs, final_probs)

class TestDTQW2D:
    def test_initialization(self):
        qw = DTQW2D(Nx=5, Ny=5, coin_type='grover')
        assert qw.Nx == 5
        assert qw.Ny == 5
        assert qw.coin.shape == (4, 4)
    
    def test_probability_conservation(self):
        qw = DTQW2D(Nx=5, Ny=5, coin_type='grover')
        qw.initialize_state(init_pos=(0, 0))
        qw.evolve(10)
        probs = qw.get_probabilities()
        assert np.isclose(np.sum(probs), 1.0)

class TestDTQW3D:
    def test_initialization(self):
        qw = DTQW3D(Nx=3, Ny=3, Nz=3, coin_type='fourier')
        assert qw.Nx == 3
        assert qw.Ny == 3
        assert qw.Nz == 3
        assert qw.coin.shape == (6, 6)