# app.py
import gradio as gr
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from PIL import Image
import io

from transformers import T5Tokenizer, T5ForConditionalGeneration

# Load trained model
model_path = './t5_summarizer'
tokenizer = T5Tokenizer.from_pretrained(model_path)
model = T5ForConditionalGeneration.from_pretrained(model_path)

# --- Local summarizer ---
def summarize_with_local_model(prompt):
    try:
        inputs = tokenizer("summarize: " + prompt, return_tensors="pt", max_length=512, truncation=True)
        summary_ids = model.generate(inputs["input_ids"], max_length=128)
        return tokenizer.decode(summary_ids[0], skip_special_tokens=True)
    except Exception as e:
        return f"Summary Failed: {str(e)}"

# --- Chart generator ---
def generate_static_chart(df, chart_type, x, y=None):
    try:
        plt.figure(figsize=(10, 6))
        if chart_type == "bar":
            if y:
                sns.barplot(data=df, x=x, y=y)
            else:
                df[x].value_counts().plot(kind='bar')
        elif chart_type == "pie":
            df[x].value_counts().plot.pie(autopct='%1.1f%%')
        elif chart_type == "line" and y:
            sns.lineplot(data=df, x=x, y=y)
        elif chart_type == "scatter" and y:
            sns.scatterplot(data=df, x=x, y=y)
        elif chart_type == "hist":
            sns.histplot(df[x], bins=10, kde=True)
        elif chart_type == "correlation heatmap":
            corr = df.select_dtypes(include='number').corr()
            sns.heatmap(corr, annot=True, cmap='coolwarm', fmt=".2f")
        else:
            return None

        plt.title(f"{chart_type.capitalize()} of {x}" + (f" vs {y}" if y else ""))
        if y:
            plt.legend([y])
        plt.tight_layout()
        buf = io.BytesIO()
        plt.savefig(buf, format='png')
        buf.seek(0)
        return Image.open(buf)

    except Exception as e:
        print("Chart error:", e)
        return None

# --- Main logic ---
def visualize(csv_file, chart_type, x_col, y_col):
    if csv_file is None:
        raise gr.Error("Please upload a CSV file.")
    if not chart_type:
        raise gr.Error("Please select a chart type.")

    try:
        df = pd.read_csv(csv_file.name)
    except Exception as e:
        return f"Error reading CSV: {e}", None, "Error loading CSV."

    if x_col not in df.columns or (y_col and y_col not in df.columns):
        return (
            f"Invalid column names. Available columns: {', '.join(df.columns)}",
            None,
            df.to_markdown()
        )

    numeric_required = ["line", "scatter", "hist", "correlation heatmap"]
    if chart_type in numeric_required:
        if not pd.api.types.is_numeric_dtype(df[x_col]) or (y_col and not pd.api.types.is_numeric_dtype(df[y_col])):
            return (
                f"The chart type '{chart_type}' requires numeric columns.",
                None,
                df.to_markdown()
            )

    sample = df[[x_col] + ([y_col] if y_col else [])].head(10).to_csv(index=False)
    prompt = f"Analyze this {chart_type} chart with sample data:\n{sample}\nIdentify trends or insights."
    summary = summarize_with_local_model(prompt)

    chart = generate_static_chart(df, chart_type, x_col, y_col)
    if chart is None:
        return "Chart generation failed. Try different columns.", None, df.to_markdown()

    return summary, chart, df.to_markdown()

# --- Column selector helper ---
def get_columns(csv_file):
    try:
        df = pd.read_csv(csv_file.name)
        cols = list(df.columns)
        has_numeric = not df.select_dtypes(include='number').empty
        chart_choices = ["bar", "pie"] + (["line", "scatter", "hist", "correlation heatmap"] if has_numeric else [])
        return (
            gr.update(choices=cols, visible=True),
            gr.update(choices=cols, visible=True),
            ", ".join(cols),
            gr.update(choices=chart_choices)
        )
    except Exception as e:
        return (
            gr.update(visible=False),
            gr.update(visible=False),
            f"Error: {e}",
            gr.update(choices=[])
        )

def reset_outputs():
    return "", None, "", gr.update(visible=False), gr.update(visible=False), gr.update(choices=[])

# --- Gradio Interface ---
demo = gr.Blocks()

with demo:
    gr.Markdown("## 📊 Smart Data Visualizer with Local AI Summary (T5)")
    gr.Markdown("Upload a CSV and select columns to visualize. AI will describe trends.")

    with gr.Row():
        csv_file = gr.File(label="Upload CSV", file_types=[".csv"])
        chart_type = gr.Dropdown([], label="Chart Type")

    with gr.Row():
        x_input = gr.Dropdown(label="X-axis Column", visible=False)
        y_input = gr.Dropdown(label="Y-axis Column (optional)", visible=False)

    available_cols = gr.Markdown("")

    with gr.Row():
        summary_output = gr.Textbox(label="📄 AI Summary")
        chart_output = gr.Image(type="pil", label="📊 Chart")

    with gr.Accordion("📋 Full CSV Preview", open=False):
        csv_preview = gr.Markdown("")

    run_btn = gr.Button("Generate Visualization")

    run_btn.click(
        fn=visualize,
        inputs=[csv_file, chart_type, x_input, y_input],
        outputs=[summary_output, chart_output, csv_preview]
    )

    csv_file.change(
        fn=get_columns,
        inputs=csv_file,
        outputs=[x_input, y_input, available_cols, chart_type]
    )

    csv_file.clear(
        fn=reset_outputs,
        inputs=None,
        outputs=[summary_output, chart_output, csv_preview, x_input, y_input, chart_type]
    )

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