app.py

"""ML Data Analysis Studio — Gradio app for Hugging Face Spaces.

Chat about ML models, upload data in any common format, preprocess it,
train regression/classification models, inspect the generated code,
and download PDF/HTML reports plus the cleaned dataset.
"""

import os
import traceback

import gradio as gr
import pandas as pd

import chat_engine
import ml_models
from data_processor import (
    CLEANED_PATH,
    generate_preprocessing_code,
    load_data,
    preprocess_data,
    profile_data,
    profile_text,
)
from report_generator import generate_html_report, generate_pdf_report

# ---------------------------------------------------------------- state ----
STATE = {
    "raw_df": None,
    "clean_df": None,
    "file_name": None,
    "profile": None,
    "steps": [],
    "result": None,
}


# ----------------------------------------------------------------- chat ----
def chat_fn(message, history):
    if not message.strip():
        return ""
    return chat_engine.respond(message, history, STATE["file_name"])


# --------------------------------------------------------------- upload ----
def upload_fn(file):
    if file is None:
        return "Upload a file to begin.", None, gr.update(choices=[]), ""
    file_path = file if isinstance(file, str) else file.name
    try:
        df = load_data(file_path)
    except Exception as e:
        return f"❌ Could not read file: {e}", None, gr.update(choices=[]), ""
    STATE["raw_df"] = df
    STATE["file_name"] = os.path.basename(file_path)
    STATE["profile"] = profile_data(df)
    STATE["clean_df"] = None
    STATE["result"] = None
    msg = (
        f"✅ Loaded **{STATE['file_name']}**. Do you want me to analyse this data? "
        f"Review the profile below, adjust the cleaning options, and click "
        f"**Run Preprocessing**.\n\n" + profile_text(STATE["profile"])
    )
    cols = df.columns.tolist()
    return msg, df.head(20), gr.update(choices=cols, value=cols[-1]), ""


# ----------------------------------------------------------- preprocess ----
def preprocess_fn(missing_strategy, drop_dups, encode_cats, scaling, outliers, target):
    if STATE["raw_df"] is None:
        return "⚠️ Upload a dataset first.", None, "", gr.update(choices=[])
    try:
        clean_df, steps = preprocess_data(
            STATE["raw_df"],
            missing_strategy=missing_strategy,
            drop_duplicates=drop_dups,
            encode_categoricals=encode_cats,
            scaling=scaling,
            clip_outliers=outliers,
            target_column=target or None,
        )
    except Exception as e:
        return f"❌ Preprocessing failed: {e}", None, "", gr.update(choices=[])
    STATE["clean_df"] = clean_df
    STATE["steps"] = steps
    code = generate_preprocessing_code(
        missing_strategy, drop_dups, encode_cats, scaling, outliers, target or None
    )
    msg = "### ✅ Preprocessing complete\n\n" + "\n".join(f"- {s}" for s in steps)
    cols = clean_df.columns.tolist()
    default_target = target if target in cols else (cols[-1] if cols else None)
    return msg, clean_df.head(20), code, gr.update(choices=cols, value=default_target)


# ---------------------------------------------------------------- train ----
def suggest_task_fn(target):
    df = STATE["clean_df"] if STATE["clean_df"] is not None else STATE["raw_df"]
    if df is None or not target or target not in df.columns:
        return gr.update()
    task = ml_models.suggest_task(df, target)
    return gr.update(
        value=task,
        info=f"Suggested: {task} (based on the '{target}' column)",
    )


def task_models_fn(task):
    models = (
        list(ml_models.CLASSIFICATION_MODELS)
        if task == "Classification"
        else list(ml_models.REGRESSION_MODELS)
    )
    return gr.update(choices=models, value=models[0])


def train_fn(target, task, model_name, test_size):
    empty = (None,) * 4
    if STATE["clean_df"] is None:
        if STATE["raw_df"] is None:
            return ("⚠️ Upload a dataset first (Data & Preprocessing tab).", *empty)
        return ("⚠️ Run preprocessing first (Data & Preprocessing tab).", *empty)
    if not target:
        return ("⚠️ Select a target column.", *empty)
    valid = (
        ml_models.CLASSIFICATION_MODELS
        if task == "Classification"
        else ml_models.REGRESSION_MODELS
    )
    if model_name not in valid:
        model_name = list(valid)[0]
    try:
        result = ml_models.train_model(
            STATE["clean_df"], target, model_name, task, test_size
        )
    except Exception as e:
        traceback.print_exc()
        return (f"❌ Training failed: {e}", *empty)
    STATE["result"] = result

    metrics_md = "\n".join(f"| {k} | {v} |" for k, v in result["metrics"].items())
    msg = (
        f"### ✅ {model_name} trained\n\n"
        f"**Task:** {task} | **Target:** `{target}` | "
        f"**Train/test:** {result['n_train']:,}/{result['n_test']:,} | "
        f"**Features used:** {len(result['features'])}\n\n"
        f"| Metric | Value |\n|---|---|\n{metrics_md}"
    )
    code = ml_models.generate_model_code(model_name, task, target, test_size)
    return (
        msg,
        result["plot_path"],
        result["importance_path"],
        code,
        pd.DataFrame([result["metrics"]]),
    )


# ------------------------------------------------------------- download ----
def download_fn(formats):
    if STATE["result"] is None:
        return "⚠️ Train a model first — the report includes its results.", []
    files = []
    try:
        if "PDF report" in formats:
            files.append(generate_pdf_report(STATE["profile"], STATE["steps"], STATE["result"]))
        if "HTML report" in formats:
            files.append(generate_html_report(STATE["profile"], STATE["steps"], STATE["result"]))
        if "Cleaned data (CSV)" in formats and os.path.exists(CLEANED_PATH):
            files.append(CLEANED_PATH)
    except Exception as e:
        traceback.print_exc()
        return f"❌ Report generation failed: {e}", []
    if not files:
        return "⚠️ Select at least one format.", []
    return f"✅ Generated {len(files)} file(s) — download below.", files


# ------------------------------------------------------------------ UI ----
CSS = """
.gradio-container { max-width: 1100px !important; }
footer { display: none !important; }
"""

with gr.Blocks(title="ML Data Analysis Studio", css=CSS, theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        "# 🤖 ML Data Analysis Studio\n"
        "Chat about machine learning, upload your data, clean it, train models, "
        "inspect the code, and download reports."
    )

    with gr.Tabs():
        # ---------------------------------------------------- Tab 1: Chat
        with gr.Tab("💬 ML Assistant"):
            gr.ChatInterface(
                fn=chat_fn,
                chatbot=gr.Chatbot(
                    value=[{"role": "assistant", "content": chat_engine.WELCOME}],
                    height=420,
                    type="messages",
                ),
                type="messages",
                cache_examples=False,
                examples=[
                    "What do you want to analyse today?",
                    "Which model should I use?",
                    "What is a random forest?",
                    "Explain overfitting",
                    "What do precision and recall mean?",
                ],
            )

        # -------------------------------------- Tab 2: Data & Preprocessing
        with gr.Tab("📊 Data & Preprocessing"):
            gr.Markdown("### Step 1 — Upload your dataset (CSV, TSV, Excel, JSON, or Parquet)")
            file_input = gr.File(
                label="Upload data",
                file_types=[".csv", ".tsv", ".xlsx", ".xls", ".json", ".parquet", ".txt"],
            )
            upload_status = gr.Markdown("Upload a file to begin.")
            data_preview = gr.Dataframe(label="Data preview (first 20 rows)")

            gr.Markdown("### Step 2 — Preprocessing & cleaning options")
            with gr.Row():
                missing_dd = gr.Dropdown(
                    [
                        "Impute (mean/mode)",
                        "Impute (median/mode)",
                        "Drop rows with missing values",
                    ],
                    value="Impute (mean/mode)",
                    label="Missing values",
                )
                scale_dd = gr.Dropdown(
                    ["None", "Standard (z-score)", "Min-Max (0-1)"],
                    value="None",
                    label="Scale numeric features",
                )
                target_dd_pre = gr.Dropdown(
                    [], label="Target column (kept out of encoding/scaling)"
                )
            with gr.Row():
                dups_cb = gr.Checkbox(True, label="Remove duplicate rows")
                encode_cb = gr.Checkbox(True, label="One-hot encode categoricals")
                outliers_cb = gr.Checkbox(False, label="Clip outliers (1.5×IQR)")
            gr.Markdown(
                "*Always applied: trims whitespace, converts placeholder values "
                "('N/A', '?', '-') to missing, parses numeric-looking text "
                "('$1,234', '45%'), and drops empty or constant columns.*"
            )
            preprocess_btn = gr.Button("🧹 Run Preprocessing", variant="primary")
            preprocess_status = gr.Markdown()
            clean_preview = gr.Dataframe(label="Cleaned data preview")
            with gr.Accordion("👨‍💻 View preprocessing code", open=False):
                preprocess_code = gr.Code(language="python")

        # ------------------------------------------- Tab 3: Model Training
        with gr.Tab("🧠 Model Training"):
            gr.Markdown("### Step 3 — Configure and train a model")
            with gr.Row():
                target_dd = gr.Dropdown([], label="Target column (what to predict)")
                task_radio = gr.Radio(
                    ["Regression", "Classification"], value="Regression", label="Task"
                )
            with gr.Row():
                model_dd = gr.Dropdown(
                    list(ml_models.REGRESSION_MODELS),
                    value="Linear Regression",
                    label="Model",
                )
                test_slider = gr.Slider(0.1, 0.4, 0.2, step=0.05, label="Test set fraction")
            train_btn = gr.Button("🚀 Train Model", variant="primary")
            train_status = gr.Markdown()
            with gr.Row():
                result_plot = gr.Image(label="Result plot", type="filepath")
                importance_plot = gr.Image(label="Feature importance", type="filepath")
            metrics_df = gr.Dataframe(label="Metrics", visible=True)
            with gr.Accordion("👨‍💻 View model training code", open=False):
                model_code = gr.Code(language="python")

        # ----------------------------------------------- Tab 4: Downloads
        with gr.Tab("📥 Reports & Downloads"):
            gr.Markdown(
                "### Step 4 — Export your analysis\n"
                "Generates a report with the dataset profile, preprocessing steps, "
                "model results, and charts."
            )
            formats_cg = gr.CheckboxGroup(
                ["PDF report", "HTML report", "Cleaned data (CSV)"],
                value=["PDF report", "HTML report", "Cleaned data (CSV)"],
                label="Formats",
            )
            report_btn = gr.Button("📄 Generate Downloads", variant="primary")
            report_status = gr.Markdown()
            report_files = gr.Files(label="Your downloads")

    # ------------------------------------------------------------ wiring ----
    file_input.change(
        upload_fn,
        inputs=file_input,
        outputs=[upload_status, data_preview, target_dd_pre, preprocess_status],
    ).then(lambda t: gr.update(choices=STATE["raw_df"].columns.tolist() if STATE["raw_df"] is not None else [], value=t),
           inputs=target_dd_pre, outputs=target_dd)

    preprocess_btn.click(
        preprocess_fn,
        inputs=[missing_dd, dups_cb, encode_cb, scale_dd, outliers_cb, target_dd_pre],
        outputs=[preprocess_status, clean_preview, preprocess_code, target_dd],
    )

    target_dd.change(suggest_task_fn, inputs=target_dd, outputs=task_radio)
    task_radio.change(task_models_fn, inputs=task_radio, outputs=model_dd)

    train_btn.click(
        train_fn,
        inputs=[target_dd, task_radio, model_dd, test_slider],
        outputs=[train_status, result_plot, importance_plot, model_code, metrics_df],
    )

    report_btn.click(download_fn, inputs=formats_cg, outputs=[report_status, report_files])


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