import os
import gradio as gr
import pandas as pd
import numpy as np
import re
import plotly.express as px
from collections import Counter

def load_data(file_obj):
    """Safely loads CSV, Excel, or TXT file into a Pandas DataFrame."""
    if file_obj is None:
        return None, gr.update(choices=[], visible=False), "Please upload a file."
    
    file_path = file_obj.name
    ext = os.path.splitext(file_path)[1].lower()
    
    try:
        if ext == '.csv':
            df = pd.read_csv(file_path)
        elif ext in ['.xls', '.xlsx']:
            df = pd.read_excel(file_path)
        elif ext == '.txt':
            with open(file_path, 'r', encoding='utf-8') as f:
                content = f.read()
            df = pd.DataFrame({'text': [content]})
        else:
            return None, gr.update(choices=[], visible=False), "Unsupported file format. Please upload .csv, .xlsx, or .txt."
        
        string_cols = [col for col in df.columns if df[col].dtype == 'object' or df[col].astype(str).str.len().mean() > 5]
        if not string_cols:
            string_cols = list(df.columns)
            
        return df, gr.update(choices=string_cols, value=string_cols[0], visible=True), f"Successfully loaded dataset with {len(df)} rows."
    except Exception as e:
        return None, gr.update(choices=[], visible=False), f"Error loading file: {str(e)}"

def calculate_collocations(docs, n_gram_type, metric, min_freq, top_n):
    """Calculates n-gram collocations using Raw Frequency or PMI."""
    # Combine texts and tokenize
    words = []
    stopwords = {'the', 'and', 'for', 'that', 'with', 'this', 'have', 'from', 'your', 'will', 'not', 'are', 'was', 'were', 'but', 'how', 'they', 'our', 'them', 'their', 'she', 'him', 'her', 'his', 'has', 'had', 'been', 'would', 'could', 'should'}
    
    for doc in docs:
        cleaned = re.sub(r'[^\w\s]', ' ', doc.lower())
        doc_words = [w.strip() for w in cleaned.split() if w.strip() and w.strip() not in stopwords and len(w.strip()) > 2]
        words.extend(doc_words)
        
    if len(words) < 5:
        return pd.DataFrame()
        
    n = 2 if n_gram_type == "Bigrams (2-word pairs)" else 3
    
    # Generate n-grams
    ngrams_list = []
    for i in range(len(words) - n + 1):
        ngram = tuple(words[i:i+n])
        ngrams_list.append(ngram)
        
    ngram_counts = Counter(ngrams_list)
    
    # Filter by minimum frequency
    filtered_ngrams = {k: v for k, v in ngram_counts.items() if v >= min_freq}
    
    if not filtered_ngrams:
        return pd.DataFrame()
        
    results = []
    
    if metric == "Raw Joint Frequency":
        for ngram, count in Counter(filtered_ngrams).most_common(top_n):
            results.append({
                "Word Phrase": " ".join(ngram),
                "Score": float(count),
                "Measure": "Frequency"
            })
    else:
        # Pointwise Mutual Information (PMI)
        # PMI(x, y) = log2( P(x,y) / (P(x)*P(y)) )
        word_counts = Counter(words)
        total_words = len(words)
        total_ngrams = len(ngrams_list)
        
        pmi_scores = {}
        for ngram, count in filtered_ngrams.items():
            # Joint probability
            p_joint = count / total_ngrams
            
            # Marginal probabilities product
            p_marginals = 1.0
            for word in ngram:
                p_marginals *= (word_counts[word] / total_words)
                
            pmi = np.log2(p_joint / p_marginals)
            pmi_scores[ngram] = pmi
            
        sorted_pmi = sorted(pmi_scores.items(), key=lambda x: x[1], reverse=True)[:top_n]
        for ngram, score in sorted_pmi:
            results.append({
                "Word Phrase": " ".join(ngram),
                "Score": round(float(score), 4),
                "Measure": "PMI Score"
            })
            
    return pd.DataFrame(results)

def run_analysis(file_obj, text_col, n_gram_type, metric, min_freq, top_n):
    if file_obj is None:
        return None, None, None, "Please upload a dataset first."
        
    # Re-load data
    df, _, _ = load_data(file_obj)
    if df is None:
        return None, None, None, "Failed to parse the file."
        
    docs = df[text_col].astype(str).fillna("").tolist()
    if not docs:
        return None, None, None, "No text documents found in the selected column."
        
    try:
        df_res = calculate_collocations(docs, n_gram_type, metric, min_freq, top_n)
        
        if df_res.empty:
            return None, None, None, "No collocations met the minimum frequency filter. Try lowering 'Min Word Co-occurrences'."
            
        # Plotly Bar Chart
        fig = px.bar(
            df_res,
            x="Score",
            y="Word Phrase",
            orientation="h",
            color="Score",
            title=f"Top Collocations via {metric}",
            template="plotly_dark",
            color_continuous_scale="Cividis"
        )
        fig.update_layout(yaxis={'categoryorder': 'total ascending'}, height=450, margin=dict(l=20, r=20, t=40, b=20))
        
        # Export CSV
        csv_path = "collocations_report.csv"
        df_res.to_csv(csv_path, index=False)
        
        status_md = f"Successfully calculated top **{len(df_res)}** collocations using {metric}."
        
        return df_res, fig, csv_path, status_md
    except Exception as e:
        return None, None, None, f"Analysis failed: {str(e)}"

custom_css = """
body {
    background-color: #0b0f19;
    color: #f3f4f6;
}
.gradio-container {
    font-family: 'Inter', sans-serif !important;
}
h1, h2 {
    color: #6366f1 !important;
}
"""

with gr.Blocks(theme=gr.themes.Default(primary_hue="indigo", secondary_hue="slate"), css=custom_css) as demo:
    df_state = gr.State()
    
    gr.HTML("""
    <div style="text-align: center; margin-bottom: 2rem;">
        <h1 style="font-size: 2.5rem; font-weight: 700; margin-bottom: 0.5rem; background: linear-gradient(to right, #6366f1, #a855f7); -webkit-background-clip: text; -webkit-text-fill-color: transparent;">Interactive Collocation Analyzer</h1>
        <p style="font-size: 1.1rem; color: #94a3b8; max-width: 800px; margin: 0 auto;">
            Extract and analyze recurring word pairs (Bigrams) or triplets (Trigrams) that frequently co-occur in your text. 
            Toggle between raw joint counts and Pointwise Mutual Information (PMI) to reveal locked idioms and idioms.
        </p>
    </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 1. Upload Dataset")
            file_input = gr.File(label="Upload Dataset (.csv, .xlsx, .txt)", file_types=[".csv", ".xlsx", ".txt"])
            text_column_selector = gr.Dropdown(
                label="Target Text Column", 
                choices=[], 
                visible=False,
                interactive=True
            )
            status_text = gr.Markdown("No file uploaded yet.")
            
            gr.Markdown("### 2. Configure Collocations")
            n_gram_type = gr.Radio(
                choices=["Bigrams (2-word pairs)", "Trigrams (3-word triplets)"],
                value="Bigrams (2-word pairs)",
                label="Co-occurrence Type"
            )
            metric_selector = gr.Radio(
                choices=["Raw Joint Frequency", "Pointwise Mutual Information (PMI)"],
                value="Raw Joint Frequency",
                label="Association Metric"
            )
            with gr.Row():
                min_freq = gr.Slider(minimum=1, maximum=50, value=3, step=1, label="Min Word Co-occurrences")
                top_n = gr.Slider(minimum=5, maximum=40, value=15, step=1, label="Phrases to Display")
                
            run_btn = gr.Button("Analyze Collocations", variant="primary")
            
        with gr.Column(scale=2):
            gr.Markdown("### 3. Collocations Results")
            status_markdown = gr.Markdown("Configure settings and click 'Analyze Collocations' to run.")
            
            with gr.Tabs():
                with gr.TabItem("Collocations Plot"):
                    chart_output = gr.Plot(label="Collocation Strength Plot")
                with gr.TabItem("Collocations Table"):
                    table_output = gr.Dataframe(
                        headers=["Word Phrase", "Score", "Measure"],
                        datatype=["str", "number", "str"],
                        interactive=False,
                        wrap=True
                    )
                    
            gr.Markdown("### 4. Export")
            download_csv = gr.File(label="Download Collocations Report (CSV)")

    file_input.change(
        fn=load_data,
        inputs=file_input,
        outputs=[df_state, text_column_selector, status_text]
    )
    
    run_btn.click(
        fn=run_analysis,
        inputs=[file_input, text_column_selector, n_gram_type, metric_selector, min_freq, top_n],
        outputs=[table_output, chart_output, download_csv, status_markdown]
    )

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