src/mlp_regression.py

import pandas as pd
import numpy as np
from sklearn.datasets import fetch_california_housing, make_regression
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
from plotly.subplots import make_subplots
import plotly.graph_objects as go
import torch
import torch.nn as nn
import torch.optim as optim
import time
import threading

_model_lock = threading.RLock()
_current_model = None
_current_scaler = None

class MLP(nn.Module):
    def __init__(self, input_dim, hidden_layers_config):
        super(MLP, self).__init__()
        layers_list = []
        prev_dim = input_dim
        
        for i, layer_config in enumerate(hidden_layers_config):
            neurons = layer_config['neurons']
            activation = layer_config.get('activation', 'relu')
            
            layers_list.append(nn.Linear(prev_dim, neurons))
            
            if activation.lower() == 'relu':
                layers_list.append(nn.ReLU())
            elif activation.lower() == 'sigmoid':
                layers_list.append(nn.Sigmoid())
            elif activation.lower() == 'tanh':
                layers_list.append(nn.Tanh())
            elif activation.lower() in ['leakyrelu', 'leaky_relu']:
                layers_list.append(nn.LeakyReLU(0.01))
            else:
                layers_list.append(nn.ReLU())
            
            prev_dim = neurons
        
        layers_list.append(nn.Linear(prev_dim, 1))
        
        self.network = nn.Sequential(*layers_list)
    
    def forward(self, x):
        return self.network(x)

def load_data(file_obj=None, dataset_choice="California Housing"):
    if file_obj is not None:
        if file_obj.name.endswith(".csv"):
            encodings = ["utf-8", "latin-1", "iso-8859-1", "cp1252"]
            for encoding in encodings:
                try:
                    return pd.read_csv(file_obj.name, encoding=encoding)
                except UnicodeDecodeError:
                    continue
            return pd.read_csv(file_obj.name, encoding="utf-8", errors="replace")
        elif file_obj.name.endswith((".xlsx", ".xls")):
            return pd.read_excel(file_obj.name)
        else:
            raise ValueError("Unsupported format. Upload CSV or Excel files.")
    
    datasets = {
        "California Housing": lambda: _california_housing_to_df(),
        "Synthetic": lambda: _synthetic_regression(),
    }
    if dataset_choice not in datasets:
        raise ValueError(f"Unknown dataset: {dataset_choice}")
    return datasets[dataset_choice]()

def _california_housing_to_df():
    data = fetch_california_housing()
    df = pd.DataFrame(data.data, columns=data.feature_names)
    df["target"] = data.target
    return df

def _synthetic_regression():
    X, y = make_regression(n_samples=1000, n_features=20, n_informative=15, 
                          noise=10.0, random_state=42)
    df = pd.DataFrame(X, columns=[f"feature_{i}" for i in range(X.shape[1])])
    df["target"] = y
    return df

def create_input_components(df, target_col):
    feature_cols = [c for c in df.columns if c != target_col]
    components = []
    for col in feature_cols:
        data = df[col]
        val = pd.to_numeric(data, errors="coerce").dropna().mean()
        val = 0.0 if pd.isna(val) else float(val)
        components.append({
            "name": col,
            "type": "number",
            "value": round(val, 3),
            "minimum": None,
            "maximum": None,
        })
    return components

def preprocess_data(df, target_col, new_point_dict):
    feature_cols = [c for c in df.columns if c != target_col]
    X = df[feature_cols].copy()
    y = df[target_col].copy()

    for col in feature_cols:
        X[col] = pd.to_numeric(X[col], errors="coerce").fillna(0.0)
    
    y = pd.to_numeric(y, errors="coerce").fillna(0.0)

    new_point = []
    for col in feature_cols:
        if col in new_point_dict:
            try:
                new_point.append(float(new_point_dict[col]))
            except Exception:
                new_point.append(0.0)
        else:
            new_point.append(0.0)
    
    new_point = np.array(new_point, dtype=float).reshape(1, -1)
    
    if new_point.shape[1] != X.shape[1]:
        if new_point.shape[1] < X.shape[1]:
            padding = np.zeros((1, X.shape[1] - new_point.shape[1]))
            new_point = np.hstack([new_point, padding])
        else:
            new_point = new_point[:, :X.shape[1]]

    return X.values, np.array(y, dtype=float), new_point, feature_cols

def build_mlp_model(input_dim, hidden_layers_config):
    if not hidden_layers_config or len(hidden_layers_config) == 0:
        raise ValueError("At least one hidden layer is required")
    
    model = MLP(input_dim, hidden_layers_config)
    return model

def train_mlp_with_validation(X_train, y_train, X_val, y_val, hidden_layers_config, 
                               epochs, learning_rate, batch_size, optimizer_name, 
                               reg_type, reg_rate, device='cpu'):
    scaler_X = StandardScaler()
    scaler_y = StandardScaler()
    
    X_train_norm = scaler_X.fit_transform(X_train)
    X_val_norm = scaler_X.transform(X_val)
    
    y_train_norm = scaler_y.fit_transform(y_train.reshape(-1, 1)).flatten()
    y_val_norm = scaler_y.transform(y_val.reshape(-1, 1)).flatten()
    
    input_dim = X_train_norm.shape[1]
    model = build_mlp_model(input_dim, hidden_layers_config)
    model = model.to(device)
    
    if batch_size is None or batch_size <= 0:
        batch_size = len(X_train_norm)
    
    criterion = nn.MSELoss()
    
    if optimizer_name.lower() == 'adam':
        optimizer = optim.Adam(model.parameters(), lr=learning_rate, weight_decay=reg_rate if reg_type == 'l2' else 0)
    elif optimizer_name.lower() == 'sgd':
        optimizer = optim.SGD(model.parameters(), lr=learning_rate, weight_decay=reg_rate if reg_type == 'l2' else 0)
    elif optimizer_name.lower() == 'rmsprop':
        optimizer = optim.RMSprop(model.parameters(), lr=learning_rate, weight_decay=reg_rate if reg_type == 'l2' else 0)
    else:
        optimizer = optim.Adam(model.parameters(), lr=learning_rate)
    
    if reg_type == 'l1' and reg_rate > 0:
        l1_reg = lambda: sum(p.abs().sum() for p in model.parameters())
    else:
        l1_reg = None
    
    X_train_tensor = torch.FloatTensor(X_train_norm).to(device)
    y_train_tensor = torch.FloatTensor(y_train_norm.reshape(-1, 1)).to(device)
    X_val_tensor = torch.FloatTensor(X_val_norm).to(device)
    y_val_tensor = torch.FloatTensor(y_val_norm.reshape(-1, 1)).to(device)
    
    train_losses = []
    val_losses = []
    train_maes = []
    val_maes = []
    train_r2s = []
    val_r2s = []
    
    with _model_lock:
        for epoch in range(epochs):
            model.train()
            train_loss = 0.0
            
            indices = torch.randperm(len(X_train_tensor))
            for i in range(0, len(X_train_tensor), batch_size):
                batch_indices = indices[i:i+batch_size]
                X_batch = X_train_tensor[batch_indices]
                y_batch = y_train_tensor[batch_indices]
                
                optimizer.zero_grad()
                outputs = model(X_batch)
                loss = criterion(outputs, y_batch)
                
                if l1_reg:
                    loss = loss + reg_rate * l1_reg()
                
                loss.backward()
                optimizer.step()
                
                train_loss += loss.item()
            
            model.eval()
            with torch.no_grad():
                train_outputs = model(X_train_tensor)
                val_outputs = model(X_val_tensor)
                
                train_loss_norm = criterion(train_outputs, y_train_tensor).item()
                val_loss_norm = criterion(val_outputs, y_val_tensor).item()
                
                train_pred_denorm = scaler_y.inverse_transform(train_outputs.cpu().numpy()).flatten()
                val_pred_denorm = scaler_y.inverse_transform(val_outputs.cpu().numpy()).flatten()
                
                train_mae = mean_absolute_error(y_train, train_pred_denorm)
                val_mae = mean_absolute_error(y_val, val_pred_denorm)
                
                train_r2 = r2_score(y_train, train_pred_denorm)
                val_r2 = r2_score(y_val, val_pred_denorm)
            
            train_losses.append(train_loss_norm)
            val_losses.append(val_loss_norm)
            train_maes.append(train_mae)
            val_maes.append(val_mae)
            train_r2s.append(train_r2)
            val_r2s.append(val_r2)
    
    model = model.cpu()
    return model, scaler_X, scaler_y, train_losses, val_losses, train_maes, val_maes, train_r2s, val_r2s

def create_training_loss_chart(train_losses, train_maes):
    if not train_losses or len(train_losses) == 0:
        return None
    
    epochs = list(range(1, len(train_losses) + 1))
    valid_losses = [loss if not (np.isinf(loss) or np.isnan(loss)) else None for loss in train_losses]
    
    fig = make_subplots(
        rows=2, cols=1,
        subplot_titles=("Training Loss (MSE)", "Training MAE"),
        vertical_spacing=0.15,
        row_heights=[0.5, 0.5]
    )
    
    fig.add_trace(
        go.Scatter(
            x=epochs,
            y=valid_losses,
            mode='lines+markers',
            name='Training Loss (MSE)',
            line=dict(color='#1976D2', width=3),
            marker=dict(size=6),
            showlegend=True
        ),
        row=1, col=1
    )
    
    if train_maes and len(train_maes) == len(train_losses):
        valid_maes = [mae if not (np.isinf(mae) or np.isnan(mae)) else None for mae in train_maes]
        fig.add_trace(
            go.Scatter(
                x=epochs,
                y=valid_maes,
                mode='lines+markers',
                name='Training MAE',
                line=dict(color='#42A5F5', width=3),
                marker=dict(size=6),
                showlegend=True
            ),
            row=2, col=1
        )
    
    fig.update_xaxes(title_text="Epoch", row=1, col=1, showgrid=True, gridwidth=1, gridcolor='lightgray')
    fig.update_yaxes(title_text="MSE", row=1, col=1, showgrid=True, gridwidth=1, gridcolor='lightgray')
    fig.update_xaxes(title_text="Epoch", row=2, col=1, showgrid=True, gridwidth=1, gridcolor='lightgray')
    fig.update_yaxes(title_text="MAE", row=2, col=1, showgrid=True, gridwidth=1, gridcolor='lightgray')
    
    fig.update_layout(
        title="Training Metrics Over Epochs",
        plot_bgcolor="white",
        height=600,
        margin=dict(l=40, r=40, t=80, b=40)
    )
    
    return fig

def create_validation_loss_chart(val_losses, val_maes):
    if not val_losses or len(val_losses) == 0:
        return None
    
    epochs = list(range(1, len(val_losses) + 1))
    valid_losses = [loss if not (np.isinf(loss) or np.isnan(loss)) else None for loss in val_losses]
    
    fig = make_subplots(
        rows=2, cols=1,
        subplot_titles=("Validation Loss (MSE)", "Validation MAE"),
        vertical_spacing=0.15,
        row_heights=[0.5, 0.5]
    )
    
    fig.add_trace(
        go.Scatter(
            x=epochs,
            y=valid_losses,
            mode='lines+markers',
            name='Validation Loss (MSE)',
            line=dict(color='#7B1FA2', width=3),
            marker=dict(size=6),
            showlegend=True
        ),
        row=1, col=1
    )
    
    if val_maes and len(val_maes) == len(val_losses):
        valid_maes = [mae if not (np.isinf(mae) or np.isnan(mae)) else None for mae in val_maes]
        fig.add_trace(
            go.Scatter(
                x=epochs,
                y=valid_maes,
                mode='lines+markers',
                name='Validation MAE',
                line=dict(color='#BA68C8', width=3),
                marker=dict(size=6),
                showlegend=True
            ),
            row=2, col=1
        )
    
    fig.update_xaxes(title_text="Epoch", row=1, col=1, showgrid=True, gridwidth=1, gridcolor='lightgray')
    fig.update_yaxes(title_text="MSE", row=1, col=1, showgrid=True, gridwidth=1, gridcolor='lightgray')
    fig.update_xaxes(title_text="Epoch", row=2, col=1, showgrid=True, gridwidth=1, gridcolor='lightgray')
    fig.update_yaxes(title_text="MAE", row=2, col=1, showgrid=True, gridwidth=1, gridcolor='lightgray')
    
    fig.update_layout(
        title="Validation Metrics Over Epochs",
        plot_bgcolor="white",
        height=600,
        margin=dict(l=40, r=40, t=80, b=40)
    )
    
    return fig

def create_results_display(model, prediction_value, feature_cols, 
                          epochs, learning_rate, hidden_layers_config,
                          optimizer_name, reg_type, reg_rate, split_info):
    input_dim = len(feature_cols)
    arch_desc = f"{input_dim} → "
    if hidden_layers_config:
        arch_desc += " → ".join([str(layer['neurons']) for layer in hidden_layers_config])
        arch_desc += " → "
    arch_desc += "1"
    
    activations = []
    for layer in hidden_layers_config:
        act = layer.get('activation', 'relu')
        if act.lower() in ['leakyrelu', 'leaky_relu']:
            activations.append('LeakyReLU')
        else:
            activations.append(act.upper())
    activation_desc = ", ".join(activations) if activations else "None"
    
    reg_desc = f"{reg_type.upper()}(λ={reg_rate})" if reg_type != 'none' and reg_rate > 0 else 'None'
    
    total_params = sum(p.numel() for p in model.parameters())
    
    html_content = f"""
    <div style='background:#E3F2FD;border-left:6px solid #1976D2;padding:14px 16px;border-radius:10px;'>
        <strong style='color:#0D47A1;'>🧠 MLP (Multi-Layer Perceptron) Regression Results</strong><br><br>
        
        <div style='margin:8px 0;'>
            <strong style='color:#1976D2;'>🏗️ Model Architecture:</strong><br>
            • Architecture: {arch_desc}<br>
            • Hidden Layers: {len(hidden_layers_config)}<br>
            • Activation Functions: {activation_desc}<br>
            • Output Activation: Linear (Regression)<br>
        </div>
        
        <div style='margin:8px 0;'>
            <strong style='color:#1976D2;'>🔧 Training Configuration:</strong><br>
            • Epochs: {epochs} | Learning Rate: {learning_rate}<br>
            • Optimizer: {optimizer_name.upper()}<br>
            • Batch Size: {split_info.get('batch_size', 'Full Batch')} | Features: {len(feature_cols)}<br>
            • Regularization: {reg_desc}<br>
            • Normalization: Standardized | Loss: Mean Squared Error (MSE)<br>
        </div>
        
        <div style='margin:8px 0;'>
            <strong style='color:#1976D2;'>📊 Data Split:</strong><br>
            • Training: {split_info['train_size']} samples ({split_info['train_ratio']:.1%})<br>
            • Validation: {split_info['val_size']} samples ({split_info['val_ratio']:.1%})<br>
        </div>
        
        <div style='margin:8px 0;'>
            <strong style='color:#1976D2;'>📈 Performance Metrics:</strong><br>
            • Training MSE: <span style='background:#BBDEFB;padding:2px 6px;border-radius:4px;'><strong>{split_info['train_mse']:.4f}</strong></span><br>
            • Validation MSE: <span style='background:#C5CAE9;padding:2px 6px;border-radius:4px;'><strong>{split_info['val_mse']:.4f}</strong></span><br>
            • Training MAE: <span style='background:#BBDEFB;padding:2px 6px;border-radius:4px;'><strong>{split_info['train_mae']:.4f}</strong></span><br>
            • Validation MAE: <span style='background:#C5CAE9;padding:2px 6px;border-radius:4px;'><strong>{split_info['val_mae']:.4f}</strong></span><br>
            • Training R²: <span style='background:#BBDEFB;padding:2px 6px;border-radius:4px;'><strong>{split_info['train_r2']:.4f}</strong></span><br>
            • Validation R²: <span style='background:#C5CAE9;padding:2px 6px;border-radius:4px;'><strong>{split_info['val_r2']:.4f}</strong></span><br>
            • Training Time: <span style='background:#E1BEE7;padding:2px 6px;border-radius:4px;'><strong>{split_info['training_time']:.4f}s</strong></span><br>
        </div>
        
        <div style='margin:8px 0;'>
            <strong style='color:#1976D2;'>🎯 Model Parameters:</strong><br>
            • Total Parameters: <code style='background:#F3E5F5;padding:2px 6px;border-radius:4px;'>{total_params:,}</code><br>
            • Trainable Parameters: {total_params:,}<br>
        </div>
        
        <div style='margin:8px 0;'>
            <strong style='color:#1976D2;'>🔮 Prediction:</strong><br>
            • Predicted Value: <span style='background:#DCEDC8;padding:2px 6px;border-radius:4px;'><strong>{prediction_value:.4f}</strong></span><br>
            <em style='font-size:0.9em;color:#424242;'>* The model predicts a continuous numerical value for the target variable</em><br>
        </div>
    </div>
    """
    
    return html_content

def run_mlp_and_visualize(df, target_col, new_point_dict, hidden_layers_config,
                          epochs, learning_rate, batch_size_str="Full Batch", 
                          train_test_split_ratio=0.8,
                          optimizer_name="adam", reg_type="none", reg_rate=0.001):
    try:
        X, y, new_point, feature_cols = preprocess_data(df, target_col, new_point_dict)
    except Exception as e:
        return None, None, f"<div style='background:#FFEBEE;border-left:6px solid #C62828;padding:14px 16px;border-radius:10px;'><strong>🧠 MLP Error</strong><br><br>❌ Data preprocessing error: {str(e)}</div>", None
    
    if epochs < 1:
        return None, None, f"<div style='background:#FFEBEE;border-left:6px solid #C62828;padding:14px 16px;border-radius:10px;'><strong>🧠 MLP Error</strong><br><br>❌ Number of epochs must be ≥ 1.</div>", None
    if learning_rate <= 0:
        return None, None, f"<div style='background:#FFEBEE;border-left:6px solid #C62828;padding:14px 16px;border-radius:10px;'><strong>🧠 MLP Error</strong><br><br>❌ Learning rate must be > 0.</div>", None
    if len(hidden_layers_config) == 0:
        return None, None, f"<div style='background:#FFEBEE;border-left:6px solid #C62828;padding:14px 16px;border-radius:10px;'><strong>🧠 MLP Error</strong><br><br>❌ At least one hidden layer is required.</div>", None
    
    for i, layer in enumerate(hidden_layers_config):
        if layer.get('neurons', 0) < 1:
            return None, None, f"<div style='background:#FFEBEE;border-left:6px solid #C62828;padding:14px 16px;border-radius:10px;'><strong>🧠 MLP Error</strong><br><br>❌ Layer {i+1} must have at least 1 neuron.</div>", None
    
    test_size = 1.0 - train_test_split_ratio
    X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=test_size, random_state=42)
    
    if batch_size_str == "Full Batch" or batch_size_str is None or batch_size_str == "":
        batch_size = None
    else:
        try:
            batch_size = int(batch_size_str)
            if batch_size <= 0:
                batch_size = None
            if batch_size is not None and batch_size > len(X_train):
                batch_size = len(X_train)
        except (ValueError, TypeError):
            batch_size = None
    
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    
    try:
        start_time = time.time()
        model, scaler_X, scaler_y, train_losses, val_losses, train_maes, val_maes, train_r2s, val_r2s = train_mlp_with_validation(
            X_train, y_train, X_val, y_val, hidden_layers_config,
            epochs, learning_rate, batch_size, optimizer_name, reg_type, reg_rate, device
        )
        training_time = time.time() - start_time
    except Exception as e:
        import traceback
        error_msg = str(e)
        traceback.print_exc()
        return None, None, f"<div style='background:#FFEBEE;border-left:6px solid #C62828;padding:14px 16px;border-radius:10px;'><strong>🧠 MLP Training Error</strong><br><br>❌ {error_msg}</div>", None
    
    _set_current_model(model, (scaler_X, scaler_y))
    
    X_train_norm = scaler_X.transform(X_train)
    X_val_norm = scaler_X.transform(X_val)
    new_point_norm = scaler_X.transform(new_point)
    
    try:
        model.eval()
        with torch.no_grad():
            train_pred_norm = model(torch.FloatTensor(X_train_norm)).numpy()
            val_pred_norm = model(torch.FloatTensor(X_val_norm)).numpy()
            prediction_norm = model(torch.FloatTensor(new_point_norm)).numpy()[0][0]
        
        train_pred = scaler_y.inverse_transform(train_pred_norm).flatten()
        val_pred = scaler_y.inverse_transform(val_pred_norm).flatten()
        prediction_value = float(scaler_y.inverse_transform([[prediction_norm]])[0][0])
    except Exception as e:
        return None, None, f"<div style='background:#FFEBEE;border-left:6px solid #C62828;padding:14px 16px;border-radius:10px;'><strong>🧠 MLP Prediction Error</strong><br><br>❌ {str(e)}</div>", None
    
    train_mse = mean_squared_error(y_train, train_pred)
    val_mse = mean_squared_error(y_val, val_pred)
    train_mae = mean_absolute_error(y_train, train_pred)
    val_mae = mean_absolute_error(y_val, val_pred)
    train_r2 = r2_score(y_train, train_pred)
    val_r2 = r2_score(y_val, val_pred)
    
    final_train_loss = train_losses[-1] if train_losses and len(train_losses) > 0 else 0.0
    final_val_loss = val_losses[-1] if val_losses and len(val_losses) > 0 else 0.0
    final_train_mae = train_maes[-1] if train_maes and len(train_maes) > 0 else 0.0
    final_val_mae = val_maes[-1] if val_maes and len(val_maes) > 0 else 0.0
    
    train_loss_fig = create_training_loss_chart(train_losses, train_maes)
    val_loss_fig = create_validation_loss_chart(val_losses, val_maes)
    
    results_display = create_results_display(
        model, prediction_value, feature_cols, epochs, 
        learning_rate, hidden_layers_config, optimizer_name, 
        reg_type, reg_rate,
        split_info={
            "train_size": len(X_train),
            "val_size": len(X_val),
            "train_ratio": train_test_split_ratio,
            "val_ratio": 1.0 - train_test_split_ratio,
            "train_mse": train_mse,
            "val_mse": val_mse,
            "train_mae": train_mae,
            "val_mae": val_mae,
            "train_r2": train_r2,
            "val_r2": val_r2,
            "batch_size": batch_size_str if batch_size_str else "Full Batch",
            "training_time": training_time
        }
    )
    
    return train_loss_fig, val_loss_fig, results_display, prediction_value

def _get_current_model():
    return _current_model

def _set_current_model(model, scalers):
    global _current_model, _current_scaler
    _current_model = model
    _current_scaler = scalers