import math
import sys
from PIL import Image

import numpy as np
import torch
from torch.utils.data import Dataset
from sklearn.gaussian_process.kernels import RBF
from sklearn.gaussian_process import GaussianProcessRegressor

sys.path.append("..")

from skssl.utils.helpers import rescale_range

__all__ = ["SineDataset", "GPDataset"]


class GPDataset(Dataset):
    """
    Dataset of functions generated by a gaussian process.

    Parameters
    ----------
    kernel : sklearn.gaussian_process.kernels
        The kernel specifying the covariance function of the GP. If None is
        passed, the kernel "1.0 * RBF(1.0)" is used as default.

    n_samples : int, optional
        Number of sampled functios contained in dataset.

    n_points : int, optional
        Number of points at which to evaluate f(x) for x in min_max.

    min_max : tuple of floats, optional
        Min and max point at which to evaluate the function (bounds).
    """

    def __init__(self,
                 kernel=1. * RBF(length_scale=1.),
                 min_max=(-5, 5),
                 n_samples=1000,
                 n_points=100):

        self.n_samples = n_samples
        self.n_points = n_points
        self.min_max = min_max
        self.gp = GaussianProcessRegressor(kernel=kernel)
        self.data, self.targets = self.precompute_data()

    def __len__(self):
        return self.n_samples

    def __getitem__(self, index):
        # doesn't use index because randomly gnerated in any case => sample
        # in order which enables to know when epoch is finished and regenerate
        # new functions
        self.counter += 1
        if self.counter == self.n_samples:
            self.data, self.targets = self.precompute_data()
        return self.data[self.counter], self.targets[self.counter]

    def precompute_data(self, min_max=None, n_samples=None, ):
        self.counter = 0
        return self._precompute_helper(self.min_max, self.n_samples, self.n_points)

    def _precompute_helper(self, min_max, n_samples, n_points):
        # sample from a grid
        X = np.linspace(*min_max, n_points)
        # add noise (with standard deviation of "stepsize") to not be on a grid
        X += np.random.randn(*X.shape) * (min_max[1] - min_max[0]) / n_points
        # make sure that still in bound
        X = X.clip(min=min_max[0], max=min_max[1])
        # sort which is convenient for plotting
        X.sort()

        targets = self.gp.sample_y(X[:, np.newaxis], n_samples).transpose(1, 0)
        targets = torch.from_numpy(targets)
        targets = targets.view(n_samples, n_points, 1).float()

        X = torch.from_numpy(X)
        X = X.view(1, -1, 1).expand(n_samples, n_points, 1).float()
        # rescale features to [-1,1]
        # uses `self.min_max` like that possible to precompute extrapolation data
        X = rescale_range(X, self.min_max, (-1, 1))

        return X, targets

    def extrapolation_samples(self, n_samples=1, test_min_max=None, n_points=None):
        """Return a batch of extrapolation

        Parameters
        ----------
        n_samples : int, optional
            Number of sampled function (i.e. batch size).

        test_min_max : float, optional
            Testing range. If `None` uses training one.

        n_points : int, optional
            Number of points at which to evaluate f(x) for x in min_max. If None
            uses `self.n_points`.
        """
        if test_min_max is None:
            test_min_max = self.min_max
        n_points = n_points if n_points is not None else self.n_points
        return self._precompute_helper(test_min_max, n_samples, n_points)


class SineDataset(Dataset):
    """
    Dataset of functions f(x) = a * sin(x - b) where a and b are randomly
    sampled.

    Notes
    -----
    - modified from: https://github.com/EmilienDupont/neural-processes/

    Parameters
    ----------
    amplitude_range : tuple of float, optional
        Defines the range from which the amplitude (i.e. a) of the sine function
        is sampled.

    shift_range : tuple of float, optional
        Defines the range from which the shift (i.e. b) of the sine function is
        sampled.

    n_samples : int, optional
        Number of sampled functios contained in dataset.

    n_points : int, optional
        Number of points at which to evaluate f(x) for x in min_max.

    std_noise : float, optional
        Standard deviation of the noise to add

    min_max : tuple of floats, optional
        Min and max point at which to evaluate the function (bounds).
    """

    def __init__(self,
                 amplitude_range=(-1., 1.),
                 shift_range=(-.5, .5),
                 n_samples=1000,
                 n_points=100,
                 std_noise=0.01,
                 min_max=(-math.pi, math.pi)):

        self.amplitude_range = amplitude_range
        self.shift_range = shift_range
        self.n_samples = n_samples
        self.n_points = n_points
        self.x_dim = 1  # x and y dim are fixed for this dataset.
        self.y_dim = 1
        self.std_noise = std_noise
        self.min_max = min_max

        # Generate data
        a_min, a_max = amplitude_range
        b_min, b_max = shift_range
        # Sample random amplitude
        a = (a_max - a_min) * torch.rand(n_samples, 1, self.y_dim) + a_min
        # Sample random shift
        b = (b_max - b_min) * torch.rand(n_samples, 1, self.y_dim) + b_min
        # Shape (n_samples, n_points, x_dim)
        self.data = torch.linspace(*self.min_max, n_points)
        self.data = self.data.view(1, -1, 1).expand(self.n_samples, self.n_points, self.x_dim)
        # Shape (n_samples, n_points, y_dim)
        self.targets = (a * torch.sin(self.data - b)
                        ) + torch.randn_like(self.data) * self.std_noise

        # rescale features to [-1,1]
        self.data = rescale_range(self.data, self.min_max, (-1, 1))

    def __getitem__(self, index):
        return self.data[index], self.targets[index]

    def __len__(self):
        return self.n_samples
