import gradio as gr
import pandas as pd
import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.model_selection import train_test_split
from sklearn.naive_bayes import MultinomialNB
from sklearn.linear_model import LogisticRegression
from sklearn.svm import LinearSVC
from sklearn.metrics import classification_report, confusion_matrix, accuracy_score
import plotly.graph_objects as go

# Global variables to store the trained model, vectorizer, and categories
global_vectorizer = None
global_model = None
global_classes = None

def train_classifier(file_obj, algorithm):
    global global_vectorizer, global_model, global_classes
    
    if file_obj is None:
        return "Please upload a CSV or Excel labeled training file.", None, None, 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)
        
    # Standardize column headers
    text_col, label_col = None, None
    for col in df.columns:
        if col.lower() in ['text', 'document', 'content', 'body', 'sentence']:
            text_col = col
        elif col.lower() in ['label', 'category', 'class', 'target', 'topic']:
            label_col = col
            
    if not text_col or not label_col:
        # Fallbacks
        string_cols = df.select_dtypes(include=['object']).columns
        if len(string_cols) >= 2:
            text_col = string_cols[0]
            label_col = string_cols[1]
        else:
            return "Could not find 'Text' and 'Label' columns. Make sure your sheet has at least two columns.", None, None, gr.update(visible=False)
            
    df = df.dropna(subset=[text_col, label_col])
    
    if len(df) < 10:
        return "Training dataset is too small. Please provide at least 10 labeled rows.", None, None, gr.update(visible=False)
        
    texts = df[text_col].astype(str).tolist()
    labels = df[label_col].astype(str).tolist()
    
    # Split
    X_train, X_test, y_train, y_test = train_test_split(texts, labels, test_size=0.25, random_state=42)
    
    # Vectorizer
    vectorizer = TfidfVectorizer(stop_words='english', max_features=2000)
    X_train_vec = vectorizer.fit_transform(X_train)
    X_test_vec = vectorizer.transform(X_test)
    
    # Model select
    if algorithm == "Naive Bayes":
        model = MultinomialNB()
    elif algorithm == "Logistic Regression":
        model = LogisticRegression(random_state=42, max_iter=1000)
    else:  # Linear SVM
        model = LinearSVC(random_state=42)
        
    model.fit(X_train_vec, y_train)
    preds = model.predict(X_test_vec)
    
    # Metrics
    acc = accuracy_score(y_test, preds)
    classes = sorted(list(set(labels)))
    report = classification_report(y_test, preds, output_dict=True)
    report_df = pd.DataFrame(report).transpose().round(3).reset_index().rename(columns={"index": "Metric Class"})
    
    # Save globals for real-time inference
    global_vectorizer = vectorizer
    global_model = model
    global_classes = classes
    
    # 4. Generate Visual Plotly Confusion Matrix
    cm = confusion_matrix(y_test, preds, labels=classes)
    
    fig = go.Figure(data=go.Heatmap(
        z=cm,
        x=classes,
        y=classes,
        colorscale='Oranges',
        text=cm,
        texttemplate="%{text}",
        hoverinfo='z'
    ))
    
    fig.update_layout(
        title=f"Confusion Matrix (Test Accuracy: {acc:.2%})",
        paper_bgcolor='#16100c',
        plot_bgcolor='#16100c',
        font_color='#f4eee6',
        xaxis=dict(title="Predicted label", gridcolor='rgba(255,255,255,0.05)'),
        yaxis=dict(title="True label", gridcolor='rgba(255,255,255,0.05)'),
        margin=dict(l=40, r=40, t=50, b=40)
    )
    
    metrics_summary_html = f"""
    <div style='display: grid; grid-template-columns: repeat(2, 1fr); gap: 1rem; margin-bottom: 1.5rem;'>
        <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;'>Model Testing Accuracy</div>
            <div style='font-size: 2rem; font-weight: bold; margin-top: 0.5rem;'>{acc:.2%}</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;'>Number of Target Classes</div>
            <div style='font-size: 2rem; font-weight: bold; margin-top: 0.5rem;'>{len(classes)}</div>
        </div>
    </div>
    """
    
    return "", metrics_summary_html, fig, report_df, gr.update(visible=True)

def classify_new_text(new_text):
    global global_vectorizer, global_model, global_classes
    
    if global_model is None or global_vectorizer is None:
        return "Please train a classification model first using the panel on the left.", None
        
    if not new_text or len(new_text.strip()) < 3:
        return "Please enter a valid text to classify.", None
        
    # Vectorize
    vec = global_vectorizer.transform([new_text])
    
    # Predict
    if hasattr(global_model, "predict_proba"):
        probs = global_model.predict_proba(vec)[0]
    else:  # LinearSVC uses decision function
        decision = global_model.decision_function(vec)[0]
        # Map decision scores to pseudo-probabilities via softmax or sigmoid
        if len(global_classes) == 2:
            # For binary LinearSVC, decision is a single float
            probs = np.array([1 / (1 + np.exp(decision)), 1 / (1 + np.exp(-decision))])
        else:
            exp_scores = np.exp(decision - np.max(decision))
            probs = exp_scores / exp_scores.sum()
            
    pred_idx = np.argmax(probs)
    predicted_label = global_classes[pred_idx]
    confidence = probs[pred_idx]
    
    # Generate horizontal Plotly bar chart
    fig = go.Figure(go.Bar(
        x=probs,
        y=global_classes,
        orientation='h',
        marker=dict(color='#ff7043', line=dict(width=1, color='#16100c')),
        text=[f"{p:.1%}" for p in probs],
        textposition='auto'
    ))
    
    fig.update_layout(
        title="Class Probability Distribution",
        paper_bgcolor='#16100c',
        plot_bgcolor='#16100c',
        font_color='#f4eee6',
        xaxis=dict(showgrid=True, gridcolor='rgba(255,255,255,0.05)', range=[0, 1]),
        yaxis=dict(gridcolor='rgba(255,255,255,0.05)'),
        margin=dict(l=40, r=40, t=50, b=40)
    )
    
    result_html = f"""
    <div style='background: rgba(255, 112, 67, 0.05); border-left: 4px solid #ff7043; border-radius: 4px; padding: 1.5rem; margin-bottom: 1rem;'>
        <div style='font-size: 0.8rem; text-transform: uppercase; color: #f4eee6; letter-spacing: 0.1em;'>Predicted Category</div>
        <div style='font-size: 2.2rem; font-weight: bold; color: #ff7043; margin-top: 0.5rem;'>{predicted_label}</div>
        <div style='font-size: 0.95rem; margin-top: 0.5rem; opacity: 0.8;'>Confidence Score: <strong>{confidence:.2%}</strong></div>
    </div>
    """
    
    return result_html, fig

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="Text Classifier Studio") as demo:
    gr.Markdown(
        """
        # 🏷️ Custom Text Classification Studio
        ### Upload a labeled training sheet (CSV containing Text and Category labels) to train a custom machine learning classifier locally. Test it instantly with live texts!
        """
    )
    
    error_msg = gr.Markdown("", visible=False)
    
    with gr.Row():
        with gr.Column(scale=1):
            file_obj = gr.File(label="Upload Training CSV or Excel", file_types=[".csv", ".xlsx"])
            gr.Markdown("💡 **Tip**: Make sure your sheet has a **Text** column and a **Label** column (e.g., 'Politics', 'Sports', 'Art').")
            
            algorithm = gr.Radio(
                choices=["Naive Bayes", "Logistic Regression", "Linear Support Vector (SVM)"],
                value="Naive Bayes",
                label="Classification Algorithm"
            )
            
            train_btn = gr.Button("Train Custom Classifier", variant="primary")
            
        with gr.Column(scale=2):
            stats_box = gr.HTML()
            
            with gr.Tabs():
                with gr.TabItem("Validation & Diagnostics"):
                    plot_cm = gr.Plot()
                    table_report = gr.Dataframe(headers=["Metric Class", "precision", "recall", "f1-score", "support"])
                    
                with gr.TabItem("Live Model Playground"):
                    with gr.Group(visible=False) as inference_group:
                        new_text_input = gr.Textbox(
                            label="Enter New Text to Classify",
                            placeholder="Write or paste any paragraph here to test the trained model in real-time...",
                            lines=5
                        )
                        predict_btn = gr.Button("Predict Category", variant="secondary")
                        prediction_result = gr.HTML()
                        plot_probs = gr.Plot()
                    
                    no_model_warning = gr.Markdown(
                        "⚠️ **No Model Trained Yet**: Upload a training dataset on the left and click 'Train Custom Classifier' to unlock the live playground!",
                        visible=True
                    )

    def on_train_success(file_obj, algo):
        err, stats, plot, report, update_group = train_classifier(file_obj, algo)
        if err:
            return gr.update(value=err, visible=True), "", None, None, gr.update(visible=False), gr.update(visible=True)
        return gr.update(visible=False), stats, plot, report, update_group, gr.update(visible=False)
        
    train_btn.click(
        on_train_success,
        inputs=[file_obj, algorithm],
        outputs=[error_msg, stats_box, plot_cm, table_report, inference_group, no_model_warning]
    )
    
    predict_btn.click(
        classify_new_text,
        inputs=[new_text_input],
        outputs=[prediction_result, plot_probs]
    )

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