import torch
import torch.nn.functional as F
import torch.nn as nn

from TargetDistributions.base import BaseTargetDistribution
import copy




class NN(nn.Module):
    """Instantiate a neural net with weights, w ~ prior N(w; 0,1) """
    def __init__(self, x_dim=2, y_dim=2, n_hidden_layers=2, layer_width=10,
                 linear_activations=False, fixed_variance=False, use_bias=True):
        super(NN, self).__init__()
        self.fixed_variance = fixed_variance
        self.use_bias = use_bias
        self.x_dim = x_dim
        self.y_dim = y_dim
        if linear_activations:
            self.activation = lambda x: x
        else:
            self.activation = F.elu
        self.hidden_layers = nn.ModuleList()
        in_dim = x_dim
        for i in range(n_hidden_layers):
            assert layer_width > 0
            hidden_layer = nn.Linear(in_features=in_dim, out_features=layer_width, bias=use_bias)
            nn.init.normal_(hidden_layer.weight, mean=0.0, std=1.0)
            self.hidden_layers.append(hidden_layer)
            in_dim = layer_width

        self.output_layer_means = nn.Linear(in_dim, y_dim, bias=use_bias)
        nn.init.normal_(self.output_layer_means.weight, mean=0.0, std=1.0)
        if fixed_variance is False:
            self.output_layer_log_stds = nn.Linear(in_dim, y_dim, bias=use_bias)
            nn.init.normal_(self.output_layer_log_stds.weight, mean=0.0, std=1.0)
        else:
            self.register_buffer("log_stds", torch.ones(y_dim))


    def forward(self, x):
        return self.sample_posterior_y_given_x(x)

    def sample_posterior_y_given_x(self, x):
        posterior_distribution = self.posterior_y_given_x(x)
        return posterior_distribution.sample()


    def posterior_y_given_x_log_prob(self, y, x):
        posterior_distribution = self.posterior_y_given_x(x)
        if self.y_dim > 1:
            return torch.sum(posterior_distribution.log_prob(y), dim=-1)
        else:
            return posterior_distribution.log_prob(y)


    def posterior_y_given_x(self, x):
        for hidden_layer in self.hidden_layers:
            x = self.activation(hidden_layer(x))
        means = self.output_layer_means(x)
        # reparameterise to keep std resonably high, so that we have resonable density of most of y
        # but not ridiculously high or low, as this gives us nans
        if self.fixed_variance is False:
            log_std = self.output_layer_log_stds(x)/20 + 1
        else:
            log_std = self.log_stds/20 + 1
        stds = torch.exp(log_std)
        return torch.distributions.normal.Normal(loc=means, scale=stds)

    def set_parameters(self, flat_parameter_tensor):
        assert flat_parameter_tensor.numel() == self.n_parameters
        new_state_dict = {}
        param_counter = 0
        # keys = list(dict(model.state_dict()).keys())
        # key = keys[1]
        new_state_dict = self.state_dict()
        for name, parameter in self.named_parameters():
            if parameter.requires_grad:
                n_param = parameter.numel()
                tensor = flat_parameter_tensor[param_counter:n_param + param_counter]
                tensor = tensor.reshape(parameter.shape)
                new_state_dict[name] = tensor
                param_counter += n_param
        self.load_state_dict(new_state_dict)

    @property
    def n_parameters(self):
        # number of trainable parameters
        return sum([tensor.numel() for tensor in self.parameters() if tensor.requires_grad])


class Target(nn.Module):
    """y = f_theta(x) where theta ~ N(0,1) (theta sampled once during initialisation), and x ~ N(0,1)
    Thus a target dataset is generated by sampling x and computing y
    The goal is to use this provide x & y data that we can use to get the posterior over weights of a BNN
    """
    def __init__(self, x_dim=2, y_dim=2, n_hidden_layers=2, layer_width=10,
                 fixed_variance=False, linear_activations=False, use_bias=True):
        super(Target, self).__init__()
        self.model = NN(x_dim, y_dim, n_hidden_layers, layer_width, fixed_variance=fixed_variance,
                        linear_activations=linear_activations, use_bias=use_bias)
        self.register_buffer("prior_loc", torch.zeros(x_dim))
        self.register_buffer("prior_covariance", torch.eye(x_dim))


    @property
    def prior(self):
        return torch.distributions.multivariate_normal.MultivariateNormal(loc=self.prior_loc,
                                                                       covariance_matrix=self.prior_covariance)

    @torch.no_grad()
    def sample(self, n_points=100):
        x = self.prior.sample((n_points,))
        y = self.model(x)
        return x, y

class PosteriorBNN(BaseTargetDistribution, nn.Module):
    """
     p(w | X, Y) proportional to p(w) p(Y | X, w)
     where we generate X, Y datasets using the Target class
     if we set n_hidden_layer=0, x_dim=1, y_dim=1 we should be able to visualise p(w | X, Y) in 3D
    """
    def __init__(self, n_datapoints=100, x_dim=2, y_dim=2, n_hidden_layers=1, layer_width=5,
                 linear_activations=False, fixed_variance=False, use_bias=True):
        super(PosteriorBNN, self).__init__()
        self.model = NN(x_dim, y_dim, n_hidden_layers, layer_width,
                        linear_activations=linear_activations, fixed_variance=fixed_variance, use_bias=use_bias)
        self.n_parameters = self.model.n_parameters
        self.register_buffer("prior_loc", torch.zeros(self.n_parameters))
        self.register_buffer("prior_covariance", torch.eye(self.n_parameters))
        self.target = Target(x_dim, y_dim, n_hidden_layers, layer_width,
                             linear_activations=linear_activations, fixed_variance=fixed_variance, use_bias=use_bias)
        X, Y = self.target.sample(n_datapoints)
        self.register_buffer("X", X)
        self.register_buffer("Y", Y)


    def log_prob(self, w):
        if len(w.shape) > 1:
            log_probs = list(map(self.log_prob_single_w, torch.split(w, 1, dim=0)))
            return torch.stack(log_probs)
        else:
            return self.log_prob_single_w(w)


    def log_prob_single_w(self, w):
        """p(w | X, Y) proportional to p(w) p(Y | X, w)"""
        w = torch.squeeze(w)
        model = copy.deepcopy(self.model)
        model.set_parameters(w)
        # keys = list(dict(model.state_dict()).keys())
        # key = keys[1]
        # model.state_dict()[key] == self.model.state_dict()[key] # want these to be false
        log_p_x = self.prior_w.log_prob(w)
        log_p_Y_given_w_X = torch.sum(model.posterior_y_given_x_log_prob(self.Y, self.X))
        return log_p_x + log_p_Y_given_w_X

    @property
    def prior_w(self):
        return torch.distributions.multivariate_normal.MultivariateNormal(
            loc=self.prior_loc, covariance_matrix=self.prior_covariance)



if __name__ == '__main__':
    #"""
    posterior_bnn = PosteriorBNN(n_datapoints=10, x_dim=2, y_dim=2, n_hidden_layers=1, layer_width=3,
                                 fixed_variance=True).to("cuda")
    for _ in range(5):
        samples_w = torch.randn(posterior_bnn.model.n_parameters).to("cuda")
        print(posterior_bnn.log_prob(samples_w))
    """
    from ImportanceSampling.VanillaImportanceSampler import VanillaImportanceSampling
    from FittedModels.Models.FlowModel import FlowModel
    from FittedModels.train import LearntDistributionManager
    torch.manual_seed(0)
    posterior_bnn = PosteriorBNN(n_datapoints=10, x_dim=2, y_dim=2, n_hidden_layers=1, layer_width=2)
    x = torch.zeros(posterior_bnn.n_parameters)
    print(posterior_bnn.log_prob(x))
    learnt_sampler = FlowModel(x_dim=posterior_bnn.n_parameters, n_flow_steps=3, scaling_factor=5)
    tester = LearntDistributionManager(posterior_bnn, learnt_sampler, VanillaImportanceSampling, loss_type="DReG")
    expectation_before, info_before = tester.estimate_expectation(int(1e4), lambda x: torch.sum(x))
    print(expectation_before, info_before)
    """

    """
    posterior_bnn = PosteriorBNN(n_datapoints=2, x_dim=1, y_dim=1, n_hidden_layers=0, layer_width=0
                                 ,linear_activations=True, fixed_variance=True, use_bias=True)
    
    assert posterior_bnn.n_parameters == 2
    from utils import plot_distribution
    import matplotlib.pyplot as plt
    plot_distribution(posterior_bnn, n_points=100)
    plt.show()
    """




