import torch
from torch.nn import functional as F
import numpy as np
from utils.data_io import PlaceHolder

def sum_except_batch(x):
    return x.reshape(x.size(0), -1).sum(dim=-1)

def assert_correctly_masked(variable, node_mask):
    assert (variable * (1 - node_mask.long())).abs().max().item() < 1e-4, \
        'Variables not masked properly.'

def linear_beta_schedule_discrete(timesteps, beta_start=1e-4, beta_end=0.02):
    """Linear schedule from beta_start to beta_end over timesteps."""
    return np.linspace(beta_start, beta_end, timesteps + 1)

def cosine_beta_schedule_discrete(timesteps, s=0.008):
    """ Cosine schedule as proposed in https://openreview.net/forum?id=-NEXDKk8gZ. """
    steps = timesteps + 2
    x = np.linspace(0, steps, steps)

    alphas_cumprod = np.cos(0.5 * np.pi * ((x / steps) + s) / (1 + s)) ** 2
    alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
    alphas = (alphas_cumprod[1:] / alphas_cumprod[:-1])
    betas = 1 - alphas
    return betas.squeeze()

def custom_beta_schedule_discrete(timesteps, average_num_nodes=30, s=0.008):
    """ Cosine schedule as proposed in https://openreview.net/forum?id=-NEXDKk8gZ. """
    steps = timesteps + 2
    x = np.linspace(0, steps, steps)

    alphas_cumprod = np.cos(0.5 * np.pi * ((x / steps) + s) / (1 + s)) ** 2
    alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
    alphas = (alphas_cumprod[1:] / alphas_cumprod[:-1])
    betas = 1 - alphas

    assert timesteps >= 100

    p = 4 / 5       # 1 - 1 / num_edge_classes
    num_edges = average_num_nodes * (average_num_nodes - 1) / 2

    # First 100 steps: only a few updates per graph
    updates_per_graph = 1.2
    beta_first = updates_per_graph / (p * num_edges)

    betas[betas < beta_first] = beta_first
    return np.array(betas)


def check_mask_correct(variables, node_mask):
    for i, variable in enumerate(variables):
        if len(variable) > 0:
            assert_correctly_masked(variable, node_mask)


def check_tensor_same_size(*args):
    for i, arg in enumerate(args):
        if i == 0:
            continue
        assert args[0].size() == arg.size()



def reverse_tensor(x):
    return x[torch.arange(x.size(0) - 1, -1, -1)]


def sample_discrete_features(probX, probE, prob_pos, node_mask, deterministic=False, exclude_null_pos=False):
    ''' Sample features from multinomial distribution with given probabilities (probX, probE, proby)
        :param probX: bs, n, dx_out        node features
        :param probE: bs, n, n, de_out     edge features
        :param proby: bs, dy_out           global features.
    '''
    bs, n, _ = probX.shape

    # Noise X
    # The masked rows should define probability distributions as well
    probX[~node_mask] = 1 / probX.shape[-1]
    # Flatten the probability tensor to sample with multinomial
    probX = probX.reshape(bs * n, -1)       # (bs * n, dx_out)

    # Sample X
    probX = probX + 1e-12
    probX = probX / probX.sum(dim=-1, keepdim=True)
    if deterministic:
        X_t = probX.argmax(dim=-1)
    else:
        X_t = probX.multinomial(1)      # (bs * n, 1)
    X_t = X_t.reshape(bs, n)        # (bs, n)

    # Noise E
    # The masked rows should define probability distributions as well
    inverse_edge_mask = ~(node_mask.unsqueeze(1) * node_mask.unsqueeze(2))
    diag_mask = torch.eye(n).unsqueeze(0).expand(bs, -1, -1)
    probE[inverse_edge_mask] = 1 / probE.shape[-1]
    probE[diag_mask.bool()] = 1 / probE.shape[-1]
    probE = probE.reshape(bs * n * n, -1)           # (bs * n * n, de_out)

    probE = probE + 1e-12
    probE = probE / probE.sum(dim=-1, keepdim=True)
    # Sample E
    if deterministic:
        E_t = probE.argmax(dim=-1).reshape(bs, n, n)
    else:
        E_t = probE.multinomial(1).reshape(bs, n, n)    # (bs, n, n)
    E_t = torch.triu(E_t, diagonal=1)
    E_t = (E_t + torch.transpose(E_t, 1, 2))

    # Noise pos, the same as Noies E but non-symmetrical
    prob_pos[inverse_edge_mask] = 1 / prob_pos.shape[-1]
    prob_pos[diag_mask.bool()] = 1 / prob_pos.shape[-1]
    prob_pos = prob_pos.reshape(bs * n * n, -1)           # (bs * n * n, dpos_out)
    if exclude_null_pos:
        prob_pos[:, 0] = 0
    prob_pos = prob_pos + 1e-12
    prob_pos = prob_pos / prob_pos.sum(dim=-1, keepdim=True)
    if deterministic:
        pos_t = prob_pos.argmax(dim=-1).reshape(bs, n, n)
    else:
        pos_t = prob_pos.multinomial(1).reshape(bs, n, n)    # (bs, n, n)

    return PlaceHolder(X=X_t, E=E_t, pos=pos_t)
    # return PlaceHolder(X=X_t, E=E_t, y=torch.zeros(bs, 0).type_as(X_t))

def reverse_diffusion(predX_0, X_t, Qt, Qsb, Qtb):
    """ M: X or E
        Compute xt @ Qt.T * x0 @ Qsb / x0 @ Qtb @ xt.T for each possible value of x0
        X_t: bs, n, dt          or bs, n, n, dt
        Qt: bs, d_t-1, dt
        Qsb: bs, d0, d_t-1
        Qtb: bs, d0, dt.
    """                                  
    Qt_T = Qt.transpose(-1, -2)                             # bs, N, dt
    assert Qt.dim() == 3
    left_term = X_t @ Qt_T                                  # bs, N, d_t-1
    right_term = predX_0 @ Qsb                         
    numerator = left_term * right_term                      # bs, N, d_t-1
    
    denominator = Qtb @ X_t.transpose(-1, -2)                 # bs, d0, N
    denominator = denominator.transpose(-1, -2)               # bs, N, d0
    return numerator / denominator.clamp_min(1e-5)

def mask_distributions(true_X, true_E, pred_X, pred_E, node_mask):
    # Add a small value everywhere to avoid nans
    pred_X = pred_X.clamp_min(1e-18)
    pred_X = pred_X / torch.sum(pred_X, dim=-1, keepdim=True)

    pred_E = pred_E.clamp_min(1e-18)
    pred_E = pred_E / torch.sum(pred_E, dim=-1, keepdim=True)

    # Set masked rows to arbitrary distributions, so it doesn't contribute to loss
    row_X = torch.ones(true_X.size(-1), dtype=true_X.dtype, device=true_X.device)
    row_E = torch.zeros(true_E.size(-1), dtype=true_E.dtype, device=true_E.device).clamp_min(1e-18)
    row_E[0] = 1.

    diag_mask = ~torch.eye(node_mask.size(1), device=node_mask.device, dtype=torch.bool).unsqueeze(0)
    true_X[~node_mask] = row_X
    true_E[~(node_mask.unsqueeze(1) * node_mask.unsqueeze(2) * diag_mask), :] = row_E
    pred_X[~node_mask] = row_X.type_as(pred_X)
    pred_E[~(node_mask.unsqueeze(1) * node_mask.unsqueeze(2) * diag_mask), :] = row_E.type_as(pred_E)

    return true_X, true_E, pred_X, pred_E

def posterior_distributions(X, X_t, Qt, Qsb, Qtb, X_dim):
    bs, n, d = X.shape
    X = X.float()
    Qt_X_T = torch.transpose(Qt.X, -2, -1).float()                  # (bs, d, d)
    left_term = X_t @ Qt_X_T                                        # (bs, N, d)
    right_term = X @ Qsb.X                                          # (bs, N, d)
    
    numerator = left_term * right_term                              # (bs, N, d)
    denominator = X @ Qtb.X                                         # (bs, N, d) @ (bs, d, d) = (bs, N, d)
    denominator = denominator * X_t 
    
    num_X = numerator[:, :, :X_dim]
    num_E = numerator[:, :, X_dim:].reshape(bs, n*n, -1)

    deno_X = denominator[:, :, :X_dim]
    deno_E = denominator[:, :, X_dim:].reshape(bs, n*n, -1)

    # denominator = (denominator * X_t).sum(dim=-1)                   # (bs, N, d) * (bs, N, d) + sum = (bs, N)
    denominator = denominator.unsqueeze(-1)                         # (bs, N, 1)

    deno_X = deno_X.sum(dim=-1).unsqueeze(-1)
    deno_E = deno_E.sum(dim=-1).unsqueeze(-1)

    deno_X[deno_X == 0.] = 1
    deno_E[deno_E == 0.] = 1
    prob_X = num_X / deno_X
    prob_E = num_E / deno_E
    
    prob_E = prob_E / prob_E.sum(dim=-1, keepdim=True)
    prob_X = prob_X / prob_X.sum(dim=-1, keepdim=True)
    return PlaceHolder(X=prob_X, E=prob_E, y=None)

def sample_discrete_feature_noise(limit_dist, node_mask):
    """ Sample from the limit distribution of the diffusion process"""
    bs, n_max = node_mask.shape
    device = node_mask.device
    x_limit = limit_dist.X[None, None, :].expand(bs, n_max, -1)
    U_X = x_limit.flatten(end_dim=-2).multinomial(1).reshape(bs, n_max)
    U_X = F.one_hot(U_X.long(), num_classes=x_limit.shape[-1]).type_as(x_limit).to(device)
    # U_X = F.one_hot(U_X.long(), num_classes=x_limit.shape[-1]).to(node_mask.device)

    e_limit = limit_dist.E[None, None, None, :].expand(bs, n_max, n_max, -1)
    U_E = e_limit.flatten(end_dim=-2).multinomial(1).reshape(bs, n_max, n_max)
    U_E = F.one_hot(U_E.long(), num_classes=e_limit.shape[-1]).type_as(x_limit).to(device)
    # U_E = F.one_hot(U_E.long(), num_classes=e_limit.shape[-1]).to(node_mask.device)

    pos_limit = limit_dist.pos[None, None, None, :].expand(bs, n_max, n_max, -1)
    U_pos = pos_limit.flatten(end_dim=-2).multinomial(1).reshape(bs, n_max, n_max)
    U_pos = F.one_hot(U_pos.long(), num_classes=pos_limit.shape[-1]).type_as(x_limit).to(device)

    # Get upper triangular part of edge noise, without main diagonal
    upper_triangular_mask = torch.zeros_like(U_E).to(device)
    indices = torch.triu_indices(row=U_E.size(1), col=U_E.size(2), offset=1)
    upper_triangular_mask[:, indices[0], indices[1], :] = 1

    U_E = U_E * upper_triangular_mask
    U_E = (U_E + torch.transpose(U_E, 1, 2))

    assert (U_E == torch.transpose(U_E, 1, 2)).all()
    return PlaceHolder(X=U_X, E=U_E, pos=U_pos).mask(node_mask)

def index_QE(X, q_e, n_bond=5):
    bs, n, n_atom = X.shape
    node_indices = X.argmax(-1)  # (bs, n)

    exp_ind1 = node_indices[ :, :, None, None, None].expand(bs, n, n_atom, n_bond, n_bond)
    exp_ind2 = node_indices[ :, :, None, None, None].expand(bs, n, n, n_bond, n_bond)
    
    q_e = torch.gather(q_e, 1, exp_ind1)
    q_e = torch.gather(q_e, 2, exp_ind2) # (bs, n, n, n_bond, n_bond)

    node_mask = X.sum(-1) != 0
    no_edge = (~node_mask)[:, :, None] & (~node_mask)[:, None, :]
    q_e[no_edge] = torch.tensor([1, 0, 0, 0, 0]).type_as(q_e)

    return q_e
