import numpy as np
import pandas as pd
from tqdm import tqdm
from .graph_persistence_analysis import graph_persistence_diagram

try:
    from persim import bottleneck, wasserstein, plot_diagrams
    PERSIM_AVAILABLE = True
except ImportError:
    PERSIM_AVAILABLE = False

def _check_persim_available():
    if not PERSIM_AVAILABLE:
        raise ImportError(
            "The 'persim' library is required for TDA stability metrics. "
            "Please install it using `pip install persim`."
        )

def tgraph_persistence_diagrams(
    tgraph, 
    start=None, 
    end=None, 
    positions=None,
    maxdim=1,
    invert_method='reciprocal'
):
    """
    Compute persistence diagrams for each time step in a TGraph.
    
    Args:
        tgraph: ciga.TGraph object
        start, end: Time range (optional)
        positions: Position columns (if different from tgraph's default)
        maxdim: Max homology dimension
        invert_method: Weight inversion method
        
    Returns:
        dict: Mapping from time_key to diagrams list
    """
    # Use TGraph's position if not provided
    if positions is None:
        positions = tgraph._position
        
    # We need to iterate over time steps.
    # TGraph doesn't expose a simple "iterate all steps" easily without knowing the structure.
    # We can rely on the data's unique time points.
    
    df = tgraph.data
    pos_cols = list(positions)
    
    # Get unique time points sorted
    # Check if positions are in index (which is standard for TGraph)
    if all(col in df.index.names for col in pos_cols):
        unique_times = df.index.to_frame(index=False)[pos_cols].drop_duplicates().sort_values(by=pos_cols)
    else:
        # Fallback to columns if not in index (or mixed)
        unique_times = df[pos_cols].drop_duplicates().sort_values(by=pos_cols)
        
    time_keys = [tuple(x) for x in unique_times.values]
    
    results = {}
    
    # Iterate
    desc = "Computing Persistence Diagrams"
    for t_key in tqdm(time_keys, desc=desc):
        # We need check if t_key is within start/end range if provided
        # This is a bit complex with generic tuples, simpler if we let the user filter beforehand 
        # or we just compute all and let user slice. 
        # For now, let's compute all or trust the user passed a filtered TGraph (not implemented yet)
        # or we check bounds strictly.
        
        # Get graph snapshot
        # Note: accumulate=False is typical for "snapshot" topology
        # accumulate=True would be "evolution of cumulative topology"
        g = tgraph.get_graph(time_point=t_key, accumulate=False)
        
        if g is None or g.ecount() == 0:
            # Empty graph or only nodes
            results[t_key] = [np.zeros((0, 2)) for _ in range(maxdim + 1)]
            continue
            
        diagrams = graph_persistence_diagram(
            g, 
            maxdim=maxdim, 
            invert_method=invert_method
        )
        results[t_key] = diagrams
        
    return results


def tgraph_stability(
    tgraph,
    dimension=1,
    metric='bottleneck', # or 'wasserstein'
    delta_step=1,
    maxdim=1,
    invert_method='reciprocal'
):
    """
    Calculate topological stability (distance between consecutive diagrams).
    
    Args:
        tgraph: TGraph object
        dimension: Homology dimension to compare (0 for components, 1 for loops)
        metric: 'bottleneck' or 'wasserstein'
        delta_step: Step size for comparison (default 1 means compare t with t+1)
        maxdim: Max dimension to compute diagrams for
        invert_method: Weight inversion method
        
    Returns:
        pd.DataFrame: DataFrame with columns [*time_cols, 'stability', 'metric']
    """
    _check_persim_available()
    
    # 1. Compute all diagrams
    diagrams_map = tgraph_persistence_diagrams(
        tgraph, maxdim=maxdim, invert_method=invert_method
    )
    
    # 2. Sort keys
    sorted_keys = sorted(diagrams_map.keys())
    
    stability_results = []
    
    for i in range(len(sorted_keys) - delta_step):
        t_current = sorted_keys[i]
        t_next = sorted_keys[i + delta_step]
        
        dgm1 = diagrams_map[t_current][dimension]
        dgm2 = diagrams_map[t_next][dimension]
        
        # Calculate distance
        if metric == 'bottleneck':
            dist = bottleneck(dgm1, dgm2)
        elif metric == 'wasserstein':
            dist = wasserstein(dgm1, dgm2)
        else:
            raise ValueError("Metric must be 'bottleneck' or 'wasserstein'")
            
        # Result uses the "current" time as the index for the change happening *after* it
        # or between it and next.
        entry = {k: v for k, v in zip(tgraph._position, t_current)}
        entry['stability'] = dist
        entry['metric'] = metric
        entry['next_time'] = str(t_next)
        
        stability_results.append(entry)
        
    return pd.DataFrame(stability_results)
