import torch
from torch.distributions import Distribution, TransformedDistribution, AffineTransform


class ImplicitQuantile(Distribution):
    arg_constraints = {}
    
    def __init__(
        self,
        implicit_quantile_function,
        taus,
        nn_output,
        predicted_quantiles,
        validate_args=None,
    ):
        self.predicted_quantiles = predicted_quantiles[0]
        self.taus = taus
        self.quantile_function = implicit_quantile_function
        self.input_data = nn_output

        super(ImplicitQuantile, self).__init__(
            batch_shape=self.predicted_quantiles.shape, validate_args=validate_args
        )

    @torch.no_grad()
    def sample(self, sample_shape=torch.Size()):
        """See arXiv: 1806.06923
        Once the model has learned how to predict a given quantile tau, one can sample from the
        distribution of the target, by sampling tau values.
        """
        if len(sample_shape) == 0:
            num_parallel_samples = 1
        else:
            num_parallel_samples = sample_shape[0]
        input_data = torch.repeat_interleave(
            self.input_data, repeats=num_parallel_samples, dim=0
        )
        batch_size = input_data.shape[0]
        forecast_length = input_data.shape[1]
        device = input_data.device

        taus = torch.rand(size=(batch_size, forecast_length), device=device)
        samples = self.quantile_function(input_data, taus)
        if len(sample_shape) == 0:
            return samples
        else:
            return samples.reshape((num_parallel_samples, -1, forecast_length))

    def log_prob(self, value):
        # Assumes same distribution for all steps in the future, conditionally on the input data
        return -self.quantile_loss(self.predicted_quantiles, value, self.taus)

    @staticmethod
    def quantile_loss(quantile_forecast, target, tau):
        return torch.abs(
            (quantile_forecast - target) * ((target <= quantile_forecast).float() - tau)
        )


class TransformedImplicitQuantile(TransformedDistribution):
    def __init__(self, base_distribution, transforms):
        super().__init__(base_distribution, transforms)

    def log_prob(self, x):
        scale = 1.0
        for transform in reversed(self.transforms):
            assert isinstance(transform, AffineTransform), "Not an AffineTransform"
            x = transform.inv(x)
            scale *= transform.scale
        p = self.base_dist.log_prob(x)
        return p * scale