app.py

# pip install seaborn mlxtend
# app.py
# app.py
import os
import io
import tempfile
import json
import math
import traceback
from typing import Optional, List, Tuple, Dict, Any, Union
import pandas as pd
import numpy as np
import gradio as gr
import seaborn as sns
from scipy import stats
import warnings
warnings.filterwarnings('ignore')
import plotly.express as px
import plotly.graph_objects as go
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split, GridSearchCV, cross_val_score
from sklearn.preprocessing import StandardScaler, MinMaxScaler, OneHotEncoder, LabelEncoder
from sklearn.compose import ColumnTransformer
from sklearn.pipeline import Pipeline
from sklearn.impute import SimpleImputer
from sklearn.feature_selection import RFECV
from sklearn.linear_model import LinearRegression, Ridge, Lasso, LogisticRegression
from sklearn.tree import DecisionTreeRegressor, DecisionTreeClassifier
from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier, GradientBoostingRegressor, GradientBoostingClassifier
from sklearn.svm import SVR, SVC
from sklearn.neighbors import KNeighborsClassifier, KNeighborsRegressor
from sklearn.naive_bayes import GaussianNB
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score, accuracy_score, classification_report, confusion_matrix, roc_auc_score, roc_curve, silhouette_score, davies_bouldin_score, calinski_harabasz_score
from sklearn.cluster import KMeans, DBSCAN
from sklearn.decomposition import PCA
from sklearn.exceptions import NotFittedError
import scipy.stats as stats
import matplotlib
matplotlib.use("Agg")  # for headless plotting
import matplotlib.pyplot as plt

# --- lightweight password hashing (stdlib) ---
import os as _os
import hashlib, binascii, hmac as _hmac

def generate_password_hash(password: str) -> str:
    salt = _os.urandom(16)
    dk = hashlib.pbkdf2_hmac('sha256', password.encode('utf-8'), salt, 100000)
    return binascii.hexlify(salt).decode() + ':' + binascii.hexlify(dk).decode()

def check_password_hash(stored_hash: str, password: str) -> bool:
    try:
        salt_hex, dk_hex = stored_hash.split(':')
    except ValueError:
        return False
    salt = binascii.unhexlify(salt_hex)
    dk = binascii.unhexlify(dk_hex)
    new_dk = hashlib.pbkdf2_hmac('sha256', password.encode('utf-8'), salt, 100000)
    return _hmac.compare_digest(new_dk, dk)

# ------------------------------------------------

# -----------------------------
# Simple in-repo auth store (file-backed)
# -----------------------------
AUTH_STORE = "auth_store.json"
def _ensure_auth():
    if not os.path.exists(AUTH_STORE):
        admin_pass = generate_password_hash("password")
        with open(AUTH_STORE, "w") as f:
            json.dump({"admin": admin_pass}, f)

def authenticate(username: str, password: str) -> bool:
    _ensure_auth()
    with open(AUTH_STORE, "r") as f:
        data = json.load(f)
    stored = data.get(username)
    if not stored:
        return False
    return check_password_hash(stored, password)

# seed auth
_ensure_auth()

# -----------------------------
# App constants & helpers
# -----------------------------
HIGH_CARD_THRESHOLD_COUNT = 50
HIGH_CARD_THRESHOLD_RATIO = 0.10
LOGO_PATH = "DataSynth.png"

# -----------------------------
# Small utility functions
# -----------------------------
def read_file_to_df(uploaded) -> pd.DataFrame:
    if uploaded is None:
        return pd.DataFrame()
    try:
        if hasattr(uploaded, "name") and os.path.exists(uploaded.name):
            path = uploaded.name
        elif isinstance(uploaded, dict) and "name" in uploaded:
            path = uploaded["name"]
        else:
            content = uploaded.read()
            return _read_bytes_content(content)

        return _read_file_path(path)
        
    except Exception as e:
        print(f"read_file_to_df error: {e}")
        try:
          if 'path' in locals():
            return pd.read_csv(path, encoding_errors='ignore', engine='python')
        except Exception:
          pass
        raise

def _read_bytes_content(content) -> pd.DataFrame:
  formats_to_try = [
      ('csv', lambda: pd.read_csv(io.BytesIO(content))),
      ('csv_utf8', lambda: pd.read_csv(io.BytesIO(content), encoding='utf-8')),
      ('csv_latin1', lambda: pd.read_csv(io.BytesIO(content), encoding='latin1')),
      ('csv_ignore', lambda: pd.read_csv(io.BytesIO(content), encoding_errors='ignore')),
      ('excel', lambda: pd.read_excel(io.BytesIO(content))),
      ('json', lambda: pd.read_json(io.BytesIO(content))),
  ]
  for format_name, reader_func in formats_to_try:
    try:
      return reader_func()
    except Exception as e:
      print(f"Failed to read as {format_name}: {e}")
      continue
  raise ValueError("Could not read file with any supported format")

def _read_file_path(path) -> pd.DataFrame:
    file_ext = os.path.splitext(path)[1].lower()
    
    if file_ext in ['.csv', '.txt', '.tsv']:
        encodings = [None, 'utf-8', 'latin1', 'cp1252', 'iso-8859-1']
        for encoding in encodings:
            try:
                if file_ext == '.tsv':
                    return pd.read_csv(path, sep='\t', encoding=encoding, encoding_errors='ignore')
                else:
                    return pd.read_csv(path, encoding=encoding, encoding_errors='ignore')
            except Exception:
                continue
    
    elif file_ext in ['.xlsx', '.xls', '.xlsm', '.xlsb']:
        return pd.read_excel(path)
    
    elif file_ext == '.json':
        return pd.read_json(path)
    
    elif file_ext in ['.parquet']:
        return pd.read_parquet(path)
    
    elif file_ext in ['.feather']:
        return pd.read_feather(path)
    
    try:
        return pd.read_csv(path, encoding_errors='ignore', engine='python')
    except Exception:
        return pd.read_excel(path)

def comprehensive_data_profile(df: pd.DataFrame) -> Dict[str, Any]:
    if df is None or df.empty:
        return {}
    
    profile = {}
    profile["rows"], profile["columns"] = df.shape
    profile["memory_usage"] = df.memory_usage(deep=True).sum() / 1024**2
    
    dtypes = df.dtypes.apply(lambda x: x.name).to_dict()
    profile["dtypes"] = dtypes
    
    nulls = df.isnull().sum().to_dict()
    profile["nulls"] = nulls
    profile["null_pct"] = {k: (v / len(df)) for k, v in nulls.items()}
    
    unique_counts = df.nunique(dropna=False).to_dict()
    profile["unique"] = unique_counts
    profile["duplicate_rows"] = df.duplicated().sum()
    profile["duplicate_pct"] = profile["duplicate_rows"] / len(df)
    
    numeric_cols = df.select_dtypes(include=[np.number]).columns.tolist()
    categorical_cols = df.select_dtypes(include=["object", "category"]).columns.tolist()
    datetime_cols = df.select_dtypes(include=["datetime64"]).columns.tolist()
    
    profile["numeric_cols"] = numeric_cols
    profile["categorical_cols"] = categorical_cols
    profile["datetime_cols"] = datetime_cols
    
    high_cardinality = []
    for col, cnt in unique_counts.items():
        if col in categorical_cols and (cnt > HIGH_CARD_THRESHOLD_COUNT or 
                                      (cnt / len(df) > HIGH_CARD_THRESHOLD_RATIO)):
            high_cardinality.append(col)
    profile["high_cardinality"] = high_cardinality
    
    quantitative_stats = {}
    for col in numeric_cols:
        col_data = df[col].dropna()
        if len(col_data) > 0:
            quantitative_stats[col] = {
                "mean": col_data.mean(),
                "median": col_data.median(),
                "std": col_data.std(),
                "min": col_data.min(),
                "max": col_data.max(),
                "q1": col_data.quantile(0.25),
                "q3": col_data.quantile(0.75),
                "skew": col_data.skew(),
                "kurtosis": col_data.kurtosis(),
                "zeros": (col_data == 0).sum(),
                "zeros_pct": (col_data == 0).sum() / len(col_data),
                "outliers": detect_outliers_iqr(col_data)
            }
    profile["quantitative_stats"] = quantitative_stats
    
    qualitative_stats = {}
    for col in categorical_cols:
        col_data = df[col].dropna()
        if len(col_data) > 0:
            value_counts = col_data.value_counts()
            qualitative_stats[col] = {
                "top_value": value_counts.index[0] if len(value_counts) > 0 else None,
                "top_freq": value_counts.iloc[0] if len(value_counts) > 0 else 0,
                "top_freq_pct": value_counts.iloc[0] / len(col_data) if len(value_counts) > 0 else 0,
                "unique_values": len(value_counts),
                "entropy": stats.entropy(value_counts.values) if len(value_counts) > 0 else 0
            }
    profile["qualitative_stats"] = qualitative_stats
    
    quality_metrics = {
        "completeness": (len(df) - df.isnull().sum().sum()) / (len(df) * len(df.columns)),
        "uniqueness": 1 - (profile["duplicate_pct"]),
        "validity": {}
    }
    profile["quality_metrics"] = quality_metrics
    
    if len(numeric_cols) > 1:
        profile["correlation_matrix"] = df[numeric_cols].corr().round(3).to_dict()
    else:
        profile["correlation_matrix"] = {}
    
    profile["head"] = df.head(10).to_dict(orient="records")
    
    return profile

def detect_outliers_iqr(series):
    Q1 = series.quantile(0.25)
    Q3 = series.quantile(0.75)
    IQR = Q3 - Q1
    lower_bound = Q1 - 1.5 * IQR
    upper_bound = Q3 + 1.5 * IQR
    return ((series < lower_bound) | (series > upper_bound)).sum()

def profile_to_enhanced_markdown(profile: Dict[str, Any]) -> str:
    if not profile:
        return "No data loaded."
    
    md = []
    md.append("# 📊 Comprehensive Data Profile Report")
    md.append("---")
    
    md.append("## 🎯 Dataset Overview")
    with io.StringIO() as overview:
        overview.write(f"**Rows:** {profile['rows']:,} | ")
        overview.write(f"**Columns:** {profile['columns']} | ")
        overview.write(f"**Memory:** {profile['memory_usage']:.2f} MB | ")
        overview.write(f"**Duplicates:** {profile['duplicate_rows']} ({profile['duplicate_pct']:.2%})")
        md.append(overview.getvalue())
    md.append("")
    
    md.append("## 📈 Data Quality Scorecard")
    quality = profile["quality_metrics"]
    md.append(f"**Completeness:** {quality['completeness']:.2%} | ")
    md.append(f"**Uniqueness:** {quality['uniqueness']:.2%}")
    md.append("")
    
    md.append("## 🗂️ Column Type Summary")
    md.append(f"**Numeric:** {len(profile['numeric_cols'])} | ")
    md.append(f"**Categorical:** {len(profile['categorical_cols'])} | ")
    md.append(f"**DateTime:** {len(profile['datetime_cols'])}")
    md.append("")
    
    md.append("### 🔍 Detailed Column Analysis")
    md.append("| Column | Type | Nulls | Null % | Unique | Completeness | Issues |")
    md.append("|--------|------|-------|---------|--------|--------------|--------|")
    
    for col in profile["dtypes"].keys():
        dtype = profile["dtypes"][col]
        nulls = profile["nulls"].get(col, 0)
        null_pct = profile["null_pct"].get(col, 0)
        uniq = profile["unique"].get(col, 0)
        completeness = 1 - null_pct
        
        issues = []
        if null_pct > 0.5:
            issues.append("🔴 High nulls")
        elif null_pct > 0.2:
            issues.append("🟡 Medium nulls")
        
        if col in profile["high_cardinality"]:
            issues.append("🔵 High cardinality")
            
        if col in profile["numeric_cols"]:
            stats = profile["quantitative_stats"].get(col, {})
            outliers = stats.get("outliers", 0)
            if outliers > 0:
                issues.append("⚫ Outliers")
        
        issues_str = ", ".join(issues) if issues else "✅ Good"
        
        md.append(f"| {col} | {dtype} | {nulls} | {null_pct:.2%} | {uniq} | {completeness:.2%} | {issues_str} |")
    md.append("")
    
    if profile["quantitative_stats"]:
        md.append("### 📈 Quantitative Columns Analysis")
        md.append("| Column | Mean | Std | Min | Max | Skew | Outliers | Zeros |")
        md.append("|--------|------|-----|-----|-----|------|----------|-------|")
        
        for col, stats in profile["quantitative_stats"].items():
            md.append(f"| {col} | {stats['mean']:.2f} | {stats['std']:.2f} | {stats['min']:.2f} | {stats['max']:.2f} | {stats['skew']:.2f} | {stats['outliers']} | {stats['zeros']} |")
        md.append("")
    
    if profile["qualitative_stats"]:
        md.append("### 📊 Qualitative Columns Analysis")
        md.append("| Column | Top Value | Top Freq | Top % | Unique | Entropy |")
        md.append("|--------|-----------|----------|-------|--------|---------|")
        
        for col, stats in profile["qualitative_stats"].items():
            top_value = str(stats["top_value"])[:20] + "..." if len(str(stats["top_value"])) > 20 else str(stats["top_value"])
            md.append(f"| {col} | {top_value} | {stats['top_freq']} | {stats['top_freq_pct']:.2%} | {stats['unique_values']} | {stats['entropy']:.2f} |")
        md.append("")
    
    if profile["high_cardinality"]:
        md.append("### ⚠️ High Cardinality Columns")
        md.append("The following columns have high cardinality (may impact modeling):")
        for col in profile["high_cardinality"]:
            md.append(f"- **{col}**: {profile['unique'][col]} unique values")
        md.append("")
    
    if profile["correlation_matrix"]:
        md.append("### 🔗 Correlation Highlights")
        corr_matrix = profile["correlation_matrix"]
        numeric_cols = list(corr_matrix.keys())
        
        strong_corrs = []
        for i, col1 in enumerate(numeric_cols):
            for j, col2 in enumerate(numeric_cols):
                if i < j:
                    corr = abs(corr_matrix[col1][col2])
                    if corr > 0.7:
                        strong_corrs.append((col1, col2, corr_matrix[col1][col2]))
        
        if strong_corrs:
            md.append("**Strong Correlations (|r| > 0.7):**")
            for col1, col2, corr in sorted(strong_corrs, key=lambda x: abs(x[2]), reverse=True):
                md.append(f"- {col1} ↔ {col2}: {corr:.3f}")
        else:
            md.append("No strong correlations found among numeric columns.")
        md.append("")
    
    return "\n".join(md)

def create_distribution_plots(df: pd.DataFrame, profile: Dict[str, Any]):
    numeric_cols = profile.get("numeric_cols", [])
    categorical_cols = profile.get("categorical_cols", [])
    
    plots = {}
    
    for col in numeric_cols[:4]:
        try:
            fig = px.histogram(df, x=col, title=f"Distribution of {col}", 
                             marginal="box", nbins=50)
            plots[f"num_{col}"] = fig
        except Exception as e:
            print(f"Plot error for {col}: {e}")
    
    for col in categorical_cols[:3]:
        try:
            value_counts = df[col].value_counts().head(10)
            fig = px.bar(x=value_counts.index, y=value_counts.values,
                        title=f"Top 10 Values in {col}")
            fig.update_layout(xaxis_title=col, yaxis_title="Count")
            plots[f"cat_{col}"] = fig
        except Exception as e:
            print(f"Plot error for {col}: {e}")
    
    return plots

def create_correlation_plot(df: pd.DataFrame, profile: Dict[str, Any]):
    numeric_cols = profile.get("numeric_cols", [])
    if len(numeric_cols) < 2:
        return None
    
    try:
        corr_matrix = df[numeric_cols].corr()
        fig = px.imshow(corr_matrix, 
                       title="Correlation Matrix",
                       color_continuous_scale="RdBu_r",
                       aspect="auto")
        return fig
    except Exception as e:
        print(f"Correlation plot error: {e}")
        return None

# -----------------------------
# Data cleaning & feature engineering helpers
# -----------------------------
def drop_high_cardinality(df: pd.DataFrame, threshold_count=HIGH_CARD_THRESHOLD_COUNT, threshold_ratio=HIGH_CARD_THRESHOLD_RATIO):
    n = len(df)
    cols_to_drop = []
    for c in df.columns:
        if df[c].nunique(dropna=False) > threshold_count or (df[c].nunique(dropna=False)/max(1,n) > threshold_ratio):
            if df[c].dtype == "object" or str(df[c].dtype).startswith("category"):
                cols_to_drop.append(c)
    return df.drop(columns=cols_to_drop, errors='ignore'), cols_to_drop

def impute_df(df: pd.DataFrame, numeric_strategy="mean", categorical_strategy="most_frequent", fill_value: Optional[str]=None):
    df = df.copy()
    num_cols = df.select_dtypes(include=[np.number]).columns.tolist()
    cat_cols = df.select_dtypes(include=["object", "category"]).columns.tolist()
    if num_cols:
        imp = SimpleImputer(strategy=numeric_strategy)
        df[num_cols] = imp.fit_transform(df[num_cols])
    if cat_cols:
        if categorical_strategy == "constant" and fill_value is not None:
            imp2 = SimpleImputer(strategy="constant", fill_value=fill_value)
        else:
            imp2 = SimpleImputer(strategy=categorical_strategy)
        df[cat_cols] = imp2.fit_transform(df[cat_cols])
    return df

def treat_outliers_iqr(df: pd.DataFrame, cols: List[str], method="cap"):
    df = df.copy()
    for c in cols:
        if c not in df.columns:
            continue
        if not np.issubdtype(df[c].dtype, np.number):
            continue
        q1 = df[c].quantile(0.25)
        q3 = df[c].quantile(0.75)
        iqr = q3 - q1
        lower = q1 - 1.5 * iqr
        upper = q3 + 1.5 * iqr
        if method == "remove":
            df = df[(df[c] >= lower) & (df[c] <= upper)]
        elif method == "cap":
            df[c] = np.where(df[c] < lower, lower, df[c])
            df[c] = np.where(df[c] > upper, upper, df[c])
    return df

def parse_dates(df: pd.DataFrame, col: str, fmt: Optional[str]=None):
    df = df.copy()
    try:
        if fmt:
            df[col] = pd.to_datetime(df[col], format=fmt, errors="coerce")
        else:
            df[col] = pd.to_datetime(df[col], errors="coerce", infer_datetime_format=True)
    except Exception as e:
        print("parse_dates", e)
    return df

def text_clean(df: pd.DataFrame, cols: List[str], lower=True, strip=True):
    df = df.copy()
    for c in cols:
        if c not in df.columns:
            continue
        df[c] = df[c].astype(str)
        if strip:
            df[c] = df[c].str.strip()
        if lower:
            df[c] = df[c].str.lower()
    return df

def transform_cols(df: pd.DataFrame, cols: List[str], method="log"):
    df = df.copy()
    for c in cols:
        if c in df.columns and np.issubdtype(df[c].dtype, np.number):
            if method == "log":
                df[c] = df[c].apply(lambda x: np.log(x) if x>0 else x)
            elif method == "sqrt":
                df[c] = df[c].apply(lambda x: np.sqrt(x) if x>=0 else x)
    return df

# -----------------------------
# Enhanced Data Preparation Functions
# -----------------------------
def update_column_lists(df):
    def get_columns_safe(df):
        if df is None:
            return []
        cols = list(df.columns)
        if hasattr(df.columns, "levels") and hasattr(df.columns, "names"):
            try:
                cols = ['_'.join(map(str, c)) if isinstance(c, (list, tuple)) else str(c) for c in cols]
            except Exception:
                cols = [str(c) for c in cols]
        else:
            cols = [str(c) for c in cols]
        return cols

    cols = get_columns_safe(df)
    single_value = cols[0] if cols else None
    multi_value = cols if cols else []

    return (
        gr.update(choices=cols, value=multi_value),
        gr.update(choices=cols, value=single_value),
        gr.update(choices=cols, value=single_value),
        gr.update(choices=cols, value=multi_value),
        gr.update(choices=cols, value=multi_value),
        gr.update(choices=cols, value=multi_value),
        gr.update(choices=cols, value=multi_value),
        gr.update(choices=cols, value=single_value),
        gr.update(choices=cols, value=single_value),
        gr.update(choices=cols, value=multi_value),
        gr.update(choices=cols, value=multi_value),
        gr.update(choices=cols, value=multi_value),
        gr.update(choices=cols, value=multi_value),
        gr.update(choices=cols, value=multi_value),
        gr.update(choices=cols, value=single_value)
    )

def apply_column_operations(df, selected_cols, rename_mapping, dtype_conversions):
    df = df.copy()
    
    if selected_cols:
        df = df[selected_cols]
    
    for old_name, new_name in rename_mapping.items():
        if old_name in df.columns:
            df = df.rename(columns={old_name: new_name})
    
    for col, target_dtype in dtype_conversions.items():
        if col in df.columns:
            try:
                if target_dtype == "numeric":
                    df[col] = pd.to_numeric(df[col], errors='coerce')
                elif target_dtype == "integer":
                    df[col] = pd.to_numeric(df[col], errors='coerce').astype('Int64')
                elif target_dtype == "float":
                    df[col] = pd.to_numeric(df[col], errors='coerce').astype(float)
                elif target_dtype == "datetime":
                    df[col] = pd.to_datetime(df[col], errors='coerce')
                elif target_dtype == "category":
                    df[col] = df[col].astype('category')
                elif target_dtype == "boolean":
                    df[col] = df[col].astype(bool)
            except Exception as e:
                print(f"Error converting {col} to {target_dtype}: {e}")
    
    return df

def detect_high_cardinality_cols(df, threshold):
    high_card_cols = []
    for col in df.columns:
        if df[col].dtype in ['object', 'category']:
            unique_count = df[col].nunique()
            if unique_count > threshold:
                high_card_cols.append((col, unique_count))
    return high_card_cols

def advanced_imputation(df, num_cols, num_method, num_custom, cat_cols, cat_method, cat_custom):
    df = df.copy()
    
    if num_cols:
        for col in num_cols:
            if col in df.columns:
                if num_method == "Mean":
                    df[col].fillna(df[col].mean(), inplace=True)
                elif num_method == "Median":
                    df[col].fillna(df[col].median(), inplace=True)
                elif num_method == "Mode":
                    df[col].fillna(df[col].mode()[0] if not df[col].mode().empty else 0, inplace=True)
                elif num_method == "Zero":
                    df[col].fillna(0, inplace=True)
                elif num_method == "Custom Value" and num_custom:
                    try:
                        custom_val = float(num_custom)
                        df[col].fillna(custom_val, inplace=True)
                    except ValueError:
                        df[col].fillna(0, inplace=True)
                elif num_method in ["Forward Fill", "LOCF"]:
                    df[col].fillna(method='ffill', inplace=True)
                elif num_method in ["Backward Fill", "NOCB"]:
                    df[col].fillna(method='bfill', inplace=True)
                elif num_method == "Linear Interpolation":
                    df[col].interpolate(method='linear', inplace=True)
    
    if cat_cols:
        for col in cat_cols:
            if col in df.columns:
                if cat_method == "Mode":
                    df[col].fillna(df[col].mode()[0] if not df[col].mode().empty else "Unknown", inplace=True)
                elif cat_method == "Most Frequent":
                    df[col].fillna(df[col].value_counts().index[0] if len(df[col].value_counts()) > 0 else "Unknown", inplace=True)
                elif cat_method == "Arbitrary ('Unknown')":
                    df[col].fillna("Unknown", inplace=True)
                elif cat_method == "Constant Value" and cat_custom:
                    df[col].fillna(cat_custom, inplace=True)
    
    return df

def apply_text_operations(df, text_cols, operations):
    df = df.copy()
    
    for col in text_cols:
        if col in df.columns:
            df[col] = df[col].astype(str)
            
            for op in operations:
                if op == "Remove leading/trailing whitespace":
                    df[col] = df[col].str.strip()
                elif op == "Convert to lowercase":
                    df[col] = df[col].str.lower()
                elif op == "Convert to uppercase":
                    df[col] = df[col].str.upper()
                elif op == "Remove special characters":
                    df[col] = df[col].str.replace(r'[^\w\s]', '', regex=True)
                elif op == "Remove numbers":
                    df[col] = df[col].str.replace(r'\d+', '', regex=True)
                elif op == "Remove extra spaces":
                    df[col] = df[col].str.replace(r'\s+', ' ', regex=True)
    
    return df

def split_text_column(df, col, delimiter):
    if col not in df.columns:
        return df, "Column not found"
    
    df = df.copy()
    try:
        split_df = df[col].str.split(delimiter, expand=True)
        new_col_names = [f"{col}_{i+1}" for i in range(split_df.shape[1])]
        split_df.columns = new_col_names
        df = pd.concat([df, split_df], axis=1)
        return df, f"Successfully split {col} into {len(new_col_names)} columns"
    except Exception as e:
        return df, f"Error splitting column: {str(e)}"

def create_formula_column(df, formula, new_col_name):
    df = df.copy()
    
    try:
        if "=" in formula:
            formula = formula.split("=")[1].strip()
        
        formula = formula.replace("CURRENT_YEAR", str(pd.Timestamp.now().year))
        
        try:
            df[new_col_name] = df.eval(formula)
            return df, f"Successfully created column '{new_col_name}'"
        except:
            df[new_col_name] = formula
            return df, f"Created column '{new_col_name}' with constant value"
            
    except Exception as e:
        return df, f"Error creating formula column: {str(e)}"

# -----------------------------
# Visualization NLP
# -----------------------------
def nlp_to_chart_instruction(query: str, df: pd.DataFrame):
    q = query.lower()
    numeric = df.select_dtypes(include=[np.number]).columns.tolist()
    categorical = df.select_dtypes(include=["object", "category"]).columns.tolist()
    tokens = q.split()
    
    if "hist" in q or "histogram" in q or "distribution" in q:
        for c in numeric:
            if c.lower() in q:
                return ("hist", [c])
        if numeric:
            return ("hist", [numeric[0]])
    
    if "scatter" in q or "vs" in q or "versus" in q:
        for c in numeric:
            if c.lower() in q:
                x = c
                for d in numeric:
                    if d!=c and d.lower() in q:
                        return ("scatter", [x, d])
                if len(numeric)>1:
                    return ("scatter", [numeric[0], numeric[1]])
        if len(numeric)>=2:
            return ("scatter", [numeric[0], numeric[1]])
    
    if "bar" in q or "count" in q or "counts" in q or "value counts" in q:
        for c in categorical:
            if c.lower() in q:
                return ("bar", [c])
        if categorical:
            return ("bar", [categorical[0]])
    
    if "box" in q or "outlier" in q:
        for c in numeric:
            if c.lower() in q:
                return ("box", [c])
        if numeric:
            return ("box", [numeric[0]])
    
    return ("table", [])

def render_chart_from_instruction(instruction: Tuple[str, List[str]], df: pd.DataFrame):
    typ, cols = instruction
    if typ == "hist":
        c = cols[0]
        fig = px.histogram(df, x=c, title=f"Distribution of {c}")
        return fig
    if typ == "scatter":
        x, y = cols[:2]
        fig = px.scatter(df, x=x, y=y, title=f"{y} vs {x}")
        return fig
    if typ == "bar":
        c = cols[0]
        vc = df[c].value_counts().reset_index()
        vc.columns = [c, "count"]
        fig = px.bar(vc, x=c, y="count", title=f"Counts of {c}")
        return fig
    if typ == "box":
        c = cols[0]
        fig = px.box(df, y=c, title=f"Box plot of {c}")
        return fig
    return None

# -----------------------------
# Modeling - Regression & Classification
# -----------------------------
REGRESSION_MODELS = {
    "Linear Regression": LinearRegression,
    "Ridge Regression": Ridge,
    "Lasso Regression": Lasso,
    "Decision Tree Regressor": DecisionTreeRegressor,
    "Random Forest Regressor": RandomForestRegressor,
    "Gradient Boosting Regressor": GradientBoostingRegressor,
    "Support Vector Regressor": SVR,
    "K-Neighbors Regressor": KNeighborsRegressor
}

CLASSIFICATION_MODELS = {
    "Logistic Regression": LogisticRegression,
    "Decision Tree Classifier": DecisionTreeClassifier,
    "Random Forest Classifier": RandomForestClassifier,
    "Gradient Boosting Classifier": GradientBoostingClassifier,
    "Support Vector Classifier": SVC,
    "Gaussian Naive Bayes": GaussianNB,
    "K-Neighbors Classifier": KNeighborsClassifier
}

# Model hyperparameter grids for tuning
MODEL_PARAM_GRIDS = {
    "Random Forest Regressor": {
        'model__n_estimators': [50, 100, 200],
        'model__max_depth': [None, 10, 20],
        'model__min_samples_split': [2, 5, 10]
    },
    "Random Forest Classifier": {
        'model__n_estimators': [50, 100, 200],
        'model__max_depth': [None, 10, 20],
        'model__min_samples_split': [2, 5, 10]
    },
    "Gradient Boosting Regressor": {
        'model__n_estimators': [50, 100],
        'model__learning_rate': [0.01, 0.1, 0.2],
        'model__max_depth': [3, 5, 7]
    },
    "Gradient Boosting Classifier": {
        'model__n_estimators': [50, 100],
        'model__learning_rate': [0.01, 0.1, 0.2],
        'model__max_depth': [3, 5, 7]
    },
    "Logistic Regression": {
        'model__C': [0.1, 1, 10],
        'model__penalty': ['l2', 'none']
    }
}

def prepare_features_targets(df: pd.DataFrame, target: str, drop_cols: List[str]=None, drop_high_card=True):
    df = df.copy()
    if drop_cols:
        df = df.drop(columns=drop_cols, errors='ignore')
    if drop_high_card:
        df, dropped = drop_high_cardinality(df)
    if target not in df.columns:
        raise ValueError("Target column not in dataframe")
    X = df.drop(columns=[target])
    y = df[target]
    
    # For classification, check if y needs encoding
    if y.dtype == 'object' or y.dtype.name == 'category':
        le = LabelEncoder()
        y = le.fit_transform(y)
    
    return X, y

def auto_build_preprocessor(X: pd.DataFrame, scaler_choice: str="standard", onehot=True):
    num_cols = X.select_dtypes(include=[np.number]).columns.tolist()
    cat_cols = X.select_dtypes(include=["object", "category"]).columns.tolist()
    transformers = []
    
    if num_cols:
        if scaler_choice == "standard":
            num_pipeline = Pipeline([("imputer", SimpleImputer(strategy="mean")), ("scaler", StandardScaler())])
        elif scaler_choice == "minmax":
            num_pipeline = Pipeline([("imputer", SimpleImputer(strategy="mean")), ("scaler", MinMaxScaler())])
        else:
            num_pipeline = Pipeline([("imputer", SimpleImputer(strategy="mean"))])
        transformers.append(("num", num_pipeline, num_cols))
    
    if cat_cols and onehot:
        cat_pipeline = Pipeline([("imputer", SimpleImputer(strategy="most_frequent")), 
                               ("onehot", OneHotEncoder(drop='first', sparse=False, handle_unknown='ignore'))])
        transformers.append(("cat", cat_pipeline, cat_cols))
    
    preprocessor = ColumnTransformer(transformers=transformers, remainder='drop')
    return preprocessor

def run_regression_classification(task: str, model_name: str, X, y, scaler_choice="standard", 
                                 test_size=0.2, random_state=42, do_rfecv=False, 
                                 do_hyperparameter_tuning=False):
    
    if task == "regression":
        ModelClass = REGRESSION_MODELS.get(model_name)
    else:
        ModelClass = CLASSIFICATION_MODELS.get(model_name)

    if ModelClass is None:
        raise ValueError(f"Model {model_name} not found for task {task}")

    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=test_size, random_state=random_state)
    preprocessor = auto_build_preprocessor(X_train, scaler_choice=scaler_choice, onehot=True)
    model = ModelClass()
    
    # Create pipeline
    pipe = Pipeline([("pre", preprocessor), ("model", model)])
    
    # Hyperparameter tuning if requested
    if do_hyperparameter_tuning and model_name in MODEL_PARAM_GRIDS:
        param_grid = MODEL_PARAM_GRIDS[model_name]
        gs = GridSearchCV(pipe, param_grid, cv=3, n_jobs=1, scoring='r2' if task=='regression' else 'accuracy')
        gs.fit(X_train, y_train)
        trained = gs.best_estimator_
        y_pred = trained.predict(X_test)
        best_params = gs.best_params_
    else:
        pipe.fit(X_train, y_train)
        trained = pipe
        y_pred = pipe.predict(X_test)
        best_params = None
    
    results = {}
    
    # Calculate metrics
    if task == "regression":
        results["mse"] = mean_squared_error(y_test, y_pred)
        results["rmse"] = math.sqrt(results["mse"])
        results["mae"] = mean_absolute_error(y_test, y_pred)
        results["r2"] = r2_score(y_test, y_pred)
        
        # Cross-validation scores
        cv_scores = cross_val_score(pipe, X, y, cv=5, scoring='r2')
        results["cv_r2_mean"] = cv_scores.mean()
        results["cv_r2_std"] = cv_scores.std()
        
        # Residuals
        resid = y_test - y_pred
        results["residuals"] = resid
        
    else:  # classification
        results["accuracy"] = accuracy_score(y_test, y_pred)
        results["report"] = classification_report(y_test, y_pred, output_dict=True)
        results["confusion_matrix"] = confusion_matrix(y_test, y_pred).tolist()
        
        # Cross-validation scores
        cv_scores = cross_val_score(pipe, X, y, cv=5, scoring='accuracy')
        results["cv_accuracy_mean"] = cv_scores.mean()
        results["cv_accuracy_std"] = cv_scores.std()
        
        # ROC AUC if applicable
        try:
            if hasattr(trained.named_steps['model'], "predict_proba"):
                y_score = trained.predict_proba(X_test)[:,1]
                results["roc_auc"] = roc_auc_score(y_test, y_score)
                fpr, tpr, _ = roc_curve(y_test, y_score)
                results["roc_curve"] = (fpr.tolist(), tpr.tolist())
        except:
            pass
    
    # Feature importance
    feat_importance = None
    try:
        model_obj = trained.named_steps['model']
        if hasattr(model_obj, "feature_importances_"):
            feature_names = []
            pre = trained.named_steps['pre']
            if hasattr(pre, 'transformers_'):
                for name, trans, cols in pre.transformers_:
                    if name == "num":
                        feature_names += cols
                    elif name == "cat":
                        ohe = trans.named_steps['onehot']
                        if hasattr(ohe, 'get_feature_names_out'):
                            names = list(ohe.get_feature_names_out(cols))
                            feature_names += names
                        else:
                            feature_names += cols
            importances = model_obj.feature_importances_
            feat_importance = list(zip(feature_names, importances))
            feat_importance.sort(key=lambda x: x[1], reverse=True)
    except:
        pass
    
    results["feature_importance"] = feat_importance
    results["trained_model"] = trained
    results["best_params"] = best_params
    results["y_test"] = y_test
    results["y_pred"] = y_pred
    
    return results

# -----------------------------
# Clustering Functions
# -----------------------------
def run_clustering(clustering_method: str, X, n_clusters=3, eps=0.5, min_samples=5):
    X_scaled = StandardScaler().fit_transform(X.select_dtypes(include=[np.number]))
    
    if clustering_method == "K-Means":
        model = KMeans(n_clusters=n_clusters, random_state=42, n_init=10)
        labels = model.fit_predict(X_scaled)
        centroids = model.cluster_centers_
        
        # Calculate metrics
        try:
            silhouette = silhouette_score(X_scaled, labels)
            davies_bouldin = davies_bouldin_score(X_scaled, labels)
            calinski_harabasz = calinski_harabasz_score(X_scaled, labels)
        except:
            silhouette = None
            davies_bouldin = None
            calinski_harabasz = None
        
        results = {
            "method": "K-Means",
            "labels": labels,
            "centroids": centroids,
            "inertia": model.inertia_,
            "silhouette_score": silhouette,
            "davies_bouldin_score": davies_bouldin,
            "calinski_harabasz_score": calinski_harabasz,
            "model": model
        }
        
    elif clustering_method == "DBSCAN":
        model = DBSCAN(eps=eps, min_samples=min_samples)
        labels = model.fit_predict(X_scaled)
        
        # Calculate metrics (only if more than 1 cluster found)
        n_clusters_found = len(set(labels)) - (1 if -1 in labels else 0)
        
        if n_clusters_found > 1:
            try:
                silhouette = silhouette_score(X_scaled, labels)
                davies_bouldin = davies_bouldin_score(X_scaled, labels)
                calinski_harabasz = calinski_harabasz_score(X_scaled, labels)
            except:
                silhouette = None
                davies_bouldin = None
                calinski_harabasz = None
        else:
            silhouette = None
            davies_bouldin = None
            calinski_harabasz = None
        
        results = {
            "method": "DBSCAN",
            "labels": labels,
            "core_sample_indices": model.core_sample_indices_,
            "n_clusters_found": n_clusters_found,
            "silhouette_score": silhouette,
            "davies_bouldin_score": davies_bouldin,
            "calinski_harabasz_score": calinski_harabasz,
            "model": model
        }
    
    return results

# -----------------------------
# Rule-Based (Apriori) Functions
# -----------------------------
def prepare_transaction_data(df, transaction_col=None, threshold=0.5):
    """
    Prepare data for association rule mining.
    For categorical columns, create binary indicators.
    For numeric columns, bin them into categories.
    """
    df_prep = df.copy()
    
    # Convert numeric columns to categorical by binning
    for col in df_prep.select_dtypes(include=[np.number]).columns:
        if len(df_prep[col].unique()) > 10:  # Only bin if many unique values
            # Create 5 equal-width bins
            df_prep[col] = pd.qcut(df_prep[col], q=5, duplicates='drop')
            df_prep[col] = df_prep[col].astype(str)
    
    # Convert all to string for consistency
    for col in df_prep.columns:
        df_prep[col] = df_prep[col].astype(str)
    
    # If transaction column is specified, use it
    if transaction_col and transaction_col in df_prep.columns:
        transactions = []
        for items in df_prep[transaction_col]:
            if pd.isna(items):
                transactions.append([])
            else:
                # Split by common delimiters
                if isinstance(items, str):
                    split_items = [item.strip() for item in str(items).replace(';', ',').split(',')]
                    transactions.append([f"{transaction_col}={item}" for item in split_items if item])
                else:
                    transactions.append([f"{transaction_col}={items}"])
    else:
        # Create transactions from entire dataframe (one-hot like)
        transactions = []
        for _, row in df_prep.iterrows():
            transaction = []
            for col in df_prep.columns:
                if col == transaction_col:
                    continue
                value = str(row[col])
                if value and value.lower() not in ['nan', 'null', 'none', '']:
                    transaction.append(f"{col}={value}")
            if transaction:
                transactions.append(transaction)
    
    return transactions


def simple_apriori(transactions, min_support=0.1):
    """
    A simplified Apriori algorithm implementation.
    Returns frequent itemsets and their support.
    """
    from collections import defaultdict
    import itertools
    
    # Count item frequencies
    item_counts = defaultdict(int)
    for transaction in transactions:
        for item in set(transaction):  # Use set to avoid duplicate items in same transaction
            item_counts[item] += 1
    
    total_transactions = len(transactions)
    
    # Get frequent 1-itemsets
    frequent_1_itemsets = {}
    for item, count in item_counts.items():
        support = count / total_transactions
        if support >= min_support:
            frequent_1_itemsets[frozenset([item])] = support
    
    frequent_itemsets = {1: frequent_1_itemsets}
    
    # Generate larger itemsets (up to 3-itemsets for simplicity)
    k = 2
    while True:
        candidate_itemsets = set()
        prev_itemsets = list(frequent_itemsets[k-1].keys())
        
        # Generate candidates
        for i in range(len(prev_itemsets)):
            for j in range(i+1, len(prev_itemsets)):
                itemset1 = prev_itemsets[i]
                itemset2 = prev_itemsets[j]
                union_set = itemset1.union(itemset2)
                if len(union_set) == k:
                    candidate_itemsets.add(union_set)
        
        if not candidate_itemsets:
            break
        
        # Count support for candidates
        candidate_counts = defaultdict(int)
        for transaction in transactions:
            trans_set = set(transaction)
            for candidate in candidate_itemsets:
                if candidate.issubset(trans_set):
                    candidate_counts[candidate] += 1
        
        # Filter by minimum support
        frequent_k_itemsets = {}
        for itemset, count in candidate_counts.items():
            support = count / total_transactions
            if support >= min_support:
                frequent_k_itemsets[itemset] = support
        
        if not frequent_k_itemsets:
            break
        
        frequent_itemsets[k] = frequent_k_itemsets
        k += 1
        if k > 3:  # Limit to 3-itemsets for performance
            break
    
    return frequent_itemsets

def generate_association_rules(frequent_itemsets, min_confidence=0.5, min_lift=1.0):
    """Generate association rules from frequent itemsets."""
    rules = []
    
    for k, itemsets in frequent_itemsets.items():
        if k < 2:
            continue
        
        for itemset, support_itemset in itemsets.items():
            itemset_list = list(itemset)
            
            # Generate all non-empty proper subsets
            for i in range(1, k):
                for antecedent in itertools.combinations(itemset_list, i):
                    antecedent_set = frozenset(antecedent)
                    consequent_set = itemset - antecedent_set
                    
                    if not consequent_set:
                        continue
                    
                    # Find support of antecedent
                    antecedent_support = None
                    for size in range(1, k):
                        if size in frequent_itemsets and antecedent_set in frequent_itemsets[size]:
                            antecedent_support = frequent_itemsets[size][antecedent_set]
                            break
                    
                    if antecedent_support is None or antecedent_support == 0:
                        continue
                    
                    # Calculate confidence and lift
                    confidence = support_itemset / antecedent_support
                    consequent_support = None
                    
                    # Find support of consequent
                    for size in range(1, k):
                        if size in frequent_itemsets and consequent_set in frequent_itemsets[size]:
                            consequent_support = frequent_itemsets[size][consequent_set]
                            break
                    
                    if consequent_support is None or consequent_support == 0:
                        continue
                    
                    lift = confidence / consequent_support
                    
                    if confidence >= min_confidence and lift >= min_lift:
                        rules.append({
                            'antecedents': set(antecedent),
                            'consequents': set(consequent_set),
                            'support': support_itemset,
                            'confidence': confidence,
                            'lift': lift
                        })
    
    # Sort by confidence, then lift
    rules.sort(key=lambda x: (x['confidence'], x['lift']), reverse=True)
    return rules

def run_association_mining(transactions, min_support=0.1, min_confidence=0.5, min_lift=1.0):
    """Run association rule mining using our simple implementation."""
    frequent_itemsets = simple_apriori(transactions, min_support)
    rules = generate_association_rules(frequent_itemsets, min_confidence, min_lift)
    return frequent_itemsets, rules

def run_frequency_based_rules(df, categorical_cols=None, min_frequency=0.1, min_cooccurrence=0.5):
    """
    An alternative rule discovery method based on frequency and co-occurrence.
    Simpler but effective for many use cases.
    """
    if categorical_cols is None:
        categorical_cols = df.select_dtypes(include=['object', 'category']).columns.tolist()
    
    rules = []
    
    for col1 in categorical_cols:
        for col2 in categorical_cols:
            if col1 == col2:
                continue
            
            # Calculate contingency table
            contingency = pd.crosstab(df[col1], df[col2], normalize='all')
            
            # Find strong associations
            for val1 in contingency.index:
                for val2 in contingency.columns:
                    p_val1_val2 = contingency.loc[val1, val2]
                    p_val1 = df[col1].value_counts(normalize=True).get(val1, 0)
                    p_val2 = df[col2].value_counts(normalize=True).get(val2, 0)
                    
                    if p_val1 > 0 and p_val2 > 0 and p_val1_val2 > 0:
                        confidence = p_val1_val2 / p_val1
                        lift = p_val1_val2 / (p_val1 * p_val2)
                        
                        if p_val1_val2 >= min_frequency and confidence >= min_cooccurrence:
                            rules.append({
                                'rule': f"If {col1} = {val1} then {col2} = {val2}",
                                'support': p_val1_val2,
                                'confidence': confidence,
                                'lift': lift,
                                'antecedent': f"{col1}={val1}",
                                'consequent': f"{col2}={val2}"
                            })
    
    # Sort by confidence and lift
    rules.sort(key=lambda x: (x['confidence'], x['lift']), reverse=True)
    return rules

def run_apriori(df_encoded, min_support=0.1, min_confidence=0.5, min_lift=1.0, max_length=4):
    # Find frequent itemsets
    frequent_itemsets = apriori(df_encoded, min_support=min_support, use_colnames=True, max_len=max_length)
    
    if frequent_itemsets.empty:
        return None, None
    
    # Generate association rules
    rules = association_rules(frequent_itemsets, metric="confidence", min_threshold=min_confidence)
    rules = rules[rules['lift'] >= min_lift]
    
    # Sort by confidence and lift
    rules = rules.sort_values(['confidence', 'lift'], ascending=[False, False])
    
    return frequent_itemsets, rules

def run_fp_growth(df_encoded, min_support=0.1, min_confidence=0.5, min_lift=1.0, max_length=4):
    # Find frequent itemsets using FP-Growth
    frequent_itemsets = fpgrowth(df_encoded, min_support=min_support, use_colnames=True, max_len=max_length)
    
    if frequent_itemsets.empty:
        return None, None
    
    # Generate association rules
    rules = association_rules(frequent_itemsets, metric="confidence", min_threshold=min_confidence)
    rules = rules[rules['lift'] >= min_lift]
    
    # Sort by confidence and lift
    rules = rules.sort_values(['confidence', 'lift'], ascending=[False, False])
    
    return frequent_itemsets, rules

# -----------------------------
# Gradio UI with Enhanced Model Tab
# -----------------------------
css = """
/* Vanta background container */
#vanta-bg {
  width: 100%;
  height: 380px;
  position: relative;
  overflow: hidden;
  border-radius: 12px;
  margin-bottom: 8px;
}

.login-card {
  position: absolute;
  left: 50%;
  top: 50%;
  transform: translate(-50%, -50%);
  width: 420px;
  max-width: calc(100% - 24px);
  background: linear-gradient(180deg, rgba(255,255,255,0.98), rgba(245,245,255,0.95));
  border-radius: 14px;
  box-shadow: 0 12px 36px rgba(0,0,0,0.18);
  padding: 22px;
  z-index: 999;
  border: 1px solid rgba(0,0,0,0.06);
}

.login-logo {
  display:flex;
  align-items:center;
  gap:12px;
  margin-bottom:8px;
}

.brand-title {
  font-weight:700;
  font-size:18px;
  color:#2b2b6b;
}

.btn-animate {
  transition: transform 0.12s ease-in-out, box-shadow 0.12s;
}
.btn-animate:active {
  transform: translateY(2px) scale(0.995);
  box-shadow: 0 6px 18px rgba(0,0,0,0.12) inset;
}

.app-desc {
  font-size: 13px;
  color: #444;
  margin-top: 8px;
  text-align: center;
}

#profile-report table {
    width: 100%;
    border-collapse: collapse;
    font-size: 14px;
    table-layout: fixed;
}

#profile-report th, #profile-report td {
    padding: 8px 12px;
    border: 1px solid #ddd;
    text-align: left;
    word-wrap: break-word;
    white-space: nowrap;
    overflow: hidden;
    text-overflow: ellipsis;
}

#profile-report th {
    background-color: #f5f5f5;
    font-weight: 600;
}

#profile-report tr:nth-child(even) {
    background-color: #f9f9f9;
}

.markdown-body {
    max-width: 100% !important;
}

.gr-markdown {
    max-height: 70vh;
    overflow-y: auto;
}

.dataframe-container {
    max-height: 500px;
    overflow-y: auto;
    border: 1px solid #ddd;
    border-radius: 8px;
    padding: 10px;
}

.model-section {
    padding: 15px;
    border: 1px solid #e0e0e0;
    border-radius: 8px;
    margin-bottom: 15px;
    background-color: #f9f9f9;
}

.model-metrics {
    display: grid;
    grid-template-columns: repeat(auto-fit, minmax(200px, 1fr));
    gap: 10px;
    margin: 15px 0;
}

.metric-card {
    background: white;
    padding: 12px;
    border-radius: 6px;
    border-left: 4px solid #2b2b6b;
    box-shadow: 0 2px 4px rgba(0,0,0,0.1);
}

.metric-value {
    font-size: 24px;
    font-weight: bold;
    color: #2b2b6b;
}

.metric-label {
    font-size: 12px;
    color: #666;
    text-transform: uppercase;
}
"""

vanta_html = """
<div id="vanta-bg" style="width:100%;height:380px;border-radius:12px;position:relative;">
  <div class="login-card" role="region" aria-label="Login card">
    <div class="login-logo">
      <img src="{LOGO}" alt="logo" style="height:48px;width:48px;border-radius:8px;"/>
      <div>
        <div class="brand-title">DataSynth — Analytics Hub</div>
        <div style="font-size:12px;color:#666;">Fast, modular data profiling & model building</div>
      </div>
    </div>
    <div style="margin-top:8px;">
      <div style="font-size:13px;color:#333;margin-bottom:6px;">Sign in to continue</div>
    </div>
    <div class="app-desc">Upload your dataset, prepare it, visualize with NL, and build ML models — all in one place.</div>
  </div>
</div>

<script src="https://cdnjs.cloudflare.com/ajax/libs/three.js/r121/three.min.js"></script>
<script src="https://cdn.jsdelivr.net/npm/vanta@0.5.21/dist/vanta.net.min.js"></script>
<script>
  (function(){
    try {
      if (typeof VANTA !== 'undefined') {
        VANTA.NET({
          el: "#vanta-bg",
          color: 0x2b2b6b,
          backgroundColor: 0xffffff,
          points: 10.00,
          maxDistance: 26.00
        });
      }
    } catch(e) {
      console.warn("Vanta failed to initialize", e);
    }
  })();
</script>
"""

vanta_html = vanta_html.replace("{LOGO}", LOGO_PATH)

with gr.Blocks(css=css, title="DataSynth — Analytics Hub") as demo:
    with gr.Column():
        v_html = gr.HTML(vanta_html)
        with gr.Row():
            username_in = gr.Textbox(label="Username", placeholder="username", interactive=True)
            password_in = gr.Textbox(label="Password", placeholder="password", type="password")
        with gr.Row():
            login_btn = gr.Button("Log in", elem_classes="btn-animate")
            login_msg = gr.Label(value="")
        raw_state = gr.State(value=None)
        df_state = gr.State(value=None)
        clean_state = gr.State(value=None)
        profile_state = gr.State(value=None)
        model_results_state = gr.State(value=None)

    with gr.Column(visible=False) as main_area:
        gr.Markdown("# 🚀 DataSynth Analytics Workspace")
        with gr.Tabs() as main_tabs:
            with gr.TabItem("📊 Data Profiling"):
                gr.Markdown("### 📁 Upload & Analyze Your Dataset")
                upload = gr.File(label="Upload Dataset", file_types=[".csv", ".xlsx", ".xls", ".xlsm", ".xlsb", ".json", ".parquet", ".feather", ".txt", ".tsv"],
                                 type="filepath")
                with gr.Row():
                    profile_btn = gr.Button("🚀 Run Comprehensive Data Profiling", variant="primary", size="lg")
                with gr.Tabs() as data_tabs:
                    with gr.TabItem("📊 Profile Report"):
                        profile_md = gr.Markdown(
                            "No dataset loaded. Upload your data and click 'Run Comprehensive Data Profiling'.",
                            elem_id="profile-report"
                        )

                    with gr.TabItem("👁️ Data Preview"):
                        gr.HTML("<div class='dataframe-container'>")
                        sample_table = gr.Dataframe(
                            interactive=False,
                            label="Sample Data (First 100 rows)",
                        )
                        gr.HTML("</div>")

                    with gr.TabItem("📈 Visualizations"):
                        gr.Markdown("### Correlation Analysis")
                        corr_plot = gr.Plot(label="Correlation Heatmap")

            with gr.TabItem("🔧 Prepare"):
                gr.Markdown("## 🛠️ Advanced Data Preparation Studio")
                with gr.Tabs() as prep_tabs:
                    with gr.TabItem("📋 Column Management"):
                        with gr.Row():
                            with gr.Column(scale=1):
                                gr.Markdown("### Column Selection & Operations")
                                available_cols = gr.CheckboxGroup(
                                    label="Select Columns to Keep",
                                    choices=[],
                                    interactive=True
                                )
                                select_all_cols = gr.Button("Select All")
                                deselect_all_cols = gr.Button("Deselect All")

                                gr.Markdown("### Date Operations")
                                date_col_selected = gr.Dropdown(
                                    label="Date Column",
                                    choices=[],
                                    interactive=True
                                )

                                # Add other missing components as needed:
                                text_ops_cols = gr.CheckboxGroup(
                                    label="Select Text Columns",
                                    choices=[],
                                    interactive=True
                                )

                                groupby_cols = gr.Dropdown(
                                    label="Group By Columns",
                                    choices=[],
                                    interactive=True,
                                    multiselect=True
                                )

                                agg_columns = gr.Dropdown(
                                    label="Columns to Aggregate",
                                    choices=[],
                                    interactive=True,
                                    multiselect=True
                                )

                                pivot_columns = gr.Dropdown(
                                    label="Columns to Pivot",
                                    choices=[],
                                    interactive=True,
                                    multiselect=True
                                )

                                pivot_values = gr.Dropdown(
                                    label="Values Columns",
                                    choices=[],
                                    interactive=True,
                                    multiselect=True
                                )

                                join_key = gr.Dropdown(
                                    label="Join Key Column",
                                    choices=[],
                                    interactive=True
                                )
                                
                                gr.Markdown("### Column Renaming")
                                rename_col_old = gr.Dropdown(
                                    label="Column to Rename",
                                    choices=[],
                                    interactive=True
                                )
                                rename_col_new = gr.Textbox(
                                    label="New Column Name",
                                    placeholder="Enter new column name"
                                )
                                rename_btn = gr.Button("Rename Column")
                                
                                gr.Markdown("### Data Type Conversion")
                                dtype_col = gr.Dropdown(
                                    label="Column to Convert",
                                    choices=[],
                                    interactive=True
                                )
                                dtype_target = gr.Dropdown(
                                    label="Target Data Type",
                                    choices=["string", "numeric", "integer", "float", "datetime", "category", "boolean"],
                                    value="string",
                                    interactive=True
                                )
                                convert_dtype_btn = gr.Button("Convert Data Type")
                                
                            with gr.Column(scale=1):
                                gr.Markdown("### High Cardinality Management")
                                high_card_threshold = gr.Slider(
                                    minimum=1,
                                    maximum=100,
                                    value=50,
                                    step=1,
                                    label="High Cardinality Threshold (unique values)"
                                )
                                high_card_action = gr.Radio(
                                    choices=["Show only", "Drop columns"],
                                    value="Show only",
                                    label="Action for High Cardinality Columns"
                                )
                                high_card_btn = gr.Button("Apply High Cardinality Filter")
                                high_card_results = gr.Markdown("High cardinality columns will appear here")
                    
                    with gr.TabItem("🎯 Missing Values"):
                        with gr.Row():
                            with gr.Column(scale=1):
                                gr.Markdown("### Numeric Imputation")
                                num_impute_cols = gr.CheckboxGroup(
                                    label="Select Numeric Columns",
                                    choices=[],
                                    interactive=True
                                )
                                num_impute_method = gr.Dropdown(
                                    choices=[
                                        "Mean", "Median", "Mode", "KNN", "MICE", 
                                        "LOCF", "NOCB", "Linear Interpolation", 
                                        "Forward Fill", "Backward Fill", "Zero", "Custom Value"
                                    ],
                                    value="Median",
                                    label="Imputation Method"
                                )
                                num_custom_value = gr.Textbox(
                                    label="Custom Value (if selected)",
                                    visible=False
                                )
                                
                                gr.Markdown("### Categorical Imputation")
                                cat_impute_cols = gr.CheckboxGroup(
                                    label="Select Categorical Columns",
                                    choices=[],
                                    interactive=True
                                )
                                cat_impute_method = gr.Dropdown(
                                    choices=[
                                        "Mode", "Most Frequent", "Arbitrary ('Unknown')", 
                                        "KNN", "MICE", "Constant Value"
                                    ],
                                    value="Mode",
                                    label="Imputation Method"
                                )
                                cat_custom_value = gr.Textbox(
                                    label="Custom Value (if selected)",
                                    value="Unknown",
                                    visible=False
                                )
                                
                                impute_btn = gr.Button("Apply Imputation")
                                
                            with gr.Column(scale=1):
                                gr.Markdown("### Missing Value Analysis")
                                missing_summary = gr.Markdown("Missing value summary will appear here")
                                missing_heatmap = gr.Plot(label="Missing Value Heatmap")
                    
                    with gr.TabItem("⚙️ Feature Engineering"):
                        with gr.Row():
                            with gr.Column(scale=1):
                                gr.Markdown("### Text Operations")
                                text_ops_cols = gr.CheckboxGroup(
                                    label="Select Text Columns",
                                    choices=[],
                                    interactive=True
                                )
                                text_operations = gr.CheckboxGroup(
                                    choices=[
                                        "Remove leading/trailing whitespace",
                                        "Convert to lowercase", 
                                        "Convert to uppercase",
                                        "Remove special characters",
                                        "Remove numbers",
                                        "Remove extra spaces"
                                    ],
                                    label="Text Operations"
                                )
                                
                                gr.Markdown("### Value Replacement")
                                replace_col = gr.Dropdown(
                                    label="Column for Value Replacement",
                                    choices=[],
                                    interactive=True
                                )
                                replace_old = gr.Textbox(
                                    label="Value to Replace",
                                    placeholder="Enter value to find"
                                )
                                replace_new = gr.Textbox(
                                    label="Replacement Value", 
                                    placeholder="Enter new value"
                                )
                                replace_btn = gr.Button("Replace Values")
                                
                                gr.Markdown("### Text to Columns")
                                split_col = gr.Dropdown(
                                    label="Column to Split",
                                    choices=[],
                                    interactive=True
                                )
                                split_delimiter = gr.Textbox(
                                    label="Delimiter",
                                    value=",",
                                    placeholder="Enter delimiter (e.g., ',', ';', ' ')",
                                    max_lines=1
                                )
                                split_btn = gr.Button("Split Column")
                                
                            with gr.Column(scale=1):
                                gr.Markdown("### Formula-Based Columns")
                                formula_expr = gr.Textbox(
                                    label="Formula Expression",
                                    placeholder="Example: Sales * Quantity, or LEFT(ProductName, 3)",
                                    lines=2
                                )
                                formula_new_col = gr.Textbox(
                                    label="New Column Name",
                                    placeholder="Enter name for new column"
                                )
                                formula_examples = gr.Markdown("""
                                **Formula Examples:**
                                - `Revenue = Sales * Price`
                                - `FullName = FirstName + ' ' + LastName`  
                                - `Profit = Revenue - Cost`
                                - `Category = LEFT(ProductCode, 3)`
                                - `Age = CURRENT_YEAR - BirthYear`
                                """)
                                formula_btn = gr.Button("Create New Column")
                
                with gr.Row():
                    apply_all_btn = gr.Button("🚀 Apply All Changes", variant="primary", size="lg")
                    download_clean_btn = gr.Button("📥 Download Clean CSV", variant="secondary")
                
                prep_output = gr.Markdown("Preparation status will appear here")
                
            with gr.TabItem("📊 Visualize (NL)"):
                gr.Markdown("Type a natural-language request to create a plot (e.g., 'histogram of age', 'scatter income vs age', 'bar of country').")
                nl_input = gr.Textbox(label="Describe chart")
                nl_btn = gr.Button("Create Chart")
                nl_plot = gr.Plot()
                
            with gr.TabItem("🤖 Model"):
                gr.Markdown("## 🤖 Machine Learning Studio")
                
                with gr.Row():
                    with gr.Column(scale=1):
                        gr.Markdown("### 🎯 Select Task Type")
                        task_select = gr.Radio(
                            choices=["Regression", "Classification", "Clustering", "Rule-Based (Association)"],
                            value="Regression",
                            label="ML Task Type"
                        )
                        
                        gr.Markdown("### 📊 Dataset Selection")
                        dataset_select = gr.Radio(
                            choices=["Use Original Data", "Use Cleaned/Transformed Data"],
                            value="Use Original Data",
                            label="Select Dataset"
                        )
                        
                        # Common for Regression/Classification
                        with gr.Column(visible=True) as reg_class_params:
                            gr.Markdown("### 📈 Regression/Classification Settings")
                            target_col = gr.Dropdown(
                                choices=[],
                                label="Target Variable",
                                interactive=True
                            )
                            scaler_choice = gr.Dropdown(
                                choices=["StandardScaler", "MinMaxScaler", "None"],
                                value="StandardScaler",
                                label="Feature Scaling"
                            )
                            model_select = gr.Dropdown(
                                choices=list(REGRESSION_MODELS.keys()),
                                value="Random Forest Regressor",
                                label="Select Model"
                            )
                            rfecv_opt = gr.Checkbox(
                                label="Run RFECV (Feature Selection)",
                                value=False
                            )
                            hyperparam_tuning = gr.Checkbox(
                                label="Hyperparameter Tuning",
                                value=False
                            )
                        
                        # Clustering specific
                        with gr.Column(visible=False) as clustering_params:
                            gr.Markdown("### 🎪 Clustering Settings")
                            clustering_method = gr.Radio(
                                choices=["K-Means", "DBSCAN"],
                                value="K-Means",
                                label="Clustering Algorithm"
                            )
                            
                            with gr.Row():
                                n_clusters = gr.Slider(
                                    minimum=2,
                                    maximum=20,
                                    value=3,
                                    step=1,
                                    label="Number of Clusters (K-Means)"
                                )
                            
                            with gr.Row(visible=False) as dbscan_params:
                                eps_value = gr.Slider(
                                    minimum=0.1,
                                    maximum=5.0,
                                    value=0.5,
                                    step=0.1,
                                    label="Epsilon (DBSCAN)"
                                )
                                min_samples_value = gr.Slider(
                                    minimum=2,
                                    maximum=20,
                                    value=5,
                                    step=1,
                                    label="Min Samples (DBSCAN)"
                                )
                        
                        # Rule-Based specific
                        with gr.Column(visible=False) as rule_based_params:
                            gr.Markdown("### 🔗 Rule-Based Mining Settings")
                            rule_method = gr.Radio(
                                choices=["Apriori (Simple)", "Frequency-Based Rules"],
                                value="Apriori (Simple)",
                                label="Rule Discovery Method"
                            )
                            
                            transaction_col = gr.Dropdown(
                                choices=[],
                                label="Transaction Column (optional - for Apriori)",
                                interactive=True
                            )
                            
                            with gr.Row():
                                min_support = gr.Slider(
                                    minimum=0.01,
                                    maximum=1.0,
                                    value=0.1,
                                    step=0.01,
                                    label="Minimum Support"
                                )
                                min_confidence = gr.Slider(
                                    minimum=0.01,
                                    maximum=1.0,
                                    value=0.5,
                                    step=0.01,
                                    label="Minimum Confidence"
                                )
                            
                            with gr.Row():
                                min_lift = gr.Slider(
                                    minimum=0.1,
                                    maximum=10.0,
                                    value=1.0,
                                    step=0.1,
                                    label="Minimum Lift"
                                )

                            max_rule_length = gr.Slider(
                                minimum=2,
                                maximum=10,
                                value=4,
                                step=1,
                                label="Maximum Rule Length"
                            )
                    
                    with gr.Column(scale=2):
                        gr.Markdown("### 🚀 Model Training & Evaluation")
                        train_btn = gr.Button("Train Model", variant="primary", size="lg")
                        
                        # Results display
                        model_results = gr.Markdown("Model results will appear here...")
                        
                        # Performance metrics display
                        with gr.Column(visible=False) as perf_metrics:
                            gr.Markdown("### 📊 Performance Metrics")
                            metrics_display = gr.HTML("")
                        
                        # Feature importance/Cluster visualization
                        model_plot = gr.Plot(label="Visualization")
                        
                        # Detailed results
                        with gr.Tabs():
                            with gr.TabItem("📋 Detailed Results"):
                                detailed_results = gr.Dataframe(
                                    label="Detailed Results",
                                    interactive=False,
                                    wrap=True
                                )
                            
                            with gr.TabItem("📈 Performance Graphs"):
                                performance_plots = gr.Plot(label="Performance Analysis")
                            
                            with gr.TabItem("💾 Export Results"):
                                export_results = gr.File(label="Download Results")
                                export_btn = gr.Button("Generate Export")
                
                # Model comparison section
                with gr.Row(visible=False) as model_comparison:
                    gr.Markdown("## 🏆 Top 5 Models Comparison")
                    compare_results = gr.Dataframe(label="Model Comparison")
                    compare_plot = gr.Plot(label="Comparison Visualization")
                
                # Model improvement section
                with gr.Row(visible=False) as model_improvement:
                    gr.Markdown("## 🛠️ Improve Model Performance")
                    
                    with gr.Column():
                        improvement_method = gr.Radio(
                            choices=["Hyperparameter Tuning", "Feature Engineering", "Ensemble Methods", "Cross-Validation"],
                            value="Hyperparameter Tuning",
                            label="Improvement Method"
                        )
                        
                        with gr.Column(visible=True) as hyperparam_tuning_options:
                            gr.Markdown("### Hyperparameter Tuning Options")
                            tuning_method = gr.Radio(
                                choices=["Grid Search", "Random Search", "Bayesian Optimization"],
                                value="Grid Search",
                                label="Tuning Method"
                            )
                            cv_folds = gr.Slider(
                                minimum=3,
                                maximum=10,
                                value=5,
                                step=1,
                                label="CV Folds"
                            )
                        
                        improve_btn = gr.Button("Apply Improvement", variant="primary")
                        improvement_results = gr.Markdown("Improvement results will appear here...")
                        
                    with gr.Column():
                        gr.Markdown("### Model Interpretation")
                        interpretability_method = gr.Radio(
                            choices=["Feature Importance", "SHAP Values", "Partial Dependence Plots", "LIME"],
                            value="Feature Importance",
                            label="Interpretability Method"
                        )
                        interpret_btn = gr.Button("Generate Interpretation")
                        interpretation_plot = gr.Plot(label="Model Interpretation")
                
            with gr.TabItem("📄 Report"):
                gr.Markdown("Generate a comprehensive report summarizing all analyses.")
                report_btn = gr.Button("Generate Comprehensive Report")
                report_download = gr.File(label="Download Report (.md)")

    # --- Callbacks ---
    def _do_login(username, password):
        if not username or not password:
            return gr.update(value="Enter username and password"), gr.update(visible=False)
        ok = authenticate(username.strip(), password.strip())
        if ok:
            return gr.update(value=f"Welcome — {username}"), gr.update(visible=True)
        else:
            return gr.update(value="Invalid credentials"), gr.update(visible=False)

    login_btn.click(fn=_do_login, inputs=[username_in, password_in], outputs=[login_msg, main_area])

    available_cols = gr.CheckboxGroup(choices=[], interactive=True, label="Select Columns to Keep")
    rename_col_old = gr.Dropdown(choices=[], interactive=True, label="Column to Rename")
    dtype_col = gr.Dropdown(choices=[], interactive=True, label="Column to Convert")
    num_impute_cols = gr.CheckboxGroup(choices=[], interactive=True, label="Select Numeric Columns")
    cat_impute_cols = gr.CheckboxGroup(choices=[], interactive=True, label="Select Categorical Columns")
    replace_col = gr.Dropdown(choices=[], interactive=True, label="Column for Value Replacement")
    split_col = gr.Dropdown(choices=[], interactive=True, label="Column to Split")
    date_col_selected = gr.Dropdown(choices=[], interactive=True, label="Date Column")
    pivot_index = gr.Dropdown(choices=[], interactive=True, label="Index Columns", multiselect=True)
    text_ops_cols = gr.CheckboxGroup(choices=[], interactive=True, label="Select Text Columns")
    groupby_cols = gr.Dropdown(choices=[], interactive=True, label="Group By Columns", multiselect=True)
    agg_columns = gr.Dropdown(choices=[], interactive=True, label="Columns to Aggregate", multiselect=True)
    pivot_columns = gr.Dropdown(choices=[], interactive=True, label="Columns to Pivot", multiselect=True)
    pivot_values = gr.Dropdown(choices=[], interactive=True, label="Values Columns", multiselect=True)
    join_key = gr.Dropdown(choices=[], interactive=True, label="Join Key Column")

    prepare_widget_outputs = [
        available_cols, rename_col_old, dtype_col, num_impute_cols,
        cat_impute_cols, replace_col, split_col, date_col_selected, pivot_index,
        text_ops_cols, groupby_cols, agg_columns, pivot_columns, pivot_values, join_key
    ]

    # Load & profile dataset
    def _load_and_profile(uploaded):
        try:
            if uploaded is None:
                base_fail = (gr.update(value="No file uploaded."),
                             pd.DataFrame(), None, None, None)
                return base_fail + tuple([gr.update()]*len(prepare_widget_outputs))

            df = read_file_to_df(uploaded)
            prof = comprehensive_data_profile(df)
            md = profile_to_enhanced_markdown(prof)
            corr_plot_fig = create_correlation_plot(df, prof)

            tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".csv")
            df.to_csv(tmp.name, index=False)

            prepare_updates = update_column_lists(df)
            
            if prepare_updates is None:
                prepare_updates = ()
            elif not isinstance(prepare_updates, (list, tuple)):
                prepare_updates = (prepare_updates,)
            else:
                prepare_updates = tuple(prepare_updates)

            needed = len(prepare_widget_outputs)
            if len(prepare_updates) < needed:
                prepare_updates = prepare_updates + tuple([gr.update()] * (needed - len(prepare_updates)))
            elif len(prepare_updates) > needed:
                prepare_updates = prepare_updates[:needed]

            return (md, df.head(100), tmp.name, prof, corr_plot_fig, *prepare_updates)

        except Exception as e:
            traceback.print_exc()
            base_fail = (gr.update(value=f"Error loading file: {e}"),
                         pd.DataFrame(), None, None, None)
            return base_fail + tuple([gr.update()] * len(prepare_widget_outputs))

    profile_btn.click(
        fn=_load_and_profile,
        inputs=[upload],
        outputs=[
            profile_md, sample_table, df_state, profile_state, corr_plot,
            *prepare_widget_outputs
        ]
    )

    # Update task-specific UI based on task selection
    def update_task_ui(task_type):
        if task_type == "Regression":
            return (
                gr.Column.update(visible=True),   # reg_class_params
                gr.Column.update(visible=False),  # clustering_params
                gr.Column.update(visible=False),  # rule_based_params
                gr.Dropdown.update(choices=list(REGRESSION_MODELS.keys()), value="Random Forest Regressor"),  # model_select
                gr.Row.update(visible=False),     # model_comparison
                gr.Row.update(visible=False)      # model_improvement
            )
        elif task_type == "Classification":
            return (
                gr.Column.update(visible=True),
                gr.Column.update(visible=False),
                gr.Column.update(visible=False),
                gr.Dropdown.update(choices=list(CLASSIFICATION_MODELS.keys()), value="Random Forest Classifier"),
                gr.Row.update(visible=False),
                gr.Row.update(visible=False)
            )
        elif task_type == "Clustering":
            return (
                gr.Column.update(visible=False),
                gr.Column.update(visible=True),
                gr.Column.update(visible=False),
                gr.Dropdown.update(choices=[], value=None),
                gr.Row.update(visible=True),
                gr.Row.update(visible=True)
            )
        else:  # Rule-Based
            return (
                gr.Column.update(visible=False),
                gr.Column.update(visible=False),
                gr.Column.update(visible=True),
                gr.Dropdown.update(choices=[], value=None),
                gr.Row.update(visible=False),
                gr.Row.update(visible=False)
            )

    task_select.change(
        fn=update_task_ui,
        inputs=[task_select],
        outputs=[
            reg_class_params, clustering_params, rule_based_params,
            model_select, model_comparison, model_improvement
        ]
    )

    # Update DBSCAN parameters visibility
    def update_clustering_ui(clustering_method):
        if clustering_method == "DBSCAN":
            return gr.Row.update(visible=True)
        else:
            return gr.Row.update(visible=False)

    clustering_method.change(
        fn=update_clustering_ui,
        inputs=[clustering_method],
        outputs=[dbscan_params]
    )

    # Train model based on task type
    def train_model_wrapper(task_type, model_name, dataset_choice, target_var, 
                           scaler_type, do_rfecv, do_tuning,
                           clustering_algo, n_clusters, eps, min_samples,
                           rule_algo, trans_col, min_sup, min_conf, min_lft, max_len,
                           df_path, clean_path):
        try:
            # Select dataset
            if dataset_choice == "Use Cleaned/Transformed Data" and clean_path:
                df = pd.read_csv(clean_path)
            elif df_path:
                df = pd.read_csv(df_path)
            else:
                return "No dataset available.", None, None, None
            
            results_summary = ""
            plot_fig = None
            detailed_df = pd.DataFrame()
            
            if task_type in ["Regression", "Classification"]:
                if not target_var or target_var not in df.columns:
                    return "Please select a valid target variable.", None, None, None
                
                X, y = prepare_features_targets(df, target_var)
                
                # Train multiple models for comparison
                if task_type == "Regression":
                    model_list = ["Random Forest Regressor", "Gradient Boosting Regressor", 
                                 "Linear Regression", "Ridge Regression", "Decision Tree Regressor"]
                else:
                    model_list = ["Random Forest Classifier", "Gradient Boosting Classifier", 
                                 "Logistic Regression", "Decision Tree Classifier", "Gaussian Naive Bayes"]
                
                all_results = []
                comparison_data = []
                
                for model in model_list[:5]:  # Top 5 models
                    try:
                        task_lower = task_type.lower()
                        model_results = run_regression_classification(
                            task=task_lower,
                            model_name=model,
                            X=X,
                            y=y,
                            scaler_choice=scaler_type,
                            do_rfecv=do_rfecv,
                            do_hyperparameter_tuning=do_tuning
                        )
                        
                        if task_type == "Regression":
                            score = model_results.get("r2", 0)
                        else:
                            score = model_results.get("accuracy", 0)
                        
                        comparison_data.append({
                            "Model": model,
                            "Score": round(score, 4),
                            "Best Params": str(model_results.get("best_params", "Default"))
                        })
                        
                        all_results.append((model, model_results))
                        
                    except Exception as e:
                        print(f"Error training {model}: {e}")
                        continue
                
                # Select best model
                if all_results:
                    best_model, best_results = max(all_results, key=lambda x: 
                        x[1]["r2"] if task_type=="Regression" else x[1]["accuracy"])
                    
                    # Generate summary
                    results_summary = f"## 🏆 Best Model: {best_model}\n\n"
                    
                    if task_type == "Regression":
                        results_summary += f"""
                        ### 📊 Regression Metrics
                        - **R² Score:** {best_results['r2']:.4f}
                        - **RMSE:** {best_results['rmse']:.4f}
                        - **MAE:** {best_results['mae']:.4f}
                        - **Cross-Validation R²:** {best_results.get('cv_r2_mean', 0):.4f} (±{best_results.get('cv_r2_std', 0):.4f})
                        """
                        
                        # Residual plot
                        fig, ax = plt.subplots(1, 2, figsize=(12, 5))
                        
                        # Actual vs Predicted
                        ax[0].scatter(best_results['y_test'], best_results['y_pred'], alpha=0.5)
                        ax[0].plot([best_results['y_test'].min(), best_results['y_test'].max()],
                                 [best_results['y_test'].min(), best_results['y_test'].max()], 'r--', lw=2)
                        ax[0].set_xlabel('Actual')
                        ax[0].set_ylabel('Predicted')
                        ax[0].set_title('Actual vs Predicted')
                        
                        # Residuals histogram
                        residuals = best_results['y_test'] - best_results['y_pred']
                        ax[1].hist(residuals, bins=30, edgecolor='black')
                        ax[1].axvline(x=0, color='r', linestyle='--')
                        ax[1].set_xlabel('Residuals')
                        ax[1].set_ylabel('Frequency')
                        ax[1].set_title('Residuals Distribution')
                        
                        plt.tight_layout()
                        plot_fig = fig
                        
                    else:  # Classification
                        results_summary += f"""
                        ### 📊 Classification Metrics
                        - **Accuracy:** {best_results['accuracy']:.4f}
                        - **Cross-Validation Accuracy:** {best_results.get('cv_accuracy_mean', 0):.4f} (±{best_results.get('cv_accuracy_std', 0):.4f})
                        """
                        
                        if 'roc_auc' in best_results:
                            results_summary += f"- **ROC AUC:** {best_results['roc_auc']:.4f}\n"
                        
                        # Confusion matrix plot
                        cm = best_results['confusion_matrix']
                        fig = go.Figure(data=go.Heatmap(
                            z=cm,
                            x=["Predicted 0", "Predicted 1"],
                            y=["Actual 0", "Actual 1"],
                            colorscale="Blues",
                            text=cm,
                            texttemplate="%{text}",
                            textfont={"size": 16}
                        ))
                        fig.update_layout(
                            title="Confusion Matrix",
                            xaxis_title="Predicted",
                            yaxis_title="Actual"
                        )
                        plot_fig = fig
                    
                    # Feature importance if available
                    if best_results.get('feature_importance'):
                        fi_df = pd.DataFrame(best_results['feature_importance'], 
                                           columns=['Feature', 'Importance'])
                        fi_df = fi_df.head(10)  # Top 10 features
                        
                        fig2 = px.bar(fi_df, x='Importance', y='Feature', 
                                    orientation='h',
                                    title='Top 10 Feature Importances')
                        plot_fig = fig2
                        detailed_df = fi_df
                    
                    # Model comparison table
                    comparison_df = pd.DataFrame(comparison_data)
                    results_summary += f"\n### 📈 Top 5 Models Comparison\n{comparison_df.to_markdown(index=False)}"
                    
            elif task_type == "Clustering":
                # Prepare data for clustering (numeric only)
                X = df.select_dtypes(include=[np.number])
                if X.empty:
                    return "No numeric columns found for clustering.", None, None, None
                
                if clustering_algo == "K-Means":
                    results = run_clustering("K-Means", X, n_clusters=n_clusters)
                else:  # DBSCAN
                    results = run_clustering("DBSCAN", X, eps=eps, min_samples=min_samples)
                
                # Add cluster labels to dataframe
                df_clustered = df.copy()
                df_clustered['Cluster'] = results['labels']
                
                # Generate summary
                results_summary = f"## 🎪 {results['method']} Clustering Results\n\n"
                results_summary += f"- **Number of clusters found:** {len(set(results['labels'])) - (1 if -1 in results['labels'] else 0)}\n"
                
                if results.get('silhouette_score'):
                    results_summary += f"- **Silhouette Score:** {results['silhouette_score']:.4f}\n"
                if results.get('davies_bouldin_score'):
                    results_summary += f"- **Davies-Bouldin Index:** {results['davies_bouldin_score']:.4f}\n"
                if results.get('calinski_harabasz_score'):
                    results_summary += f"- **Calinski-Harabasz Index:** {results['calinski_harabasz_score']:.4f}\n"
                
                if results['method'] == "K-Means":
                    results_summary += f"- **Inertia:** {results['inertia']:.2f}\n"
                
                # Cluster visualization with PCA
                try:
                    pca = PCA(n_components=2)
                    X_pca = pca.fit_transform(StandardScaler().fit_transform(X))
                    
                    cluster_df = pd.DataFrame({
                        'PC1': X_pca[:, 0],
                        'PC2': X_pca[:, 1],
                        'Cluster': results['labels']
                    })
                    
                    fig = px.scatter(cluster_df, x='PC1', y='PC2', color='Cluster',
                                   title=f'{results["method"]} Clustering (PCA Visualization)',
                                   color_continuous_scale='viridis')
                    plot_fig = fig
                    
                except Exception as e:
                    print(f"PCA visualization error: {e}")
                
                # Cluster statistics
                cluster_stats = df_clustered.groupby('Cluster').agg(['mean', 'std', 'count']).round(2)
                detailed_df = cluster_stats
                
            else:  # Rule-Based
                try:
                    if rule_algo == "Apriori (Simple)":
                        # Prepare transaction data
                        transactions = prepare_transaction_data(df, trans_col if trans_col else None)
                        
                        if not transactions:
                            return "No valid transactions found in the data.", None, None
                        
                        # Run Apriori
                        frequent_itemsets, rules = run_association_mining(
                            transactions, 
                            min_support=min_sup,
                            min_confidence=min_conf,
                            min_lift=min_lft
                        )
                        
                        results_summary = f"## 🔗 Simple Apriori Association Rules\n\n"
                        
                    else:  # Frequency-Based Rules
                        # Identify categorical columns
                        cat_cols = df.select_dtypes(include=['object', 'category']).columns.tolist()
                        if not cat_cols:
                            # Try to use all columns as categorical
                            cat_cols = df.columns.tolist()
                        
                        rules_list = run_frequency_based_rules(
                            df, 
                            categorical_cols=cat_cols,
                            min_frequency=min_sup,
                            min_cooccurrence=min_conf
                        )
                        
                        # Convert to rules format similar to Apriori
                        rules = []
                        for rule in rules_list:
                            rules.append({
                                'antecedents': {rule['antecedent']},
                                'consequents': {rule['consequent']},
                                'support': rule['support'],
                                'confidence': rule['confidence'],
                                'lift': rule['lift'],
                                'rule_str': rule['rule']
                            })
                        
                        results_summary = f"## 🔗 Frequency-Based Association Rules\n\n"
                    
                    if not rules:
                        return "No association rules found with the given parameters.", None, None
                    
                    results_summary += f"- **Total rules found:** {len(rules)}\n"
                    results_summary += f"- **Minimum support:** {min_sup}\n"
                    results_summary += f"- **Minimum confidence:** {min_conf}\n"
                    results_summary += f"- **Minimum lift:** {min_lft}\n\n"
                    
                    # Top 10 rules by confidence
                    top_rules = rules[:10]
                    results_summary += "### 🏆 Top 10 Rules by Confidence\n\n"
                    for idx, rule in enumerate(top_rules):
                        if 'rule_str' in rule:
                            rule_text = rule['rule_str']
                        else:
                            antecedents = ', '.join(list(rule['antecedents']))
                            consequents = ', '.join(list(rule['consequents']))
                            rule_text = f"IF {antecedents} THEN {consequents}"
                        
                        results_summary += f"{idx+1}. **{rule_text}**  \n"
                        results_summary += f"   Support: {rule['support']:.3f}, Confidence: {rule['confidence']:.3f}, Lift: {rule['lift']:.3f}\n\n"
                    
                    # Prepare rules for visualization
                    plot_data = []
                    for rule in rules[:20]:  # Limit to 20 for plotting
                        if 'rule_str' in rule:
                            rule_name = rule['rule_str'][:50] + "..." if len(rule['rule_str']) > 50 else rule['rule_str']
                        else:
                            antecedents = ', '.join(list(rule['antecedents']))[:20]
                            consequents = ', '.join(list(rule['consequents']))[:20]
                            rule_name = f"{antecedents}→{consequents}"
                        
                        plot_data.append({
                            'rule': rule_name,
                            'support': rule['support'],
                            'confidence': rule['confidence'],
                            'lift': rule['lift']
                        })
                    
                    if plot_data:
                        plot_df = pd.DataFrame(plot_data)
                        fig = px.scatter(plot_df, x='support', y='confidence',
                                       size='lift', color='lift',
                                       hover_name='rule',
                                       title=f'{rule_algo} Association Rules',
                                       labels={'support': 'Support', 'confidence': 'Confidence'})
                        plot_fig = fig
                    
                    # Prepare detailed results
                    detailed_data = []
                    for rule in rules[:20]:
                        if 'rule_str' in rule:
                            rule_str = rule['rule_str']
                        else:
                            antecedents = ', '.join(list(rule['antecedents']))
                            consequents = ', '.join(list(rule['consequents']))
                            rule_str = f"{antecedents} → {consequents}"
                        
                        detailed_data.append({
                            'Rule': rule_str,
                            'Support': f"{rule['support']:.3f}",
                            'Confidence': f"{rule['confidence']:.3f}",
                            'Lift': f"{rule['lift']:.3f}"
                        })
                    
                    detailed_df = pd.DataFrame(detailed_data)
                    
                except Exception as e:
                    return f"Error in rule-based mining: {str(e)}", None, None
            
            # Return successful results
            return results_summary, plot_fig, detailed_df
            
        except Exception as e:
            traceback.print_exc()
            return f"Error in model training: {str(e)}", None, None

    train_btn.click(
        fn=train_model_wrapper,
        inputs=[
            task_select, model_select, dataset_select, target_col,
            scaler_choice, rfecv_opt, hyperparam_tuning,
            clustering_method, n_clusters, eps_value, min_samples_value,
            rule_method, transaction_col, min_support, min_confidence, min_lift, max_rule_length,
            df_state, clean_state
        ],
        outputs=[model_results, model_plot, detailed_results]
    )

    # Update column lists for model tab
    def update_model_columns(df_path, clean_path, profile):
        try:
            if clean_path:
                df = pd.read_csv(clean_path)
            elif df_path:
                df = pd.read_csv(df_path)
            else:
                return gr.Dropdown.update(choices=[]), gr.Dropdown.update(choices=[])
            
            cols = list(df.columns)
            
            # For rule-based, also update transaction column
            transaction_choices = cols + [None]
            
            return (
                gr.Dropdown.update(choices=cols),
                gr.Dropdown.update(choices=transaction_choices, value=None)
            )
            
        except:
            return gr.Dropdown.update(choices=[]), gr.Dropdown.update(choices=[])

    # Update when dataset changes
    df_state.change(
        fn=update_model_columns,
        inputs=[df_state, clean_state, profile_state],
        outputs=[target_col, transaction_col]
    )
    
    clean_state.change(
        fn=update_model_columns,
        inputs=[df_state, clean_state, profile_state],
        outputs=[target_col, transaction_col]
    )

    # Apply preparation changes
    def _apply_all_preparations(df_path, selected_cols, high_card_threshold, high_card_action):
        try:
            if not df_path:
                return "No dataset loaded", None
            
            df = pd.read_csv(df_path)
            
            if selected_cols:
                df = df[selected_cols]
            
            high_card_cols = detect_high_cardinality_cols(df, high_card_threshold)
            if high_card_action == "Drop columns" and high_card_cols:
                cols_to_drop = [col for col, count in high_card_cols]
                df = df.drop(columns=cols_to_drop)
            
            tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".csv")
            df.to_csv(tmp.name, index=False)
            
            summary = f"""
            **Preparation Summary:**
            - Final dataset: {len(df)} rows, {len(df.columns)} columns
            - High cardinality columns found: {len(high_card_cols)}
            - Columns kept: {', '.join(df.columns.tolist())}
            """
            
            return summary, tmp.name
            
        except Exception as e:
            return f"Preparation failed: {str(e)}", None

    apply_all_btn.click(
        fn=_apply_all_preparations,
        inputs=[df_state, available_cols, high_card_threshold, high_card_action],
        outputs=[prep_output, clean_state]
    )

    # Helper functions for prepare tab
    def select_all_callback(df_path):
        try:
            if not df_path:
                return gr.update(choices=[], value=[])
            df = pd.read_csv(df_path)
            cols = [str(c) for c in df.columns]
            return gr.update(choices=cols, value=cols)
        except:
            return gr.update(choices=[], value=[])

    def deselect_all_callback(df_path):
        try:
            if not df_path:
                return gr.update(choices=[], value=[])
            df = pd.read_csv(df_path)
            cols = [str(c) for c in df.columns]
            return gr.update(choices=cols, value=[])
        except:
            return gr.update(choices=[], value=[])

    select_all_cols.click(fn=select_all_callback, inputs=[df_state], outputs=[available_cols])
    deselect_all_cols.click(fn=deselect_all_callback, inputs=[df_state], outputs=[available_cols])

    # NLP visualization
    def _nl_visualize(query, clean_path, df_path):
        try:
            df = None
            if clean_path:
                df = pd.read_csv(clean_path)
            elif df_path:
                df = pd.read_csv(df_path)
            if df is None or df.empty:
                return None
            instr = nlp_to_chart_instruction(query, df)
            fig = render_chart_from_instruction(instr, df)
            return fig
        except Exception as e:
            print("nlp visualize error", e)
            return None

    nl_btn.click(fn=_nl_visualize, inputs=[nl_input, clean_state, df_state], outputs=[nl_plot])

    # Report generation
    def _generate_report(profile, prep_summary, model_summary, feature_imp_df):
        try:
            lines = []
            lines.append("# 📊 DataSynth Executive Report")
            lines.append("---\n")
            
            lines.append("## 📈 Data Profile Summary")
            if profile:
                lines.append(f"- **Rows:** {profile.get('rows', 0):,}")
                lines.append(f"- **Columns:** {profile.get('columns', 0)}")
                lines.append(f"- **Numeric Columns:** {len(profile.get('numeric_cols', []))}")
                lines.append(f"- **Categorical Columns:** {len(profile.get('categorical_cols', []))}")
            else:
                lines.append("No profile data available.")
            
            lines.append("\n## 🔧 Data Preparation")
            lines.append(prep_summary or "No preparation performed.")
            
            lines.append("\n## 🤖 Model Results")
            lines.append(model_summary or "No model training performed.")
            
            if isinstance(feature_imp_df, pd.DataFrame) and not feature_imp_df.empty:
                lines.append("\n## 📊 Feature Importance")
                lines.append(feature_imp_df.to_markdown(index=False))
            
            # Save markdown
            tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".md", mode="w", encoding="utf-8")
            tmp.write("\n".join(lines))
            tmp.flush()
            return tmp.name
        except Exception as e:
            return None

    report_btn.click(
        fn=_generate_report,
        inputs=[profile_state, prep_output, model_results, detailed_results],
        outputs=[report_download]
    )

    demo.load(lambda: None, outputs=[])

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)