import numpy as np
from scipy.stats import multivariate_normal
import torch
from torch import nn
from Bibi2018PLDNN.torch.gaussian.relu import zero_mean_covariance
from Wu2019MP.bayes_utils import relu_covariance as relu_covariance_wu
import torch.nn.functional as F


def set_seed(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)


def generate_covariance_variance_and_mean(batch_size, input_size):
    A = torch.randn(batch_size, input_size, input_size)
    cov_matrices = torch.matmul(A.transpose(1, 2), A)
    variances = torch.diagonal(cov_matrices, dim1=1, dim2=2)
    variances = variances.unsqueeze(1)
    means = torch.randn(batch_size, 1, input_size)
    return cov_matrices.cuda(), variances.cuda(), means.cuda()


dtype = torch.float32
EPSILON = 1e-9


def standard_gaussian_pdf(x):
    return (1./np.sqrt(2*np.pi)) * torch.exp(-x*x / 2.0)


def standard_gaussian_cdf(x):
    return 0.5 * (1.0 + torch.erf(x * (1./np.sqrt(2))))


def gelu_plus(x):
    return x * standard_gaussian_cdf(x) + standard_gaussian_pdf(x)


def gelu_plus_multiply(x):
    return x * gelu_plus(x) + standard_gaussian_cdf(x)


def fill_diagonal(cov, var):
    input_size = cov.shape[1]
    diag_indices = torch.arange(input_size)
    cov = cov.clone()
    cov[:, diag_indices, diag_indices] = var.squeeze(1)
    return cov


def cal_mean_var_tp(p, q):
    mean = 1 - 2 * p - q
    var = 1 - q - mean ** 2
    return mean, var


def bivariate_standard_gaussian_cdf(x, rho):
    cov = [[1, rho], [rho, 1]]
    return multivariate_normal.cdf(x, cov=cov)


def bivariate_standard_gaussian_cdf_normalization(mean_y_batch, std_y_batch, rho_y_batch):
    batch_size = mean_y_batch.shape[0]
    input_size = mean_y_batch.shape[2]

    cov_list = []

    for k in range(batch_size):
        mean_y = mean_y_batch[k:k + 1]
        std_y = std_y_batch[k:k + 1]
        rho_y = rho_y_batch[k]
        ratio = (mean_y / std_y).cpu().numpy().flatten()
        rho_y = rho_y.cpu().numpy()

        cov = np.zeros((input_size, input_size))

        for i in range(1, input_size):
            for j in range(i):
                rho = rho_y[i][j]
                x = np.array((ratio[i], ratio[j]))
                cdf = bivariate_standard_gaussian_cdf(x, rho)
                cov[i, j] = cdf
                cov[j, i] = cdf
        cov_list.append(cov)

    cov_matrix_batch = np.stack(cov_list)

    return torch.tensor(cov_matrix_batch, dtype=torch.float32).cuda()


def bivariate_standard_gaussian_cdf_normalization_approximation(mean_y_batch, std_y_batch, rho_y_batch):
    ratio = mean_y_batch / std_y_batch
    ratio_transpose = torch.transpose(ratio, 1, 2)
    t1 = standard_gaussian_cdf(ratio_transpose * ratio)
    t2 = standard_gaussian_pdf(ratio_transpose * ratio)
    return t1 + t2 * rho_y_batch


def moment_propagation_x2y_simplified(mean_x, var_x, cov_x, layer):
    if not isinstance(layer, nn.Linear):
        raise TypeError(f"Layer {layer} is not an instance of nn.Linear()")

    output_size = layer.weight.shape[0]

    mean_w = layer.mean
    var_w = layer.var
    b = layer.bias.unsqueeze(0).unsqueeze(0)

    mean_y = mean_w * torch.sum(mean_x, dim=2, keepdim=True).expand(-1, -1, output_size) + b

    t1 = torch.sum(mean_x ** 2, dim=2, keepdim=True)
    t2 = torch.sum(var_x, dim=2, keepdim=True)
    t3 = mean_w ** 2 * torch.sum(cov_x, dim=(1, 2), keepdim=True)

    var_y = (var_w * (t1 + t2) + t3).expand(-1, -1, output_size)

    cov_y_without_diag = t3.expand(-1, output_size, output_size)
    cov_y = fill_diagonal(cov_y_without_diag, var_y)

    return mean_y, var_y, cov_y


def moment_propagation_y2z(mean_y, var_y, cov_y, layer, approx=False):
    std_y = torch.sqrt(var_y)
    std_y_transpose = torch.transpose(std_y, 1, 2)
    s1s2 = std_y_transpose * std_y
    rho_y = cov_y / s1s2
    rho_y_clamped = torch.clamp(rho_y, -1/(1+EPSILON), 1/(1+EPSILON))
    ratio = mean_y / std_y
    ratio_transpose = torch.transpose(ratio, 1, 2)

    denominator = torch.sqrt(1 - rho_y_clamped * rho_y_clamped)
    w_12 = (ratio_transpose - rho_y * ratio) / denominator
    w_21 = (- rho_y * ratio_transpose + ratio) / denominator

    mean_z = std_y * gelu_plus(ratio)
    var_z = var_y * gelu_plus_multiply(ratio) - mean_z * mean_z

    t1 = ratio_transpose * ratio
    t2 = rho_y
    t31 = (standard_gaussian_pdf(ratio_transpose) * ratio) * standard_gaussian_cdf(w_21)
    t32 = (ratio_transpose * standard_gaussian_pdf(ratio)) * standard_gaussian_cdf(w_12)
    t33 = (denominator * standard_gaussian_pdf(ratio)) * standard_gaussian_pdf(w_12)
    t3 = t31 + t32 + t33
    t4 = gelu_plus(ratio_transpose) * gelu_plus(ratio)
    if approx:
        phi2 = bivariate_standard_gaussian_cdf_normalization_approximation(mean_y, std_y, rho_y)
    else:
        phi2 = bivariate_standard_gaussian_cdf_normalization(mean_y, std_y, rho_y)
    cov_without_diag = s1s2 * ((t1 + t2) * phi2 + t3 - t4)
    cov_z = fill_diagonal(cov_without_diag, var_z)
    return mean_z, var_z, cov_z


def propagate_moments(mean, var, cov, layers, input_shape=None, approx=False):
    with torch.no_grad():
        if input_shape:
            mean = mean.reshape(-1, *input_shape)
        for layer in layers:
            if isinstance(layer, nn.Linear):
                mean, var, cov = moment_propagation_x2y_simplified(mean, var, cov, layer)
            elif isinstance(layer, nn.ReLU):
                mean, var, cov = moment_propagation_y2z(mean, var, cov, layer, approx)
    return mean, var, cov


def mc_stats_fp(mc_iteration, x, mu_1, sigma_1, mu_2, sigma_2, model, device):
    ys = []
    with torch.no_grad():
        for _ in range(mc_iteration):
            w1 = torch.normal(mu_1, sigma_1, size=model.fc1.weight.size()).to(device)
            w2 = torch.normal(mu_2, sigma_2, size=model.fc2.weight.size()).to(device)

            model.fc1.weight.data = w1
            model.fc2.weight.data = w2

            y = model(x)
            ys.append(y)

    ys = torch.stack(ys, dim=0).squeeze()

    mean_mc = torch.mean(ys, dim=0)
    var_mc = torch.var(ys, dim=0)
    cov_mc = torch.cov(ys.T)
    return mean_mc, var_mc, cov_mc


def mc_stats_fp_batch(mc_iteration, x, mu_1, sigma_1, mu_2, sigma_2, b_1, b_2, device, batch_size=1000):
    out_list = []

    for i in range(0, mc_iteration, batch_size):
        current_batch_size = min(batch_size, mc_iteration - i)

        w_1 = torch.normal(
            mu_1.unsqueeze(0).expand(current_batch_size, *mu_1.shape),
            sigma_1.unsqueeze(0).expand(current_batch_size, *sigma_1.shape)
        ).to(device)

        w_2 = torch.normal(
            mu_2.unsqueeze(0).expand(current_batch_size, *mu_2.shape),
            sigma_2.unsqueeze(0).expand(current_batch_size, *sigma_2.shape)
        ).to(device)

        out_batch = forward_fixed_input_batch_weights(x, w_1, w_2, b_1, b_2)
        out_list.append(out_batch)

    out_all = torch.cat(out_list, dim=0)

    mean_mc = torch.mean(out_all, dim=0)
    var_mc = torch.var(out_all, dim=0)
    cov_mc = torch.cov(out_all.T)

    return mean_mc, var_mc, cov_mc


def forward_fixed_input_batch_weights(x_0, weight1, weight2, bias1, bias2):
    hidden = torch.einsum('bhd,d->bh', weight1, x_0) + bias1
    hidden = F.relu(hidden)

    output = torch.einsum('bkh,bh->bk', weight2, hidden) + bias2
    return output


def mc_stats_tp(mc_iteration, x, p_1, q_1, p_2, q_2, model, device):
    ys = []
    values = torch.tensor([-1., 0., 1.], device=device)
    probs_1 = torch.tensor([p_1, q_1, 1 - p_1 - q_1], device=device)
    probs_2 = torch.tensor([p_2, q_2, 1 - p_2 - q_2], device=device)
    weight_dist_1 = torch.distributions.Categorical(probs=probs_1)
    weight_dist_2 = torch.distributions.Categorical(probs=probs_2)

    with torch.no_grad():
        for _ in range(mc_iteration):
            indices_1 = weight_dist_1.sample(model.fc1.weight.size())
            w1 = values[indices_1].to(device)
            indices_2 = weight_dist_2.sample(model.fc2.weight.size())
            w2 = values[indices_2].to(device)

            model.fc1.weight.data = w1
            model.fc2.weight.data = w2

            y = model(x)
            ys.append(y)

    ys = torch.stack(ys, dim=0).squeeze()
    mean_mc = torch.mean(ys, dim=0)
    var_mc = torch.var(ys, dim=0)
    cov_mc = torch.cov(ys.T)
    return mean_mc, var_mc, cov_mc


def mc_stats_tp_batch(mc_iteration, x, p_1, q_1, p_2, q_2, b_1, b_2, device, batch_size=1000):
    values = torch.tensor([-1., 0., 1.], device=device)
    probs_1 = torch.tensor([p_1, q_1, 1 - p_1 - q_1], device=device)
    probs_2 = torch.tensor([p_2, q_2, 1 - p_2 - q_2], device=device)
    weight_dist_1 = torch.distributions.Categorical(probs=probs_1)
    weight_dist_2 = torch.distributions.Categorical(probs=probs_2)

    out_list = []

    for i in range(0, mc_iteration, batch_size):
        current_batch_size = min(batch_size, mc_iteration - i)

        indices_1 = weight_dist_1.sample((current_batch_size, b_1.shape[0], x.shape[0]))
        w_1 = values[indices_1].to(device)
        indices_2 = weight_dist_2.sample((current_batch_size, b_2.shape[0], b_1.shape[0]))
        w_2 = values[indices_2].to(device)

        out_batch = forward_fixed_input_batch_weights(x, w_1, w_2, b_1, b_2)
        out_list.append(out_batch)

    out_all = torch.cat(out_list, dim=0)

    mean_mc = torch.mean(out_all, dim=0)
    var_mc = torch.var(out_all, dim=0)
    cov_mc = torch.cov(out_all.T)

    return mean_mc, var_mc, cov_mc


def register_parameter_2layer_mlp(model, mean_1, var_1, mean_2, var_2):
    model.fc1.register_buffer("mean", torch.tensor(mean_1))
    model.fc1.register_buffer("var", torch.tensor(var_1))
    model.fc2.register_buffer("mean", torch.tensor(mean_2))
    model.fc2.register_buffer("var", torch.tensor(var_2))
    return model


def moment_propagation_y2z_bibi2018(mean_y, var_y, cov_y, layer):
    std_y = torch.sqrt(var_y)
    ratio = mean_y / std_y

    mean_z = std_y * gelu_plus(ratio)
    var_z = var_y * gelu_plus_multiply(ratio) - mean_z * mean_z

    cov_z = zero_mean_covariance(cov_y)
    return mean_z, var_z, cov_z


def propagate_moments_bibi2018(mean, var, cov, layers, input_shape=None):
    with torch.no_grad():
        if input_shape:
            mean = mean.reshape(-1, *input_shape)
        for layer in layers:
            if isinstance(layer, nn.Linear):
                mean, var, cov = moment_propagation_x2y_simplified(mean, var, cov, layer)
            elif isinstance(layer, nn.ReLU):
                mean, var, cov = moment_propagation_y2z_bibi2018(mean, var, cov, layer)
    return mean, var, cov


def moment_propagation_y2z_wu2019(mean_y, var_y, cov_y, layer):
    std_y = torch.sqrt(var_y)
    ratio = mean_y / std_y

    mean_z = std_y * gelu_plus(ratio)
    var_z = var_y * gelu_plus_multiply(ratio) - mean_z * mean_z

    var_y = torch.squeeze(var_y, 1)
    mean_y = torch.squeeze(mean_y, 1)
    cov_z = relu_covariance_wu(cov_y, var_y, mean_y)
    return mean_z, var_z, cov_z


def propagate_moments_wu2019(mean, var, cov, layers, input_shape=None):
    with torch.no_grad():
        if input_shape:
            mean = mean.reshape(-1, *input_shape)
        for layer in layers:
            if isinstance(layer, nn.Linear):
                mean, var, cov = moment_propagation_x2y_simplified(mean, var, cov, layer)
            elif isinstance(layer, nn.ReLU):
                mean, var, cov = moment_propagation_y2z_wu2019(mean, var, cov, layer)
    return mean, var, cov

