app.py

import gradio as gr
import pandas as pd
import numpy as np
import sklearn
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler, LabelEncoder, OneHotEncoder
from sklearn.impute import SimpleImputer
from sklearn.compose import ColumnTransformer
from sklearn.pipeline import Pipeline
# Scikit-learn Models
from sklearn.linear_model import LogisticRegression, LinearRegression
from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
from sklearn.svm import SVC, SVR
# Metrics
from sklearn.metrics import accuracy_score, classification_report, mean_squared_error, r2_score
# Dataset generators
from sklearn.datasets import make_classification, make_regression

import joblib
import os
import time
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import TensorDataset, DataLoader
import torchvision # For transforms, even if data is basic
import torchvision.transforms as T

# ONNX specific imports
import skl2onnx
from skl2onnx import convert_sklearn
from skl2onnx.common.data_types import FloatTensorType, Int64TensorType, StringTensorType
import onnxruntime as rt

import traceback
import tempfile
import json
import math
import collections.abc # For Gradio issue with new Python versions

# --- Global Variables / Constants ---
TEMP_DIR = "temp_outputs"
os.makedirs(TEMP_DIR, exist_ok=True)
MAX_DATASET_ROWS_WARN = 30000 # Reduced slightly due to increased complexity
MAX_GENERATED_ROWS = 50000 # Max rows for generation
MAX_GENERATED_COLS = 100   # Max cols for generation

# --- Helper Functions ---
def count_sklearn_parameters(model):
    if hasattr(model, 'coef_'):
        return model.coef_.size + (model.intercept_.size if hasattr(model, 'intercept_') else 0)
    if hasattr(model, 'support_vectors_'):
        return model.support_vectors_.size
    if isinstance(model, (RandomForestClassifier, RandomForestRegressor)):
        try:
            return sum(tree.tree_.node_count for tree in model.estimators_)
        except: return "N/A (Complex Ensemble)"
    return "N/A"

def count_pytorch_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

def get_temp_filepath(filename_base, extension):
    # Ensure extension does not start with a dot if it's passed with one
    clean_extension = extension.lstrip('.')
    return os.path.join(TEMP_DIR, f"{filename_base}_{time.strftime('%Y%m%d-%H%M%S')}.{clean_extension}")


# --- PyTorch Model Definitions ---
class SimpleMLP(nn.Module):
    def __init__(self, input_dim, hidden_layers_str, output_dim, activation_fn_str="relu", task_type="classification"):
        super(SimpleMLP, self).__init__()
        layers = []
        if not isinstance(input_dim, int) or input_dim <= 0:
            raise ValueError(f"Input dimension must be a positive integer, got {input_dim}")

        hidden_units_list = []
        if hidden_layers_str and isinstance(hidden_layers_str, str) and hidden_layers_str.strip():
            try:
                hidden_units_list = [int(x.strip()) for x in hidden_layers_str.split(',') if x.strip()]
                if any(h_units <= 0 for h_units in hidden_units_list):
                    raise ValueError("Hidden layer units must be positive integers.")
            except ValueError as e:
                raise ValueError(f"Invalid hidden layer string '{hidden_layers_str}'. Error: {e}")
        
        current_dim = input_dim
        for h_units in hidden_units_list:
            layers.append(nn.Linear(current_dim, h_units))
            if activation_fn_str.lower() == "relu": layers.append(nn.ReLU())
            elif activation_fn_str.lower() == "tanh": layers.append(nn.Tanh())
            elif activation_fn_str.lower() == "sigmoid": layers.append(nn.Sigmoid())
            else: layers.append(nn.ReLU())
            current_dim = h_units
        
        layers.append(nn.Linear(current_dim, output_dim))
        
        if task_type == "classification":
            if output_dim == 1: layers.append(nn.Sigmoid()) # Binary
            elif output_dim > 1: layers.append(nn.Softmax(dim=-1)) # Multi-class
        self.network = nn.Sequential(*layers)
    def forward(self, x): return self.network(x)

class SimpleCNN(nn.Module):
    def __init__(self, input_channels, img_size_wh, num_classes=10,
                 c_out1=16, k1=3, s1=1, p1=1, pool1_k=2, pool1_s=2,
                 c_out2=32, k2=3, s2=1, p2=1, pool2_k=2, pool2_s=2,
                 fc_hidden=128):
        super(SimpleCNN, self).__init__()
        self.input_channels = input_channels
        self.img_h, self.img_w = img_size_wh
        self.num_classes = num_classes

        self.conv1 = nn.Conv2d(self.input_channels, c_out1, kernel_size=k1, stride=s1, padding=p1)
        self.relu1 = nn.ReLU()
        self.pool1 = nn.MaxPool2d(kernel_size=pool1_k, stride=pool1_s)

        h_out_conv1 = (self.img_h - k1 + 2 * p1) // s1 + 1
        w_out_conv1 = (self.img_w - k1 + 2 * p1) // s1 + 1
        h_pool1 = (h_out_conv1 - pool1_k) // pool1_s + 1
        w_pool1 = (w_out_conv1 - pool1_k) // pool1_s + 1
        
        self.conv2 = nn.Conv2d(c_out1, c_out2, kernel_size=k2, stride=s2, padding=p2)
        self.relu2 = nn.ReLU()
        self.pool2 = nn.MaxPool2d(kernel_size=pool2_k, stride=pool2_s)

        h_out_conv2 = (h_pool1 - k2 + 2 * p2) // s2 + 1
        w_out_conv2 = (w_pool1 - k2 + 2 * p2) // s2 + 1
        h_pool2 = (h_out_conv2 - pool2_k) // pool2_s + 1
        w_pool2 = (w_out_conv2 - pool2_k) // pool2_s + 1
        
        self.flattened_size = c_out2 * h_pool2 * w_pool2
        if self.flattened_size <= 0:
            raise ValueError(f"Calculated flattened size is {self.flattened_size}. Check CNN params and image size. Conv1_out:({h_out_conv1},{w_out_conv1}), Pool1_out:({h_pool1},{w_pool1}), Conv2_out:({h_out_conv2},{w_out_conv2}), Pool2_out:({h_pool2},{w_pool2})")
        
        self.fc1 = nn.Linear(self.flattened_size, fc_hidden)
        self.relu3 = nn.ReLU()
        self.fc2 = nn.Linear(fc_hidden, num_classes)
        
        if num_classes > 1 or (num_classes == 1 and task_type=="classification"): # Adapt for binary vs regression
            self.final_activation = nn.Softmax(dim=1) if num_classes > 1 else nn.Sigmoid()
        else: # Regression output from fc2
            self.final_activation = nn.Identity()


    def forward(self, x):
        x = self.pool1(self.relu1(self.conv1(x)))
        x = self.pool2(self.relu2(self.conv2(x)))
        x = x.view(-1, self.flattened_size)
        x = self.relu3(self.fc1(x))
        x = self.fc2(x)
        x = self.final_activation(x)
        return x

# --- Parameter Target Helpers ---
PARAM_RANGES = collections.OrderedDict([ # Ordered for consistent UI
    ("Tiny (<10k)", (0, 10000)),
    ("Small (10k-50k)", (10000, 50000)),
    ("Medium (50k-250k)", (50000, 250000)),
    ("Large (250k-1M)", (250000, 1000000)),
])

def suggest_mlp_layers_for_range(input_dim, output_dim, target_range_str, current_logs=""):
    logs = current_logs
    if not target_range_str or target_range_str not in PARAM_RANGES:
        logs += "Invalid parameter range selected for MLP suggestion.\n"; return "", logs
    min_p, max_p = PARAM_RANGES[target_range_str]
    target_p_avg = (min_p + max_p) // 2
    suggested_layers_str = ""
    if input_dim <=0 or output_dim <=0:
        logs += "Input/Output dims must be positive for MLP suggestion.\n"; return "", logs

    h1_candidate = max(1, int(target_p_avg / (input_dim + output_dim + 1e-6)))
    params_1_layer = (input_dim * h1_candidate + h1_candidate) + (h1_candidate * output_dim + output_dim)
    if min_p <= params_1_layer <= max_p and h1_candidate > 0:
        suggested_layers_str = str(h1_candidate)
        logs += f"Suggested 1 hidden layer: {h1_candidate} units (Est. Params: {params_1_layer})\n"
    else:
        h_base = max(1, int(math.sqrt(target_p_avg / 2.0)))
        h1 = min(2048, max(1, int(h_base * (input_dim / (input_dim + output_dim + 1e-6)) * 2 + h_base / 2)))
        h2 = min(2048, max(1, int(h_base * (output_dim / (input_dim + output_dim + 1e-6)) * 2 + h_base / 2)))
        params_2_layers = (input_dim * h1 + h1) + (h1 * h2 + h2) + (h2 * output_dim + output_dim)
        if min_p <= params_2_layers <= max_p and h1 > 0 and h2 > 0:
            suggested_layers_str = f"{h1},{h2}"
            logs += f"Suggested 2 hidden layers: {h1},{h2} units (Est. Params: {params_2_layers})\n"
        else:
            if target_p_avg < 50000: suggested_layers_str = str(max(1, int(target_p_avg / (input_dim + output_dim + 100)))) or "32"
            elif target_p_avg < 250000: h = max(1,int(math.sqrt(target_p_avg/1.5))); suggested_layers_str=f"{h},{h//2}" if h>0 and h//2 >0 else "128,64"
            else: h = max(1,int(math.sqrt(target_p_avg/2.0))); suggested_layers_str=f"{h},{h},{h//2}" if h>0 and h//2 >0 else "256,256,128"
            logs += f"Fallback suggestion: {suggested_layers_str} (Verify params).\n"
    if not suggested_layers_str: suggested_layers_str = "64"; logs += "Defaulting to '64'.\n"
    return suggested_layers_str, logs

def estimate_current_mlp_params(input_dim_str, hidden_layers_str, output_dim_str, current_logs=""):
    logs = current_logs
    try:
        input_dim = int(input_dim_str); output_dim = int(output_dim_str)
        if input_dim <= 0 or output_dim <= 0: return "Input/Output dims must be > 0", logs
        temp_mlp = SimpleMLP(input_dim, hidden_layers_str, output_dim)
        params = count_pytorch_parameters(temp_mlp); del temp_mlp
        return f"{params:,}", logs
    except Exception as e: logs += f"Error estimating MLP params: {e}\n"; return "Error", logs

def estimate_cnn_params(img_h_str, img_w_str, num_classes_str, current_logs=""):
    logs = current_logs
    try:
        img_h, img_w, num_classes = int(img_h_str), int(img_w_str), int(num_classes_str)
        if not (img_h > 0 and img_w > 0 and num_classes > 0): return "Image dims/classes must be > 0", logs
        # Using default SimpleCNN params here. A real app would pass them.
        temp_cnn = SimpleCNN(input_channels=1, img_size_wh=(img_h, img_w), num_classes=num_classes)
        params = count_pytorch_parameters(temp_cnn); del temp_cnn
        return f"{params:,}", logs
    except Exception as e: logs += f"Error estimating CNN params: {e}\n"; return "Error", logs

# --- Dataset and Preprocessing ---
def generate_dataset_backend(task_type, n_samples_str, n_features_str,
                             n_classes_or_informative_str, dataset_format,
                             ai_suggest_ds_shape, target_param_range_str, model_type_selection,
                             current_logs=""):
    logs = current_logs + "\n--- Generating Dataset ---\n"
    try:
        n_samples = int(n_samples_str); n_features = int(n_features_str); n_classes_or_informative = int(n_classes_or_informative_str)
    except ValueError: logs += "Invalid numbers for dataset generation.\n"; return None, "Error", logs, None

    if ai_suggest_ds_shape:
        n_samples_sugg, n_features_sugg, n_classes_or_informative_sugg = 5000, 10, 2
        if task_type == "Tabular Regression": n_classes_or_informative_sugg = min(n_features_sugg // 2, 5)
        elif task_type == "Basic Image Classification": n_samples_sugg, n_features_sugg = 500, 0 # features not tabular
        
        is_nn = "Network" in model_type_selection
        if is_nn and target_param_range_str in PARAM_RANGES:
            min_p, max_p = PARAM_RANGES[target_param_range_str]; avg_p = (min_p + max_p) / 2
            if avg_p > 200000: n_samples_sugg = min(MAX_GENERATED_ROWS, n_samples_sugg * 2); n_features_sugg = min(MAX_GENERATED_COLS, n_features_sugg * 2) if task_type.startswith("Tabular") else n_features_sugg
            elif avg_p < 50000: n_samples_sugg = max(100, n_samples_sugg // 2); n_features_sugg = max(3, n_features_sugg // 2) if task_type.startswith("Tabular") else n_features_sugg
        
        n_samples, n_features, n_classes_or_informative = n_samples_sugg, n_features_sugg, n_classes_or_informative_sugg
        logs += f"AI Suggested Dataset: Samples={n_samples}, Feats={n_features}, Classes/Informative={n_classes_or_informative}\n"
    
    n_samples = max(10, min(n_samples, MAX_GENERATED_ROWS))
    if task_type.startswith("Tabular"): n_features = max(1, min(n_features, MAX_GENERATED_COLS))
    if n_samples > MAX_DATASET_ROWS_WARN: logs += f"Warning: Generating {n_samples} rows. May be slow.\n"

    df = None; X_data=None; y_data=None # Init X_data, y_data
    try:
        if task_type == "Tabular Classification":
            n_cls = max(2, n_classes_or_informative)
            n_inf = max(1, min(n_features, n_classes_or_informative if n_classes_or_informative > n_cls else n_features // 2))
            X_data, y_data = make_classification(n_samples=n_samples, n_features=n_features, n_informative=n_inf,
                                       n_redundant=max(0,n_features - n_inf)//2, n_classes=n_cls, flip_y=0.05, random_state=42)
            df = pd.DataFrame(X_data, columns=[f'feature_{i}' for i in range(n_features)]); df['target'] = y_data
        elif task_type == "Tabular Regression":
            n_inf = max(1, min(n_features, n_classes_or_informative))
            X_data, y_data = make_regression(n_samples=n_samples, n_features=n_features, n_informative=n_inf, noise=10, random_state=42)
            df = pd.DataFrame(X_data, columns=[f'feature_{i}' for i in range(n_features)]); df['target'] = y_data
        elif task_type == "Basic Image Classification":
            # For SimpleCNN, let's generate 28x28 "images" (random noise)
            img_h, img_w = 28, 28 
            num_pixels = img_h * img_w
            X_data = np.random.randint(0, 256, size=(n_samples, num_pixels), dtype=np.uint8)
            y_data = np.random.randint(0, max(2, n_classes_or_informative), n_samples)
            df = pd.DataFrame(X_data, columns=[f'pixel_{i}' for i in range(num_pixels)]); df['target'] = y_data
            logs += f"Generated {img_h}x{img_w} Image placeholder data.\n"
        else: logs += f"Dataset generation for '{task_type}' not fully implemented.\n"; return None, "Task not implemented", logs, None
        
        logs += f"Generated data: {df.shape if df is not None else (X_data.shape, y_data.shape)}\n"
        file_path = get_temp_filepath("generated_dataset", dataset_format)
        if df is not None: # Save if DataFrame was created
            if dataset_format == ".csv": df.to_csv(file_path, index=False)
            elif dataset_format == ".json": df.to_json(file_path, orient='records', lines=True)
            elif dataset_format == ".parquet": df.to_parquet(file_path, index=False)
            else: logs += f"Unsupported format {dataset_format}. Defaulting to CSV.\n"; file_path=get_temp_filepath("generated_dataset","csv"); df.to_csv(file_path, index=False)
            logs += f"Dataset saved to {file_path}\n"
            return df.head(), df, logs, file_path # Return DataFrame for sklearn
        else: # Case where df might not be created (though current logic does)
            logs += "Dataset generated as numpy arrays. No file saved directly by this part of function.\n"
            # This branch needs more thought if we don't always make a df
            return pd.DataFrame(X_data[:5]), (X_data, y_data), logs, None # Return numpy arrays for PyTorch image case


    except Exception as e: error_msg=f"Error generating dataset: {traceback.format_exc()}"; logs+=error_msg+"\n"; return None, error_msg, logs, None

def preprocess_tabular_data(df_or_X, y_if_X_is_numpy, target_column_name, task_type, current_logs=""):
    logs = current_logs
    if isinstance(df_or_X, pd.DataFrame):
        df = df_or_X
        if target_column_name not in df.columns: raise ValueError(f"Target column '{target_column_name}' not found.")
        X_df = df.drop(target_column_name, axis=1)
        y_series = df[target_column_name]
    elif isinstance(df_or_X, np.ndarray) and y_if_X_is_numpy is not None: # If X,y are numpy
        X_df = pd.DataFrame(df_or_X, columns=[f'feature_{i}' for i in range(df_or_X.shape[1])]) # Temp DF for pipeline
        y_series = pd.Series(y_if_X_is_numpy)
    else: raise ValueError("Invalid input for preprocess_tabular_data.")

    numerical_features = X_df.select_dtypes(include=np.number).columns.tolist()
    categorical_features = X_df.select_dtypes(include='object').columns.tolist()
    logs += f"Numerical: {numerical_features}, Categorical: {categorical_features}\n"

    preprocessor = ColumnTransformer(transformers=[
        ('num', Pipeline([('imputer', SimpleImputer(strategy='mean')), ('scaler', StandardScaler())]), numerical_features),
        ('cat', Pipeline([('imputer', SimpleImputer(strategy='most_frequent')), ('onehot', OneHotEncoder(handle_unknown='ignore', sparse_output=False))]), categorical_features) # sparse_output=False for easier handling
    ], remainder='passthrough') # passthrough to keep unhandled columns if any

    X_processed_np = preprocessor.fit_transform(X_df)
    
    try: feature_names_out = preprocessor.get_feature_names_out()
    except AttributeError: # Older sklearn
        cat_encoder = preprocessor.named_transformers_['cat'].named_steps['onehot']
        if hasattr(cat_encoder, 'get_feature_names_out'):
            cat_feature_names = cat_encoder.get_feature_names_out(categorical_features)
        elif hasattr(cat_encoder, 'get_feature_names'): # even older
             cat_feature_names = cat_encoder.get_feature_names(categorical_features)
        else: cat_feature_names = [f"cat_feat_{i}" for i in range(X_processed_np.shape[1] - len(numerical_features))] # Fallback
        feature_names_out = numerical_features + list(cat_feature_names)

    processed_input_dim = X_processed_np.shape[1]
    logs += f"Tabular data preprocessed. X shape: {X_processed_np.shape}, Processed input dim: {processed_input_dim}\n"

    if task_type.endswith("Classification"):
        le = LabelEncoder()
        y_processed_np = le.fit_transform(y_series)
        num_classes = len(le.classes_)
        logs += f"Target encoded. Classes: {num_classes} ({le.classes_})\n"
        # For binary classification with PyTorch, often output 1 neuron with Sigmoid or BCEWithLogitsLoss
        # If num_classes is 2, some PyTorch setups expect output_dim=1.
        # Scikit-learn handles this internally.
        output_dim_nn = 1 if num_classes == 2 else num_classes
    else: # Regression
        y_processed_np = y_series.astype(float).values
        num_classes = 1 # Output dim for regression for NN
        output_dim_nn = 1
    
    return X_processed_np, y_processed_np, preprocessor, logs, processed_input_dim, output_dim_nn, feature_names_out


# --- Training Functions ---
def train_model_sklearn(data_input_obj, target_column, task_type, model_name, model_output_format, current_logs=""):
    logs = current_logs + f"\n--- Training Scikit-learn Model: {model_name} ---\n"
    model_path_out, metrics_out, model_params_out = None, "Training failed.", "N/A"

    df = None
    if isinstance(data_input_obj, str): # Filepath
        try:
            if data_input_obj.endswith('.csv'): df = pd.read_csv(data_input_obj)
            elif data_input_obj.endswith('.json'): df = pd.read_json(data_input_obj, lines=True)
            elif data_input_obj.endswith('.parquet'): df = pd.read_parquet(data_input_obj)
            else: logs += f"Unsupported file: {data_input_obj}\n"; return logs, "Error: Unsupported file.", None, "N/A"
        except Exception as e: logs += f"Error reading {data_input_obj}: {e}\n"; return logs, f"Error reading: {e}", None, "N/A"
    elif isinstance(data_input_obj, pd.DataFrame): df = data_input_obj
    else: logs += "Invalid data for training.\n"; return logs, "Error: Invalid data.", None, "N/A"

    if target_column not in df.columns:
        logs += f"Target '{target_column}' not found.\n"; return logs, f"Error: Target '{target_column}' not found.", None, "N/A"

    try:
        X_processed_np, y_processed_np, preprocessor, logs, _, _, feature_names = preprocess_tabular_data(df, None, target_column, task_type, logs)
    except ValueError as e: logs += f"Preprocessing error: {e}\n"; return logs, f"Error: {e}", None, "N/A"

    X_train, X_test, y_train, y_test = train_test_split(X_processed_np, y_processed_np, test_size=0.2, random_state=42)
    logs += f"Train/Test split. Train: {X_train.shape}, Test: {X_test.shape}\n"

    model = None
    if task_type == "Tabular Classification":
        if model_name == "Logistic Regression": model = LogisticRegression(max_iter=1000, random_state=42)
        elif model_name == "Random Forest Classifier": model = RandomForestClassifier(random_state=42)
        elif model_name == "Support Vector Machine (SVM) Classifier": model = SVC(random_state=42, probability=True) # probability=True for ONNX if it needs predict_proba
    elif task_type == "Tabular Regression":
        if model_name == "Linear Regression": model = LinearRegression()
        elif model_name == "Random Forest Regressor": model = RandomForestRegressor(random_state=42)
        elif model_name == "Support Vector Machine (SVR) Regressor": model = SVR()
    if model is None: logs += f"Model {model_name} or task {task_type} not supported.\n"; return logs, "Model/Task Error", None, "N/A"

    try:
        logs += f"Starting training for {model_name}...\n"; start_time = time.time()
        model.fit(X_train, y_train)
        logs += f"Training completed in {time.time() - start_time:.2f}s.\n"
        model_params_out = str(count_sklearn_parameters(model))
        logs += f"Est. Model Params: {model_params_out}\n"
        y_pred = model.predict(X_test)

        if task_type == "Tabular Classification":
            acc = accuracy_score(y_test, y_pred)
            report = classification_report(y_test, y_pred, zero_division=0)
            metrics_out = f"Accuracy: {acc:.4f}\n\nClassification Report:\n{report}"
        elif task_type == "Tabular Regression":
            mse = mean_squared_error(y_test, y_pred)
            r2 = r2_score(y_test, y_pred)
            metrics_out = f"Mean Squared Error: {mse:.4f}\nR2 Score: {r2:.4f}"
        logs += "\n--- Evaluation Metrics ---\n" + metrics_out + "\n"

        # Full pipeline for inference: preprocessor + model
        full_pipeline_for_saving = Pipeline([('preprocessor', preprocessor), ('model', model)])
        model_filename_base = f"sklearn_{model_name.replace(' ', '_').lower()}"

        if model_output_format == ".pkl (Scikit-learn)":
            model_path_out = get_temp_filepath(model_filename_base, "pkl")
            joblib.dump(full_pipeline_for_saving, model_path_out)
            logs += f"Model (with preprocessor) saved to {model_path_out} as PKL.\n"
        
        elif model_output_format == ".onnx (ONNX)":
            model_path_out = get_temp_filepath(model_filename_base, "onnx")
            
            # Define initial types for ONNX conversion based on preprocessed input
            # The preprocessor converts all to numerical. Shape is (batch_size, num_processed_features)
            # num_processed_features = X_train.shape[1]
            initial_type = [('float_input', FloatTensorType([None, X_train.shape[1]]))] # None for batch size
            
            # For models with string inputs *before* preprocessing, it's more complex.
            # Here, we assume the `full_pipeline_for_saving` takes the raw DataFrame structure as input.
            # So, we need to define initial_types based on the *original* DataFrame features.
            
            # Re-create initial types based on the *original* df structure, before preprocessing
            # This is complex because ColumnTransformer input spec is not trivial for skl2onnx for mixed types.
            # The EASIEST way for skl2onnx with ColumnTransformer is to convert the *fitted preprocessor separately*
            # OR, provide initial types that match the *input to the preprocessor*.

            # Let's try providing initial types for the raw input to the preprocessor
            raw_X_for_types = df.drop(target_column, axis=1).infer_objects() # Infer object dtypes to str for ONNX
            onnx_initial_types = []
            for col_name in raw_X_for_types.columns:
                col_dtype = raw_X_for_types[col_name].dtype
                if pd.api.types.is_numeric_dtype(col_dtype):
                    # Forcing float32 for ONNX compatibility
                    onnx_initial_types.append((col_name, FloatTensorType([None, 1])))
                elif pd.api.types.is_string_dtype(col_dtype) or col_dtype == 'object':
                     onnx_initial_types.append((col_name, StringTensorType([None, 1])))
                else:
                    logs += f"Warning: Unsupported dtype {col_dtype} for column {col_name} in ONNX conversion. Skipping.\n"
            
            if not onnx_initial_types:
                logs += "Error: Could not determine ONNX initial types for raw input. Aborting ONNX export.\n"
                raise ValueError("ONNX initial types failed.")

            try:
                options = {id(full_pipeline_for_saving): {'zipmap': False}} # Disable zipmap for classifier output
                onnx_model = convert_sklearn(full_pipeline_for_saving, initial_types=onnx_initial_types,
                                             target_opset=12, options=options) # Target opset can be important
                with open(model_path_out, "wb") as f:
                    f.write(onnx_model.SerializeToString())
                logs += f"Model (with preprocessor) saved to {model_path_out} as ONNX.\n"

                # Optional: Verify ONNX model
                sess = rt.InferenceSession(model_path_out, providers=rt.get_available_providers())
                logs += f"ONNX model loaded successfully with ONNX Runtime. Input names: {[inp.name for inp in sess.get_inputs()]}\n"
            except Exception as onnx_e:
                logs += f"Error during ONNX conversion/saving: {traceback.format_exc()}\n"
                model_path_out = None # Clear path if saving failed
                metrics_out += "\nONNX EXPORT FAILED."

        else:
            logs += f"Unsupported format '{model_output_format}'. Saving as .pkl\n"
            model_path_out = get_temp_filepath(model_filename_base, "pkl")
            joblib.dump(full_pipeline_for_saving, model_path_out)

    except Exception as e:
        error_msg = f"Error during sklearn training/eval: {traceback.format_exc()}"; logs += error_msg + "\n"; metrics_out = error_msg
    return logs, metrics_out, model_path_out, model_params_out


def train_model_pytorch(data_input_obj, target_column, task_type, model_type_pt,
                        mlp_hidden_layers_str, mlp_activation,
                        # CNN specific (using defaults in SimpleCNN for now)
                        # cnn_img_h_str, cnn_img_w_str, # Now derived from data
                        epochs_str, batch_size_str, lr_str,
                        model_output_format, current_logs=""):
    logs = current_logs + f"\n--- Training PyTorch Model: {model_type_pt} ---\n"
    model_path_out, metrics_out, model_params_out, plot_out = None, "Training failed.", "N/A", None

    df_for_pytorch = None; X_numpy_for_pytorch=None; y_numpy_for_pytorch=None # For flexibility
    if isinstance(data_input_obj, str): # Filepath
        try:
            # For PyTorch, we might want to handle data differently, esp images
            if data_input_obj.endswith('.csv'): df_for_pytorch = pd.read_csv(data_input_obj)
            elif data_input_obj.endswith('.json'): df_for_pytorch = pd.read_json(data_input_obj, lines=True)
            elif data_input_obj.endswith('.parquet'): df_for_pytorch = pd.read_parquet(data_input_obj)
            else: logs += f"Unsupported file: {data_input_obj}\n"; return logs, "Error", None, "N/A", None
        except Exception as e: logs += f"Error reading {data_input_obj}: {e}\n"; return logs, f"Error: {e}", None, "N/A", None
    elif isinstance(data_input_obj, pd.DataFrame): df_for_pytorch = data_input_obj
    elif isinstance(data_input_obj, tuple) and len(data_input_obj) == 2 and \
         isinstance(data_input_obj[0], np.ndarray) and isinstance(data_input_obj[1], np.ndarray):
        X_numpy_for_pytorch, y_numpy_for_pytorch = data_input_obj # If data was (X,y) from generation
    else: logs += "Invalid data for PyTorch training.\n"; return logs, "Error", None, "N/A", None

    try:
        epochs = int(epochs_str); batch_size = int(batch_size_str); lr = float(lr_str)
        if not (epochs > 0 and batch_size > 0 and lr > 0): raise ValueError("Params must be >0.")
    except ValueError as e: logs += f"Invalid training params: {e}\n"; return logs, f"Error: {e}", None, "N/A", None

    processed_input_dim_actual = -1; nn_output_dim_actual = -1; preprocessor_pipeline = None
    X_processed_np = None; y_processed_np = None

    if model_type_pt == "Simple Neural Network (MLP)":
        if not task_type.startswith("Tabular"):
            logs += "MLP requires Tabular task.\n"; return logs, "MLP Task Error", None, "N/A", None
        try:
            # Pass df_for_pytorch or (X_numpy_for_pytorch, y_numpy_for_pytorch)
            data_arg1 = df_for_pytorch if df_for_pytorch is not None else X_numpy_for_pytorch
            data_arg2 = y_numpy_for_pytorch if df_for_pytorch is None else None
            X_processed_np, y_processed_np, preprocessor_pipeline, logs, processed_input_dim_actual, nn_output_dim_actual, _ = \
                preprocess_tabular_data(data_arg1, data_arg2, target_column, task_type, logs)
        except ValueError as e: logs+=f"MLP Preprocessing error: {e}\n"; return logs,f"Error: {e}",None,"N/A",None

    elif model_type_pt == "Simple Convolutional Network (CNN)":
        if task_type != "Basic Image Classification":
            logs += "Warning: CNN selected, but task is not Basic Image Classification. Output may be unexpected.\n"
        
        if df_for_pytorch is not None:
            if target_column not in df_for_pytorch.columns:
                 logs += f"Target '{target_column}' not found for CNN.\n"; return logs, "CNN Target Error", None, "N/A", None
            X_raw = df_for_pytorch.drop(target_column, axis=1).values
            y_raw = df_for_pytorch[target_column].values
        elif X_numpy_for_pytorch is not None and y_numpy_for_pytorch is not None:
            X_raw = X_numpy_for_pytorch
            y_raw = y_numpy_for_pytorch
        else:
            logs += "No valid data found for CNN.\n"; return logs, "CNN Data Error", None, "N/A", None

        le = LabelEncoder(); y_processed_np = le.fit_transform(y_raw)
        nn_output_dim_actual = len(le.classes_)
        if nn_output_dim_actual == 2: nn_output_dim_actual = 1 # Binary output for NN

        pixels_per_sample = X_raw.shape[1]
        img_dim_approx = int(math.sqrt(pixels_per_sample))
        img_h, img_w, input_channels = (28,28,1) # Default
        if img_dim_approx * img_dim_approx == pixels_per_sample:
            img_h, img_w = img_dim_approx, img_dim_approx
        else: logs += f"Warning: Cannot infer square image from {pixels_per_sample} pixels. Defaulting to 28x28 for CNN.\n"
        
        # Reshape and normalize (basic)
        X_processed_np = X_raw.reshape(-1, input_channels, img_h, img_w).astype(np.float32) / 255.0
        processed_input_dim_actual = (input_channels, img_h, img_w) # For CNN constructor
        logs += f"CNN Data: X reshaped to {X_processed_np.shape}, y: {y_processed_np.shape}, NN Output Dim: {nn_output_dim_actual}\n"
    else: logs += f"Unknown PyTorch model: {model_type_pt}\n"; return logs, "Unknown PyTorch model", None, "N/A", None

    X_tensor = torch.tensor(X_processed_np, dtype=torch.float32)
    # Adjust y_tensor dtype based on loss function expectations
    y_dtype = torch.float32 if (nn_output_dim_actual == 1 and task_type.endswith("Regression")) or \
                              (nn_output_dim_actual == 1 and task_type.endswith("Classification")) \
                           else torch.long # MSELoss/BCELoss with float, CrossEntropy with long
    y_tensor = torch.tensor(y_processed_np, dtype=y_dtype)
    if nn_output_dim_actual == 1 and task_type.endswith("Classification"): y_tensor = y_tensor.unsqueeze(1) # For BCE based loss
    if task_type.endswith("Regression"):  y_tensor = y_tensor.unsqueeze(1) # MSELoss expects [N,1]

    dataset = TensorDataset(X_tensor, y_tensor)
    # Use num_workers=0 on free tier to avoid issues with multiprocessing
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=0) 

    pytorch_model = None
    try:
        if model_type_pt == "Simple Neural Network (MLP)":
            pytorch_model = SimpleMLP(input_dim=processed_input_dim_actual, hidden_layers_str=mlp_hidden_layers_str,
                                    output_dim=nn_output_dim_actual, activation_fn_str=mlp_activation,
                                    task_type="classification" if task_type.endswith("Classification") else "regression")
        elif model_type_pt == "Simple Convolutional Network (CNN)":
            channels, h, w = processed_input_dim_actual
            pytorch_model = SimpleCNN(input_channels=channels, img_size_wh=(h,w), num_classes=nn_output_dim_actual)
    except Exception as model_e:
        logs += f"Error creating PyTorch model: {traceback.format_exc()}\n"; return logs, f"Model Creation Error: {model_e}", None, "N/A", None

    if pytorch_model is None: logs += "Failed to instantiate PyTorch model.\n"; return logs, "Model instantiate fail", None, "N/A", None
    model_params_val = count_pytorch_parameters(pytorch_model)
    model_params_out = f"{model_params_val:,}"
    logs += f"PyTorch Model: {model_params_out} params.\n"
    if model_params_val > 500000: logs += "Warning: >500k params on CPU will be SLOW.\n"

    is_classification_task = task_type.endswith("Classification") or model_type_pt == "Simple Convolutional Network (CNN)" # Treat CNN as classification here
    if is_classification_task:
        criterion = nn.BCELoss() if nn_output_dim_actual == 1 else nn.CrossEntropyLoss()
    else: # Regression
        criterion = nn.MSELoss()
    optimizer = optim.Adam(pytorch_model.parameters(), lr=lr)

    logs += f"Starting PyTorch training for {epochs} epochs...\n"; start_time = time.time()
    epoch_losses = []
    pytorch_model.train()
    for epoch in range(epochs):
        epoch_loss_sum = 0.0; num_batches = 0
        for batch_X, batch_y in dataloader:
            optimizer.zero_grad()
            outputs = pytorch_model(batch_X)
            loss = criterion(outputs, batch_y)
            loss.backward(); optimizer.step()
            epoch_loss_sum += loss.item(); num_batches += 1
        avg_epoch_loss = epoch_loss_sum / num_batches if num_batches > 0 else 0
        epoch_losses.append(avg_epoch_loss)
        logs += f"Epoch {epoch+1}/{epochs}, Avg Loss: {avg_epoch_loss:.4f}\n"
        # yield logs, metrics_out, model_path_out, model_params_out, None # For streaming, but makes UI complex
    
    training_time = time.time() - start_time
    logs += f"PyTorch training completed in {training_time:.2f} seconds.\n"
    
    # Basic evaluation (on last batch for simplicity, or could do full test set)
    # A proper eval loop on a test set would be better here.
    pytorch_model.eval()
    with torch.no_grad():
        # For simplicity, let's just report final training loss.
        # A full evaluation on a test split would be needed for proper metrics.
        if is_classification_task:
            # This is a very rough accuracy on the last training batch for demo
            if dataloader.dataset: # Check if dataset is not empty
                try:
                    last_batch_X, last_batch_y = next(iter(dataloader)) # Get one batch
                    outputs = pytorch_model(last_batch_X)
                    if nn_output_dim_actual == 1: # Binary
                        predicted = (outputs > 0.5).float()
                    else: # Multi-class
                        _, predicted = torch.max(outputs.data, 1)
                    correct = (predicted == last_batch_y.view_as(predicted)).sum().item()
                    total = last_batch_y.size(0)
                    acc = correct / total if total > 0 else 0
                    metrics_out = f"Final Training Loss: {avg_epoch_loss:.4f}\nApprox. Accuracy on a batch: {acc*100:.2f}% (Note: Proper eval needs a test set)"
                except StopIteration: # Dataloader was empty
                    metrics_out = f"Final Training Loss: {avg_epoch_loss:.4f}\n (Dataloader empty, cannot get batch accuracy)"

            else:
                metrics_out = f"Final Training Loss: {avg_epoch_loss:.4f}\n (No data for batch accuracy)"
        else: # Regression
            metrics_out = f"Final Training Loss (MSE): {avg_epoch_loss:.4f}"
    logs += "\n--- PyTorch Metrics (Simplified) ---\n" + metrics_out + "\n"

    # Loss plot
    if epoch_losses:
        import matplotlib.pyplot as plt
        fig, ax = plt.subplots()
        ax.plot(range(1, epochs + 1), epoch_losses, marker='o')
        ax.set_xlabel("Epoch")
        ax.set_ylabel("Average Loss")
        ax.set_title("Training Loss Curve")
        plot_out = fig # Gradio can display matplotlib figures
        logs += "Loss curve generated.\n"


    # Save model (and preprocessor if MLP)
    model_filename_base = f"pytorch_{model_type_pt.replace(' ', '_').lower()}"
    if model_output_format == ".pt (PyTorch)":
        model_path_out = get_temp_filepath(model_filename_base, "pt")
        if model_type_pt == "Simple Neural Network (MLP)" and preprocessor_pipeline:
            torch.save({
                'model_state_dict': pytorch_model.state_dict(),
                'preprocessor': preprocessor_pipeline,
                'input_dim': processed_input_dim_actual, # From preprocessing
                'output_dim': nn_output_dim_actual, # From preprocessing
                'hidden_layers_str': mlp_hidden_layers_str,
                'activation_fn': mlp_activation,
                'task_type': task_type
            }, model_path_out)
            logs += f"PyTorch MLP (model + preprocessor) saved to {model_path_out}\n"
        else: # CNN or MLP without preprocessor explicitly bundled (less common)
            torch.save(pytorch_model.state_dict(), model_path_out)
            logs += f"PyTorch {model_type_pt} (model state_dict) saved to {model_path_out}\n"
    # Add ONNX export for PyTorch later if needed (torch.onnx.export)
    else:
        logs += f"Unsupported format '{model_output_format}' for PyTorch. Saving as .pt\n"
        model_path_out = get_temp_filepath(model_filename_base, "pt")
        torch.save(pytorch_model.state_dict(), model_path_out) # Fallback to state_dict

    return logs, metrics_out, model_path_out, model_params_out, plot_out


# --- Gradio UI Definition ---
# Define choices
TASK_CHOICES = ["Tabular Classification", "Tabular Regression", "Basic Image Classification"] # Simple Text removed for focus
MODEL_FAMILIES = ["Scikit-learn (Classical ML)", "PyTorch (Neural Networks)"]
SKLEARN_MODELS_CLASSIFICATION = ["Logistic Regression", "Random Forest Classifier", "Support Vector Machine (SVM) Classifier"]
SKLEARN_MODELS_REGRESSION = ["Linear Regression", "Random Forest Regressor", "Support Vector Machine (SVR) Regressor"]
PYTORCH_MODELS = ["Simple Neural Network (MLP)", "Simple Convolutional Network (CNN)"]
DATASET_FORMATS = [".csv", ".json", ".parquet"]
MODEL_OUTPUT_FORMATS_SKLEARN = [".pkl (Scikit-learn)", ".onnx (ONNX)"]
MODEL_OUTPUT_FORMATS_PYTORCH = [".pt (PyTorch)"] # ".onnx (ONNX)" can be added later
MLP_ACTIVATIONS = ["relu", "tanh", "sigmoid"]

CLONE_GUIDE_TEXT = """
## How to Clone & Upgrade This Space for More Power:
(Instructions as provided in previous response - omitted here for brevity but should be included)
"""

def update_model_options(task_choice, model_family_choice):
    if model_family_choice == "Scikit-learn (Classical ML)":
        if task_choice == "Tabular Classification": return gr.update(choices=SKLEARN_MODELS_CLASSIFICATION, value=SKLEARN_MODELS_CLASSIFICATION[0], visible=True)
        elif task_choice == "Tabular Regression": return gr.update(choices=SKLEARN_MODELS_REGRESSION, value=SKLEARN_MODELS_REGRESSION[0], visible=True)
        else: return gr.update(choices=[], value=None, visible=False) # Sklearn not for image task here
    elif model_family_choice == "PyTorch (Neural Networks)":
        if task_choice.startswith("Tabular"): return gr.update(choices=[PYTORCH_MODELS[0]], value=PYTORCH_MODELS[0], visible=True) # Only MLP for tabular
        elif task_choice == "Basic Image Classification": return gr.update(choices=[PYTORCH_MODELS[1]], value=PYTORCH_MODELS[1], visible=True) # Only CNN for image
        else: return gr.update(choices=[], value=None, visible=False)
    return gr.update(choices=[], value=None, visible=False)

def update_param_range_visibility(model_family_choice):
    return gr.update(visible=(model_family_choice == "PyTorch (Neural Networks)"))

def update_pytorch_specific_options_visibility(model_choice_pytorch):
    is_mlp = model_choice_pytorch == "Simple Neural Network (MLP)"
    is_cnn = model_choice_pytorch == "Simple Convolutional Network (CNN)"
    return gr.update(visible=is_mlp), gr.update(visible=is_cnn) # MLP Group, CNN Group

def update_model_output_formats(model_family_choice):
    if model_family_choice == "Scikit-learn (Classical ML)":
        return gr.update(choices=MODEL_OUTPUT_FORMATS_SKLEARN, value=MODEL_OUTPUT_FORMATS_SKLEARN[0])
    elif model_family_choice == "PyTorch (Neural Networks)":
        return gr.update(choices=MODEL_OUTPUT_FORMATS_PYTORCH, value=MODEL_OUTPUT_FORMATS_PYTORCH[0])
    return gr.update(choices=[], value=None)


css = """
.gradio-container { font-family: 'IBM Plex Sans', sans-serif; }
.gr-button { color: white; border-color: black; background: black; }
.gr-input { border-radius: 8px; }
.gr-output { border-radius: 8px; }
"""

with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="orange"), css=css) as demo:
    gr.Markdown("# 🧠 Universal AI Model Trainer (CPU Edition)")
    gr.Markdown("Create, train, and download AI models. Optimized for CPU - expect longer training for complex models.")
    
    # Global state to store generated data path or df
    # This helps pass data between dataset generation and training without re-upload
    # For DataFrames, it's better to pass them directly if possible, or save/load paths.
    generated_data_state = gr.State(None) 
    current_logs_state = gr.State("") # To accumulate logs

    with gr.Tabs():
        with gr.TabItem("1. Define Task & Model"):
            with gr.Row():
                task_type_dd = gr.Dropdown(TASK_CHOICES, label="Select Task Type", value=TASK_CHOICES[0])
                model_family_dd = gr.Dropdown(MODEL_FAMILIES, label="Select Model Family", value=MODEL_FAMILIES[0])
            
            model_specific_dd = gr.Dropdown(label="Select Specific Model", interactive=True) # Populated by callback
            
            # PyTorch Parameter Range (only visible for PyTorch)
            pytorch_param_range_dd = gr.Dropdown(list(PARAM_RANGES.keys()), label="Target Parameter Range (for NNs)", 
                                                 info="Guides NN architecture suggestions. Training >250k params on CPU is slow.",
                                                 value=list(PARAM_RANGES.keys())[1], visible=False)
            
            # PyTorch MLP Specifics (only visible for MLP)
            with gr.Group(visible=False) as pt_mlp_specific_group:
                gr.Markdown("#### MLP Configuration")
                # Input dim will be determined after data preprocessing for MLP. User doesn't set it here.
                # Output dim also determined by data (num_classes or 1 for regression)
                pt_mlp_hidden_layers_txt = gr.Textbox(label="Hidden Layer Sizes (comma-separated, e.g., 128,64)", value="64,32")
                pt_mlp_activation_dd = gr.Dropdown(MLP_ACTIVATIONS, label="Activation Function", value="relu")
                pt_mlp_suggest_btn = gr.Button("Suggest MLP Layers for Target Range")
                pt_mlp_param_count_txt = gr.Textbox(label="Estimated MLP Parameters", interactive=False)
                # For MLP param estimation, we'd need #input_features and #output_classes from data step
                # This means estimation might be better placed *after* dataset is defined.
                # For now, placeholder or user has to guess input/output dims.
                # Simplified: we'll show actual params *after* training or with a dedicated button post-data.

            # PyTorch CNN Specifics (Placeholder - visible for CNN)
            with gr.Group(visible=False) as pt_cnn_specific_group:
                gr.Markdown("#### CNN Configuration (Simplified for Demo)")
                gr.Markdown("SimpleCNN uses fixed architecture for now (2 conv layers, 1 FC). Parameters mainly come from image size/classes.")
                # For CNN param estimation, we need image H, W, num_classes from data step.
                # cnn_img_h_param_est = gr.Number(label="Est. Image Height (for param count)", value=28, visible=False) # Hidden, used by callback
                # cnn_img_w_param_est = gr.Number(label="Est. Image Width (for param count)", value=28, visible=False)
                # cnn_num_classes_param_est = gr.Number(label="Est. Num Classes (for param count)", value=10, visible=False)
                pt_cnn_param_count_txt = gr.Textbox(label="Estimated CNN Parameters", interactive=False)
                # Actual CNN param count shown after training or with dedicated button post-data.


        with gr.TabItem("2. Configure Dataset"):
            dataset_source_rb = gr.Radio(["Generate new dataset", "Upload my own dataset (CSV, JSON, Parquet)"], 
                                         label="Dataset Source", value="Generate new dataset")
            
            with gr.Group(visible=True) as generate_dataset_group: # Visible by default
                gr.Markdown("#### Generate Synthetic Dataset")
                with gr.Row():
                    ds_gen_samples_num = gr.Number(label="Number of Rows (Samples)", value=1000)
                    ds_gen_features_num = gr.Number(label="Number of Features (Columns, if tabular)", value=10)
                ds_gen_classes_informative_num = gr.Number(label="Num Classes (for Classification) / Num Informative Features (for Regression)", value=2)
                ds_gen_ai_suggest_cb = gr.Checkbox(label="Let AI suggest optimal rows/columns based on model type & param range?", value=False)
                ds_gen_format_dd = gr.Dropdown(DATASET_FORMATS, label="Generated Dataset Download Format", value=".csv")
                generate_dataset_btn = gr.Button("Generate & Preview Dataset", variant="secondary")
            
            with gr.Group(visible=False) as upload_dataset_group:
                gr.Markdown("#### Upload Dataset")
                ds_upload_file = gr.File(label="Upload your dataset file", file_types=[".csv", ".json", ".parquet"])
            
            target_column_name_txt = gr.Textbox(label="Target Column Name (Case-Sensitive)", placeholder="e.g., 'target' or 'label'")
            dataset_preview_df = gr.DataFrame(label="Dataset Preview (First 5 Rows)", interactive=False)
            generated_dataset_download_file = gr.File(label="Download Generated Dataset", interactive=False)

        with gr.TabItem("3. Train Model & Get Results"):
            gr.Markdown("Ensure Model and Dataset are configured before training.")
            with gr.Row():
                # Training Hyperparameters (Common for PyTorch)
                # For Scikit-learn, HPs are mostly defaults or need more complex UI
                # These are mainly for PyTorch NNs
                train_epochs_num = gr.Number(label="Epochs (for NNs)", value=10)
                train_batch_size_num = gr.Number(label="Batch Size (for NNs)", value=32)
                train_learning_rate_num = gr.Number(label="Learning Rate (for NNs)", value=0.001)
            
            model_output_format_dd = gr.Dropdown(label="Select Model Output Format", choices=MODEL_OUTPUT_FORMATS_SKLEARN, value=MODEL_OUTPUT_FORMATS_SKLEARN[0]) # Default to sklearn
            train_model_btn = gr.Button("🚀 Train Model", variant="primary")
            
            gr.Markdown("---")
            gr.Markdown("### Training Progress & Results")
            training_log_txt = gr.Textbox(label="Training Log & Status", lines=15, interactive=False, max_lines=50)
            model_param_count_output_txt = gr.Textbox(label="Actual Trained Model Parameters", interactive=False)
            evaluation_metrics_txt = gr.Textbox(label="Evaluation Metrics", lines=7, interactive=False)
            loss_plot_img = gr.Plot(label="Training Loss Curve (PyTorch NNs)")
            download_trained_model_file = gr.File(label="Download Trained Model", interactive=False)

        with gr.TabItem("ℹ️ Guide & Info"):
            gr.Markdown("### Using This Space")
            gr.Markdown("- **Free CPU Tier:** Training large or complex models will be slow. Memory is also limited (around 15GB RAM).")
            gr.Markdown("- **Workflow:** 1. Define Task/Model -> 2. Configure Dataset -> 3. Train.")
            gr.Markdown("- **Dataset Generation:** For 'Basic Image Classification', random pixel data is generated (not real images).")
            gr.Markdown("- **Parameters:** For Neural Networks, the 'Target Parameter Range' helps suggest architectures. 1M params is already large for CPU training.")
            gr.Markdown("- **ONNX Export (Scikit-learn):** Converts Scikit-learn pipelines (preprocessor + model) to ONNX. Input to the ONNX model should be raw data matching the original training DataFrame structure.")
            gr.Markdown(CLONE_GUIDE_TEXT)
            
    # --- Event Handlers ---
    # Update model choices based on task and family
    task_type_dd.change(fn=update_model_options, inputs=[task_type_dd, model_family_dd], outputs=model_specific_dd)
    model_family_dd.change(fn=update_model_options, inputs=[task_type_dd, model_family_dd], outputs=model_specific_dd)
    
    # Show/hide PyTorch parameter range dropdown
    model_family_dd.change(fn=update_param_range_visibility, inputs=model_family_dd, outputs=pytorch_param_range_dd)
    
    # Show/hide PyTorch MLP/CNN specific groups
    # This needs model_specific_dd as input, which is tricky if it's dynamically populated.
    # Let's assume model_specific_dd is the PyTorch model dropdown for this context.
    # This means model_specific_dd must *only* be active/relevant when model_family_dd is PyTorch.
    def combined_pytorch_ui_update(model_family_choice, pytorch_model_choice):
        param_range_visible = (model_family_choice == "PyTorch (Neural Networks)")
        if not param_range_visible: # If not PyTorch, hide all PyTorch specific groups
            return gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
        
        is_mlp = (pytorch_model_choice == "Simple Neural Network (MLP)")
        is_cnn = (pytorch_model_choice == "Simple Convolutional Network (CNN)")
        return gr.update(visible=param_range_visible), gr.update(visible=is_mlp), gr.update(visible=is_cnn)

    model_family_dd.change(fn=combined_pytorch_ui_update, 
                           inputs=[model_family_dd, model_specific_dd], 
                           outputs=[pytorch_param_range_dd, pt_mlp_specific_group, pt_cnn_specific_group])
    model_specific_dd.change(fn=combined_pytorch_ui_update, # Also trigger when specific PyTorch model changes
                             inputs=[model_family_dd, model_specific_dd],
                             outputs=[pytorch_param_range_dd, pt_mlp_specific_group, pt_cnn_specific_group])

    # Suggest MLP Layers
    def mlp_suggest_proxy(target_range_str, current_logs, dataset_preview_df, target_col_name, task_type):
        logs = current_logs
        input_dim_est = 10 # default if no data
        output_dim_est = 2 if task_type.endswith("Classification") else 1 # default
        
        if dataset_preview_df is not None and isinstance(dataset_preview_df, pd.DataFrame) and not dataset_preview_df.empty and target_col_name:
            try:
                # Attempt to get processed input dim. This is a simplified estimation.
                # A full preprocessing run is too heavy here.
                temp_X = dataset_preview_df.drop(target_col_name, axis=1, errors='ignore')
                num_cols = len(temp_X.select_dtypes(include=np.number).columns)
                cat_cols = temp_X.select_dtypes(include='object').columns
                # Rough estimate of one-hot encoded features
                one_hot_est = sum(min(10, dataset_preview_df[col].nunique()) for col in cat_cols) # cap nunique
                input_dim_est = num_cols + one_hot_est
                input_dim_est = max(1, input_dim_est) # Ensure > 0

                if task_type.endswith("Classification"):
                    output_dim_est = max(1, dataset_preview_df[target_col_name].nunique())
                    if output_dim_est == 2: output_dim_est = 1 # For binary an output of 1 is common in NNs
                logs += f"Estimated input_dim: {input_dim_est}, output_dim: {output_dim_est} for MLP suggestion.\n"
            except Exception as e:
                logs += f"Could not estimate dims from preview for MLP suggestion: {e}. Using defaults.\n"
        else:
            logs += "Dataset preview not available for MLP dimension estimation. Using defaults.\n"
            
        suggested_str, logs = suggest_mlp_layers_for_range(input_dim_est, output_dim_est, target_range_str, logs)
        
        # Also estimate params for the suggestion
        param_count_str = "Error"
        if suggested_str:
            param_count_str, logs = estimate_current_mlp_params(str(input_dim_est), suggested_str, str(output_dim_est), logs)

        return suggested_str, logs, param_count_str

    pt_mlp_suggest_btn.click(
        fn=mlp_suggest_proxy,
        inputs=[pytorch_param_range_dd, current_logs_state, dataset_preview_df, target_column_name_txt, task_type_dd],
        outputs=[pt_mlp_hidden_layers_txt, training_log_txt, pt_mlp_param_count_txt] # Use training_log_txt for logs from suggestion
    )
    
    # Estimate MLP params when hidden layers text changes (might be too slow if hooked to .change)
    # A button is safer for this. For now, rely on suggestion button or post-training report.
    # We can add an "Estimate Current MLP Params" button if needed.

    # Show/hide dataset generation/upload groups
    def toggle_dataset_source_groups(source_choice):
        return gr.update(visible=(source_choice == "Generate new dataset")), \
               gr.update(visible=(source_choice == "Upload my own dataset (CSV, JSON, Parquet)"))
    dataset_source_rb.change(fn=toggle_dataset_source_groups, inputs=dataset_source_rb, 
                             outputs=[generate_dataset_group, upload_dataset_group])

    # Update model output formats based on family
    model_family_dd.change(fn=update_model_output_formats, inputs=model_family_dd, outputs=model_output_format_dd)

    # Dataset Generation Button
    def generate_dataset_wrapper(task_type, n_samples, n_features, n_classes_info, ds_format, ai_sugg, param_range, model_type, logs_in):
        preview, data_obj, logs_out, file_out = generate_dataset_backend(
            task_type, n_samples, n_features, n_classes_info, ds_format, ai_sugg, param_range, model_type, logs_in
        )
        # Store the actual data (DataFrame or (X,y) tuple) in state if generation was successful
        # If it's a filepath (from upload), store the path.
        # For generated data, store the df or (X,y) tuple to avoid disk I/O if not necessary before training.
        stored_data = data_obj if data_obj is not None else None
        return preview, stored_data, logs_out, file_out

    generate_dataset_btn.click(
        fn=generate_dataset_wrapper,
        inputs=[task_type_dd, ds_gen_samples_num, ds_gen_features_num, ds_gen_classes_informative_num, 
                ds_gen_format_dd, ds_gen_ai_suggest_cb, pytorch_param_range_dd, model_specific_dd, current_logs_state],
        outputs=[dataset_preview_df, generated_data_state, training_log_txt, generated_dataset_download_file]
    )

    # Handle dataset upload
    def process_uploaded_file(file_obj, logs_in):
        logs = logs_in
        if file_obj is None:
            return None, logs, "Please upload a file first.", None
        logs += f"Uploaded file: {file_obj.name}\n"
        
        # For preview, try to read a few lines
        df_preview = None
        try:
            if file_obj.name.endswith(".csv"):
                df_preview = pd.read_csv(file_obj.name, nrows=5)
            elif file_obj.name.endswith(".json"): # Assuming JSONL
                df_preview = pd.read_json(file_obj.name, lines=True, nrows=5)
            elif file_obj.name.endswith(".parquet"):
                # Reading only 5 rows from parquet is not straightforward without loading more.
                # For simplicity, load full and take head, or skip preview.
                temp_df = pd.read_parquet(file_obj.name)
                df_preview = temp_df.head()
            logs += "Preview generated for uploaded file.\n"
        except Exception as e:
            logs += f"Could not generate preview for {file_obj.name}: {e}\n"
            return None, logs, f"Error previewing: {e}", file_obj.name # Return path even if preview fails
            
        return df_preview, logs, "File ready for training.", file_obj.name # Store path in generated_data_state

    ds_upload_file.upload(
        fn=process_uploaded_file,
        inputs=[ds_upload_file, current_logs_state],
        outputs=[dataset_preview_df, training_log_txt, training_log_txt, generated_data_state] # Use training_log for status, then store path
    )


    # Train Model Button
    def train_model_wrapper(data_state_val, # This will be DataFrame, (X,y) tuple, or filepath string
                            target_col, task_type, model_family, model_name, # Common params
                            # Sklearn specific (none for now beyond model_name)
                            # PyTorch specific
                            pt_model_type, pt_mlp_hidden, pt_mlp_activ, #pt_cnn_params (later)
                            epochs, batch_size, lr,
                            model_out_format,
                            logs_in): # Accumulate logs

        current_logs = logs_in + "\n--- Initiating Training ---\n"
        current_logs += f"Data state type: {type(data_state_val)}\n"

        if data_state_val is None:
            current_logs += "Error: No dataset loaded or generated. Please go to Tab 2.\n"
            return current_logs, "No data available.", None, "N/A", None, None # logs, metrics, model_file, params, plot, download_btn_update

        if not target_col and (task_type.startswith("Tabular") or (isinstance(data_state_val, pd.DataFrame) and model_type_pt != "Simple Convolutional Network (CNN)")) : # Target col needed for tabular
             current_logs += "Error: Target column name is required for this task/data.\n"
             return current_logs, "Target column needed.", None, "N/A", None, None
        
        # Ensure logs are passed and returned correctly by train functions
        if model_family == "Scikit-learn (Classical ML)":
            logs, metrics, model_file, params = train_model_sklearn(
                data_state_val, target_col, task_type, model_name, model_out_format, current_logs
            )
            return logs, metrics, model_file, params, None, model_file # No plot for sklearn here
        
        elif model_family == "PyTorch (Neural Networks)":
            # model_name here is the PyTorch model type (MLP or CNN)
            logs, metrics, model_file, params, plot = train_model_pytorch(
                data_state_val, target_col, task_type, model_name, 
                pt_mlp_hidden, pt_mlp_activ,
                epochs, batch_size, lr,
                model_out_format, current_logs
            )
            return logs, metrics, model_file, params, plot, model_file
        else:
            current_logs += f"Unknown model family: {model_family}\n"
            return current_logs, "Unknown model family.", None, "N/A", None, None

    train_model_btn.click(
        fn=train_model_wrapper,
        inputs=[
            generated_data_state, target_column_name_txt, task_type_dd, model_family_dd, model_specific_dd,
            # PyTorch specific inputs (will be None if not PyTorch family, but passed)
            model_specific_dd, # This is pt_model_type if family is PyTorch
            pt_mlp_hidden_layers_txt, pt_mlp_activation_dd,
            train_epochs_num, train_batch_size_num, train_learning_rate_num,
            model_output_format_dd,
            training_log_txt # Pass current log content to append
        ],
        outputs=[
            training_log_txt, evaluation_metrics_txt, download_trained_model_file, 
            model_param_count_output_txt, loss_plot_img,
            download_trained_model_file # This seems redundant, download_trained_model_file is already an output
        ]
    )

    # Clear logs button (optional)
    # clear_logs_btn = gr.Button("Clear Logs")
    # def clear_logs_func(): return "", "" # Clears current_logs_state and training_log_txt
    # clear_logs_btn.click(clear_logs_func, [], [current_logs_state, training_log_txt])


demo.queue().launch(debug=True, show_error=True) # Enable queue for longer tasks, debug for local testing