from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import Optional, List
import pandas as pd
import numpy as np
import io
import os
import datetime
from fastapi.responses import Response

from preprocessing import (
    analyze_dataset,
    preprocess_missing_values,
    remove_outliers,
    solve_imbalance,
    handle_duplicates,
    clean_data_types,
    encode_categorical,
    scale_features,
    feature_selection,
    split_dataset,
    drop_columns,
)
from elm_model import ELMClassifier, ELMRegressor
from sklearn.model_selection import KFold, StratifiedKFold, train_test_split
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, confusion_matrix, roc_auc_score, roc_curve, precision_recall_curve
from cv_pipeline import apply_pipeline_cv
import time
import json
from sklearn.preprocessing import StandardScaler, MinMaxScaler, LabelEncoder
from PIL import Image
import os

# CPU thread pinning for HF Spaces free tier (2 vCPU).
# Must be set BEFORE importing torch to take effect.
os.environ.setdefault("OMP_NUM_THREADS", "2")
os.environ.setdefault("MKL_NUM_THREADS", "2")

import torch
import torch.nn.functional as F
import torchvision.transforms as transforms
from torchvision.models import efficientnet_v2_s, EfficientNet_V2_S_Weights
import base64
import threading

torch.set_num_threads(2)
torch.set_num_interop_threads(1)

app = FastAPI(title="ML Research App")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

original_dataset: Optional[pd.DataFrame] = None
current_dataset: Optional[pd.DataFrame] = None
pipeline_history: List[dict] = []
dataset_snapshots: List[pd.DataFrame] = []  # for undo

# Persistent training state for Module 2 real-time prediction
last_trained_info = {
    "model": None,
    "scaler": None,
    "label_encoder": None,
    "features": [],
    "target": None,
    "problem_type": None,
    "activation": "sigmoid"
}


class PreprocessRequest(BaseModel):
    action: str
    params: dict = {}


class TrainELMRequest(BaseModel):
    target_column: str
    problem_type: str = "classification"
    features: List[str] = []
    
    # Split config
    split_strategy: str = "kfold"  # holdout, kfold, stratified_kfold
    num_folds: int = 5
    test_size: float = 0.2
    shuffle: bool = True
    random_seed: int = 42
    
    # ELM Hyperparams
    hidden_nodes: int = 100
    activation: str = "sigmoid"
    bias: bool = True
    repeats: int = 1


class PredictRequest(BaseModel):
    data: dict


class ImagePredictResponse(BaseModel):
    prediction: str
    confidence: float
    all_scores: dict = {}


def _require_dataset():
    if current_dataset is None:
        raise HTTPException(status_code=404, detail="No dataset uploaded.")


# ─── Upload ───
@app.post("/upload")
async def upload_dataset(file: UploadFile = File(...)):
    global original_dataset, current_dataset, pipeline_history, dataset_snapshots
    try:
        contents = await file.read()
        if file.filename.endswith('.csv'):
            df = pd.read_csv(io.StringIO(contents.decode('utf-8')))
        elif file.filename.endswith(('.xls', '.xlsx')):
            df = pd.read_excel(io.BytesIO(contents))
        else:
            raise HTTPException(status_code=400, detail="Unsupported file format. Upload CSV or Excel.")

        original_dataset = df.copy()
        current_dataset = df.copy()
        pipeline_history = []
        dataset_snapshots = []

        return {
            "message": "File uploaded successfully",
            "filename": file.filename,
            "analysis": analyze_dataset(current_dataset),
        }
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# ─── Demo Dataset ───
@app.get("/demo/{dataset_name}")
def load_demo(dataset_name: str):
    global original_dataset, current_dataset, pipeline_history, dataset_snapshots
    file_path = f"sample_data/{dataset_name}.csv"
    if not os.path.exists(file_path):
        raise HTTPException(status_code=404, detail="Demo dataset not found.")
    
    try:
        df = pd.read_csv(file_path)
        original_dataset = df.copy()
        current_dataset = df.copy()
        pipeline_history = []
        dataset_snapshots = []
        return {"message": f"Demo {dataset_name} loaded", "analysis": analyze_dataset(current_dataset)}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# ─── Analyze ───
@app.get("/analyze")
def analyze():
    _require_dataset()
    return analyze_dataset(current_dataset)


# ─── Apply preprocessing ───
def _apply_action(df: pd.DataFrame, action: str, params: dict) -> pd.DataFrame:
    """Apply a single preprocessing action and return the new DataFrame."""
    print(f"DEBUG: Applying action '{action}' with params: {params}")

    if action == "missing":
        strategy = params.get("strategy", "mean")
        columns = params.get("columns")  # optional: list of column names
        print(f"DEBUG: Missing Value Strategy: {strategy}, Columns: {columns}")
        result = preprocess_missing_values(df, strategy, columns)
        print(f"DEBUG: Missing Values Before: {df.isnull().sum().sum()}, After: {result.isnull().sum().sum()}")
        return result

    elif action == "duplicates":
        keep = params.get("keep", "first")
        return handle_duplicates(df, keep)

    elif action == "outliers":
        method = params.get("method", "iqr")
        treatment = params.get("treatment", "remove")
        columns = params.get("columns")
        threshold = float(params.get("threshold", 1.5 if method == "iqr" else 3.0))
        return remove_outliers(df, method, threshold, treatment, columns)

    elif action == "type_cleaning":
        column = params.get("column")
        sub_action = params.get("sub_action", "trim")
        return clean_data_types(df, column, sub_action)

    elif action == "encoding":
        column = params.get("column")
        method = params.get("method", "label")
        target_column = params.get("target_column")
        return encode_categorical(df, column, method, target_column)

    elif action == "scaling":
        columns = params.get("columns")
        method = params.get("method", "standard")
        return scale_features(df, columns, method)

    elif action == "imbalance":
        target_col = params.get("target_column")
        method = params.get("method", "smote")
        if not target_col:
            raise ValueError("Target column required for imbalance handling.")
        return solve_imbalance(df, target_col, method)

    elif action == "feature_selection":
        method = params.get("method", "variance")
        threshold = float(params.get("threshold", 0.0))
        target_col = params.get("target_column")
        k_features = int(params.get("k", 10))
        return feature_selection(df, method, threshold, target_col, k_features)

    elif action == "split":
        test_size = float(params.get("test_size", 0.2))
        stratify_col = params.get("stratify_col")
        random_state = int(params.get("random_state", 42))
        return split_dataset(df, test_size, stratify_col, random_state)

    elif action == "drop_columns":
        cols = params.get("columns", [])
        return drop_columns(df, cols)

    else:
        raise ValueError(f"Unknown action: {action}")


@app.post("/preprocess")
def preprocess(req: PreprocessRequest):
    global current_dataset
    _require_dataset()

    try:
        # Save snapshot for undo
        dataset_snapshots.append(current_dataset.copy())

        new_df = _apply_action(current_dataset, req.action, req.params)
        current_dataset = new_df

        # Record step
        step = {
            "step": len(pipeline_history) + 1,
            "action": req.action,
            "params": req.params,
            "timestamp": datetime.datetime.now().isoformat(),
            "rows_before": dataset_snapshots[-1].shape[0],
            "cols_before": dataset_snapshots[-1].shape[1],
            "rows_after": current_dataset.shape[0],
            "cols_after": current_dataset.shape[1],
        }
        pipeline_history.append(step)

        return {
            "message": f"Step {step['step']}: {req.action} applied successfully.",
            "step": step,
            "analysis": analyze_dataset(current_dataset),
            "pipeline": pipeline_history,
        }
    except Exception as e:
        # Rollback snapshot
        if dataset_snapshots:
            dataset_snapshots.pop()
        raise HTTPException(status_code=500, detail=str(e))


# ─── Preview (dry-run) ───
@app.post("/preprocess/preview")
def preprocess_preview(req: PreprocessRequest):
    _require_dataset()
    try:
        preview_df = _apply_action(current_dataset.copy(), req.action, req.params)
        return {
            "before": {
                "rows": current_dataset.shape[0],
                "cols": current_dataset.shape[1],
                "missing": int(current_dataset.isnull().sum().sum()),
                "analysis": analyze_dataset(current_dataset),
            },
            "after": {
                "rows": preview_df.shape[0],
                "cols": preview_df.shape[1],
                "missing": int(preview_df.isnull().sum().sum()),
                "analysis": analyze_dataset(preview_df),
            },
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
# ─── Python Pipeline Export Helper ───
def generate_python_script(pipeline):
    code = [
        "import pandas as pd",
        "import numpy as np",
        "from scipy import stats",
        "",
        "def preprocess_data(file_path):",
        "    df = pd.read_csv(file_path)",
        "    print('Original shape:', df.shape)",
        ""
    ]
    
    for i, step in enumerate(pipeline):
        action = step['action']
        p = step['params']
        code.append(f"    # Step {i+1}: {action}")
        
        if action == 'missing':
            strat = p.get('strategy')
            cols = p.get('columns')
            c_str = f"{cols}" if cols else "df.columns"
            if strat == 'drop': code.append(f"    df = df.dropna(subset={c_str})")
            elif strat == 'drop_cols': code.append(f"    df = df.drop(columns=[c for c in {c_str} if df[c].isnull().any()])")
            elif strat == 'mean': 
                code.append(f"    for c in {c_str}:\n        if pd.api.types.is_numeric_dtype(df[c]): df[c] = df[c].fillna(df[c].mean())")
            elif strat == 'median':
                code.append(f"    for c in {c_str}:\n        if pd.api.types.is_numeric_dtype(df[c]): df[c] = df[c].fillna(df[c].median())")
            elif strat == 'mode':
                code.append(f"    for c in {c_str}:\n        df[c] = df[c].fillna(df[c].mode()[0])")
            elif strat == 'constant':
                code.append(f"    for c in {c_str}:\n        df[c] = df[c].fillna(0)")
                
        elif action == 'duplicates':
            keep = p.get('keep', 'first')
            if keep == 'false': keep = False
            else: keep = f"'{keep}'"
            code.append(f"    df = df.drop_duplicates(keep={keep})")
            
        elif action == 'outliers':
            method = p.get('method')
            treat = p.get('treatment')
            thresh = p.get('threshold')
            cols = p.get('columns')
            if not cols: code.append(f"    num_cols = df.select_dtypes(include=[np.number]).columns")
            else: code.append(f"    num_cols = {cols}")
            
            if method == 'zscore':
                code.append(f"    for c in num_cols:")
                code.append(f"        z = np.abs(stats.zscore(df[c].dropna()))")
                code.append(f"        outliers = df[c].dropna().index[z >= {thresh}]")
                if treat == 'remove': code.append(f"        df = df.drop(index=outliers)")
                elif treat == 'cap':
                    code.append(f"        m, s = df[c].mean(), df[c].std()")
                    code.append(f"        df[c] = df[c].clip(lower=m - {thresh}*s, upper=m + {thresh}*s)")
                elif treat == 'null': code.append(f"        df.loc[outliers, c] = np.nan")
            elif method == 'iqr':
                code.append(f"    for c in num_cols:")
                code.append(f"        Q1, Q3 = df[c].quantile(0.25), df[c].quantile(0.75)")
                code.append(f"        IQR = Q3 - Q1")
                code.append(f"        lower, upper = Q1 - {thresh}*IQR, Q3 + {thresh}*IQR")
                if treat == 'remove':
                    code.append(f"        df = df[~((df[c] < lower) | (df[c] > upper)) | df[c].isna()]")
                elif treat == 'cap': code.append(f"        df[c] = df[c].clip(lower=lower, upper=upper)")
                elif treat == 'null': code.append(f"        df.loc[(df[c] < lower) | (df[c] > upper), c] = np.nan")
                
        elif action == 'type_cleaning':
            col = p.get('column')
            act = p.get('sub_action')
            if act == 'trim': code.append(f"    if df['{col}'].dtype == 'object': df['{col}'] = df['{col}'].str.strip()")
            elif act == 'lowercase': code.append(f"    if df['{col}'].dtype == 'object': df['{col}'] = df['{col}'].str.lower()")
            elif act == 'uppercase': code.append(f"    if df['{col}'].dtype == 'object': df['{col}'] = df['{col}'].str.upper()")
            elif act == 'to_numeric': code.append(f"    df['{col}'] = pd.to_numeric(df['{col}'], errors='coerce')")
            elif act == 'to_datetime': code.append(f"    df['{col}'] = pd.to_datetime(df['{col}'], errors='coerce')")
            elif act == 'remove_special': code.append(f"    if df['{col}'].dtype == 'object': df['{col}'] = df['{col}'].str.replace(r'[^a-zA-Z0-9\\s]', '', regex=True)")
            
        elif action == 'encoding':
            col = p.get('column')
            meth = p.get('method')
            tgt = p.get('target_column')
            if meth in ['label', 'ordinal']:
                code.append(f"    df['{col}'] = df['{col}'].astype('category').cat.codes")
            elif meth == 'onehot':
                code.append(f"    df = pd.concat([df.drop(columns=['{col}']), pd.get_dummies(df['{col}'], prefix='{col}')], axis=1)")
            elif meth == 'frequency':
                code.append(f"    df['{col}'] = df['{col}'].map(df['{col}'].value_counts(normalize=True))")
            elif meth == 'target':
                code.append(f"    tgt = df['{tgt}']")
                code.append(f"    if not pd.api.types.is_numeric_dtype(tgt): tgt = tgt.astype('category').cat.codes")
                code.append(f"    df['{col}'] = df['{col}'].map(tgt.groupby(df['{col}']).mean()).fillna(tgt.mean())")
                
        elif action == 'scaling':
            cols = p.get('columns')
            meth = p.get('method')
            c_str = f"{cols}" if cols else "df.select_dtypes(include=[np.number]).columns"
            if meth == 'standard':
                code.append(f"    for c in {c_str}:\n        if df[c].std() > 0: df[c] = (df[c] - df[c].mean()) / df[c].std()")
            elif meth == 'minmax':
                code.append(f"    for c in {c_str}:\n        if df[c].max() > df[c].min(): df[c] = (df[c] - df[c].min()) / (df[c].max() - df[c].min())")
            elif meth == 'robust':
                code.append(f"    for c in {c_str}:\n        iqr = df[c].quantile(0.75) - df[c].quantile(0.25)\n        if iqr > 0: df[c] = (df[c] - df[c].median()) / iqr")
            elif meth == 'log':
                code.append(f"    for c in {c_str}:\n        df[c] = np.log1p(df[c]) if (df[c].dropna() >= 0).all() else df[c]")
            elif meth == 'sqrt':
                code.append(f"    for c in {c_str}:\n        if (df[c].dropna() >= 0).all(): df[c] = np.sqrt(df[c])")
            elif meth == 'boxcox':
                code.append(f"    for c in {c_str}:\n        if (df[c].dropna() > 0).all(): df.loc[df[c].notnull(), c], _ = stats.boxcox(df[c].dropna())")
            elif meth == 'yeojohnson':
                code.append(f"    for c in {c_str}:\n        df.loc[df[c].notnull(), c], _ = stats.yeojohnson(df[c].dropna())")
            
        elif action == 'feature_selection':
            meth = p.get('method')
            thresh = p.get('threshold')
            tgt = p.get('target_column')
            code.append(f"    num_cols = df.select_dtypes(include=[np.number]).columns.tolist()")
            if meth == 'variance':
                code.append(f"    from sklearn.feature_selection import VarianceThreshold")
                code.append(f"    vt = VarianceThreshold(threshold={thresh})")
                code.append(f"    X = df[num_cols].fillna(df[num_cols].mean()).fillna(0)")
                code.append(f"    vt.fit(X)")
                code.append(f"    drop_cols = [num_cols[i] for i in range(len(num_cols)) if vt.variances_[i] <= {thresh}]")
                if tgt: code.append(f"    if '{tgt}' in drop_cols: drop_cols.remove('{tgt}')")
                code.append(f"    df = df.drop(columns=drop_cols)")
            elif meth == 'correlation':
                code.append(f"    corr = df[num_cols].corr().abs()")
                code.append(f"    upper = corr.where(np.triu(np.ones(corr.shape), k=1).astype(bool))")
                code.append(f"    drop_cols = [c for c in upper.columns if any(upper[c] > {thresh})]")
                if tgt: code.append(f"    if '{tgt}' in drop_cols: drop_cols.remove('{tgt}')")
                code.append(f"    df = df.drop(columns=drop_cols)")
            elif meth in ['kbest', 'mutual_info']:
                k = p.get('k')
                code.append(f"    from sklearn.feature_selection import SelectKBest, f_classif, f_regression, mutual_info_classif, mutual_info_regression")
                code.append(f"    tgt = '{tgt}'")
                code.append(f"    if tgt in num_cols: num_cols.remove(tgt)")
                code.append(f"    X = df[num_cols].fillna(df[num_cols].mean()).fillna(0)")
                code.append(f"    y = df[tgt]")
                code.append(f"    is_class = not pd.api.types.is_numeric_dtype(y) or y.nunique() <= 20")
                code.append(f"    y_enc = y.astype('category').cat.codes if is_class else y")
                if meth == 'kbest':
                    code.append(f"    score_func = f_classif if is_class else f_regression")
                else:
                    code.append(f"    score_func = mutual_info_classif if is_class else mutual_info_regression")
                code.append(f"    selector = SelectKBest(score_func=score_func, k=min({k}, len(num_cols)))")
                code.append(f"    selector.fit(X, y_enc)")
                code.append(f"    keep_cols = [num_cols[i] for i in range(len(num_cols)) if selector.get_support()[i]]")
                code.append(f"    df = df.drop(columns=[c for c in num_cols if c not in keep_cols])")
                
        elif action == 'imbalance':
            meth = p.get('method')
            tgt = p.get('target_column')
            code.append(f"    from imblearn.over_sampling import SMOTE, RandomOverSampler")
            code.append(f"    from imblearn.under_sampling import RandomUnderSampler")
            code.append(f"    from imblearn.combine import SMOTEENN")
            code.append(f"    X = df.drop(columns=['{tgt}'])")
            code.append(f"    y = df['{tgt}']")
            code.append(f"    for c in X.select_dtypes(include=['object', 'category']).columns: X[c] = X[c].astype('category').cat.codes")
            if meth == 'smote': code.append(f"    sampler = SMOTE(random_state=42)")
            elif meth == 'random_over': code.append(f"    sampler = RandomOverSampler(random_state=42)")
            elif meth == 'random_under': code.append(f"    sampler = RandomUnderSampler(random_state=42)")
            elif meth == 'smote_enn': code.append(f"    sampler = SMOTEENN(random_state=42)")
            code.append(f"    X_res, y_res = sampler.fit_resample(X, y.astype('category').cat.codes)")
            code.append(f"    df = pd.DataFrame(X_res, columns=X.columns)")
            code.append(f"    df['{tgt}'] = y_res")
            
        code.append("")
        
    code.append("    print('Final shape:', df.shape)")
    code.append("    return df")
    code.append("")
    code.append("if __name__ == '__main__':")
    code.append("    # df_clean = preprocess_data('your_dataset.csv')")
    code.append("    # df_clean.to_csv('cleaned_data.csv', index=False)")
    code.append("")
    
    return "\n".join(code)


# ─── Undo ───
@app.post("/undo")
def undo():
    global current_dataset
    _require_dataset()
    if not dataset_snapshots:
        raise HTTPException(status_code=400, detail="Nothing to undo.")

    current_dataset = dataset_snapshots.pop()
    removed = pipeline_history.pop() if pipeline_history else None

    return {
        "message": f"Undid step: {removed['action']}" if removed else "Undo done.",
        "analysis": analyze_dataset(current_dataset),
        "pipeline": pipeline_history,
    }


# ─── Reset ───
@app.post("/reset")
def reset():
    global current_dataset, pipeline_history, dataset_snapshots
    if original_dataset is None:
        raise HTTPException(status_code=404, detail="No dataset uploaded.")

    current_dataset = original_dataset.copy()
    pipeline_history = []
    dataset_snapshots = []

    return {
        "message": "Dataset reset to original.",
        "analysis": analyze_dataset(current_dataset),
        "pipeline": [],
    }


# ─── Pipeline History ───
@app.get("/pipeline")
def get_pipeline():
    return {"pipeline": pipeline_history}


# ─── Named dataset snapshots ────────────────────────────────────────────
#
# Why this exists:
# The backend keeps a single session (`original_dataset`, `current_dataset`,
# `pipeline_history`) and the frontend has two independent contexts:
# Module 1 (Data Forensic) and Module 2 (ELM Studio). When a user goes
# Module 1 → Module 2 and then uploads a NEW dataset in Module 2, the
# backend's globals get overwritten — but the frontend's Module 1 context
# still believes its dataset is loaded. Clicking "Import from Module 1"
# then trains against a server session that no longer contains Module 1's
# columns, producing 400 "Target column not found" errors.
#
# The snapshot endpoints let the frontend save the current backend state
# under a name (e.g. "forensic") before a destructive upload, then restore
# it later when "Import from Module 1" is clicked. The snapshot deep-copies
# the dataframes so subsequent edits do not mutate the saved state.
# ────────────────────────────────────────────────────────────────────────

# {name: {"original_dataset": df, "current_dataset": df, "pipeline_history": [...]}}
module_snapshots: dict = {}


@app.post("/dataset/snapshot/save")
def save_dataset_snapshot(name: str):
    """Save the current dataset + pipeline state under a named slot.

    The frontend calls this from Module 2 before a fresh upload to preserve
    whatever Module 1 left behind, so a later "Import from Module 1" can
    restore it.
    """
    if current_dataset is None or original_dataset is None:
        raise HTTPException(status_code=400, detail="No dataset loaded to snapshot.")
    module_snapshots[name] = {
        "original_dataset": original_dataset.copy(),
        "current_dataset": current_dataset.copy(),
        "pipeline_history": list(pipeline_history),
    }
    return {
        "saved": name,
        "shape": list(current_dataset.shape),
        "pipeline_steps": len(pipeline_history),
    }


@app.post("/dataset/snapshot/restore")
def restore_dataset_snapshot(name: str):
    """Restore a previously-saved named snapshot into the live session."""
    global original_dataset, current_dataset, pipeline_history, dataset_snapshots
    snap = module_snapshots.get(name)
    if snap is None:
        raise HTTPException(status_code=404, detail=f"Snapshot '{name}' not found.")
    original_dataset = snap["original_dataset"].copy()
    current_dataset = snap["current_dataset"].copy()
    pipeline_history = list(snap["pipeline_history"])
    dataset_snapshots = []  # undo stack belongs to the previous session
    return {
        "restored": name,
        "shape": list(current_dataset.shape),
        "analysis": analyze_dataset(current_dataset),
        "pipeline": pipeline_history,
    }


@app.get("/dataset/snapshot/list")
def list_dataset_snapshots():
    """List which snapshots are currently saved (for debugging / UI hints)."""
    return {
        "snapshots": [
            {"name": name, "shape": list(snap["current_dataset"].shape)}
            for name, snap in module_snapshots.items()
        ]
    }


# ─── Export Dataset ───
@app.get("/export/dataset/{export_format}")
def export_dataset(export_format: str):
    _require_dataset()
    
    if export_format.lower() == "csv":
        stream = io.StringIO()
        current_dataset.to_csv(stream, index=False)
        response = Response(content=stream.getvalue(), media_type="text/csv")
        response.headers["Content-Disposition"] = "attachment; filename=cleaned_dataset.csv"
        return response
    elif export_format.lower() in ["xls", "xlsx", "excel"]:
        stream = io.BytesIO()
        try:
            with pd.ExcelWriter(stream, engine='openpyxl') as writer:
                current_dataset.to_excel(writer, index=False, sheet_name='CleanedData')
        except Exception as e:
            # Catch both ImportError and any other writing errors
            error_msg = str(e)
            if "openpyxl" in error_msg.lower() or isinstance(e, ImportError):
                raise HTTPException(status_code=500, detail="openpyxl is missing. Please run: python -m pip install openpyxl and restart backend.")
            raise HTTPException(status_code=500, detail=f"Excel Export Error: {error_msg}")
        response = Response(content=stream.getvalue(), media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet")
        response.headers["Content-Disposition"] = f"attachment; filename=cleaned_dataset.xlsx"
        return response
    elif export_format.lower() == "python":
        script_content = generate_python_script(pipeline_history)
        response = Response(content=script_content, media_type="text/x-python")
        response.headers["Content-Disposition"] = "attachment; filename=pipeline.py"
        return response
    else:
        raise HTTPException(status_code=400, detail="Unsupported export format.")


# ─── ELM Training ───
@app.post("/train")
def train_model(req: TrainELMRequest):
    _require_dataset()
    if req.target_column not in original_dataset.columns:
        raise HTTPException(status_code=400, detail=f"Target column '{req.target_column}' not found.")

    try:
        # Separate Global vs Fold-specific actions to prevent leakage and improve performance
        global_actions = ["duplicates", "drop_columns", "type_cleaning"]
        global_steps = [s for s in pipeline_history if s['action'] in global_actions]
        cv_history = [s for s in pipeline_history if s['action'] not in global_actions]

        # 1. Apply Global Preprocessing (Stateless)
        df_prepared = original_dataset.copy()
        for step in global_steps:
            df_prepared = _apply_action(df_prepared, step['action'], step['params'])
            
        df_raw = df_prepared
        y_raw = df_raw[req.target_column]
        
        fold_results = []
        is_classification = req.problem_type == "classification"
        
        for repeat in range(req.repeats):
            seed = req.random_seed + repeat
            
            splits = []
            if req.split_strategy == "holdout":
                train_idx, test_idx = train_test_split(
                    np.arange(len(df_raw)), 
                    test_size=req.test_size, 
                    random_state=seed, 
                    shuffle=req.shuffle,
                    stratify=y_raw if (req.shuffle and is_classification) else None
                )
                splits.append((train_idx, test_idx))
            else:
                if req.split_strategy == "stratified_kfold" and is_classification:
                    cv = StratifiedKFold(n_splits=req.num_folds, shuffle=req.shuffle, random_state=seed if req.shuffle else None)
                    splits = list(cv.split(df_raw, y_raw))
                else:
                    cv = KFold(n_splits=req.num_folds, shuffle=req.shuffle, random_state=seed if req.shuffle else None)
                    splits = list(cv.split(df_raw))
            
            for fold_idx, (train_idx, test_idx) in enumerate(splits, 1):
                start_time = time.time()
                
                df_train_raw = df_raw.iloc[train_idx]
                df_test_raw = df_raw.iloc[test_idx]
                
                # Apply preprocessing WITHOUT LEAKAGE (Only stateful steps learned per fold)
                df_train_proc, df_test_proc = apply_pipeline_cv(df_train_raw, df_test_raw, cv_history)
                
                # Drop NAs
                df_train_proc = df_train_proc.dropna(subset=[req.target_column])
                df_test_proc = df_test_proc.dropna(subset=[req.target_column])
                
                X_train = df_train_proc.drop(columns=[req.target_column])
                y_train = df_train_proc[req.target_column]
                X_test = df_test_proc.drop(columns=[req.target_column])
                y_test = df_test_proc[req.target_column]
                
                # Only use selected features if any
                if req.features:
                    X_train = X_train[[c for c in req.features if c in X_train.columns]]
                    X_test = X_test[[c for c in req.features if c in X_test.columns]]
                    
                # Encode target if classification and not numeric yet
                if is_classification:
                    from sklearn.preprocessing import LabelEncoder
                    le = LabelEncoder()
                    y_train = le.fit_transform(y_train)
                    y_test = le.transform(y_test)
                
                # Train ELM
                if is_classification:
                    elm = ELMClassifier(hidden_nodes=req.hidden_nodes, activation=req.activation, random_state=seed)
                else:
                    elm = ELMRegressor(hidden_nodes=req.hidden_nodes, activation=req.activation, random_state=seed)
                    
                elm.fit(X_train, y_train)
                y_pred = elm.predict(X_test)
                
                # Evaluate
                metrics = {}
                cm_list = None
                curve_data = None
                
                if is_classification:
                    is_multiclass = len(np.unique(y_train)) > 2
                    avg = 'macro' if is_multiclass else 'binary'
                    metrics['accuracy'] = float(accuracy_score(y_test, y_pred))
                    metrics['precision'] = float(precision_score(y_test, y_pred, average=avg, zero_division=0))
                    metrics['recall'] = float(recall_score(y_test, y_pred, average=avg, zero_division=0))
                    metrics['f1_score'] = float(f1_score(y_test, y_pred, average=avg, zero_division=0))
                    
                    try:
                        H_test = elm._activate(np.dot(X_test.values.astype(np.float64), elm.input_weights_) + elm.biases_)
                        y_pred_proba = np.dot(H_test, elm.output_weights_)
                        exp_p = np.exp(y_pred_proba - np.max(y_pred_proba, axis=1, keepdims=True))
                        y_pred_proba = exp_p / np.sum(exp_p, axis=1, keepdims=True)
                        
                        if is_multiclass:
                            metrics['roc_auc'] = float(roc_auc_score(y_test, y_pred_proba, multi_class='ovr'))
                        else:
                            probs = y_pred_proba[:, 1] if y_pred_proba.shape[1] > 1 else y_pred_proba[:, 0]
                            metrics['roc_auc'] = float(roc_auc_score(y_test, probs))
                            # Calculate curves for binary
                            fpr, tpr, _ = roc_curve(y_test, probs)
                            prc, rec, _ = precision_recall_curve(y_test, probs)
                            # Downsample curves if too large to save bandwidth
                            if len(fpr) > 100:
                                idx = np.linspace(0, len(fpr)-1, 100).astype(int)
                                fpr, tpr = fpr[idx], tpr[idx]
                            if len(prc) > 100:
                                idx = np.linspace(0, len(prc)-1, 100).astype(int)
                                prc, rec = prc[idx], rec[idx]
                                
                            curve_data = {
                                "roc": {"fpr": fpr.tolist(), "tpr": tpr.tolist()},
                                "pr": {"precision": prc.tolist(), "recall": rec.tolist()}
                            }
                    except Exception as e:
                        metrics['roc_auc'] = None
                        
                    cm = confusion_matrix(y_test, y_pred)
                    cm_list = cm.tolist()
                else:
                    from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
                    metrics['mae'] = float(mean_absolute_error(y_test, y_pred))
                    metrics['mse'] = float(mean_squared_error(y_test, y_pred))
                    metrics['rmse'] = float(np.sqrt(metrics['mse']))
                    metrics['r2'] = float(r2_score(y_test, y_pred))
                    
                train_time = time.time() - start_time
                
                # Class distribution for viz
                if is_classification:
                    dist_train = pd.Series(y_train).value_counts().to_dict()
                    dist_test = pd.Series(y_test).value_counts().to_dict()
                    dist_train = {f"Class {k}": int(v) for k, v in dist_train.items()}
                    dist_test = {f"Class {k}": int(v) for k, v in dist_test.items()}
                else:
                    dist_train = {}
                    dist_test = {}
                
                fold_results.append({
                    "repeat": repeat + 1,
                    "fold": fold_idx,
                    "train_size": len(X_train),
                    "test_size": len(X_test),
                    "distribution": {"train": dist_train, "test": dist_test},
                    "metrics": metrics,
                    "confusion_matrix": cm_list,
                    "curves": curve_data,
                    "training_time": train_time
                })

        # Calculate Aggregates
        agg_metrics = {}
        if fold_results:
            metric_keys = fold_results[0]["metrics"].keys()
            for key in metric_keys:
                vals = [f["metrics"][key] for f in fold_results if f["metrics"][key] is not None]
                if vals:
                    agg_metrics[key] = {
                        "mean": float(np.mean(vals)),
                        "std": float(np.std(vals)),
                        "min": float(np.min(vals)),
                        "max": float(np.max(vals))
                    }
                else:
                    agg_metrics[key] = None
        
        response_data = {
            "summary": agg_metrics,
            "folds": fold_results,
            "pipeline_used": pipeline_history,
        }
        response_data.update(req.dict())
        return response_data

    except Exception as e:
        import traceback
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/train-finalize")
def train_finalize(req: TrainELMRequest):
    """Refined endpoint to train the final model for Module 2 prediction."""
    _require_dataset()
    global last_trained_info
    
    try:
        # We assume the user wants to train on the CURRENT dataset (after preprocessing)
        # Or you can re-run the entire pipeline from original_dataset. Let's use current_dataset.
        df = current_dataset.copy()
        
        # Check target
        if req.target_column not in df.columns:
            raise HTTPException(status_code=400, detail=f"Target column '{req.target_column}' not found.")
            
        # Select features
        features = req.features if req.features else [c for c in df.columns if c != req.target_column]
        features = [c for c in features if c in df.columns]
        
        # Drop rows with NAs in target or features
        df = df.dropna(subset=features + [req.target_column])
        
        X = df[features]
        y = df[req.target_column]
        
        # Scaler
        scaler = MinMaxScaler() # Using MinMax as requested for ELM
        X_scaled = scaler.fit_transform(X)
        
        # Label Encoder for classification
        le = None
        is_classification = req.problem_type == "classification"
        if is_classification:
            le = LabelEncoder()
            y = le.fit_transform(y)
            
        # Model
        if is_classification:
            model = ELMClassifier(hidden_nodes=req.hidden_nodes, activation=req.activation, random_state=req.random_seed)
        else:
            model = ELMRegressor(hidden_nodes=req.hidden_nodes, activation=req.activation, random_state=req.random_seed)
            
        model.fit(X_scaled, y)
        
        # Store for real-time prediction
        last_trained_info = {
            "model": model,
            "scaler": scaler,
            "label_encoder": le,
            "features": features,
            "target": req.target_column,
            "problem_type": req.problem_type,
            "activation": req.activation
        }
        
        return {
            "message": "Model trained and finalized for prediction.",
            "features": features,
            "target": req.target_column,
            "classes": le.classes_.tolist() if le else []
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/model/export")
def export_model():
    """Download the finalized model as a joblib bundle (model + scaler + label encoder + features).

    The frontend Verification Lab uses this to let users save the trained
    model for offline inference or sharing. Mirrors the legacy backend endpoint
    but adapted for hf-space's single-session model state.
    """
    if last_trained_info["model"] is None:
        raise HTTPException(status_code=404, detail="No finalized model. Train and finalize a model first.")

    try:
        import joblib
    except ImportError:
        raise HTTPException(status_code=500, detail="joblib is missing. pip install joblib.")

    bundle = {
        "model": last_trained_info["model"],
        "scaler": last_trained_info["scaler"],
        "label_encoder": last_trained_info["label_encoder"],
        "features": last_trained_info["features"],
        "target": last_trained_info["target"],
        "problem_type": last_trained_info["problem_type"],
        "activation": last_trained_info["activation"],
        "exported_at": datetime.datetime.utcnow().isoformat() + "Z",
    }
    buf = io.BytesIO()
    joblib.dump(bundle, buf)
    response = Response(content=buf.getvalue(), media_type="application/octet-stream")
    response.headers["Content-Disposition"] = "attachment; filename=elm_model.joblib"
    return response


@app.post("/predict")
def predict(req: PredictRequest):
    """Real-time prediction using the last trained model."""
    if last_trained_info["model"] is None:
        raise HTTPException(status_code=404, detail="No model trained yet. Please train a model first.")
        
    try:
        # Convert input dict to df with correct feature order
        input_data = req.data
        features = last_trained_info["features"]
        
        # Check if all features exist in input
        missing = [f for f in features if f not in input_data]
        if missing:
            raise HTTPException(status_code=400, detail=f"Missing input fields: {', '.join(missing)}")
            
        # Create row
        row = np.array([[float(input_data[f]) for f in features]])
        
        # Scale
        row_scaled = last_trained_info["scaler"].transform(row)
        
        # Model prediction
        model = last_trained_info["model"]
        prediction = model.predict(row_scaled)
        
        # Convert prediction to class label if needed
        result_label = str(prediction[0])
        probabilities = {}
        
        if last_trained_info["problem_type"] == "classification":
            if last_trained_info["label_encoder"]:
                result_label = str(last_trained_info["label_encoder"].inverse_transform([prediction])[0])
                
            # Try to get probabilities for display
            try:
                # Custom prob logic for ELM
                # H = activate(X * W + b)
                # Y = H * beta
                H = model._activate(np.dot(row_scaled, model.input_weights_) + model.biases_)
                y_raw = np.dot(H, model.output_weights_)
                
                # Softmax
                exp_y = np.exp(y_raw - np.max(y_raw))
                probs = exp_y / np.sum(exp_y)
                probs = probs.flatten()
                
                classes = last_trained_info["label_encoder"].classes_
                for i, cls in enumerate(classes):
                    probabilities[str(cls)] = float(probs[i])
            except:
                pass
                
        return {
            "prediction": result_label,
            "probabilities": probabilities
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# ─── Category 3: Image Classification (PyTorch + EfficientNetV2) ───
try:
    weights = EfficientNet_V2_S_Weights.DEFAULT
    cv_model = efficientnet_v2_s(weights=weights)
    cv_model.eval()
    preprocess_img = weights.transforms()
    imagenet_classes = weights.meta["categories"]
    print(f"EfficientNetV2-S loaded: {len(imagenet_classes)} classes")
except Exception as e:
    print(f"Warning: Could not load CV model: {e}")
    cv_model = None

@app.post("/category3/predict-image")
async def predict_image(file: UploadFile = File(...)):
    """Image Classification using EfficientNetV2-S (84.2% Top-1 accuracy)."""
    if cv_model is None:
        raise HTTPException(status_code=503, detail="Image classification model not initialized.")

    try:
        contents = await file.read()
        img = Image.open(io.BytesIO(contents)).convert('RGB')

        batch = preprocess_img(img).unsqueeze(0)

        with torch.inference_mode():
            prediction = cv_model(batch).squeeze(0).softmax(0)

        top5_prob, top5_catid = torch.topk(prediction, 5)

        results = {}
        for i in range(5):
            results[imagenet_classes[top5_catid[i]]] = float(top5_prob[i])

        top_class = imagenet_classes[top5_catid[0]]
        top_prob = float(top5_prob[0])

        return {
            "prediction": top_class,
            "confidence": top_prob,
            "all_scores": results
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# ═══════════════════════════════════════════════════════════════
# MODULE 3: AI MODEL HUB — Lazy-loaded Models
# ═══════════════════════════════════════════════════════════════

_model_lock = threading.Lock()
_xray_model = None
_yolo_model = None
_pose_model = None
_general_model = None
_food_pipeline = None
_bird_pipeline = None
_skin_bundle = None    # (processor, model, id2label)
_brain_bundle = None   # (processor, model, id2label)

# COCO animal class IDs (zero-indexed, 80-class COCO)
_COCO_ANIMAL_IDS = {14, 15, 16, 17, 18, 19, 20, 21, 22, 23}

# COCO 17-keypoint names (in YOLOv8-pose order)
_COCO_KEYPOINTS = [
    "nose", "left_eye", "right_eye", "left_ear", "right_ear",
    "left_shoulder", "right_shoulder", "left_elbow", "right_elbow",
    "left_wrist", "right_wrist", "left_hip", "right_hip",
    "left_knee", "right_knee", "left_ankle", "right_ankle",
]

# Cap people per pose request to keep payload size bounded
_POSE_MAX_PEOPLE = 10


def _get_xray_model():
    """Lazy-load torchxrayvision DenseNet121 trained on ensemble of 5 datasets.

    densenet121-res224-all is trained on NIH + CheXpert + MIMIC-CXR +
    PadChest + RSNA, giving pneumonia AUC ~0.86 (vs 0.78 for NIH-only).
    """
    global _xray_model
    if _xray_model is None:
        with _model_lock:
            if _xray_model is None:
                import torchxrayvision as xrv
                m = xrv.models.DenseNet(weights="densenet121-res224-all")
                m.eval()
                _xray_model = m
                print(f"torchxrayvision DenseNet121-all loaded: {len(m.pathologies)} pathologies")
    return _xray_model


def _get_yolo_model():
    """Lazy-load YOLOv8n (nano) from ultralytics."""
    global _yolo_model
    if _yolo_model is None:
        with _model_lock:
            if _yolo_model is None:
                from ultralytics import YOLO
                _yolo_model = YOLO("yolov8n.pt")
                print("YOLOv8n loaded")
    return _yolo_model


def _get_pose_model():
    """Lazy-load YOLOv8n-pose for 17-keypoint COCO pose estimation."""
    global _pose_model
    if _pose_model is None:
        with _model_lock:
            if _pose_model is None:
                from ultralytics import YOLO
                _pose_model = YOLO("yolov8n-pose.pt")
                print("YOLOv8n-pose loaded")
    return _pose_model


def _get_general_model():
    """Lazy-load YOLOv8n for full COCO 80-class detection (no animal filter).

    Reuses the same yolov8n.pt weights as `_get_yolo_model()` so the model is
    only downloaded once. Kept as a separate global handle to make the warmup
    and class-listing logic explicit per task.
    """
    global _general_model
    if _general_model is None:
        with _model_lock:
            if _general_model is None:
                from ultralytics import YOLO
                _general_model = YOLO("yolov8n.pt")
                print("YOLOv8n (general 80 classes) loaded")
    return _general_model


def _get_skin_model():
    """Lazy-load HuggingFace ViT for skin lesion / cancer classification.

    Anwarkh1/Skin_Cancer-Image_Classification — ViT-base fine-tuned on
    dermoscopic images covering ~9 classes (melanoma, basal cell carcinoma,
    melanocytic nevus, etc.). Returns (processor, model, id2label).
    """
    global _skin_bundle
    if _skin_bundle is None:
        with _model_lock:
            if _skin_bundle is None:
                from transformers import AutoImageProcessor, AutoModelForImageClassification
                repo = "Anwarkh1/Skin_Cancer-Image_Classification"
                processor = AutoImageProcessor.from_pretrained(repo)
                # Use eager attention so output_attentions=True works for rollout
                model = AutoModelForImageClassification.from_pretrained(
                    repo, attn_implementation="eager"
                )
                model.eval()
                id2label = {int(k): v for k, v in model.config.id2label.items()}
                _skin_bundle = (processor, model, id2label)
                print(f"Skin lesion model loaded: {len(id2label)} classes")
    return _skin_bundle


def _get_brain_model():
    """Lazy-load HuggingFace ViT for brain tumor MRI classification.

    Devarshi/Brain_Tumor_Classification — ViT-base fine-tuned on a 4-class
    MRI dataset (glioma, meningioma, pituitary, no_tumor).
    Returns (processor, model, id2label).
    """
    global _brain_bundle
    if _brain_bundle is None:
        with _model_lock:
            if _brain_bundle is None:
                from transformers import AutoImageProcessor, AutoModelForImageClassification
                repo = "Devarshi/Brain_Tumor_Classification"
                processor = AutoImageProcessor.from_pretrained(repo)
                # Use eager attention so output_attentions=True works for rollout
                model = AutoModelForImageClassification.from_pretrained(
                    repo, attn_implementation="eager"
                )
                model.eval()
                id2label = {int(k): v for k, v in model.config.id2label.items()}
                _brain_bundle = (processor, model, id2label)
                print(f"Brain tumor model loaded: {len(id2label)} classes")
    return _brain_bundle


def _get_food_pipeline():
    """Lazy-load HuggingFace ViT-Base fine-tuned on Food-101."""
    global _food_pipeline
    if _food_pipeline is None:
        with _model_lock:
            if _food_pipeline is None:
                from transformers import pipeline
                _food_pipeline = pipeline("image-classification", model="nateraw/food")
                print("Food Recognition (nateraw/food) loaded")
    return _food_pipeline


def _get_bird_pipeline():
    """Lazy-load HuggingFace EfficientNetB2 for 525 bird species."""
    global _bird_pipeline
    if _bird_pipeline is None:
        with _model_lock:
            if _bird_pipeline is None:
                from transformers import pipeline
                _bird_pipeline = pipeline(
                    "image-classification",
                    model="dennisjooo/Birds-Classifier-EfficientNetB2"
                )
                print("Bird Species (dennisjooo/Birds-Classifier-EfficientNetB2) loaded")
    return _bird_pipeline


# ─── Endpoint: Food Recognition ───
@app.post("/category3/predict-food")
async def predict_food(file: UploadFile = File(...)):
    """Food classification via HF ViT-Base Food-101 (89.1% top-1)."""
    try:
        pipe = _get_food_pipeline()
        contents = await file.read()
        img = Image.open(io.BytesIO(contents)).convert("RGB")

        raw = pipe(img, top_k=5)
        all_scores = {item["label"]: float(item["score"]) for item in raw}
        top = raw[0]

        return {
            "prediction": top["label"],
            "confidence": float(top["score"]),
            "all_scores": all_scores,
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# ─── Endpoint: Bird Species ───
@app.post("/category3/predict-bird")
async def predict_bird(file: UploadFile = File(...)):
    """Bird species classification via HF EfficientNetB2 (99.12% on 525 species)."""
    try:
        pipe = _get_bird_pipeline()
        contents = await file.read()
        img = Image.open(io.BytesIO(contents)).convert("RGB")

        raw = pipe(img, top_k=5)
        all_scores = {item["label"]: float(item["score"]) for item in raw}
        top = raw[0]

        return {
            "prediction": top["label"],
            "confidence": float(top["score"]),
            "all_scores": all_scores,
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# ─── Endpoint: Animal Detection (YOLOv8n) ───
@app.post("/detection/predict-animals")
async def predict_animals(file: UploadFile = File(...), conf_threshold: float = 0.25):
    """Animal object detection via YOLOv8n. Filters to COCO animal classes only."""
    try:
        model = _get_yolo_model()
        contents = await file.read()
        img = Image.open(io.BytesIO(contents)).convert("RGB")
        img_w, img_h = img.size

        results = model.predict(img, conf=conf_threshold, verbose=False)
        r = results[0]
        boxes = r.boxes
        names = r.names

        detections = []
        classes_detected = set()
        for i in range(len(boxes)):
            cls_id = int(boxes.cls[i])
            if cls_id not in _COCO_ANIMAL_IDS:
                continue
            xyxy = boxes.xyxy[i].tolist()
            conf = float(boxes.conf[i])
            class_name = names[cls_id]
            detections.append({
                "box": [float(xyxy[0]), float(xyxy[1]), float(xyxy[2]), float(xyxy[3])],
                "confidence": conf,
                "class_name": class_name,
                "class_id": cls_id,
            })
            classes_detected.add(class_name)

        return {
            "image_width": img_w,
            "image_height": img_h,
            "detections": detections,
            "count": len(detections),
            "classes_detected": sorted(classes_detected),
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# ─── Endpoint: Pose Estimation (YOLOv8n-pose) ───
@app.post("/detection/predict-pose")
async def predict_pose(file: UploadFile = File(...), conf_threshold: float = 0.25):
    """Multi-person pose estimation via YOLOv8n-pose.

    Returns bounding boxes for each detected person plus 17 COCO keypoints
    (nose, eyes, ears, shoulders, elbows, wrists, hips, knees, ankles).
    Capped at the top-N highest-confidence people to bound payload size.
    """
    try:
        model = _get_pose_model()
        contents = await file.read()
        img = Image.open(io.BytesIO(contents)).convert("RGB")
        img_w, img_h = img.size

        results = model.predict(img, conf=conf_threshold, verbose=False)
        r = results[0]
        boxes = r.boxes
        kpts = r.keypoints  # shape (N, 17, 3) where last dim is (x, y, conf)
        names = r.names

        # Sort persons by confidence and cap
        person_indices = list(range(len(boxes)))
        person_indices.sort(key=lambda i: float(boxes.conf[i]), reverse=True)
        person_indices = person_indices[:_POSE_MAX_PEOPLE]

        detections = []
        for i in person_indices:
            cls_id = int(boxes.cls[i])
            class_name = names.get(cls_id, "person") if isinstance(names, dict) else names[cls_id]
            xyxy = boxes.xyxy[i].tolist()
            conf = float(boxes.conf[i])

            # Extract this person's 17 keypoints
            person_kpts = []
            if kpts is not None and kpts.data is not None and i < len(kpts.data):
                kp_data = kpts.data[i].tolist()  # list of [x, y, conf]
                for idx, (kx, ky, kc) in enumerate(kp_data):
                    person_kpts.append({
                        "x": float(kx),
                        "y": float(ky),
                        "conf": float(kc),
                        "name": _COCO_KEYPOINTS[idx] if idx < len(_COCO_KEYPOINTS) else f"kp_{idx}",
                    })

            detections.append({
                "box": [float(xyxy[0]), float(xyxy[1]), float(xyxy[2]), float(xyxy[3])],
                "confidence": conf,
                "class_name": class_name,
                "class_id": cls_id,
                "keypoints": person_kpts,
            })

        return {
            "image_width": img_w,
            "image_height": img_h,
            "detections": detections,
            "count": len(detections),
            "model": "yolov8n-pose",
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# ─── Endpoint: General Object Detection (YOLOv8n, full COCO 80) ───
@app.post("/detection/predict-general")
async def predict_general(file: UploadFile = File(...), conf_threshold: float = 0.25):
    """Generic object detection via YOLOv8n on full COCO 80 classes (no filter)."""
    try:
        model = _get_general_model()
        contents = await file.read()
        img = Image.open(io.BytesIO(contents)).convert("RGB")
        img_w, img_h = img.size

        results = model.predict(img, conf=conf_threshold, verbose=False)
        r = results[0]
        boxes = r.boxes
        names = r.names

        detections = []
        classes_detected = set()
        for i in range(len(boxes)):
            cls_id = int(boxes.cls[i])
            xyxy = boxes.xyxy[i].tolist()
            conf = float(boxes.conf[i])
            class_name = names[cls_id]
            detections.append({
                "box": [float(xyxy[0]), float(xyxy[1]), float(xyxy[2]), float(xyxy[3])],
                "confidence": conf,
                "class_name": class_name,
                "class_id": cls_id,
            })
            classes_detected.add(class_name)

        return {
            "image_width": img_w,
            "image_height": img_h,
            "detections": detections,
            "count": len(detections),
            "classes_detected": sorted(classes_detected),
            "model": "yolov8n",
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# ─── Endpoint: Chest X-ray Pneumonia + Grad-CAM ───
def _compute_xray_gradcam(model, img_tensor: torch.Tensor, target_idx: int):
    """Compute Grad-CAM heatmap for torchxrayvision DenseNet121.

    Returns a (224, 224) numpy array normalized to [0, 1].
    Slower than pseudo-CAM but captures gradient information.
    """
    import numpy as np

    target_layer = model.features.denseblock4
    saved = {"act": None, "grad": None}

    def fwd_hook(_m, _inp, out):
        saved["act"] = out
        out.register_hook(lambda g: saved.__setitem__("grad", g))

    handle = target_layer.register_forward_hook(fwd_hook)

    try:
        img_tensor = img_tensor.clone().detach().requires_grad_(True)
        output = model(img_tensor)
        score = output[0, target_idx]
        model.zero_grad()
        score.backward()

        acts = saved["act"][0]               # (C, H, W)
        grads = saved["grad"][0]             # (C, H, W)
        weights = grads.mean(dim=(1, 2))     # (C,)
        cam = (weights[:, None, None] * acts).sum(dim=0)
        cam = torch.relu(cam)
        cam = cam.unsqueeze(0).unsqueeze(0)
        cam = F.interpolate(cam, size=(224, 224), mode="bilinear", align_corners=False)
        cam = cam.squeeze().detach().cpu().numpy()

        if cam.max() > 0:
            cam = (cam - cam.min()) / (cam.max() - cam.min())
        else:
            cam = np.zeros_like(cam)
        return cam
    finally:
        handle.remove()


def _compute_xray_pseudo_cam(model, img_tensor: torch.Tensor, target_idx: int):
    """Compute gradient-free CAM for torchxrayvision DenseNet121.

    Uses the classifier weights directly instead of computing gradients.
    torchxrayvision's DenseNet has a GAP-based classifier, so the weights
    of model.classifier[target_idx] tell us which feature channels matter.

    ~2-3x faster than Grad-CAM on CPU because no backward pass is needed.
    Returns a (224, 224) numpy array normalized to [0, 1].
    """
    import numpy as np

    target_layer = model.features
    saved: dict = {"act": None}

    def fwd_hook(_m, _inp, out):
        saved["act"] = out

    handle = target_layer.register_forward_hook(fwd_hook)

    try:
        with torch.inference_mode():
            _ = model(img_tensor)

        acts = saved["act"][0]                                 # (C, H, W)
        acts = torch.relu(acts)                                # DenseNet features are post-ReLU already

        # torchxrayvision DenseNet uses model.classifier which is a Linear layer.
        # Fallback: if no classifier weight is exposed, use channel mean (proxy CAM).
        classifier_weight = None
        if hasattr(model, "classifier") and hasattr(model.classifier, "weight"):
            classifier_weight = model.classifier.weight   # (num_classes, C)

        if classifier_weight is not None and target_idx < classifier_weight.shape[0]:
            class_weights = classifier_weight[target_idx]  # (C,)
            cam = (class_weights[:, None, None] * acts).sum(dim=0)
        else:
            cam = acts.mean(dim=0)

        cam = torch.relu(cam).unsqueeze(0).unsqueeze(0)
        cam = F.interpolate(cam, size=(224, 224), mode="bilinear", align_corners=False)
        cam = cam.squeeze().detach().cpu().numpy()

        if cam.max() > 0:
            cam = (cam - cam.min()) / (cam.max() - cam.min())
        else:
            cam = np.zeros_like(cam)
        return cam
    finally:
        handle.remove()


def _cam_to_jet_rgb(cam):
    """Map a (H, W) [0, 1] heatmap to a (H, W, 3) uint8 RGB jet colormap.

    Matches matplotlib's jet colormap:
        0.00 → dark blue   (0, 0, 128)
        0.25 → blue/cyan
        0.50 → green
        0.75 → yellow
        1.00 → dark red    (128, 0, 0)
    """
    import numpy as np

    cam = np.clip(cam, 0.0, 1.0)
    r = np.clip(1.5 - np.abs(4.0 * cam - 3.0), 0, 1)
    g = np.clip(1.5 - np.abs(4.0 * cam - 2.0), 0, 1)
    b = np.clip(1.5 - np.abs(4.0 * cam - 1.0), 0, 1)

    rgb = np.zeros((*cam.shape, 3), dtype=np.uint8)
    rgb[..., 0] = (r * 255).astype(np.uint8)
    rgb[..., 1] = (g * 255).astype(np.uint8)
    rgb[..., 2] = (b * 255).astype(np.uint8)
    return rgb


def _composite_heatmap_on_xray(
    cam,
    xray_gray: "np.ndarray",
    alpha: float = 0.5,
    gamma: float = 0.7,
) -> str:
    """Composite a jet heatmap onto a grayscale X-ray and return a base64 PNG.

    Produces the classic "medical Grad-CAM" look:
      - Full jet colormap across the image (blue→cyan→green→yellow→red)
      - Base X-ray always visible underneath via alpha blending
      - Gamma < 1 boosts contrast of mid-range attention

    Args:
        cam: (H, W) float array in [0, 1]
        xray_gray: (H, W) uint8 grayscale X-ray matching CAM dimensions
        alpha: weight of the heatmap (0 = only X-ray, 1 = only heatmap)
        gamma: gamma correction on the heatmap (<1 brightens mid values)
    """
    import numpy as np

    # Smooth and sharpen the CAM
    cam = np.clip(cam, 0.0, 1.0)
    if cam.max() > 0:
        cam = (cam - cam.min()) / (cam.max() - cam.min() + 1e-8)
    cam = np.power(cam, gamma)  # gamma correction for contrast

    # Jet colormap
    heatmap_rgb = _cam_to_jet_rgb(cam)  # (H, W, 3) uint8

    # Grayscale X-ray → RGB
    xray_rgb = np.stack([xray_gray, xray_gray, xray_gray], axis=-1).astype(np.float32)
    heatmap_f = heatmap_rgb.astype(np.float32)

    # Alpha blend: composite = alpha * heatmap + (1 - alpha) * xray
    composite = alpha * heatmap_f + (1.0 - alpha) * xray_rgb
    composite = np.clip(composite, 0, 255).astype(np.uint8)

    overlay = Image.fromarray(composite, mode="RGB")
    buf = io.BytesIO()
    overlay.save(buf, format="PNG", optimize=True)
    b64 = base64.b64encode(buf.getvalue()).decode("ascii")
    return f"data:image/png;base64,{b64}"


@app.post("/medical/predict-xray")
async def predict_xray(
    file: UploadFile = File(...),
    heatmap: bool = True,
    use_pseudo_cam: bool = True,
):
    """Chest X-ray multi-pathology classification via torchxrayvision DenseNet121-all.

    Query params:
        heatmap: if False, skip heatmap computation entirely (fastest mode)
        use_pseudo_cam: if True (default), use gradient-free CAM (~3x faster);
            if False, use original Grad-CAM (more detail but slower)
    """
    import numpy as np
    import torchxrayvision as xrv

    try:
        model = _get_xray_model()
        contents = await file.read()
        img = Image.open(io.BytesIO(contents)).convert("L")  # grayscale

        # Pre-resize large images to ~512px to save memory
        if max(img.size) > 512:
            ratio = 512 / max(img.size)
            new_size = (int(img.size[0] * ratio), int(img.size[1] * ratio))
            img = img.resize(new_size, Image.LANCZOS)

        # torchxrayvision preprocessing
        arr = np.array(img, dtype=np.float32)
        arr = xrv.datasets.normalize(arr, 255)      # to [-1024, 1024]
        arr = arr[None, ...]                         # (1, H, W)
        transform = transforms.Compose([
            xrv.datasets.XRayCenterCrop(),
            xrv.datasets.XRayResizer(224),
        ])
        arr = transform(arr)
        tensor = torch.from_numpy(arr).unsqueeze(0)  # (1, 1, 224, 224)

        # Inference — inference_mode is faster than no_grad
        with torch.inference_mode():
            output = model(tensor)
            scores = torch.sigmoid(output[0]).cpu().numpy()

        pathologies = {
            name: float(score)
            for name, score in zip(model.pathologies, scores)
            if name
        }
        top_name = max(pathologies, key=pathologies.get)
        top_score = pathologies[top_name]
        top_idx = list(model.pathologies).index(top_name)

        # Heatmap computation (optional)
        heatmap_b64 = None
        if heatmap:
            try:
                tensor_cam = torch.from_numpy(arr).unsqueeze(0)
                if use_pseudo_cam:
                    cam = _compute_xray_pseudo_cam(model, tensor_cam, top_idx)
                else:
                    cam = _compute_xray_gradcam(model, tensor_cam, top_idx)

                # Reconstruct the 224x224 grayscale X-ray for compositing.
                # torchxrayvision normalizes to [-1024, 1024]; invert the mapping.
                xray_float = (arr[0] + 1024.0) / 2048.0 * 255.0
                xray_uint8 = np.clip(xray_float, 0, 255).astype(np.uint8)

                heatmap_b64 = _composite_heatmap_on_xray(
                    cam, xray_uint8, alpha=0.5, gamma=0.7
                )
            except Exception as cam_err:
                print(f"Heatmap computation failed: {cam_err}")

        return {
            "prediction": top_name,
            "confidence": top_score,
            "pathologies": pathologies,
            "heatmap": heatmap_b64,
            "image_width": 224,
            "image_height": 224,
            "mode": "pseudo_cam" if use_pseudo_cam else "grad_cam",
            "output_type": "multi-label",
            "model": "DenseNet121-all (torchxrayvision)",
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# ─── Shared helper: ViT classification + attention rollout ───

def _run_vit_classifier(
    bundle,
    file_bytes: bytes,
    *,
    heatmap: bool,
    target_w: int = 384,
    target_h: int = 384,
):
    """Run a ViT classification bundle (processor, model, id2label) on raw
    image bytes and return a unified multi-class response dict.

    Returns a dict matching the schema documented for skin / brain endpoints:
        {
            "prediction": str,           # top class name
            "confidence": float,         # softmax of top class
            "classes": dict[str, float], # all softmax probabilities
            "heatmap": str | None,       # base64 PNG data-URL
            "output_type": "multi-class",
            "image_width": int,
            "image_height": int,
        }
    """
    import numpy as np
    from medical_cam import (
        compute_attention_rollout,
        composite_heatmap_on_rgb,
        resize_cam,
    )

    processor, model, id2label = bundle
    img = Image.open(io.BytesIO(file_bytes)).convert("RGB")

    # Resize for display (independent of model preprocessing)
    display_img = img.copy()
    if max(display_img.size) > max(target_w, target_h):
        display_img.thumbnail((target_w, target_h), Image.LANCZOS)
    display_w, display_h = display_img.size
    display_rgb = np.array(display_img, dtype=np.uint8)

    # Model-specific preprocessing
    inputs = processor(images=img, return_tensors="pt")
    pixel_values = inputs["pixel_values"]

    with torch.inference_mode():
        outputs = model(pixel_values=pixel_values)
        logits = outputs.logits[0]
        probs = torch.softmax(logits, dim=-1).cpu().numpy()

    classes = {id2label.get(i, str(i)): float(p) for i, p in enumerate(probs)}
    top_idx = int(np.argmax(probs))
    top_name = id2label.get(top_idx, str(top_idx))
    top_score = float(probs[top_idx])

    heatmap_b64 = None
    if heatmap:
        try:
            cam = compute_attention_rollout(model, pixel_values)
            cam_resized = resize_cam(cam, (display_w, display_h))
            heatmap_b64 = composite_heatmap_on_rgb(
                cam_resized, display_rgb, alpha=0.45, gamma=0.7
            )
        except Exception as cam_err:
            print(f"Attention rollout failed: {cam_err}")

    return {
        "prediction": top_name,
        "confidence": top_score,
        "classes": classes,
        "heatmap": heatmap_b64,
        "output_type": "multi-class",
        "image_width": display_w,
        "image_height": display_h,
    }


# ─── Endpoint: Skin Lesion / Cancer Classification ───
@app.post("/medical/predict-skin-lesion")
async def predict_skin_lesion(file: UploadFile = File(...), heatmap: bool = True):
    """Skin lesion multi-class classification (Anwarkh1/Skin_Cancer ViT) with
    attention-rollout heatmap. Research only — not a clinical diagnosis.
    """
    try:
        bundle = _get_skin_model()
        contents = await file.read()
        result = _run_vit_classifier(bundle, contents, heatmap=heatmap)
        result["model"] = "Anwarkh1/Skin_Cancer-Image_Classification"
        return result
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# ─── Endpoint: Brain Tumor MRI Classification ───
@app.post("/medical/predict-brain-tumor")
async def predict_brain_tumor(file: UploadFile = File(...), heatmap: bool = True):
    """Brain tumor multi-class classification (Devarshi/Brain_Tumor ViT) with
    attention-rollout heatmap. Research only — not a clinical diagnosis.
    """
    try:
        bundle = _get_brain_model()
        contents = await file.read()
        result = _run_vit_classifier(bundle, contents, heatmap=heatmap)
        result["model"] = "Devarshi/Brain_Tumor_Classification"
        return result
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# ─── Tabular prediction endpoints ───

class TabularPayload(BaseModel):
    """Arbitrary feature -> value payload for tabular prediction."""

    class Config:
        extra = "allow"


def _predict_tabular(bundle_name: str, payload: TabularPayload) -> dict:
    import tabular_models as tm

    try:
        bundle = tm.get_bundle(bundle_name)
    except KeyError:
        raise HTTPException(status_code=404, detail=f"Unknown tabular bundle: {bundle_name}")
    try:
        data = payload.model_dump() if hasattr(payload, "model_dump") else payload.dict()
    except Exception:
        data = dict(payload) if payload else {}
    try:
        return bundle.predict(data)
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Prediction failed: {e}")


@app.post("/tabular/predict-titanic")
async def predict_titanic(payload: TabularPayload):
    """Predict Titanic survival from passenger features."""
    return _predict_tabular("titanic-survival", payload)


@app.post("/tabular/predict-heart")
async def predict_heart(payload: TabularPayload):
    """Predict heart-disease risk from UCI Cleveland clinical features."""
    return _predict_tabular("heart-disease", payload)


@app.post("/tabular/predict-wine")
async def predict_wine(payload: TabularPayload):
    """Predict red-wine quality tier (low/medium/high) from chemistry."""
    return _predict_tabular("wine-quality", payload)


@app.get("/tabular/schema/{name}")
async def tabular_schema(name: str):
    """Expose baseline values + feature ordering for a tabular bundle."""
    import tabular_models as tm

    try:
        return tm.schema_summary(name)
    except KeyError:
        raise HTTPException(status_code=404, detail=f"Unknown tabular bundle: {name}")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Schema failed: {e}")


# ─── Warmup endpoint ───
@app.get("/models/warmup")
async def warmup_models(model: str = "all"):
    """Trigger lazy load of one or all models. Call after deploy to avoid cold start."""
    loaded = []
    try:
        if model in ("xray", "all"):
            _get_xray_model()
            loaded.append("xray")
        if model in ("yolo", "all"):
            _get_yolo_model()
            loaded.append("yolo")
        if model in ("pose", "all"):
            _get_pose_model()
            loaded.append("pose")
        if model in ("general", "all"):
            _get_general_model()
            loaded.append("general")
        if model in ("food", "all"):
            _get_food_pipeline()
            loaded.append("food")
        if model in ("bird", "all"):
            _get_bird_pipeline()
            loaded.append("bird")
        if model in ("skin", "all"):
            _get_skin_model()
            loaded.append("skin")
        if model in ("brain", "all"):
            _get_brain_model()
            loaded.append("brain")
        if model in ("titanic", "all"):
            import tabular_models as tm
            tm.get_bundle("titanic-survival")
            loaded.append("titanic")
        if model in ("heart", "all"):
            import tabular_models as tm
            tm.get_bundle("heart-disease")
            loaded.append("heart")
        if model in ("wine", "all"):
            import tabular_models as tm
            tm.get_bundle("wine-quality")
            loaded.append("wine")
        if model in ("imagenet", "all") and cv_model is not None:
            loaded.append("imagenet")
        return {"loaded": loaded, "status": "ok"}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Warmup failed: {e}")


# ─── Classes listing endpoint ───
@app.get("/models/classes")
async def get_model_classes(model: str):
    """Return the full list of classes a given model can predict.

    Used by the frontend ClassListSearch component to let users
    browse and search what each model supports.
    """
    try:
        if model == "imagenet":
            if not imagenet_classes:
                raise HTTPException(503, "ImageNet classes not loaded")
            labels = list(imagenet_classes)
            return {"classes": sorted(labels), "count": len(labels)}

        if model == "food":
            pipe = _get_food_pipeline()
            labels = list(pipe.model.config.id2label.values())
            return {"classes": sorted(labels), "count": len(labels)}

        if model == "bird":
            pipe = _get_bird_pipeline()
            labels = list(pipe.model.config.id2label.values())
            return {"classes": sorted(labels), "count": len(labels)}

        if model == "xray":
            m = _get_xray_model()
            labels = [p for p in m.pathologies if p]
            return {"classes": sorted(labels), "count": len(labels)}

        if model == "yolo":
            m = _get_yolo_model()
            labels = [
                m.names[i]
                for i in sorted(_COCO_ANIMAL_IDS)
                if i in m.names
            ]
            return {"classes": sorted(labels), "count": len(labels)}

        if model == "pose":
            return {"classes": list(_COCO_KEYPOINTS), "count": len(_COCO_KEYPOINTS)}

        if model == "general":
            m = _get_general_model()
            labels = list(m.names.values())
            return {"classes": sorted(labels), "count": len(labels)}

        if model == "skin":
            _, _, id2label = _get_skin_model()
            labels = list(id2label.values())
            return {"classes": sorted(labels), "count": len(labels)}

        if model == "brain":
            _, _, id2label = _get_brain_model()
            labels = list(id2label.values())
            return {"classes": sorted(labels), "count": len(labels)}

        raise HTTPException(404, f"Unknown model: {model}")
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(500, f"Failed to list classes: {e}")


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)