import pytest
import networkx as nx
import numpy as np
from quantumwalks import SzegedyQW

class TestSzegedyQW:
    def setup_method(self):
        self.graph = nx.cycle_graph(6)
    
    def test_initialization(self):
        sqw = SzegedyQW(self.graph, initial_nodes=[0])
        assert sqw.num_nodes == 6
        assert sqw.num_edges == 12  # 无向图每条边两个方向
        assert np.isclose(np.linalg.norm(sqw.initial_state), 1.0)
    
    def test_probability_conservation(self):
        sqw = SzegedyQW(self.graph, initial_nodes=[0])
        sqw.evolve(10)
        probs_first = sqw.get_probabilities(register='first')
        probs_second = sqw.get_probabilities(register='second')
        
        assert np.isclose(np.sum(probs_first), 1.0)
        assert np.isclose(np.sum(probs_second), 1.0)
    
    def test_finite_average(self):
        sqw = SzegedyQW(self.graph, initial_nodes=[0])
        avg_probs = sqw.average_probabilities_finite(10, register='first')
        assert np.isclose(np.sum(avg_probs), 1.0)
        assert len(avg_probs) == sqw.num_nodes
    
    def test_infinite_average(self):
        sqw = SzegedyQW(self.graph, initial_nodes=[0])
        avg_probs = sqw.average_probabilities_infinite(register='first')
        assert np.isclose(np.sum(avg_probs), 1.0)
        assert len(avg_probs) == sqw.num_nodes
    
    def test_different_initial_nodes(self):
        sqw1 = SzegedyQW(self.graph, initial_nodes=[0])
        sqw2 = SzegedyQW(self.graph, initial_nodes=[0, 1])
        
        # 应该有不同的初始状态
        assert not np.array_equal(sqw1.initial_state, sqw2.initial_state)