import time
import warnings
from typing import Any

import networkx as nx
import numpy as np
import torch_geometric.data as pyg
from scipy.stats import pearsonr, spearmanr
from tqdm import tqdm

from .graph_sample import GraphSample
from .graph_universe import GraphUniverse


class GraphFamilyGenerator:
    """
    Generates families of graphs sampled from a graph universe object dictating global graph building rules.
    Each graph in the family samples it's communities from the universe and uses the universe as global template for the graph structure.

    Initialize the graph family generator.

    Args:
        universe: GraphUniverse to sample from
        n_nodes_range: Tuple of (min_n_nodes, max_n_nodes) per graph
        n_graphs: Default number of graphs to generate (can be overridden if specified in generate_family. Kinf of default behaviour when using the CONFIG)
        n_communities_range: Tuple of (min_communities, max_communities) per graph (max defaults to universe.K if None)
        homophily_range: Tuple of (min_homophily, max_homophily) in graph family
        avg_degree_range: Tuple of (min_avg_degree, max_avg_degree) in graph family
        power_law_exponent_range: Range for power law exponent (min, max)
        degree_separation_range: Range for degree distribution separation (min, max)
        seed: Random seed for reproducibility
    """

    def __init__(
        self,
        universe: GraphUniverse,
        n_nodes_range: tuple[int, int],
        n_graphs: int = 100,
        n_communities_range: tuple[int, int] = (2, None),
        homophily_range: tuple[float, float] = (0.0, 0.4),
        avg_degree_range: tuple[float, float] = (1.0, 3.0),
        power_law_exponent_range: tuple[float, float] = (2.0, 3.5),
        degree_separation_range: tuple[float, float] = (0.5, 0.5),
        seed: int | None = 42,
    ):
        self.universe = universe
        self.min_n_nodes = n_nodes_range[0]
        self.max_n_nodes = n_nodes_range[1]
        self.n_graphs = n_graphs
        self.min_communities = n_communities_range[0]
        self.max_communities = n_communities_range[1] if n_communities_range[1] is not None else universe.K
        self.homophily_range = homophily_range
        self.avg_degree_range = avg_degree_range
        self.degree_separation_range = degree_separation_range
        self.power_law_exponent_range = power_law_exponent_range

        # Set random seed
        self.seed = seed
        if seed is not None:
            np.random.seed(seed)

        # Validate parameters
        self._validate_parameters()

        # Storage for generated graphs and metadata
        self.graphs: list[GraphSample] = []
        self.generation_metadata: list[dict[str, Any]] = []
        self.generation_stats: dict[str, Any] = {}

        self.family_generated = False

    def generate_family(
        self,
        n_graphs: int | None = None,
        show_progress: bool = True,
        collect_stats: bool = True,
        allowed_community_combinations: list[list[int]] | None = None,
    ) -> list[GraphSample]:
        """
        Generate a family of graphs from the universe.

        Args:
            n_graphs: Number of graphs to generate (if None, uses default n_graphs of initialization.)
            show_progress: Whether to show progress bar
            collect_stats: Whether to collect generation statistics
            allowed_community_combinations: Optional list of lists of community indices to be sampled from the universe
        Returns:
            List of generated GraphSample objects
        """
        self.n_graphs = n_graphs if n_graphs is not None else self.n_graphs
        start_time = time.time()
        self.graphs = []
        self.community_labels_per_graph = []
        self.generation_metadata = []
        self.graph_generation_times = []
        failed_graphs = 0

        # Reset seed at the start of generation to ensure reproducibility
        if self.seed is not None:
            np.random.seed(self.seed)

        # Progress bar setup
        if show_progress:
            pbar = tqdm(total=self.n_graphs, desc="Generating graph family")

        while len(self.graphs) < self.n_graphs:
            graph_start_time = time.time()

            try:
                # Get a random seed for this graph
                graph_seed = np.random.randint(0, 1000000)

                # Sample parameters for this graph
                params = self._sample_graph_parameters()

                if allowed_community_combinations is not None:
                    sampled_community_combination_index = np.random.randint(
                        0, len(allowed_community_combinations)
                    )
                    sampled_community_combination = allowed_community_combinations[
                        sampled_community_combination_index
                    ]

                graph_sample = GraphSample(
                    universe=self.universe,
                    num_communities=params["n_communities"],
                    n_nodes=params["n_nodes"],
                    target_homophily=params["target_homophily"],
                    target_average_degree=params["target_average_degree"],
                    power_law_exponent=params.get("power_law_exponent", None),
                    degree_separation=params.get("degree_separation", 0.5),
                    global_degree_params=params.get("global_degree_params", {}),
                    seed=graph_seed,
                    user_defined_communities=sampled_community_combination
                    if allowed_community_combinations is not None
                    else None,
                )

                # Store graph and metadata
                self.graphs.append(graph_sample)
                self.community_labels_per_graph.append(
                    np.unique(graph_sample.community_labels_universe_level, axis=0)
                )
                metadata = {
                    "graph_id": len(self.graphs) - 1,
                    "final_n_nodes": graph_sample.n_nodes,
                    "final_n_edges": graph_sample.graph.number_of_edges(),
                    "actual_density": graph_sample.graph.number_of_edges()
                    / (graph_sample.n_nodes * (graph_sample.n_nodes - 1) / 2)
                    if graph_sample.n_nodes > 1
                    else 0,
                    "timing_info": graph_sample.timing_info.copy()
                    if hasattr(graph_sample, "timing_info")
                    else {},
                    **params,  # Include sampled parameters
                }

                self.generation_metadata.append(metadata)

                graph_generation_time = time.time() - graph_start_time
                self.graph_generation_times.append(graph_generation_time)

                if show_progress:
                    pbar.update(1)

            except Exception as e:
                failed_graphs += 1
                warnings.warn(
                    f"Failed to generate graph {len(self.graphs) + failed_graphs}: {e}",
                    stacklevel=2,
                )
                # Continue to next graph
                continue

        # Use set of sorted tuples for uniqueness
        unique_set_of_community_combinations = {
            tuple(sorted(arr)) for arr in self.community_labels_per_graph
        }

        # Convert back to list of lists
        unique_list_of_community_combinations = [
            list(tup) for tup in unique_set_of_community_combinations
        ]
        self.seen_community_combinations = unique_list_of_community_combinations

        if show_progress:
            pbar.close()

        # Collect generation statistics
        if collect_stats:
            self._collect_generation_stats(start_time, failed_graphs, self.n_graphs)

        # Mark family as generated
        self.family_generated = True

    def to_pyg_graphs(self, task: str | None = None) -> list[pyg.data.Data]:
        """
        Convert the graphs to PyG graphs with the specified task as a property.

        Args:
            task: String of task to include as property on each graph.
                If None, include community detection

        Returns:
            List of PyG Data objects, each containing the task as a property
        """
        # Check if graphs are created
        assert self.family_generated, "Graph family has not been generated yet"

        # Check if task is valid
        if task is not None:
            if task not in ["community_detection", "triangle_counting"]:
                raise ValueError(f"Invalid task specified: {task}")
            self.task = task
        else:
            self.task = "community_detection"

        pyg_graphs = []
        for graph in tqdm(self.graphs, desc="Converting graphs to PyG graphs"):
            pyg_graphs.append(graph.to_pyg_graph(self.task))

        return pyg_graphs

    def get_uniquely_identifying_metadata(self) -> dict[str, Any]:
        """
        Get uniquely identifying metadata for the graph family.
        """
        uniquely_identifying_metadata = {
            "universe_parameters": {
                "K": self.universe.K,
                "feature_dim": self.universe.feature_dim,
                "center_variance": self.universe.center_variance,
                "cluster_variance": self.universe.cluster_variance,
                "edge_propensity_variance": self.universe.edge_propensity_variance,
                "seed": self.universe.seed,
            },
            "family_parameters": {
                "n_graphs": self.n_graphs,
                "n_nodes_range": [self.min_n_nodes, self.max_n_nodes],
                "n_communities_range": [self.min_communities, self.max_communities],
                "homophily_range": self.homophily_range,
                "avg_degree_range": self.avg_degree_range,
                "degree_separation_range": self.degree_separation_range,
                "power_law_exponent_range": self.power_law_exponent_range,
                "seed": self.seed,
            },
        }

        return uniquely_identifying_metadata

    def save_pyg_graphs_and_universe(
        self, task: str | None = None, root_dir: str = "datasets"
    ):
        """
        Save the PyG graphs and universe to a file.
        """
        from .dataset import GraphUniverseDataset

        uniquely_identifying_metadata = self.get_uniquely_identifying_metadata()

        # Convert the graphs to PyG graphs including task
        pyg_graphs = self.to_pyg_graphs(task)

        # Add task to metadata so it's included in the directory structure
        uniquely_identifying_metadata['task'] = task

        self.dataset = GraphUniverseDataset(
            graph_list=pyg_graphs, root=root_dir, parameters=uniquely_identifying_metadata
        )

    def _validate_parameters(self) -> None:
        """Validate initialization parameters."""
        if self.min_n_nodes <= 0:
            raise ValueError("min_n_nodes must be positive")
        if self.max_n_nodes < self.min_n_nodes:
            raise ValueError("max_n_nodes must be >= min_n_nodes")
        if self.min_communities < 1:
            raise ValueError("min_communities must be >= 1")
        if self.max_communities > self.universe.K:
            raise ValueError(f"max_communities cannot exceed universe size ({self.universe.K})")
        if self.max_communities < self.min_communities:
            raise ValueError("max_communities must be >= min_communities")

        # Validate ranges
        if len(self.homophily_range) != 2 or self.homophily_range[0] > self.homophily_range[1]:
            raise ValueError("homophily_range must be a tuple (min, max) with min <= max")
        if self.homophily_range[0] < 0.0 or self.homophily_range[1] > 1.0:
            raise ValueError("homophily_range values must be between 0.0 and 1.0")
        if len(self.avg_degree_range) != 2 or self.avg_degree_range[0] > self.avg_degree_range[1]:
            raise ValueError("avg_degree_range must be a tuple (min, max) with min <= max")
        if self.avg_degree_range[0] < 0.0 or self.avg_degree_range[1] > 100.0:
            raise ValueError("avg_degree_range values must be between 0.0 and 100.0")
        if (
            len(self.degree_separation_range) != 2
            or self.degree_separation_range[0] > self.degree_separation_range[1]
        ):
            raise ValueError("degree_separation_range must be a tuple (min, max) with min <= max")

        # Validate degree distribution parameter ranges
        if (
            len(self.power_law_exponent_range) != 2
            or self.power_law_exponent_range[0] > self.power_law_exponent_range[1]
        ):
            raise ValueError("power_law_exponent_range must be a tuple (min, max) with min <= max")

        if self.power_law_exponent_range[0] <= 1.0:
            raise ValueError("power_law_exponent_range values must be > 1.0")

    def _sample_graph_parameters(self) -> dict[str, Any]:
        """
        Sample parameters for a single graph within the specified ranges.

        Returns:
            Dictionary of sampled parameters
        """
        # Sample number of nodes
        n_nodes = np.random.randint(self.min_n_nodes, self.max_n_nodes + 1)

        # Sample number of communities
        n_communities = np.random.randint(self.min_communities, self.max_communities + 1)

        # Sample target homophily
        target_homophily = np.random.uniform(self.homophily_range[0], self.homophily_range[1])

        # Sample target average degree
        target_average_degree = np.random.uniform(
            self.avg_degree_range[0], self.avg_degree_range[1]
        )

        # Sample degree distribution related parameters
        degree_separation = np.random.uniform(
            self.degree_separation_range[0], self.degree_separation_range[1]
        )

        power_law_exponent = np.random.uniform(
            self.power_law_exponent_range[0], self.power_law_exponent_range[1]
        )
        distribution_params = {"exponent": power_law_exponent, "x_min": 1.0}

        degree_params = {
            "degree_separation": degree_separation,
            "global_degree_params": distribution_params,
            "power_law_exponent": power_law_exponent,
        }

        # Combine all parameters
        params = {
            "n_nodes": n_nodes,
            "target_homophily": target_homophily,
            "target_average_degree": target_average_degree,
            "n_communities": n_communities,
        }

        # Add degree distribution related parameters
        params.update(degree_params)

        return params

    def _generate_single_graph(
        self, max_attempts: int = 10
    ) -> GraphSample:
        """
        Generate a single graph with progress tracking.

        Args:
            max_attempts: Maximum attempts per graph before giving up

        Returns:
            Generated GraphSample object

        Raises:
            Exception: If graph generation fails after max_attempts
        """
        attempts = 0

        while attempts < max_attempts:
            attempts += 1

            try:
                # Get a random seed for this graph
                graph_seed = np.random.randint(0, 1000000)

                # Sample parameters for this graph
                params = self._sample_graph_parameters()

                # Debug: Print parameters being used
                print(f"Attempt {attempts}: Generating graph with params: {params}")

                # Create graph sample
                print(f"Attempt {attempts}: Starting GraphSample initialization...")

                try:
                    graph_sample = GraphSample(
                        universe=self.universe,
                        num_communities=params["n_communities"],
                        n_nodes=params["n_nodes"],
                        target_homophily=params["target_homophily"],
                        target_average_degree=params["target_average_degree"],
                        power_law_exponent=params.get("power_law_exponent", None),
                        degree_separation=params.get("degree_separation", 0.5),
                        global_degree_params=params.get("global_degree_params", {}),
                        seed=graph_seed,
                        user_defined_communities=None,
                    )
                    print(f"Attempt {attempts}: GraphSample initialization completed successfully")
                except Exception as e:
                    print(f"Attempt {attempts}: GraphSample initialization failed: {e!s}")
                    raise e

                print(f"Successfully generated graph with {graph_sample.n_nodes} nodes")
                return graph_sample

            except Exception as e:
                if attempts == max_attempts:
                    raise Exception(
                        f"Failed to generate graph after {attempts} attempts: {e}"
                    ) from e
                # Continue to next attempt

    def _collect_generation_stats(
        self, start_time: float, failed_graphs: int, n_graphs: int
    ) -> None:
        """Collect statistics about the generation process."""
        total_time = time.time() - start_time
        successful_graphs = len(self.graphs)

        # Basic stats
        self.generation_stats = {
            "total_time": total_time,
            "successful_graphs": successful_graphs,
            "failed_graphs": failed_graphs,
            "success_rate": successful_graphs / n_graphs if n_graphs > 0 else 0,
            "avg_time_per_graph": total_time / n_graphs if n_graphs > 0 else 0,
        }

        if successful_graphs > 0:
            # Get statistics from successful graphs
            successful_metadata = [
                m for m in self.generation_metadata if not m.get("failed", False)
            ]

            node_counts = [m["final_n_nodes"] for m in successful_metadata]
            edge_counts = [m["final_n_edges"] for m in successful_metadata]
            densities = [m["actual_density"] for m in successful_metadata]

            self.generation_stats.update(
                {
                    "node_count_stats": {
                        "mean": np.mean(node_counts),
                        "std": np.std(node_counts),
                        "min": np.min(node_counts),
                        "max": np.max(node_counts),
                    },
                    "edge_count_stats": {
                        "mean": np.mean(edge_counts),
                        "std": np.std(edge_counts),
                        "min": np.min(edge_counts),
                        "max": np.max(edge_counts),
                    },
                    "density_stats": {
                        "mean": np.mean(densities),
                        "std": np.std(densities),
                        "min": np.min(densities),
                        "max": np.max(densities),
                    },
                }
            )

    def save_family(self, filepath: str, n_graphs: int = 0) -> None:
        """
        Save the graph family to file.

        Args:
            filepath: Path to save file
            include_graphs: Whether to include graph objects (large files)
        """
        import pickle

        save_data = {
            "n_graphs": n_graphs,
            "n_nodes_range": [self.min_n_nodes, self.max_n_nodes],
            "n_communities_range": [self.min_communities, self.max_communities],
            "homophily_range": self.homophily_range,
            "degree_separation_range": self.degree_separation_range,
            "power_law_exponent_range": self.power_law_exponent_range,
        }

        with open(filepath, "wb") as f:
            pickle.dump(save_data, f)

    def __len__(self) -> int:
        """Return number of successfully generated graphs."""
        return len(self.graphs)

    def __getitem__(self, index: int) -> GraphSample:
        """Get a specific graph from the family."""
        return self.graphs[index]

    def __iter__(self):
        """Iterate over graphs in the family."""
        return iter(self.graphs)

    def analyze_graph_family_properties(self) -> dict[str, Any]:
        """Analyze the properties of the graph family once generated."""
        if self.graphs is None:
            raise ValueError(
                "No graphs in family. Please generate family first before analyzing properties."
            )

        properties = {
            "n_graphs": len(self.graphs),
            "node_counts": [],
            "edge_counts": [],
            "densities": [],
            "avg_degrees": [],
            "clustering_coefficients": [],
            "community_counts": [],
            "homophily_levels": [],
            "nr_of_triangles": [],
            "degree_distributions": [],
            "degree_distribution_power_law_exponents": [],
            "tail_ratio_95": [],
            "tail_ratio_99": [],
            "mean_edge_probability_deviation": [],
            "graph_generation_times": self.graph_generation_times,
            "gini_degree": [],
        }

        for graph in tqdm(
            self.graphs, desc="Analyzing graph family properties", total=len(self.graphs)
        ):
            properties["node_counts"].append(graph.n_nodes)
            properties["edge_counts"].append(graph.graph.number_of_edges())

            if graph.n_nodes > 1:
                density = graph.graph.number_of_edges() / (graph.n_nodes * (graph.n_nodes - 1) / 2)
                properties["densities"].append(density)
            else:
                properties["densities"].append(0.0)

            if graph.n_nodes > 1:
                avg_degree = sum(dict(graph.graph.degree()).values()) / (graph.n_nodes)
                properties["avg_degrees"].append(avg_degree)
            else:
                properties["avg_degrees"].append(0.0)

            try:
                clustering = nx.average_clustering(graph.graph)
                properties["clustering_coefficients"].append(clustering)
            except Exception as e:
                warnings.warn(f"Failed to calculate clustering coefficient: {e}", stacklevel=2)
                properties["clustering_coefficients"].append(0.0)

            properties["community_counts"].append(len(np.unique(graph.community_labels)))

            # Calculate average homophily level across communities
            if graph.n_nodes > 0 and graph.graph.number_of_edges() > 0:
                # Calculate homophily directly
                same_community_edges = 0
                total_edges = 0
                for u, v in graph.graph.edges():
                    total_edges += 1
                    if graph.community_labels[u] == graph.community_labels[v]:
                        same_community_edges += 1
                avg_homophily = same_community_edges / total_edges if total_edges > 0 else 0.0
                properties["homophily_levels"].append(avg_homophily)
            else:
                properties["homophily_levels"].append(0.0)

            # Calculate number of triangles
            if graph.n_nodes > 0:
                triangle_values = list(nx.triangles(graph.graph).values())
                properties["nr_of_triangles"].append(np.sum(triangle_values) / 3)
            else:
                properties["nr_of_triangles"].append(0.0)

            # Extract degree distribution and fit power law exponent
            if graph.n_nodes > 0:
                degrees = list(dict(graph.graph.degree()).values())
                # Fit power law exponent
                power_law_exponent = self._fit_power_law_exponent(degrees)
                tail_metrics = self._calculate_degree_tail_metrics(degrees)
                properties["tail_ratio_95"].append(tail_metrics["tail_ratio_95"])
                properties["tail_ratio_99"].append(tail_metrics["tail_ratio_99"])
                properties["degree_distributions"].append(degrees)
                properties["degree_distribution_power_law_exponents"].append(power_law_exponent)
            else:
                properties["degree_distributions"].append([])
                properties["degree_distribution_power_law_exponents"].append(0.0)

            # Calculate mean edge probability deviation
            if graph.n_nodes > 0:
                deviations = graph._calculate_community_deviations(
                    graph.graph, graph.community_labels, graph.P_sub
                )
                properties["mean_edge_probability_deviation"].append(deviations["mean_deviation"])
            else:
                properties["mean_edge_probability_deviation"].append(0.0)

            # Calculate Gini index of degree distribution
            if graph.n_nodes > 0:
                degrees = list(dict(graph.graph.degree()).values())
                # Sort degrees in non-decreasing order
                sorted_degrees = sorted(degrees)
                n = len(sorted_degrees)
                # Calculate Gini coefficient using discrete formula
                # Gini = Σ(2i - n - 1) * d_(i) / (n * Σd_(i))
                numerator = sum((2 * i - n - 1) * d for i, d in enumerate(sorted_degrees, 1))
                denominator = n * sum(sorted_degrees)
                gini = numerator / denominator if denominator > 0 else 0.0
                properties["gini_degree"].append(gini)
            else:
                properties["gini_degree"].append(0.0)

        # Calculate statistics and convert to native Python types
        for key in [
            "node_counts",
            "edge_counts",
            "densities",
            "avg_degrees",
            "clustering_coefficients",
            "community_counts",
            "homophily_levels",
            "nr_of_triangles",
            "graph_generation_times",
            "mean_edge_probability_deviation",
        ]:
            values = properties[key]
            if values:
                properties[f"{key}_mean"] = float(np.mean(values))
                properties[f"{key}_std"] = float(np.std(values))
                properties[f"{key}_min"] = float(np.min(values))
                properties[f"{key}_max"] = float(np.max(values))
            else:
                properties[f"{key}_mean"] = 0.0
                properties[f"{key}_std"] = 0.0
                properties[f"{key}_min"] = 0.0
                properties[f"{key}_max"] = 0.0

        return properties

    def analyze_graph_family_signals(self) -> dict[str, Any]:
        """Analyze the signals of the graph family."""
        if not self.graphs:
            raise ValueError(
                "No graphs in family. Please generate family first before analyzing signals."
            )

        # Calculate the signals
        signals = {
            "feature_signal": [],
            "degree_signal": [],
            "triangle_signal": [],
            "structure_signal": [],
        }

        for graph in tqdm(self.graphs, desc="Calculating signals", total=len(self.graphs)):
            signals["feature_signal"].append(graph.calculate_feature_signal())
            signals["degree_signal"].append(graph.calculate_degree_signal())
            signals["triangle_signal"].append(graph.calculate_triangle_community_signal())
            signals["structure_signal"].append(graph.calculate_structure_signal())

        return signals

    def analyze_graph_family_consistency(self) -> dict[str, Any]:
        """Analyze the consistency of the graph family."""
        if not self.graphs:
            raise ValueError(
                "No graphs in family. Please generate family first before analyzing consistency."
            )

        # Calculate the consistency
        results = {}

        print("Calculating consistency metrics...")

        # 1. Structural consistency (do communities RELATIVELY connect more to the communities they are supposed to connect to?)
        try:
            pattern_corrs = self._calculate_pattern_consistency()
            results["structure_consistency"] = pattern_corrs

        except Exception:
            results["structure_consistency"] = []

        # 2. Feature consistency (do communities of the same community indeed draw from a consistent feature distribution?)
        try:
            feature_consistency = self._calculate_feature_consistency()
            results["feature_consistency"] = feature_consistency
        except Exception:
            results["feature_consistency"] = []

        # 3. Degree consistency (do actual node degrees correlate with universe degree centers?)
        try:
            degree_consistency = self._calculate_degree_consistency()
            results["degree_consistency"] = degree_consistency
        except Exception:
            results["degree_consistency"] = []

        print("Consistency metrics calculated.")

        return results

    def _calculate_pattern_consistency(self) -> list[float]:
        """
        Measure whether relative community connection patterns are preserved.
        Uses rank correlation to focus on structural patterns rather than absolute values.
        """
        pattern_correlations = []

        for graph in self.graphs:
            try:
                # Extract relevant submatrix from universe
                P_sub = graph.P_sub

                # Get actual connections
                actual_edge_probabilities, _, _ = graph.calculate_actual_probability_matrix()

                # Calculate rank correlation per row comparing the expected and actual edge probabilities and average them
                if len(P_sub) > 1 and np.std(P_sub) > 0 and np.std(actual_edge_probabilities) > 0:
                    correlations = []
                    for i in range(len(P_sub)):
                        correlation, _ = spearmanr(P_sub[i], actual_edge_probabilities[i])
                        if not np.isnan(correlation):
                            correlations.append(correlation)
                    pattern_correlations.append(np.mean(correlations))
            except Exception as e:
                warnings.warn(
                    f"Error in pattern consistency calculation for graph: {e}", stacklevel=2
                )
                continue

        return pattern_correlations

    def _calculate_generation_fidelity(self) -> list[float]:
        """
        Measure how well graphs match their scaled probability targets (P_sub).
        """
        fidelity_scores = []

        for graph in self.graphs:
            try:
                # Use the graph's own scaled P_sub as reference
                expected_matrix = graph.P_sub

                # Get actual connections
                actual_analysis = graph.analyze_community_connections()
                actual_matrix = actual_analysis["actual_matrix"]

                # Calculate correlation between expected and actual
                expected_flat = expected_matrix.flatten()
                actual_flat = actual_matrix.flatten()

                if (
                    len(expected_flat) > 1
                    and np.std(expected_flat) > 0
                    and np.std(actual_flat) > 0
                    and not np.isnan(pearsonr(expected_flat, actual_flat)[0])
                ):
                    fidelity_scores.append(pearsonr(expected_flat, actual_flat)[0])

            except Exception as e:
                warnings.warn(
                    f"Error in generation fidelity calculation for graph: {e}", stacklevel=2
                )
                continue

        if not fidelity_scores:
            return []

        return fidelity_scores

    def _calculate_feature_consistency(self) -> list[float]:
        """
        Measure cross-graph feature consistency for communities.
        Returns per-graph consistency scores.
        """
        if not self.graphs or self.universe.feature_dim == 0:
            return []

        consistency_scores = []

        for graph_idx, graph in enumerate(self.graphs):
            graph_consistency = 0.0
            valid_communities = 0

            # For each universe community that appears in this graph
            for universe_comm_id in graph.communities:
                # Get nodes from this community in current graph
                current_nodes = np.where(graph.community_labels_universe_level == universe_comm_id)[
                    0
                ]
                current_features = graph.features[current_nodes]

                if len(current_features) == 0:
                    continue

                # Compare against same community in other graphs
                cross_graph_similarities = []

                for other_graph_idx, other_graph in enumerate(self.graphs):
                    if other_graph_idx == graph_idx:
                        continue

                    # Check if this community exists in other graph
                    if universe_comm_id in other_graph.communities:
                        other_nodes = np.where(
                            other_graph.community_labels_universe_level == universe_comm_id
                        )[0]
                        other_features = other_graph.features[other_nodes]

                        if len(other_features) == 0:
                            continue

                        # Calculate feature similarity between same communities across graphs
                        similarity = self._calculate_feature_similarity(
                            current_features, other_features
                        )
                        cross_graph_similarities.append(similarity)

                if cross_graph_similarities:
                    graph_consistency += np.mean(cross_graph_similarities)
                    valid_communities += 1

            if valid_communities > 0:
                consistency_scores.append(graph_consistency / valid_communities)
            else:
                consistency_scores.append(0.0)

        return consistency_scores

    def _calculate_feature_similarity(self, features1: np.ndarray, features2: np.ndarray) -> float:
        """
        Calculate similarity between two feature matrices.

        Args:
            features1: Shape (n1, feature_dim) - features from community in graph 1
            features2: Shape (n2, feature_dim) - features from same community in graph 2

        Returns:
            Similarity score in [0, 1]
        """
        # Calculate centroid vectors (mean across nodes, keeping all dimensions)
        centroid1 = np.mean(features1, axis=0)  # Shape: (feature_dim,)
        centroid2 = np.mean(features2, axis=0)  # Shape: (feature_dim,)

        # Calculate cosine similarity between the two centroid vectors
        dot_product = np.dot(centroid1, centroid2)
        norm1 = np.linalg.norm(centroid1)
        norm2 = np.linalg.norm(centroid2)

        if norm1 == 0 or norm2 == 0:
            return 0.0

        cosine_sim = dot_product / (norm1 * norm2)

        # Convert to [0, 1] range (cosine similarity is in [-1, 1])
        return (cosine_sim + 1) / 2

    def _measure_ordering_consistency(self, values_a: np.ndarray, values_b: np.ndarray) -> float:
        """
        Fallback ordering consistency: fraction of correctly ordered pairs between two vectors.
        Values that preserve pairwise ordering contribute positively.
        """
        num_items = len(values_a)
        if num_items <= 1:
            return 1.0
        correct_pairs = 0
        total_pairs = 0
        for i in range(num_items):
            for j in range(i + 1, num_items):
                total_pairs += 1
                diff_a = values_a[i] - values_a[j]
                diff_b = values_b[i] - values_b[j]
                if diff_a == 0 and diff_b == 0:
                    # Ties in both are considered consistent
                    correct_pairs += 1
                elif diff_a * diff_b >= 0:
                    correct_pairs += 1
        return (correct_pairs / total_pairs) if total_pairs > 0 else 1.0

    def _calculate_degree_consistency(self) -> list[float]:
        """
        Measure degree consistency by comparing relative rankings of communities
        that appear together in graphs.
        """
        degree_centers = self.universe.community_degree_propensity_vector
        consistency_scores = []

        for graph in self.graphs:
            if len(graph.communities) < 2:
                # Can't measure relative ranking with only one community
                consistency_scores.append(1.0)
                continue

            # Calculate average degree per community in this graph
            avg_degrees = self._get_community_avg_degrees(graph)

            # Get expected relative ordering from universe degree centers
            expected_ordering = self._get_expected_ordering(graph.communities, degree_centers)
            actual_ordering = self._get_actual_ordering(avg_degrees)

            # Calculate rank correlation between expected and actual orderings
            if len(expected_ordering) > 1:
                correlation = self._rank_correlation(expected_ordering, actual_ordering)
                consistency_scores.append(max(0.0, correlation))  # Convert [-1,1] to [0,1]
            else:
                consistency_scores.append(1.0)

        return consistency_scores

    def _get_community_avg_degrees(self, graph) -> np.ndarray:
        """Get average degree for each community in the graph."""
        num_communities = len(graph.communities)
        avg_degrees = np.zeros(num_communities)
        community_sizes = np.zeros(num_communities)

        for node_idx in range(graph.n_nodes):
            local_comm_id = graph.community_labels[node_idx]
            degree = graph.graph.degree[node_idx]
            avg_degrees[local_comm_id] += degree
            community_sizes[local_comm_id] += 1

        # Normalize by community size
        for i in range(num_communities):
            if community_sizes[i] > 0:
                avg_degrees[i] /= community_sizes[i]

        return avg_degrees

    def _get_expected_ordering(
        self, communities: list[int], degree_centers: np.ndarray
    ) -> np.ndarray:
        """Get expected relative ordering of communities based on universe degree centers."""
        community_centers = degree_centers[communities]
        return np.argsort(np.argsort(community_centers))  # Ranks from 0 to k-1

    def _get_actual_ordering(self, avg_degrees: np.ndarray) -> np.ndarray:
        """Get actual relative ordering of communities based on observed average degrees."""
        return np.argsort(np.argsort(avg_degrees))  # Ranks from 0 to k-1

    def _rank_correlation(self, expected_ranks: np.ndarray, actual_ranks: np.ndarray) -> float:
        """Calculate Spearman rank correlation."""
        from scipy.stats import spearmanr

        correlation, _ = spearmanr(expected_ranks, actual_ranks)
        return correlation if not np.isnan(correlation) else 0.0

    def _fit_power_law_exponent(self, degrees: list[int]) -> float:
        """
        Fit a power law exponent to a degree distribution using discrete MLE.
        This method is specifically designed for discrete network degrees.

        Args:
            degrees: List of node degrees

        Returns:
            float: Power law exponent alpha, or 0 if fitting fails or alpha not in [1, 20]
        """
        if not degrees or len(degrees) < 2:
            return 0.0

        # Convert to numpy array and filter out zeros
        degrees_array = np.array(degrees, dtype=int)
        degrees_array = degrees_array[degrees_array > 0]

        if len(degrees_array) < 2:
            return 0.0

        k_min = np.min(degrees_array)
        n = len(degrees_array)

        def negative_log_likelihood(alpha):
            if alpha <= 1.0:
                return np.inf
            try:
                # For discrete power law, the log-likelihood is:
                # L = -alpha * sum(log(k_i)) - n * log(zeta(alpha, k_min))
                # We approximate zeta(alpha, k_min) for computational efficiency

                # Approximation: zeta(alpha, k_min) ≈ sum_{k=k_min}^{k_max} k^(-alpha)
                k_max = max(100, np.max(degrees_array) * 2)  # Reasonable upper bound
                k_range = np.arange(k_min, k_max + 1)
                zeta_approx = np.sum(k_range ** (-alpha))

                if zeta_approx <= 0:
                    return np.inf

                log_likelihood = -alpha * np.sum(np.log(degrees_array)) - n * np.log(zeta_approx)
                return -log_likelihood  # Return negative for minimization

            except (OverflowError, ZeroDivisionError, ValueError):
                return np.inf

        try:
            # Find MLE estimate using bounded optimization
            from scipy.optimize import minimize_scalar

            result = minimize_scalar(negative_log_likelihood, bounds=(1.01, 10.0), method="bounded")
            if result.success and 1.0 < result.x < 20.0:
                return result.x
            else:
                return 0.0
        except Exception as e:
            warnings.warn(f"Failed to fit power law exponent: {e}", stacklevel=2)
            return 0.0

    def _calculate_degree_tail_metrics(self, degrees: list[int]) -> dict[str, float]:
        """Calculate tail-based degree metrics instead of power law fitting"""
        if len(degrees) == 0:
            return {
                "tail_ratio_95": 0.0,
                "tail_ratio_99": 0.0,
                "coefficient_variation": 0.0,
                "max_degree": 0.0,
                "mean_degree": 0.0
            }

        tail_95 = np.percentile(degrees, 95)
        tail_99 = np.percentile(degrees, 99)
        mean_degree = np.mean(degrees)

        return {
            "tail_ratio_95": tail_95 / mean_degree if mean_degree > 0 else 0.0,
            "tail_ratio_99": tail_99 / mean_degree if mean_degree > 0 else 0.0,
            "coefficient_variation": np.std(degrees) / mean_degree if mean_degree > 0 else 0.0,
            "max_degree": np.max(degrees),
            "mean_degree": mean_degree,
        }
