core/train_eval.py

import numpy as np
import pandas as pd
import torch
from torch import nn, optim
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score, explained_variance_score
from torch.utils.data import DataLoader, TensorDataset
import matplotlib.pyplot as plt
import os


def create_sequences(data, window_size, horizon=1):
    X, y = [], []
    for i in range(len(data) - window_size - horizon + 1):
        X.append(data[i:i + window_size])
        y.append(data[i + window_size:i + window_size + horizon].flatten())
    return np.array(X), np.array(y)


def mean_absolute_percentage_error(y_true, y_pred):
    """Calculate MAPE, avoiding division by zero."""
    y_true, y_pred = np.array(y_true), np.array(y_pred)
    non_zero = np.abs(y_true) > 0
    if np.sum(non_zero) == 0:
        return np.nan
    return np.mean(np.abs((y_true[non_zero] - y_pred[non_zero]) / y_true[non_zero])) * 100


def mean_absolute_scaled_error(y_true, y_pred, y_train):
    """Calculate MASE, using naive forecast as denominator."""
    y_true, y_pred = np.array(y_true), np.array(y_pred)
    errors = np.abs(y_true - y_pred)
    naive_errors = np.abs(y_train[1:] - y_train[:-1])
    mean_naive_error = np.mean(naive_errors) if len(naive_errors) > 0 else 1.0
    return np.mean(errors) / mean_naive_error if mean_naive_error != 0 else np.nan


def mean_directional_accuracy(y_true, y_pred):
    """Calculate MDA: percentage of correct direction predictions."""
    y_true, y_pred = np.array(y_true), np.array(y_pred)
    if len(y_true) < 2:
        return np.nan
    true_diff = np.sign(y_true[1:] - y_true[:-1])
    pred_diff = np.sign(y_pred[1:] - y_pred[:-1])
    correct = np.sum(true_diff == pred_diff)
    return (correct / (len(y_true) - 1)) * 100


def train_and_evaluate(
    df,
    future_df,
    model_cls,
    horizon=1,
    hidden=64,
    layers=1,
    epochs=50,
    lr=0.001,
    beta1=0.9,
    beta2=0.999,
    weight_decay=0.01,
    dropout=0.2,
    window=30,
    test_split=0.2,
    scheduler_factor=0.5,
    device="cuda" if torch.cuda.is_available() else "cpu",
    verbose=True
):
    result = {}

    original_values = df['value'].values.astype(np.float32)
    scaler = StandardScaler()
    scaled_data = scaler.fit_transform(original_values.reshape(-1, 1))
    X, y = create_sequences(scaled_data, window, horizon)

    print(f"X shape: {X.shape}, y shape: {y.shape}")

    split = int(len(X) * (1 - test_split))
    val_split = int(split * 0.9)
    X_train, X_val, X_test = X[:val_split], X[val_split:split], X[split:]
    y_train, y_val, y_test = y[:val_split], y[val_split:split], y[split:]

    print(f"X_train shape: {X_train.shape}, y_train shape: {y_train.shape}")
    print(f"X_val shape: {X_val.shape}, y_val shape: {y_val.shape}")
    print(f"X_test shape: {X_test.shape}, y_test shape: {y_test.shape}")

    batch_size = 32
    X_train_tensor = torch.tensor(X_train, dtype=torch.float32)
    y_train_tensor = torch.tensor(y_train, dtype=torch.float32)
    X_val_tensor = torch.tensor(X_val, dtype=torch.float32)
    y_val_tensor = torch.tensor(y_val, dtype=torch.float32)
    X_test_tensor = torch.tensor(X_test, dtype=torch.float32)
    y_test_tensor = torch.tensor(y_test, dtype=torch.float32)

    train_loader = DataLoader(TensorDataset(X_train_tensor, y_train_tensor), batch_size=batch_size, shuffle=True)
    val_loader = DataLoader(TensorDataset(X_val_tensor, y_val_tensor), batch_size=batch_size, shuffle=False)
    test_loader = DataLoader(TensorDataset(X_test_tensor, y_test_tensor), batch_size=batch_size, shuffle=False)

    input_dim = X_train.shape[2] if X_train.ndim == 3 else 1

    # Model-specific initialization
    model_name = model_cls.__name__
    if model_name == "CNNModel":
        model = model_cls(input_size=input_dim, output_size=horizon, num_filters=hidden).to(device)
        architecture_units = hidden  # num_filters for CNN
    elif model_name == "MLPModel":
        model = model_cls(input_size=input_dim * window, hidden_sizes=[hidden, hidden], output_size=horizon, dropout=dropout).to(device)
        architecture_units = [hidden, hidden]  # hidden_sizes for MLP
    elif model_name == "TransformerModel":
        model = model_cls(input_size=input_dim, d_model=hidden, num_layers=layers, output_size=horizon, dropout=dropout).to(device)
        architecture_units = hidden  # d_model for Transformer
    else:
        # RNN, LSTM, GRU, BiLSTM, Hybrid, CNN-GRU
        model = model_cls(input_size=input_dim, hidden_size=hidden, num_layers=layers, output_size=horizon, dropout=dropout).to(device)
        architecture_units = hidden  # hidden_size for other models

    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, betas=(beta1, beta2), weight_decay=weight_decay)
    loss_fn = nn.MSELoss()
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', patience=5, factor=scheduler_factor, threshold=1e-4)

    train_losses = []
    val_losses = []
    best_val_loss = float('inf')
    patience = 5
    counter = 0
    best_model_state = None
    last_lr = lr

    model.train()
    for epoch in range(epochs):
        epoch_loss = 0.0
        for xb, yb in train_loader:
            xb, yb = xb.to(device), yb.to(device)
            optimizer.zero_grad()
            out = model(xb)
            loss = loss_fn(out, yb)
            loss.backward()
            optimizer.step()
            epoch_loss += loss.item()
        train_losses.append(epoch_loss / len(train_loader))

        model.eval()
        val_loss = 0.0
        with torch.no_grad():
            for xb, yb in val_loader:
                xb, yb = xb.to(device), yb.to(device)
                out = model(xb)
                loss = loss_fn(out, yb)
                val_loss += loss.item()
        val_loss /= len(val_loader)
        val_losses.append(val_loss)

        scheduler.step(val_loss)
        current_lr = optimizer.param_groups[0]['lr']
        if current_lr != last_lr and verbose:
            print(f"Epoch {epoch+1}: Learning rate reduced to {current_lr:.6f}")
        last_lr = current_lr

        if verbose and (epoch + 1) % 10 == 0:
            print(f"Epoch {epoch+1}/{epochs} - Train Loss: {train_losses[-1]:.4f}, Val Loss: {val_losses[-1]:.4f}, LR: {current_lr:.6f}")

        if val_loss < best_val_loss:
            best_val_loss = val_loss
            counter = 0
            best_model_state = model.state_dict()
        else:
            counter += 1
            if counter >= patience:
                print(f"Early stopping at epoch {epoch+1}")
                break

    if best_model_state:
        model.load_state_dict(best_model_state)

    result["train_loss"] = train_losses
    result["val_loss"] = val_losses

    model.eval()
    preds, targets = [], []
    with torch.no_grad():
        for xb, yb in test_loader:
            xb = xb.to(device)
            out = model(xb).cpu().numpy()
            preds.append(out)
            targets.append(yb.numpy())

    preds = np.concatenate(preds, axis=0)
    targets = np.concatenate(targets, axis=0)

    print(f"Preds shape: {preds.shape}, Targets shape: {targets.shape}")

    preds_reshaped = preds.reshape(-1, 1)
    targets_reshaped = targets.reshape(-1, 1)
    preds_inv = scaler.inverse_transform(preds_reshaped).reshape(preds.shape)
    targets_inv = scaler.inverse_transform(targets_reshaped).reshape(targets.shape)

    mse = mean_squared_error(targets_inv, preds_inv)
    rmse = np.sqrt(mse)
    mae = mean_absolute_error(targets_inv, preds_inv)
    r2 = r2_score(targets_inv, preds_inv)
    mape = mean_absolute_percentage_error(targets_inv, preds_inv)
    evs = explained_variance_score(targets_inv, preds_inv)
    mase = mean_absolute_scaled_error(targets_inv, preds_inv, original_values[:len(original_values)-horizon])
    mda = mean_directional_accuracy(targets_inv, preds_inv)

    result["metrics"] = {
        "R² (%)": round(r2 * 100, 2),
        "Explained Variance (%)": round(evs * 100, 2),
        "MDA (%)": round(mda, 2) if not np.isnan(mda) else None,
        "RMSE": round(rmse, 2),
        "MAE": round(mae, 2),
        "MAPE (%)": round(mape, 2) if not np.isnan(mape) else None,
        "MASE": round(mase, 2) if not np.isnan(mase) else None
    }

    result["forecast"] = preds_inv
    result["actual"] = targets_inv
    result["predicted"] = result["forecast"]

    latest_window = scaled_data[-window:].reshape(1, window, 1)
    latest_input = torch.tensor(latest_window, dtype=torch.float32).to(device)

    with torch.no_grad():
        future_pred = model(latest_input).cpu().numpy()
    future_pred_reshaped = future_pred.reshape(-1, 1)
    future_pred_inv = scaler.inverse_transform(future_pred_reshaped).reshape(future_pred.shape)

    result["latest_prediction"] = future_pred_inv[0].tolist()

    if not future_df.empty:
        result["future_actuals"] = future_df['value'].values.tolist()[:horizon]

    # Update architecture details with model-specific parameters
    result["architecture"] = {
        "model_name": model_cls.__name__,
        "num_layers": layers,
        "hidden_units": architecture_units,  # Use model-specific units
        "dropout": dropout,
        "batch_size": batch_size,
        "input_size": input_dim,
        "output_size": horizon
    }

    return result