import gradio as gr
import networkx as nx
import matplotlib.pyplot as plt
from pathlib import Path
import spacy
import re

# Load spaCy model
try:
    nlp = spacy.load("en_core_web_sm")
except OSError:
    print("Downloading spaCy model...")
    import subprocess
    subprocess.run(["python", "-m", "spacy", "download", "en_core_web_sm"])
    nlp = spacy.load("en_core_web_sm")

# Define categories and their colors with better contrast
CATEGORIES = {
    "Main Theme": "#4287f5",    # Bright blue
    "Event": "#42f54b",         # Bright green
    "Person": "#f542aa",        # Pink
    "Law": "#f5d442",          # Yellow
    "Concept": "#f54242"       # Red
}

# Sample historical data for Unit 5 (1844-1877)
HISTORICAL_DATA = {
    "civil war": {
        "category": "Main Theme",
        "related": [
            ("Abraham Lincoln", "Person"),
            ("Emancipation Proclamation", "Law"),
            ("Confederate States", "Concept"),
            ("Union Army", "Concept"),
            ("Battle of Gettysburg", "Event"),
            ("Slavery", "Main Theme"),
            ("Robert E. Lee", "Person"),
            ("Ulysses S. Grant", "Person")
        ]
    },
    "reconstruction": {
        "category": "Main Theme",
        "related": [
            ("13th Amendment", "Law"),
            ("14th Amendment", "Law"),
            ("15th Amendment", "Law"),
            ("Freedmen's Bureau", "Concept"),
            ("Andrew Johnson", "Person"),
            ("Black Codes", "Law"),
            ("Radical Republicans", "Concept"),
            ("Carpetbaggers", "Concept")
        ]
    },
    "manifest destiny": {
        "category": "Main Theme",
        "related": [
            ("Mexican-American War", "Event"),
            ("Oregon Territory", "Concept"),
            ("California Gold Rush", "Event"),
            ("James K. Polk", "Person"),
            ("Treaty of Guadalupe Hidalgo", "Law"),
            ("Westward Expansion", "Concept"),
            ("Native American Displacement", "Event"),
            ("Mexican Cession", "Event")
        ]
    }
}

def categorize_term(term):
    """Categorize a term based on predefined data and NER."""
    term_lower = term.lower()
    
    # Check predefined categories first
    if term_lower in HISTORICAL_DATA:
        return HISTORICAL_DATA[term_lower]["category"]
    
    # Use spaCy for NER
    doc = nlp(term)
    for ent in doc.ents:
        if ent.label_ == "PERSON":
            return "Person"
        elif ent.label_ == "EVENT" or ent.label_ == "DATE":
            return "Event"
        elif ent.label_ == "LAW" or ent.label_ == "ORG":
            return "Law"
    
    # Default to Concept if no other category is found
    return "Concept"

def get_related_terms(term):
    """Get related terms for a given historical term."""
    term_lower = term.lower()
    
    if term_lower in HISTORICAL_DATA:
        return HISTORICAL_DATA[term_lower]["related"]
    
    # If term not in predefined data, return some general connections
    # based on the time period
    general_connections = [
        ("Civil War", "Main Theme"),
        ("Reconstruction", "Main Theme"),
        ("Abraham Lincoln", "Person"),
        ("Slavery", "Concept"),
        ("United States", "Concept")
    ]
    return general_connections[:5]  # Limit to 5 connections

def generate_context_map(term):
    """Generate a network visualization for the given term."""
    if not term or not term.strip():
        return None
    
    # Create graph
    G = nx.Graph()
    
    # Add main term
    main_term = term.strip()
    term_category = categorize_term(main_term)
    G.add_node(main_term, category=term_category)
    
    # Add related terms
    related_terms = get_related_terms(main_term)
    for related_term, category in related_terms:
        if related_term.lower() != main_term.lower():
            G.add_node(related_term, category=category)
            G.add_edge(main_term, related_term)
    
    # Create visualization
    plt.figure(figsize=(12, 12))
    plt.clf()
    
    # Set light background for better contrast
    plt.gca().set_facecolor('#ffffff')
    plt.gcf().set_facecolor('#ffffff')
    
    # Create layout
    pos = nx.spring_layout(G, k=1.5, iterations=50)
    
    # Draw nodes for each category
    for category, color in CATEGORIES.items():
        node_list = [node for node, attr in G.nodes(data=True) 
                    if attr.get('category') == category]
        if node_list:
            nx.draw_networkx_nodes(G, pos, 
                                 nodelist=node_list,
                                 node_color=color, 
                                 node_size=3000,
                                 alpha=0.7)
    
    # Draw edges
    nx.draw_networkx_edges(G, pos, edge_color='gray', width=2, alpha=0.6)
    
    # Add labels with better formatting
    labels = nx.draw_networkx_labels(G, pos, 
                                   font_size=10, 
                                   font_weight='bold',
                                   font_color='black')
    
    # Add title and legend
    plt.title(f"Historical Context Map for '{main_term}'", 
              fontsize=16, 
              pad=20)
    
    # Add category legend
    legend_elements = [plt.Line2D([0], [0], marker='o', color='w', 
                                 markerfacecolor=color, markersize=10, 
                                 label=category)
                      for category, color in CATEGORIES.items()]
    plt.legend(handles=legend_elements, loc='upper left', 
              bbox_to_anchor=(1, 1))
    
    plt.tight_layout()
    return plt.gcf()

# Create Gradio interface
iface = gr.Interface(
    fn=generate_context_map,
    inputs=gr.Textbox(
        label="Enter a historical term from Unit 5 (1844-1877)",
        placeholder="e.g., Civil War, Abraham Lincoln, Reconstruction"
    ),
    outputs=gr.Plot(),
    title="Historical Context Mapper",
    description="""Enter a term from Unit 5 (1844-1877) to see its historical context and connections.
                   The visualization will show how the term relates to key events, people, laws, and concepts from this period.""",
    examples=[
        ["Civil War"],
        ["Reconstruction"],
        ["Manifest Destiny"],
        ["Abraham Lincoln"],
        ["Emancipation Proclamation"]
    ],
    theme="default"
)

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