import numpy as np
import math

from dataset_specifications.dataset import Dataset

class SwirlsSet(Dataset):
    def __init__(self):
        super().__init__()
        self.name = "swirls"

        self.n_samples = {
            "train": 2000,
            "val": 1000,
            "test": 1000,
        }

        self.y_dim = 2

        # 2D heteroskedastic Gaussian mixture model with 2 components

    def sample_ys(self, xs):
        n = xs.shape[0]
        components = np.random.randint(2, size=n) # uniform 0,1
        angles = math.pi*components + (math.pi/2.)*xs[:,0] # Angles to centers
        means = np.stack((np.cos(angles), np.sin(angles)), axis=1)

        noise = np.random.randn(n, 2) # samples form 2d gaussian

        std = 0.3 - 0.2*np.abs(xs-1.)
        ys = means + std*noise
        return ys

    def sample(self, n):
        xs = np.random.uniform(low=0., high=2., size=(n,1))
        ys = self.sample_ys(xs)
        return np.concatenate((xs, ys), axis=1)



[((23, 21, 23, 49), 'numpy.random.randint', 'np.random.randint', (), '', True, 'import numpy as np\n'), ((27, 16, 27, 37), 'numpy.random.randn', 'np.random.randn', ({(27, 32, 27, 33): 'n', (27, 35, 27, 36): '2'}, {}), '(n, 2)', True, 'import numpy as np\n'), ((34, 13, 34, 59), 'numpy.random.uniform', 'np.random.uniform', (), '', True, 'import numpy as np\n'), ((36, 15, 36, 47), 'numpy.concatenate', 'np.concatenate', (), '', True, 'import numpy as np\n'), ((25, 26, 25, 40), 'numpy.cos', 'np.cos', ({(25, 33, 25, 39): 'angles'}, {}), '(angles)', True, 'import numpy as np\n'), ((25, 42, 25, 56), 'numpy.sin', 'np.sin', ({(25, 49, 25, 55): 'angles'}, {}), '(angles)', True, 'import numpy as np\n'), ((29, 24, 29, 37), 'numpy.abs', 'np.abs', ({(29, 31, 29, 36): '(xs - 1.0)'}, {}), '(xs - 1.0)', True, 'import numpy as np\n')]