import numpy as np
import sklearn.datasets
import torch

# if torch.backends.mps.is_available():
#     device = torch.device("mps")
#     x = torch.ones(1, device=device)
#     print (x)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
torch.manual_seed(0)
torch.set_default_dtype(torch.float32)

### UTILITY CODE FOR SAMPLING FROM TARGET DISTRIBUTIONS ###


class EmpiricalDist:
    """Empirical distribution object.

    Stores empirical data as Pytorch Tensor and has sample() method.
    """

    def __init__(self, data, device=device):
        self.data = data.to(device)
        self.n_points = data.shape[0]
        self.dims = data.shape[1]
        self.device = device

    def sample(self, n_samples):
        """Generate n_samples drawn with replacement from self.data."""
        if (
            n_samples == self.data.shape[0]
        ):  # If n_samples = size of dataset, just return the whole dataset deterministically
            return self.data
        else:
            indices = torch.randint(high=self.n_points, size=(n_samples,)).to(
                self.device
            )
            return self.data[indices].float().to(self.device)
        
class IsotropicGaussianDist:
    """Isotropic Gaussian distribution on R^D with mean 0 and fixed std."""

    def __init__(self, std, dim, device=device):
        #self.dist = torch.distributions.MultivariateNormal(
        #    torch.zeros(dim, device=device), std * torch.eye(dim, device=device)
        #)
        self.std = std
        self.dim = dim

    def sample(self, n_samples):
        #return self.dist.sample((n_samples,))
        return torch.randn(n_samples, self.dim, device=device) * self.std


class Toy2DDist:
    """FFJORD 2D toy distribution in the format of a Pytorch Distribution."""

    def __init__(self, dist_name, device=device):
        self.dist_name = dist_name
        self.device = device

    def sample(self, n_samples):
        return (
            torch.from_numpy(
                inf_train_gen(self.dist_name, rng=None, batch_size=n_samples)
            )
            .float()
            .to(device)
        )


# Dataset iterator for FFJORD toy 2D data
def inf_train_gen(data, rng=None, batch_size=200):
    if rng is None:
        rng = np.random.RandomState()

    if data == "swissroll":
        data = sklearn.datasets.make_swiss_roll(n_samples=batch_size, noise=1.0)[0]
        data = data.astype("float32")[:, [0, 2]]
        data /= 5
        return data
        
    elif data == "helix":
        # Generate batch_size samples from a helix embedded in 3D

        #t = np.linspace(0, 2, batch_size, endpoint=True)
        t = np.random.rand(batch_size) * 2
        data = np.zeros((batch_size, 3))
        data[:, 0] = np.cos(2*np.pi*t)
        data[:, 1] = np.sin(2*np.pi*t)
        data[:, 2] = 2 * t
        data += np.random.randn(batch_size, 3) * 0.025
        return data

    elif data == "circles":
        data = sklearn.datasets.make_circles(
            n_samples=batch_size, factor=0.5, noise=0.08
        )[0]
        data = data.astype("float32")
        data *= 3
        return data

    elif data == "rings":
        n_samples4 = n_samples3 = n_samples2 = batch_size // 4
        n_samples1 = batch_size - n_samples4 - n_samples3 - n_samples2

        # so as not to have the first point = last point, we set endpoint=False
        linspace4 = np.linspace(0, 2 * np.pi, n_samples4, endpoint=False)
        linspace3 = np.linspace(0, 2 * np.pi, n_samples3, endpoint=False)
        linspace2 = np.linspace(0, 2 * np.pi, n_samples2, endpoint=False)
        linspace1 = np.linspace(0, 2 * np.pi, n_samples1, endpoint=False)

        circ4_x = np.cos(linspace4)
        circ4_y = np.sin(linspace4)
        circ3_x = np.cos(linspace4) * 0.75
        circ3_y = np.sin(linspace3) * 0.75
        circ2_x = np.cos(linspace2) * 0.5
        circ2_y = np.sin(linspace2) * 0.5
        circ1_x = np.cos(linspace1) * 0.25
        circ1_y = np.sin(linspace1) * 0.25

        X = (
            np.vstack(
                [
                    np.hstack([circ4_x, circ3_x, circ2_x, circ1_x]),
                    np.hstack([circ4_y, circ3_y, circ2_y, circ1_y]),
                ]
            ).T
            * 3.0
        )
        # X = util_shuffle(X, random_state=rng) # NEED TO ADD UTIL_SHUFFLE

        # Add noise
        X = X + rng.normal(scale=0.08, size=X.shape)

        return X.astype("float32")

    elif data == "moons":
        data = sklearn.datasets.make_moons(n_samples=batch_size, noise=0.1)[0]
        data = data.astype("float32")
        data = data * 2 + np.array([-1, -0.2])
        return data

    elif data == "8gaussians":
        scale = 4.0
        centers = [
            (1, 0),
            (-1, 0),
            (0, 1),
            (0, -1),
            (1.0 / np.sqrt(2), 1.0 / np.sqrt(2)),
            (1.0 / np.sqrt(2), -1.0 / np.sqrt(2)),
            (-1.0 / np.sqrt(2), 1.0 / np.sqrt(2)),
            (-1.0 / np.sqrt(2), -1.0 / np.sqrt(2)),
        ]
        centers = [(scale * x, scale * y) for x, y in centers]

        dataset = []
        for i in range(batch_size):
            point = rng.randn(2) * 0.5
            idx = rng.randint(8)
            center = centers[idx]
            point[0] += center[0]
            point[1] += center[1]
            dataset.append(point)
        dataset = np.array(dataset, dtype="float32")
        dataset /= 1.414
        return dataset

    elif data == "pinwheel":
        radial_std = 0.3
        tangential_std = 0.1
        num_classes = 5
        num_per_class = batch_size // 5
        rate = 0.25
        rads = np.linspace(0, 2 * np.pi, num_classes, endpoint=False)

        features = rng.randn(num_classes * num_per_class, 2) * np.array(
            [radial_std, tangential_std]
        )
        features[:, 0] += 1.0
        labels = np.repeat(np.arange(num_classes), num_per_class)

        angles = rads[labels] + rate * np.exp(features[:, 0])
        rotations = np.stack(
            [np.cos(angles), -np.sin(angles), np.sin(angles), np.cos(angles)]
        )
        rotations = np.reshape(rotations.T, (-1, 2, 2))

        return 2 * rng.permutation(np.einsum("ti,tij->tj", features, rotations))

    elif data == "2spirals":
        n = np.sqrt(np.random.rand(batch_size // 2, 1)) * 540 * (2 * np.pi) / 360
        d1x = -np.cos(n) * n + np.random.rand(batch_size // 2, 1) * 0.5
        d1y = np.sin(n) * n + np.random.rand(batch_size // 2, 1) * 0.5
        x = np.vstack((np.hstack((d1x, d1y)), np.hstack((-d1x, -d1y)))) / 3
        x += np.random.randn(*x.shape) * 0.1
        return x

    elif data == "checkerboard":
        x1 = np.random.rand(batch_size) * 4 - 2
        x2_ = np.random.rand(batch_size) - np.random.randint(0, 2, batch_size) * 2
        x2 = x2_ + (np.floor(x1) % 2)
        return np.concatenate([x1[:, None], x2[:, None]], 1) * 2

    elif data == "line":
        x = rng.rand(batch_size) * 5 - 2.5
        y = x
        return np.stack((x, y), 1)
    elif data == "cos":
        x = rng.rand(batch_size) * 5 - 2.5
        y = np.sin(x) * 2.5
        return np.stack((x, y), 1)
    else:
        return inf_train_gen("8gaussians", rng, batch_size)