import numpy as np
import torch
from scipy.stats import multivariate_normal

from torchvision import datasets, transforms
from torch.utils.data import DataLoader

def generate_training_data(n_train_data, data_dim, task, task_param, seed_data):
    """
    Generate training data for a specified task ('xor', 'ising', or 'mnist').

    Args:
        n_train_data (int): Number of training samples to generate.
        data_dim (int): Dimensionality of each input sample.
        task (str): Task type, must be one of ['xor', 'ising', 'mnist'].
        task_param (float): Task-specific parameter (used differently depending on task).
        seed_data (int): Random seed for reproducibility.

    Returns:
        tuple:
            - X (np.ndarray): Input data matrix of shape (n_train_data, data_dim).
            - Y (np.ndarray): Binary labels, shape (n_train_data,).
    
    Raises:
        ValueError: If the specified task is not recognized.
    """
    if task == 'xor':
        X, Y = generate_training_data_xor(n_train_data, data_dim, task_param, seed_data)
    elif task == 'ising':
        X, Y = generate_training_data_ising(n_train_data, data_dim, task_param, seed_data)
    elif task == 'mnist':
        X, Y = generate_training_data_mnist(n_train_data, data_dim, task_param, seed_data)
    else:
        raise ValueError("Task must be either ising, xor or mnist.")
    return X, Y

def generate_training_data_ising(n_train_data, data_dim, delta_p, seed_data):
    """
    Generate training data for a synthetic binary classification task based on random Ising spin vectors.

    Args:
        n_train_data (int): Number of training samples to generate (should be even).
        data_dim (int): Dimensionality of each input sample.
        delta_p (float): Class separation parameter (larger means easier classification).
        seed_data (int): Random seed for reproducibility.

    Returns:
        tuple:
            - X (np.ndarray): Input data matrix of shape (n_train_data, data_dim).
            - Y (np.ndarray): Binary labels, shape (n_train_data,).
    """
    np.random.seed(seed_data)
    X1 = 2*(np.random.rand(int(n_train_data/2), data_dim) > 0.5 + delta_p) - 1.0
    X2 = 2*(np.random.rand(int(n_train_data/2), data_dim) > 0.5 - delta_p) - 1.0

    # setting of binary classification
    # labels +-1 for the two classes, examples presented
    # sorted by labels
    X = np.vstack ((X1, X2))
    Y = np.hstack ((np.ones(int(n_train_data/2)), -np.ones(int(n_train_data/2))))

    return X, Y


def generate_training_data_xor(n_train_data, data_dim, sigma, seed_data):
    """
    Generate high-dimensional XOR data for binary classification.
    See http://proceedings.mlr.press/v139/refinetti21b/refinetti21b.pdf.

    Based on four Gaussian clusters centered at ±e1 and ±e2 (first two dimensions),
    forming an XOR pattern.

    Args:
        n_train_data (int): Number of training samples to generate (should be divisible by 4).
        data_dim (int): Dimensionality of each input sample.
        sigma (float): Standard deviation of Gaussian noise added to the clusters.
        seed_data (int): Random seed for reproducibility.

    Returns:
        tuple:
            - X (np.ndarray): Input data matrix of shape (n_train_data, data_dim).
            - Y (np.ndarray): Binary labels, shape (n_train_data,).
    """
    rng = np.random.default_rng(seed_data)
    means = np.eye(data_dim)[:2]
    cov = sigma**2 * np.eye(data_dim)

    X1 = rng.multivariate_normal(mean=means[0], cov=cov, size=int(n_train_data / 4))
    X2 = rng.multivariate_normal(mean=-means[0], cov=cov, size=int(n_train_data / 4))
    X3 = rng.multivariate_normal(mean=means[1], cov=cov, size=int(n_train_data / 4))
    X4 = rng.multivariate_normal(mean=-means[1], cov=cov, size=int(n_train_data / 4))

    X = np.vstack ( (X1, X2, X3, X4) )
    Y = np.hstack ((np.ones(int(n_train_data/2)), -np.ones(int(n_train_data / 2))))

    return X, Y

def generate_training_data_mnist(n_train_data, data_dim, task_param, seed_data):
    """
    Generate binary classification data from the MNIST dataset, using digits 0 and 3.

    Digit 0 is labeled -1, digit 3 is labeled +1.

    Args:
        n_train_data (int): Number of training samples to generate (even, split between classes).
        data_dim (int): Dimensionality of each input sample (here 784).
        task_param (any): Unused (placeholder for interface consistency).
        seed_data (int): Random seed for reproducibility.

    Returns:
        tuple:
            - X (np.ndarray): Flattened MNIST images of shape (n_train_data, data_dim).
            - Y (np.ndarray): Binary labels, shape (n_train_data,).
    """
    torch.manual_seed(seed_data)

    transform = transforms.Compose([transforms.ToTensor(), transforms.Lambda(lambda x: torch.flatten(x))])

    mnist_train_digit0 = datasets.MNIST("./data", train=True, download=True, transform=transform)
    train_idx = mnist_train_digit0.targets == 0
    mnist_train_digit0.data = mnist_train_digit0.data[train_idx]
    mnist_train_digit0.targets = mnist_train_digit0.targets[train_idx]

    mnist_train_digit1 = datasets.MNIST("./data", train=True, download=True, transform=transform)
    train_idx = mnist_train_digit1.targets == 3
    mnist_train_digit1.data = mnist_train_digit1.data[train_idx]
    mnist_train_digit1.targets = mnist_train_digit1.targets[train_idx]

    train_loader_digit0 = DataLoader(mnist_train_digit0, batch_size=int(n_train_data/2), shuffle=True)
    train_loader_digit1 = DataLoader(mnist_train_digit1, batch_size=int(n_train_data/2), shuffle=True)

    X0 = (next(iter(train_loader_digit0))[0]).numpy()
    X1 = (next(iter(train_loader_digit1))[0]).numpy()

    X = np.vstack((X1, X0))
    Y = np.hstack((np.ones(int(n_train_data / 2)), -np.ones(int(n_train_data / 2))))

    return X, Y