import torch
import numpy as np
import random
import os


# =========================
# 🔁 Reproducibility
# =========================
def set_seed(seed: int = 42):
    """
    Set random seed for reproducibility
    """
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

    # đảm bảo deterministic
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False


# =========================
# 📊 Device management
# =========================
def get_device():
    """
    Get available device (cuda or cpu)
    """
    return torch.device("cuda" if torch.cuda.is_available() else "cpu")


def to_device(data, device):
    """
    Move tensor or model to device
    """
    if isinstance(data, (list, tuple)):
        return [to_device(x, device) for x in data]
    return data.to(device)


# =========================
# 💾 Model utilities
# =========================
def save_model(model, path: str):
    """
    Save model weights
    """
    os.makedirs(os.path.dirname(path), exist_ok=True)
    torch.save(model.state_dict(), path)


def load_model(model, path: str, device=None):
    """
    Load model weights
    """
    if device is None:
        device = get_device()

    state_dict = torch.load(path, map_location=device)
    model.load_state_dict(state_dict)
    return model


# =========================
# 📈 Metrics
# =========================
def accuracy(y_pred, y_true):
    """
    Compute classification accuracy
    """
    preds = torch.argmax(y_pred, dim=1)
    correct = (preds == y_true).sum().item()
    return correct / len(y_true)


def rmse(y_pred, y_true):
    """
    Root Mean Squared Error
    """
    return torch.sqrt(torch.mean((y_pred - y_true) ** 2)).item()


# =========================
# 🧮 Parameter counter
# =========================
def count_parameters(model):
    """
    Count trainable parameters
    """
    return sum(p.numel() for p in model.parameters() if p.requires_grad)


# =========================
# 📦 Mini DataLoader helper
# =========================
def batchify(X, y, batch_size=32, shuffle=True):
    """
    Simple batch generator (without torch DataLoader)
    """
    indices = list(range(len(X)))

    if shuffle:
        random.shuffle(indices)

    for i in range(0, len(X), batch_size):
        batch_idx = indices[i:i + batch_size]
        yield X[batch_idx], y[batch_idx]


# =========================
# 🔍 Debug helper
# =========================
def print_model_summary(model):
    """
    Print model summary (basic)
    """
    print("Model Summary:")
    print(model)
    print(f"Total parameters: {count_parameters(model):,}")