import torch


def obtain_samples(flow, y, nsamples, device=None, batch_size=512):
    """Draw samples from the posterior.
    Arguments:
        flow {Flow} -- NSF model
        y {array} -- strain data
        nsamples {int} -- number of samples desired
    Keyword Arguments:
        device {torch.device} -- model device (CPU or GPU) (default: {None})
        batch_size {int} -- batch size for sampling (default: {512})
    Returns:
        Tensor -- samples
    """

    with torch.no_grad():
        flow.eval()

        y = torch.from_numpy(y).unsqueeze(0).to(device)

        num_batches = nsamples // batch_size
        num_leftover = nsamples % batch_size

        samples = [flow.sample(batch_size, y) for _ in range(num_batches)]
        if num_leftover > 0:
            samples.append(flow.sample(num_leftover, y))

        # The batching in the nsf package seems screwed up, so we had to do it
        # ourselves, as above. They are concatenating on the wrong axis.

        # samples = flow.sample(nsamples, context=y, batch_size=batch_size)

        return torch.cat(samples, dim=1)[0]