import torch
import math


def generate_spirals(n_samples: int, length: int = 100):
    """

    :param n_samples: number of samples to generate
    :param length: length of the time sampling of the spirals
    :return: X, y data samples. X is a torch.Tensor of shape (n_samples, length, 2), y is of shape (n_samples) with 0/1
    labels
    """
    t = torch.linspace(0., 4 * math.pi, length)

    start = torch.rand(n_samples) * 2 * math.pi
    width = torch.rand(n_samples)
    speed = torch.rand(n_samples)
    x_pos = torch.cos(start.unsqueeze(1) + speed.unsqueeze(1)
                      * t.unsqueeze(0)) / (1 + width.unsqueeze(1) * t)

    x_pos[:(n_samples // 2)] *= -1
    y_pos = torch.sin(start.unsqueeze(1) + speed.unsqueeze(1)
                      * t.unsqueeze(0)) / (1 + width.unsqueeze(1) * t)

    X = torch.stack([x_pos, y_pos], dim=2)
    y = torch.zeros(n_samples, dtype=torch.int64)
    y[:(n_samples // 2)] = 1

    perm = torch.randperm(n_samples)

    X = X[perm]
    y = y[perm]

    return X, y