import gradio as gr
import pandas as pd
import networkx as nx
from pyvis.network import Network
import tempfile
import os
import plotly.graph_objects as go

def clean_time_column(df, time_col):
    # Try converting to integer (years) first
    try:
        df[time_col] = pd.to_numeric(df[time_col]).astype(int)
        return df, True
    except:
        # Fallback to string-based categorical if not purely numeric
        df[time_col] = df[time_col].astype(str)
        return df, False

def generate_vis_html(df_filtered):
    nodes = set(df_filtered['Source'].unique()).union(set(df_filtered['Target'].unique()))
    
    net = Network(
        height="500px", 
        width="100%", 
        bgcolor="#16100c", 
        font_color="#f4eee6", 
        notebook=False
    )
    
    net.set_options("""
    var options = {
      "nodes": {
        "borderWidth": 2,
        "color": {
          "border": "#2c1e16",
          "background": "#ff7043",
          "highlight": {
            "border": "#ff7043",
            "background": "#ffffff"
          }
        },
        "font": {
          "color": "#f4eee6",
          "size": 14,
          "face": "Inter, sans-serif"
        }
      },
      "edges": {
        "color": {
          "color": "rgba(255, 112, 67, 0.4)",
          "highlight": "#ff7043"
        },
        "smooth": {
          "type": "continuous"
        }
      },
      "physics": {
        "barnesHut": {
          "gravitationalConstant": -12000,
          "centralGravity": 0.3,
          "springLength": 120,
          "springConstant": 0.04
        }
      }
    }
    """)
    
    # Scale nodes by degree inside the filtered slice
    G = nx.from_pandas_edgelist(df_filtered, 'Source', 'Target', edge_attr='Weight' if 'Weight' in df_filtered.columns else None)
    degrees = dict(G.degree())
    
    for node in nodes:
        deg = degrees.get(node, 1)
        net.add_node(node, label=node, size=10 + (deg * 2), title=f"Connection count: {deg}")
        
    for _, row in df_filtered.iterrows():
        weight = row.get('Weight', 1.0)
        try:
            w_val = float(weight)
            if pd.isna(w_val): w_val = 1.0
        except:
            w_val = 1.0
        net.add_edge(row['Source'], row['Target'], value=w_val)
        
    temp_dir = tempfile.gettempdir()
    temp_path = os.path.join(temp_dir, next(tempfile._get_candidate_names()) + ".html")
    net.save_graph(temp_path)
    
    with open(temp_path, "r", encoding="utf-8") as f:
        html_content = f.read()
        
    try:
        os.remove(temp_path)
    except:
        pass
        
    escaped_html = html_content.replace('"', '"')
    iframe_code = f'<iframe srcdoc="{escaped_html}" style="width: 100%; height: 530px; border: 1px solid rgba(255, 255, 255, 0.08); border-radius: 8px;"></iframe>'
    return iframe_code

def analyze_dynamic_network(file_obj, min_val, max_val, cat_val, is_cumulative):
    if file_obj is None:
        return "Please upload a CSV or Excel file.", None, None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
        
    try:
        if file_obj.name.endswith('.csv'):
            df = pd.read_csv(file_obj.name)
        else:
            df = pd.read_excel(file_obj.name)
    except Exception as e:
        return f"Error reading file: {str(e)}", None, None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
        
    # Standardize headers
    rename_map = {}
    for col in df.columns:
        if col.lower() in ['source', 'from', 'node1']:
            rename_map[col] = 'Source'
        elif col.lower() in ['target', 'to', 'node2']:
            rename_map[col] = 'Target'
        elif col.lower() in ['weight', 'value', 'strength']:
            rename_map[col] = 'Weight'
        elif col.lower() in ['time', 'year', 'timestamp', 'date', 'interval']:
            rename_map[col] = 'Time'
            
    df = df.rename(columns=rename_map)
    
    if 'Source' not in df.columns or 'Target' not in df.columns or 'Time' not in df.columns:
        return "CSV/Excel must contain 'Source', 'Target', and a temporal column ('Time', 'Year', etc.).", None, None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
        
    if 'Weight' not in df.columns:
        df['Weight'] = 1.0
        
    df, is_numeric = clean_time_column(df, 'Time')
    
    # Filter operations
    if is_numeric:
        if is_cumulative:
            df_filtered = df[df['Time'] <= max_val]
        else:
            df_filtered = df[(df['Time'] >= min_val) & (df['Time'] <= max_val)]
    else:
        # Categorical slice
        if is_cumulative:
            # Cumulative categorical doesn't follow strict ordering, but we slice up to the selected value in sorting order
            unique_cats = sorted(df['Time'].unique())
            idx = unique_cats.index(cat_val) + 1 if cat_val in unique_cats else len(unique_cats)
            df_filtered = df[df['Time'].isin(unique_cats[:idx])]
        else:
            df_filtered = df[df['Time'] == cat_val]
            
    if df_filtered.empty:
        return "No network links found inside this specific timeframe.", None, None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
        
    # Calculate global stats vs time to plot evolution
    stats_data = []
    if is_numeric:
        unique_times = sorted(df['Time'].unique())
        for t in unique_times:
            if is_cumulative:
                sub_df = df[df['Time'] <= t]
            else:
                sub_df = df[df['Time'] == t]
                
            sub_G = nx.from_pandas_edgelist(sub_df, 'Source', 'Target')
            stats_data.append({
                "Time": t,
                "Nodes": sub_G.number_of_nodes(),
                "Edges": sub_G.number_of_edges(),
                "Density": nx.density(sub_G)
            })
    else:
        unique_times = sorted(df['Time'].unique())
        for idx, t in enumerate(unique_times):
            if is_cumulative:
                sub_df = df[df['Time'].isin(unique_times[:idx+1])]
            else:
                sub_df = df[df['Time'] == t]
                
            sub_G = nx.from_pandas_edgelist(sub_df, 'Source', 'Target')
            stats_data.append({
                "Time": t,
                "Nodes": sub_G.number_of_nodes(),
                "Edges": sub_G.number_of_edges(),
                "Density": nx.density(sub_G)
            })
            
    stats_df = pd.DataFrame(stats_data)
    
    # Generate Plotly trend line chart
    fig = go.Figure()
    fig.add_trace(go.Scatter(x=stats_df['Time'], y=stats_df['Nodes'], mode='lines+markers', name='Nodes', line=dict(color='#ff7043', width=3)))
    fig.add_trace(go.Scatter(x=stats_df['Time'], y=stats_df['Edges'], mode='lines+markers', name='Edges', line=dict(color='#ffffff', width=2)))
    
    fig.update_layout(
        title="Temporal Network Growth",
        paper_bgcolor='#16100c',
        plot_bgcolor='#16100c',
        font_color='#f4eee6',
        margin=dict(l=40, r=40, t=50, b=40),
        xaxis=dict(gridcolor='rgba(255,255,255,0.05)'),
        yaxis=dict(gridcolor='rgba(255,255,255,0.05)')
    )
    
    # Current stats block
    G_curr = nx.from_pandas_edgelist(df_filtered, 'Source', 'Target')
    stats_html = f"""
    <div style='display: grid; grid-template-columns: repeat(3, 1fr); gap: 1rem; margin-bottom: 1rem;'>
        <div style='background: rgba(255, 255, 255, 0.03); border: 1px solid rgba(255, 255, 255, 0.08); border-radius: 8px; padding: 1rem; text-align: center;'>
            <div style='font-size: 0.75rem; text-transform: uppercase; color: #ff7043; letter-spacing: 0.1em;'>Nodes (Active)</div>
            <div style='font-size: 2rem; font-weight: bold; margin-top: 0.5rem;'>{G_curr.number_of_nodes()}</div>
        </div>
        <div style='background: rgba(255, 255, 255, 0.03); border: 1px solid rgba(255, 255, 255, 0.08); border-radius: 8px; padding: 1rem; text-align: center;'>
            <div style='font-size: 0.75rem; text-transform: uppercase; color: #ff7043; letter-spacing: 0.1em;'>Edges (Active)</div>
            <div style='font-size: 2rem; font-weight: bold; margin-top: 0.5rem;'>{G_curr.number_of_edges()}</div>
        </div>
        <div style='background: rgba(255, 255, 255, 0.03); border: 1px solid rgba(255, 255, 255, 0.08); border-radius: 8px; padding: 1rem; text-align: center;'>
            <div style='font-size: 0.75rem; text-transform: uppercase; color: #ff7043; letter-spacing: 0.1em;'>Time-slice Density</div>
            <div style='font-size: 2rem; font-weight: bold; margin-top: 0.5rem;'>{nx.density(G_curr):.4f}</div>
        </div>
    </div>
    """
    
    # Generate interactive Vis.js HTML
    vis_html = generate_pyvis_html(df_filtered)
    
    # Create downloadable filtered edge list CSV
    out_csv = tempfile.mktemp(suffix=".csv")
    df_filtered.to_csv(out_csv, index=False)
    
    return "", stats_html, vis_html, fig, out_csv

def init_sliders(file_obj):
    if file_obj is None:
        return (
            gr.update(visible=False), 
            gr.update(visible=False), 
            gr.update(visible=False), 
            gr.update(visible=False), 
            "Please upload a file."
        )
        
    try:
        if file_obj.name.endswith('.csv'):
            df = pd.read_csv(file_obj.name)
        else:
            df = pd.read_excel(file_obj.name)
    except Exception as e:
        return (
            gr.update(visible=False), 
            gr.update(visible=False), 
            gr.update(visible=False), 
            gr.update(visible=False), 
            f"Error parsing sheet: {str(e)}"
        )
        
    # Find time col
    time_col = None
    for col in df.columns:
        if col.lower() in ['time', 'year', 'timestamp', 'date', 'interval']:
            time_col = col
            break
            
    if not time_col:
        return (
            gr.update(visible=False), 
            gr.update(visible=False), 
            gr.update(visible=False), 
            gr.update(visible=False), 
            "Error: Temporal column ('Time', 'Year', etc.) was not found in columns."
        )
        
    df, is_numeric = clean_time_column(df, time_col)
    
    if is_numeric:
        min_val = int(df['Time'].min())
        max_val = int(df['Time'].max())
        
        return (
            gr.update(minimum=min_val, maximum=max_val, value=min_val, visible=True),
            gr.update(minimum=min_val, maximum=max_val, value=max_val, visible=True),
            gr.update(visible=False),
            gr.update(visible=True),
            f"Loaded numeric temporal range: {min_val} to {max_val}."
        )
    else:
        unique_cats = sorted(df['Time'].unique())
        return (
            gr.update(visible=False),
            gr.update(visible=False),
            gr.update(choices=unique_cats, value=unique_cats[0], visible=True),
            gr.update(visible=True),
            f"Loaded {len(unique_cats)} unique categorical time intervals."
        )

theme = gr.themes.Default(
    primary_hue="orange",
    neutral_hue="stone"
).set(
    body_background_fill="#0d0907",
    body_text_color="#c4bbae",
    block_background_fill="#16100c",
    block_border_width="1px",
    block_label_text_color="#f4eee6"
)

with gr.Blocks(theme=theme, title="Dynamic Network Analyzer") as demo:
    gr.Markdown(
        """
        # ⏳ Dynamic Network Evolution Analyzer
        ### Upload time-stamped datasets (Source, Target, Year/Interval) to slide through histories and observe network change, structural growth, and relation trends.
        """
    )
    
    error_msg = gr.Markdown("", visible=False)
    
    with gr.Row():
        with gr.Column(scale=1):
            file_obj = gr.File(label="Upload Time-stamped Network CSV", file_types=[".csv", ".xlsx"])
            status_text = gr.Markdown("💡 **Tip**: Make sure your dataset contains **Source**, **Target**, and a temporal column (e.g., **Year**, **Time**, **Date**).")
            
            with gr.Group(visible=False) as controls_group:
                is_cumulative = gr.Checkbox(label="Cumulative Display", value=False, info="Show all connections up to selected window (instead of only within window).")
                
                with gr.Row():
                    min_slider = gr.Slider(minimum=0, maximum=100, step=1, label="Min Time Boundary", visible=False)
                    max_slider = gr.Slider(minimum=0, maximum=100, step=1, label="Max Time Boundary", visible=False)
                    
                cat_dropdown = gr.Dropdown(choices=[], label="Select Time Interval", visible=False)
                btn = gr.Button("Re-render Network", variant="primary")
            
        with gr.Column(scale=2):
            stats_box = gr.HTML()
            
            with gr.Tabs():
                with gr.TabItem("Interactive Graph Slice"):
                    vis_box = gr.HTML()
                    download_btn = gr.File(label="Download Filtered Edge List CSV")
                with gr.TabItem("Temporal Trend Charts"):
                    trend_plot = gr.Plot()

    # Initial file parsing
    file_obj.change(
        init_sliders,
        inputs=[file_obj],
        outputs=[min_slider, max_slider, cat_dropdown, controls_group, status_text]
    )

    def process(file_obj, min_val, max_val, cat_val, is_cumulative):
        err, stats, vis, plot, csv_path = analyze_dynamic_network(file_obj, min_val, max_val, cat_val, is_cumulative)
        if err:
            return gr.update(value=err, visible=True), "", "", None, gr.update(visible=False)
        return gr.update(visible=False), stats, vis, plot, gr.update(value=csv_path, visible=True)

    btn.click(
        process,
        inputs=[file_obj, min_slider, max_slider, cat_dropdown, is_cumulative],
        outputs=[error_msg, stats_box, vis_box, trend_plot, download_btn]
    )

if __name__ == "__main__":
    demo.launch()