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)



[((476, 504), 'numpy.random.randint', 'np.random.randint', (['(2)'], {'size': 'n'}), '(2, size=n)\n', (493, 504), True, 'import numpy as np\n'), ((682, 703), 'numpy.random.randn', 'np.random.randn', (['n', '(2)'], {}), '(n, 2)\n', (697, 703), True, 'import numpy as np\n'), ((858, 907), 'numpy.random.uniform', 'np.random.uniform', ([], {'low': '(0.0)', 'high': '(2.0)', 'size': '(n, 1)'}), '(low=0.0, high=2.0, size=(n, 1))\n', (875, 907), True, 'import numpy as np\n'), ((952, 984), 'numpy.concatenate', 'np.concatenate', (['(xs, ys)'], {'axis': '(1)'}), '((xs, ys), axis=1)\n', (966, 984), True, 'import numpy as np\n'), ((624, 638), 'numpy.cos', 'np.cos', (['angles'], {}), '(angles)\n', (630, 638), True, 'import numpy as np\n'), ((640, 654), 'numpy.sin', 'np.sin', (['angles'], {}), '(angles)\n', (646, 654), True, 'import numpy as np\n'), ((756, 772), 'numpy.abs', 'np.abs', (['(xs - 1.0)'], {}), '(xs - 1.0)\n', (762, 772), True, 'import numpy as np\n')]