"""
Phase 2: Clustering & Structure Analysis

This module provides tools for analyzing the structure of factor relationships:
- Hierarchical Clustering Analysis
- Minimum Spanning Tree (MST) Analysis
"""

from typing import Optional, Dict, Any, List, Tuple, Union
import pandas as pd
import numpy as np
from scipy.cluster.hierarchy import linkage, fcluster, dendrogram
from scipy.spatial.distance import squareform
from dataclasses import dataclass, field
import heapq


@dataclass
class ClusterInfo:
    """Information about a factor cluster"""
    cluster_id: int
    factor_ids: List[str]
    size: int
    mean_intra_correlation: float


@dataclass
class HierarchicalClusteringResult:
    """Result of hierarchical clustering analysis"""
    linkage_matrix: np.ndarray
    cluster_labels: Dict[str, int]  # factor_id -> cluster_id
    clusters: List[ClusterInfo]
    n_clusters: int
    factor_order: List[str]  # Optimal leaf ordering for visualization


@dataclass
class MSTEdge:
    """Edge in the Minimum Spanning Tree"""
    source: str
    target: str
    distance: float
    correlation: float


@dataclass
class MSTResult:
    """Result of Minimum Spanning Tree analysis"""
    edges: List[MSTEdge]
    total_distance: float
    centrality_scores: Dict[str, float]  # factor_id -> degree centrality
    central_factors: List[str]  # Most connected factors
    peripheral_factors: List[str]  # Least connected factors (unique alpha)


class ClusteringAnalyzer:
    """
    Analyzer for factor clustering and structure.
    
    Implements Phase 2 of the orthogonality analysis framework:
    1. Hierarchical Clustering Analysis
    2. Minimum Spanning Tree (MST) Analysis
    """
    
    def __init__(
        self,
        linkage_method: str = "ward",
        distance_threshold: Optional[float] = None,
        n_clusters: Optional[int] = None,
        central_percentile: float = 90,
        peripheral_percentile: float = 10
    ):
        """
        Initialize the clustering analyzer.
        
        Args:
            linkage_method: Linkage method for hierarchical clustering
                           ('ward', 'average', 'complete', 'single')
            distance_threshold: Threshold for cutting the dendrogram
            n_clusters: Number of clusters (alternative to distance_threshold)
            central_percentile: Percentile for identifying central factors
            peripheral_percentile: Percentile for identifying peripheral factors
        """
        self.linkage_method = linkage_method
        self.distance_threshold = distance_threshold
        self.n_clusters = n_clusters
        self.central_percentile = central_percentile
        self.peripheral_percentile = peripheral_percentile
    
    def perform_hierarchical_clustering(
        self,
        distance_matrix: pd.DataFrame,
        corr_matrix: pd.DataFrame
    ) -> HierarchicalClusteringResult:
        """
        Perform hierarchical clustering on factors.
        
        Uses correlation-based distance: d(F_i, F_j) = sqrt(2 * (1 - rho_ij))
        
        Args:
            distance_matrix: Distance matrix between factors
            corr_matrix: Correlation matrix (for intra-cluster correlation)
            
        Returns:
            HierarchicalClusteringResult with clustering information
        """
        factor_ids = distance_matrix.columns.tolist()
        n_factors = len(factor_ids)
        
        # Convert distance matrix to condensed form for scipy
        dist_values = distance_matrix.values.copy()
        np.fill_diagonal(dist_values, 0)
        condensed_dist = squareform(dist_values, checks=False)
        
        # Perform hierarchical clustering
        if self.linkage_method == "ward":
            # Ward requires condensed distance matrix
            Z = linkage(condensed_dist, method='ward')
        else:
            Z = linkage(condensed_dist, method=self.linkage_method)
        
        # Determine cluster labels
        if self.n_clusters is not None:
            labels = fcluster(Z, t=self.n_clusters, criterion='maxclust')
        elif self.distance_threshold is not None:
            labels = fcluster(Z, t=self.distance_threshold, criterion='distance')
        else:
            # Auto-determine using elbow method or default to sqrt(n)
            default_n = max(2, int(np.sqrt(n_factors)))
            labels = fcluster(Z, t=default_n, criterion='maxclust')
        
        # Create cluster label mapping
        cluster_labels = {factor_ids[i]: int(labels[i]) for i in range(n_factors)}
        
        # Analyze each cluster
        unique_labels = sorted(set(labels))
        clusters = []
        
        for cluster_id in unique_labels:
            cluster_factors = [f for f, c in cluster_labels.items() if c == cluster_id]
            
            # Calculate mean intra-cluster correlation
            if len(cluster_factors) > 1:
                cluster_corr = corr_matrix.loc[cluster_factors, cluster_factors]
                upper_mask = np.triu(np.ones_like(cluster_corr, dtype=bool), k=1)
                intra_corr = np.nanmean(cluster_corr.values[upper_mask])
            else:
                intra_corr = 1.0
            
            clusters.append(ClusterInfo(
                cluster_id=cluster_id,
                factor_ids=cluster_factors,
                size=len(cluster_factors),
                mean_intra_correlation=float(intra_corr)
            ))
        
        # Get optimal leaf ordering for visualization
        from scipy.cluster.hierarchy import leaves_list
        factor_order = [factor_ids[i] for i in leaves_list(Z)]
        
        return HierarchicalClusteringResult(
            linkage_matrix=Z,
            cluster_labels=cluster_labels,
            clusters=clusters,
            n_clusters=len(unique_labels),
            factor_order=factor_order
        )
    
    def build_minimum_spanning_tree(
        self,
        distance_matrix: pd.DataFrame
    ) -> MSTResult:
        """
        Build Minimum Spanning Tree from factor distance matrix.
        
        Uses Prim's algorithm for efficiency. MST reveals the factor network
        structure and identifies central vs peripheral factors.
        
        Args:
            distance_matrix: Distance matrix between factors
            
        Returns:
            MSTResult with MST edges and centrality analysis
        """
        factor_ids = distance_matrix.columns.tolist()
        n_factors = len(factor_ids)
        
        if n_factors < 2:
            return MSTResult(
                edges=[],
                total_distance=0.0,
                centrality_scores={factor_ids[0]: 0} if factor_ids else {},
                central_factors=[],
                peripheral_factors=factor_ids
            )
        
        # Build MST using Prim's algorithm
        mst_edges = []
        visited = set()
        
        # Start from the first factor
        visited.add(factor_ids[0])
        
        # Priority queue: (distance, source, target)
        edges_heap = []
        for j in range(1, n_factors):
            dist = distance_matrix.iloc[0, j]
            heapq.heappush(edges_heap, (dist, factor_ids[0], factor_ids[j]))
        
        while len(visited) < n_factors and edges_heap:
            dist, src, tgt = heapq.heappop(edges_heap)
            
            if tgt in visited:
                continue
            
            visited.add(tgt)
            
            # Convert distance back to correlation
            corr = 1 - (dist ** 2) / 2
            corr = np.clip(corr, -1, 1)
            
            mst_edges.append(MSTEdge(
                source=src,
                target=tgt,
                distance=float(dist),
                correlation=float(corr)
            ))
            
            # Add edges from new node
            tgt_idx = factor_ids.index(tgt)
            for j, factor in enumerate(factor_ids):
                if factor not in visited:
                    new_dist = distance_matrix.iloc[tgt_idx, j]
                    heapq.heappush(edges_heap, (new_dist, tgt, factor))
        
        # Calculate degree centrality
        degree_count = {f: 0 for f in factor_ids}
        for edge in mst_edges:
            degree_count[edge.source] += 1
            degree_count[edge.target] += 1
        
        # Normalize centrality
        max_degree = max(degree_count.values()) if degree_count else 1
        centrality_scores = {
            f: count / max_degree for f, count in degree_count.items()
        }
        
        # Identify central and peripheral factors
        centrality_values = list(centrality_scores.values())
        
        if centrality_values:
            central_threshold = np.percentile(centrality_values, self.central_percentile)
            peripheral_threshold = np.percentile(centrality_values, self.peripheral_percentile)
        else:
            central_threshold = 0
            peripheral_threshold = 0
        
        central_factors = [
            f for f, c in centrality_scores.items() if c >= central_threshold
        ]
        peripheral_factors = [
            f for f, c in centrality_scores.items() if c <= peripheral_threshold
        ]
        
        # Sort by centrality
        central_factors.sort(key=lambda x: centrality_scores[x], reverse=True)
        peripheral_factors.sort(key=lambda x: centrality_scores[x])
        
        total_distance = sum(e.distance for e in mst_edges)
        
        return MSTResult(
            edges=mst_edges,
            total_distance=total_distance,
            centrality_scores=centrality_scores,
            central_factors=central_factors,
            peripheral_factors=peripheral_factors
        )
    
    def get_dendrogram_data(
        self,
        clustering_result: HierarchicalClusteringResult
    ) -> Dict[str, Any]:
        """
        Get dendrogram plotting data from clustering result.
        
        Args:
            clustering_result: Result from hierarchical clustering
            
        Returns:
            Dictionary with dendrogram data for visualization
        """
        return {
            "linkage_matrix": clustering_result.linkage_matrix.tolist(),
            "factor_order": clustering_result.factor_order,
            "n_clusters": clustering_result.n_clusters,
        }
    
    def get_mst_graph_data(
        self,
        mst_result: MSTResult
    ) -> Dict[str, Any]:
        """
        Get MST graph data for visualization.
        
        Args:
            mst_result: Result from MST analysis
            
        Returns:
            Dictionary with nodes and edges for graph visualization
        """
        # Get all nodes
        nodes = set()
        for edge in mst_result.edges:
            nodes.add(edge.source)
            nodes.add(edge.target)
        
        node_list = [
            {
                "id": node,
                "centrality": mst_result.centrality_scores.get(node, 0),
                "is_central": node in mst_result.central_factors,
                "is_peripheral": node in mst_result.peripheral_factors
            }
            for node in nodes
        ]
        
        edge_list = [
            {
                "source": edge.source,
                "target": edge.target,
                "distance": edge.distance,
                "correlation": edge.correlation
            }
            for edge in mst_result.edges
        ]
        
        return {
            "nodes": node_list,
            "edges": edge_list,
            "total_distance": mst_result.total_distance,
            "central_factors": mst_result.central_factors,
            "peripheral_factors": mst_result.peripheral_factors
        }
    
    def find_representative_factors(
        self,
        clustering_result: HierarchicalClusteringResult,
        mst_result: MSTResult,
        ic_values: Optional[Dict[str, float]] = None
    ) -> Dict[int, str]:
        """
        Find representative factor for each cluster.
        
        Selection criteria:
        1. If IC values provided: highest IC factor in cluster
        2. Otherwise: most central factor in cluster (based on MST)
        
        Args:
            clustering_result: Hierarchical clustering result
            mst_result: MST analysis result
            ic_values: Optional dictionary of factor_id -> IC value
            
        Returns:
            Dictionary mapping cluster_id -> representative_factor_id
        """
        representatives = {}
        
        for cluster in clustering_result.clusters:
            cluster_factors = cluster.factor_ids
            
            if ic_values:
                # Select factor with highest absolute IC
                best_factor = max(
                    cluster_factors,
                    key=lambda f: abs(ic_values.get(f, 0))
                )
            else:
                # Select most central factor
                best_factor = max(
                    cluster_factors,
                    key=lambda f: mst_result.centrality_scores.get(f, 0)
                )
            
            representatives[cluster.cluster_id] = best_factor
        
        return representatives
    
    def full_analysis(
        self,
        distance_matrix: pd.DataFrame,
        corr_matrix: pd.DataFrame,
        ic_values: Optional[Dict[str, float]] = None
    ) -> Dict[str, Any]:
        """
        Perform full Phase 2 clustering analysis.
        
        Args:
            distance_matrix: Distance matrix between factors
            corr_matrix: Correlation matrix between factors
            ic_values: Optional IC values for representative selection
            
        Returns:
            Dictionary containing all analysis results
        """
        # Hierarchical clustering
        clustering_result = self.perform_hierarchical_clustering(
            distance_matrix, corr_matrix
        )
        
        # MST analysis
        mst_result = self.build_minimum_spanning_tree(distance_matrix)
        
        # Find representatives
        representatives = self.find_representative_factors(
            clustering_result, mst_result, ic_values
        )
        
        # Summary statistics
        cluster_sizes = [c.size for c in clustering_result.clusters]
        
        return {
            "clustering_result": clustering_result,
            "mst_result": mst_result,
            "representative_factors": representatives,
            "dendrogram_data": self.get_dendrogram_data(clustering_result),
            "mst_graph_data": self.get_mst_graph_data(mst_result),
            "summary": {
                "n_clusters": clustering_result.n_clusters,
                "cluster_size_mean": float(np.mean(cluster_sizes)),
                "cluster_size_std": float(np.std(cluster_sizes)),
                "cluster_size_max": int(np.max(cluster_sizes)),
                "cluster_size_min": int(np.min(cluster_sizes)),
                "n_central_factors": len(mst_result.central_factors),
                "n_peripheral_factors": len(mst_result.peripheral_factors),
                "mst_total_distance": mst_result.total_distance,
            }
        }
