import pytest
import networkx as nx
import numpy as np
from quantumwalks import DTQWPageRank, QuantumPageRank

class TestDTQWPageRank:
    def setup_method(self):
        self.digraph = nx.cycle_graph(6).to_directed()
    
    def test_initialization(self):
        pr = DTQWPageRank(self.digraph, num_steps=10, num_runs=5)
        assert pr.N == 6
        assert pr.num_steps == 10
        assert pr.num_runs == 5
        assert pr.A.shape == (6, 6)
    
    def test_ranking_output(self):
        pr = DTQWPageRank(self.digraph, num_steps=10, num_runs=5)
        ranking = pr.get_ranking()
        
        assert len(ranking) == 6
        assert np.isclose(np.sum(ranking), 1.0)
        assert all(score >= 0 for score in ranking)
    
    def test_sorted_ranking(self):
        pr = DTQWPageRank(self.digraph, num_steps=10, num_runs=5)
        sorted_ranking = pr.get_sorted_ranking()
        
        assert len(sorted_ranking) == 6
        # 检查是否按分数降序排列
        scores = [score for _, score in sorted_ranking]
        assert all(scores[i] >= scores[i+1] for i in range(len(scores)-1))

class TestQuantumPageRank:
    def setup_method(self):
        self.graph = nx.cycle_graph(6)
    
    def test_initialization(self):
        qpr = QuantumPageRank(self.graph, alpha=0.85)
        assert qpr.N == 6
        assert qpr.alpha == 0.85
        assert qpr.G.shape == (6, 6)
        assert qpr.U_red.shape == (12, 12)  # 2N x 2N
    
    def test_ranking_output(self):
        qpr = QuantumPageRank(self.graph, alpha=0.85)
        ranking = qpr.get_ranking()
        
        assert len(ranking) == 6
        assert np.isclose(np.sum(ranking), 1.0)
        assert all(score >= 0 for score in ranking)
    
    def test_sorted_ranking(self):
        qpr = QuantumPageRank(self.graph, alpha=0.85)
        sorted_ranking = qpr.get_sorted_ranking()
        
        assert len(sorted_ranking) == 6
        scores = [score for _, score in sorted_ranking]
        assert all(scores[i] >= scores[i+1] for i in range(len(scores)-1))
    
    def test_different_alpha(self):
        qpr1 = QuantumPageRank(self.graph, alpha=0.85)
        qpr2 = QuantumPageRank(self.graph, alpha=0.95)
        
        ranking1 = qpr1.get_ranking()
        ranking2 = qpr2.get_ranking()
        
        # 不同的alpha应该产生不同的排名（虽然可能很接近）
        assert not np.array_equal(ranking1, ranking2)

def test_pagerank_comparison():
    """测试两种量子PageRank方法都能正常工作"""
    graph = nx.erdos_renyi_graph(8, 0.4)
    
    # 对于无向图，需要转换为有向图才能用DTQWPageRank
    digraph = graph.to_directed()
    
    dtqw_pr = DTQWPageRank(digraph, num_steps=20, num_runs=5)
    quantum_pr = QuantumPageRank(graph, alpha=0.85)
    
    dtqw_scores = dtqw_pr.get_ranking()
    quantum_scores = quantum_pr.get_ranking()
    
    # 两种方法都应该产生有效的概率分布
    assert np.isclose(np.sum(dtqw_scores), 1.0)
    assert np.isclose(np.sum(quantum_scores), 1.0)
    assert all(score >= 0 for score in dtqw_scores)
    assert all(score >= 0 for score in quantum_scores)