import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.model_selection import train_test_split
import numpy as np
from utils import compute_metrics

def train_model(model, X, y, epochs=20, lr=0.001, test_size=0.2, batch_size=64, verbose=False):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.to(device)

    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=test_size, shuffle=False)
    X_train_tensor = torch.Tensor(X_train).to(device)
    y_train_tensor = torch.Tensor(y_train).to(device)
    X_test_tensor = torch.Tensor(X_test).to(device)
    y_test_tensor = torch.Tensor(y_test).to(device)

    if len(X_train_tensor.shape) == 2:
        X_train_tensor = X_train_tensor.unsqueeze(-1)
        X_test_tensor = X_test_tensor.unsqueeze(-1)

    criterion = nn.MSELoss()
    optimizer = optim.Adam(model.parameters(), lr=lr)

    losses = []
    for epoch in range(epochs):
        model.train()
        optimizer.zero_grad()
        output = model(X_train_tensor)
        loss = criterion(output, y_train_tensor)
        loss.backward()
        optimizer.step()

        model.eval()
        with torch.no_grad():
            val_output = model(X_test_tensor)
            val_loss = criterion(val_output, y_test_tensor)

        losses.append((loss.item(), val_loss.item()))

        if verbose and epoch % 5 == 0:
            print(f"Epoch {epoch} - Train Loss: {loss.item():.4f}, Test Loss: {val_loss.item():.4f}")

    model.eval()
    with torch.no_grad():
        preds = model(X_test_tensor).cpu().numpy()
        true_vals = y_test_tensor.cpu().numpy()

    metrics = compute_metrics(true_vals, preds)
    return model, metrics, preds, true_vals, losses