import torch
import numpy as np

from .dataset import Dataset
from .gaussian_covariates import GaussianCovariates
from ..config.config import PyTorchConfig

class LinearRegressionDataset(Dataset):
    """ An implementation of a dataset generator, where features are generated
    by some multivariate gaussian distribution and output is generated by
    a linear function plus gaussian noise.
    True function will be stored as class member named beta in a form of numpy
    vector. """

    def __init__(self, d, noise_std, covariates_kwargs_dict = {}):
        """ Parameters:
            d Denotes the number of features.
            noise_std The standard deviation of gaussian noise to be added to
                    the true function output.
            covariance_matrix The covariance matrix of the features. If None,
                    identity matrix will be assumed. Should be a numpy array
                    of size d x d if not None. """
        super().__init__()
        self.d = d
        self.noise_std = noise_std
        # Class member beta defines the true regression function. It should be
        # set from outside of the class before generating the datasets.
        # Should be of shape (d, 1).
        self.beta = None
        self.covariates = GaussianCovariates
        self.covariates_kwargs_dict = covariates_kwargs_dict

    def generate_data(self, n, pytorch_config = PyTorchConfig()):
        dtype = pytorch_config.dtype
        device = pytorch_config.device
        beta = torch.tensor(self.beta, dtype = dtype, device = device)
        d = self.d
        X = self.covariates.generate_covariates_matrix(
                n, d, pytorch_config, **self.covariates_kwargs_dict)
        y = torch.mm(X, beta)
        y += self.noise_std * torch.randn(n, 1, dtype = dtype, device = device)
        return X, y

