src/streamlit_app.py

# streamlit_data_analysis_app.py
# Streamlit Data Analysis App for Hugging Face Spaces
# Features:
# - Upload CSV / Excel
# - Automatic cleaning & standardization
# - Preprocessing (imputation, encoding, scaling)
# - Quick visualizations (histogram, boxplot, scatter, correlation heatmap)
# - Preview cleaned dataset
# - LLM-powered insights using Hugging Face Inference API
# - Auto fallback if model access (403) fails
# - Uses HF_TOKEN from Streamlit secrets or environment

import os
import io
import streamlit as st
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.impute import SimpleImputer
from sklearn.preprocessing import OneHotEncoder, OrdinalEncoder, StandardScaler
from sklearn.compose import ColumnTransformer
from sklearn.pipeline import Pipeline
from huggingface_hub import InferenceClient

# ---------- CONFIGURATION ----------
st.set_page_config(page_title="Data Analysis App", layout="wide")

# ✅ Safe HF_TOKEN loader (works locally + on Spaces)
try:
    HF_TOKEN = st.secrets["HF_TOKEN"]
except Exception:
    HF_TOKEN = os.getenv("HF_TOKEN")

if not HF_TOKEN:
    st.warning("⚠️ HF_TOKEN not found. Please add it to your Hugging Face Space secrets or environment.")
else:
    st.success("✅ Hugging Face token loaded successfully.")

# Default open-access models
MODEL_OPTIONS = {
    "mistralai/Mistral-7B-Instruct-v0.3": "Mistral 7B Instruct (open, strong)",
    "HuggingFaceH4/zephyr-7b-beta": "Zephyr 7B Beta (open, fluent)",
    "bigscience/bloom-3b": "Bloom 3B (lightweight, open)"
}

# ---------- UTILITY FUNCTIONS ----------

def read_file(uploaded_file) -> pd.DataFrame:
    """Reads uploaded CSV or Excel file."""
    name = uploaded_file.name.lower()
    if name.endswith(('.csv', '.txt')):
        return pd.read_csv(uploaded_file)
    elif name.endswith(('.xls', '.xlsx')):
        return pd.read_excel(uploaded_file)
    else:
        raise ValueError("Unsupported file type. Please upload CSV or Excel.")


def clean_column_name(col: str) -> str:
    col = str(col).strip().lower().replace("\n", " ").replace("\t", " ")
    col = "_".join(col.split())
    col = ''.join(c for c in col if (c.isalnum() or c == '_'))
    while '__' in col:
        col = col.replace('__', '_')
    return col


def standardize_dataframe(df: pd.DataFrame, drop_all_nan_cols: bool = True) -> pd.DataFrame:
    """Standardizes column names and cleans whitespace."""
    df = df.copy()
    for c in df.select_dtypes(include=['object']).columns:
        df[c] = df[c].apply(lambda x: x.strip() if isinstance(x, str) else x)
    df.columns = [clean_column_name(c) for c in df.columns]
    if drop_all_nan_cols:
        df.dropna(axis=1, how='all', inplace=True)
    for c in df.columns:
        if df[c].dtype == object:
            sample = df[c].dropna().astype(str).head(20)
            if not sample.empty:
                parsed = pd.to_datetime(sample, errors='coerce')
                if parsed.notna().sum() / len(sample) > 0.6:
                    df[c] = pd.to_datetime(df[c], errors='coerce')
    return df


def summarize_dataframe(df: pd.DataFrame, max_rows: int = 5):
    """Creates a structured summary of the dataframe."""
    summary = {'shape': df.shape, 'columns': [], 'preview': df.head(max_rows).to_dict(orient='records')}
    for c in df.columns:
        info = {'name': c, 'dtype': str(df[c].dtype), 'n_missing': int(df[c].isna().sum()), 'n_unique': int(df[c].nunique(dropna=True))}
        if pd.api.types.is_numeric_dtype(df[c]):
            info['summary'] = df[c].describe().to_dict()
        elif pd.api.types.is_datetime64_any_dtype(df[c]):
            info['summary'] = {'min': str(df[c].min()), 'max': str(df[c].max())}
        else:
            info['top_values'] = df[c].astype(str).value_counts().head(5).to_dict()
        summary['columns'].append(info)
    return summary


def prepare_preprocessing_pipeline(df: pd.DataFrame, impute_strategy_num='median', scale_numeric=True, encode_categorical='onehot'):
    """Build preprocessing pipeline for numeric and categorical features."""
    numeric_cols = list(df.select_dtypes(include=[np.number]).columns)
    cat_cols = list(df.select_dtypes(include=['object', 'category', 'bool']).columns)
    transformers = []
    if numeric_cols:
        num_pipe = [('imputer', SimpleImputer(strategy=impute_strategy_num))]
        if scale_numeric:
            num_pipe.append(('scaler', StandardScaler()))
        transformers.append(('num', Pipeline(num_pipe), numeric_cols))
    if cat_cols:
        if encode_categorical == 'onehot':
            cat_pipe = Pipeline([
                ('imputer', SimpleImputer(strategy='most_frequent')),
                ('onehot', OneHotEncoder(handle_unknown='ignore', sparse=False))
            ])
        else:
            cat_pipe = Pipeline([
                ('imputer', SimpleImputer(strategy='most_frequent')),
                ('ord', OrdinalEncoder())
            ])
        transformers.append(('cat', cat_pipe, cat_cols))
    return ColumnTransformer(transformers), numeric_cols + cat_cols


def apply_preprocessing(df: pd.DataFrame, preprocessor: ColumnTransformer) -> pd.DataFrame:
    """Applies preprocessing pipeline and returns processed DataFrame."""
    X = preprocessor.fit_transform(df)
    feature_names = []
    for name, trans, cols in preprocessor.transformers_:
        if name == 'num':
            feature_names += cols
        elif name == 'cat':
            try:
                ohe = trans.named_steps['onehot']
                for col, cats in zip(cols, ohe.categories_):
                    feature_names += [f"{col}__{c}" for c in cats]
            except Exception:
                feature_names += cols
    return pd.DataFrame(X, columns=feature_names)


# ---------- LLM INTEGRATION ----------

def build_dataset_prompt(summary, user_question=None):
    """Builds a prompt for dataset insights."""
    s = [f"Dataset shape: {summary['shape'][0]} rows, {summary['shape'][1]} columns."]
    for c in summary['columns']:
        s.append(f"- {c['name']} ({c['dtype']}) missing={c['n_missing']} unique={c['n_unique']}")
    s.append("Preview:")
    for row in summary['preview']:
        s.append(str(row))
    if user_question:
        s.append(f"User question: {user_question}")
    else:
        s.append("Please give a dataset summary, patterns, and visualization suggestions.")
    return "\n".join(s)


def call_llm(prompt: str, model: str, max_tokens: int = 512) -> str:
    """Calls the Hugging Face Inference API with error handling and fallback."""
    if not HF_TOKEN:
        return "⚠️ No Hugging Face token found."
    client = InferenceClient(token=HF_TOKEN)
    try:
        response = client.text_generation(model=model, inputs=prompt, max_new_tokens=max_tokens)
        if isinstance(response, dict):
            return response.get('generated_text', str(response))
        return str(response)
    except Exception as e:
        if "403" in str(e):
            fallback = "mistralai/Mistral-7B-Instruct-v0.3"
            if model != fallback:
                try:
                    st.warning(f"🚫 Access denied to {model}. Falling back to {fallback}...")
                    response = client.text_generation(model=fallback, inputs=prompt, max_new_tokens=max_tokens)
                    if isinstance(response, dict):
                        return response.get('generated_text', str(response))
                    return str(response)
                except Exception as e2:
                    return f"❌ Fallback model also failed: {e2}"
            return "🚫 Access denied (403). Try using an open-access model like Mistral or Zephyr."
        return f"❌ LLM call failed: {e}"

# ---------- STREAMLIT UI ----------

st.title("📊 Data Analysis & Cleaning App (Hugging Face + Streamlit)")
st.markdown("Upload CSV or Excel files, clean, preprocess, visualize, and generate insights using an LLM.")

with st.sidebar:
    st.header("⚙️ Options")
    model_choice = st.selectbox("Select LLM model", options=list(MODEL_OPTIONS.keys()), format_func=lambda k: MODEL_OPTIONS[k])
    max_tokens = st.slider("LLM max tokens", 128, 1024, 512, 64)
    impute_strategy_num = st.selectbox("Numeric imputation", ['mean', 'median', 'most_frequent'])
    encode_categorical = st.selectbox("Categorical encoding", ['onehot', 'ordinal'])
    scale_numeric = st.checkbox("Scale numeric features", True)
    show_raw_preview = st.checkbox("Show raw preview", True)

uploaded_file = st.file_uploader("📂 Upload your CSV or Excel file", type=['csv', 'xls', 'xlsx', 'txt'])

if uploaded_file:
    with st.spinner("Reading file..."):
        raw_df = read_file(uploaded_file)

    if show_raw_preview:
        st.subheader("Raw Data Preview")
        st.dataframe(raw_df.head())

    st.subheader("Data Cleaning & Standardization")
    cleaned_df = standardize_dataframe(raw_df)
    st.write(f"✅ Cleaned data shape: {cleaned_df.shape}")
    st.dataframe(cleaned_df.head())

    st.subheader("Summary")
    summary = summarize_dataframe(cleaned_df)
    st.write(f"Shape: {summary['shape']}")
    st.json(summary['columns'])

    st.subheader("Preprocessing")
    if st.button("Generate Preprocessing Pipeline"):
        preproc, _ = prepare_preprocessing_pipeline(cleaned_df, impute_strategy_num, scale_numeric, encode_categorical)
        processed_df = apply_preprocessing(cleaned_df, preproc)
        st.success("Preprocessing complete!")
        st.dataframe(processed_df.head())
        st.download_button("⬇️ Download Processed CSV", processed_df.to_csv(index=False), "processed_data.csv")

    st.subheader("Visualizations")
    viz_col = st.selectbox("Select column", options=cleaned_df.columns)
    viz_type = st.selectbox("Visualization type", ['Histogram', 'Boxplot', 'Bar (categorical)', 'Scatter', 'Correlation heatmap'])

    if viz_type == 'Scatter':
        second_col = st.selectbox("Second column", options=[c for c in cleaned_df.columns if c != viz_col])

    if st.button("Show Visualization"):
        fig, ax = plt.subplots(figsize=(8, 5))
        try:
            if viz_type == 'Histogram':
                sns.histplot(cleaned_df[viz_col], kde=True, ax=ax)
            elif viz_type == 'Boxplot':
                sns.boxplot(x=cleaned_df[viz_col], ax=ax)
            elif viz_type == 'Bar (categorical)':
                counts = cleaned_df[viz_col].astype(str).value_counts().head(20)
                sns.barplot(x=counts.values, y=counts.index, ax=ax)
            elif viz_type == 'Scatter':
                sns.scatterplot(x=cleaned_df[viz_col], y=cleaned_df[second_col], ax=ax)
            elif viz_type == 'Correlation heatmap':
                corr = cleaned_df.select_dtypes(include=[np.number]).corr()
                sns.heatmap(corr, annot=True, cmap='coolwarm', ax=ax)
            st.pyplot(fig)
        except Exception as e:
            st.error(f"Visualization failed: {e}")

    st.subheader("🧠 Ask the LLM for Insights")
    user_q = st.text_area("Enter your question (optional):")
    if st.button("Get Insights"):
        with st.spinner("Generating insights..."):
            prompt = build_dataset_prompt(summary, user_q if user_q else None)
            llm_resp = call_llm(prompt, model_choice, max_tokens)
            st.write(llm_resp)

else:
    st.info("📥 Upload a file to begin.")