import torch
from torch.nn import functional as F
import numpy as np

from utils import PlaceHolder


def straight_through_gumbel_softmax(logits, tau=1.0, hard=False, dim=-1):
    # Sample a soft version
    soft = F.gumbel_softmax(logits, tau=tau, hard=hard, dim=dim) 
    # Argmax to get the one-hot index
    index = soft.argmax(dim=dim, keepdim=True)  # shape (..., 1)
    _hard = torch.zeros_like(soft).scatter_(dim, index, 1.0)       # one-hot

    # Straight-through (hard in forward, soft grad in backward)
    #   forward pass: 'hard'
    #   backward pass gradient flows as if 'soft'
    return _hard.detach() - soft.detach() + soft


def st_sample_discrete_feature(probX, probE, node_mask, temperature=0.5, 
                               hard=False, masked=False, remove_mask_class=False):

    bs, n, dx_out = probX.shape
    _, _, _, de_out = probE.shape

    # To avoid log(0), clamp probabilities
    probX = probX.clamp(min=1e-9)
    logitsX = torch.log(probX)

    invalid_nodes = ~node_mask
    probX[invalid_nodes] = 1.0 / dx_out
    logitsX = torch.log(probX.clamp(min=1e-9))

    # ST Gumbel-Softmax across the last dim
    X_soft_st = straight_through_gumbel_softmax(logitsX, tau=temperature, hard=hard, dim=-1)

    inverse_edge_mask = ~(node_mask.unsqueeze(1) * node_mask.unsqueeze(2))
    diag_mask = torch.eye(n, device=probE.device, dtype=torch.bool).unsqueeze(0)

    # For invalid or diagonal edges, set uniform or small probabilities
    probE = probE.clone()
    probE[inverse_edge_mask] = 1.0 / de_out
    probE[diag_mask.expand(bs, n, n)] = 1.0 / de_out

    # Convert to logits
    probE = probE.clamp(min=1e-9)
    logitsE = torch.log(probE)

    # Flatten so we can apply ST Gumbel in one shot: shape (bs*n*n, de_out)
    logitsE_flat = logitsE.view(bs*n*n, de_out)
    E_soft_st = straight_through_gumbel_softmax(logitsE_flat, tau=temperature, hard=hard, dim=-1)
    E_soft_st = E_soft_st.view(bs, n, n, de_out)
        
    # Zero out dimensions below diagonal
    triu_mask = torch.triu(torch.ones(n, n), diagonal=1).bool()
    triu_mask = triu_mask.unsqueeze(0).expand(bs, -1, -1)
    E_soft_st[~triu_mask] = 0
    E_soft_st = E_soft_st + E_soft_st.transpose(1, 2)
    E_soft_st = E_soft_st.clone()
    E_soft_st[:, torch.arange(n), torch.arange(n), 0] = 0
    
    hard_E = E_soft_st.argmax(dim=-1)
    assert (hard_E == hard_E.transpose(1, 2)).all(), "E_soft_st is not symmetric"
    return PlaceHolder(X=X_soft_st, 
                       E=E_soft_st, 
                       y=torch.zeros(bs, 0).type_as(X_soft_st))


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 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 mask_distributions(true_X, true_E=None, pred_X=None, pred_E=None, node_mask=None):
    # 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)

    if pred_E is not None:
        pred_E = pred_E.clamp_min(1e-18)
        pred_E = pred_E / torch.sum(pred_E, dim=-1, keepdim=True)
    else:
        true_E = None
        pred_E = None

    # 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)
    if pred_E is not None:
        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
    if pred_E is not None:
        true_E[~(node_mask.unsqueeze(1) * node_mask.unsqueeze(2) * diag_mask), :] = row_E
    pred_X[~node_mask] = row_X
    if pred_E is not None:
        pred_E[~(node_mask.unsqueeze(1) * node_mask.unsqueeze(2) * diag_mask), :] = row_E

    return true_X, true_E, pred_X, pred_E



# torch.multinomial() -> non-differentiable sampling
# apply gumbel noise to the logits
# take a softmax with a temperature param.
# (Optional) STE with hard=True
def sample_discrete_features(probX, probE, node_mask, remove_mask_class=False, masked=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, dx_out = probX.shape
    _, _, _, de_out = probE.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
    if masked:
        probX[:, -1] = 0
    if remove_mask_class:
        probX = probX[:, :-1]
    probX = probX / probX.sum(dim=-1, keepdim=True)

    probX = probX.clamp_min(1e-12)

    X_t = probX.multinomial(1)      # (bs * n, 1)
    X_t = X_t.reshape(bs, n)        # (bs, n)

    # Sample 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
    # remove second last dimension
    if masked:
        probE[:, -1] = 0
    if remove_mask_class:
        probE = probE[:, :-1]
    probE = probE / probE.sum(dim=-1, keepdim=True)

    probE = probE.clamp_min(1e-12)

    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))
    return PlaceHolder(X=X_t, E=E_t, y=torch.zeros(bs, 0).type_as(X_t))


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


def sample_discrete_features_from_limits(limit_dist, node_mask, temperature=0.5, hard=False):
    bs, n_max = node_mask.shape

    x_limit = limit_dist.X[:, None, :].expand(bs, n_max, -1).to(node_mask.device)
    invalid_nodes = ~node_mask
    x_limit[invalid_nodes] = 1.0 / x_limit.shape[-1]
    
    logits_x = torch.log(x_limit.clamp(min=1e-9))
    X_sample = F.gumbel_softmax(logits_x, tau=temperature, hard=hard, dim=-1)
    
    e_limit = limit_dist.E[:, None, None, :].expand(bs, n_max, n_max, -1).to(node_mask.device)
    
    logits_e = torch.log(e_limit.clamp(min=1e-9))
    logits_e_flat = logits_e.view(bs * n_max * n_max, -1)
    E_sample_flat = F.gumbel_softmax(logits_e_flat, tau=temperature, hard=hard, dim=-1)
    E_sample = E_sample_flat.view(bs, n_max, n_max, -1)

    # Force symmetry for edge features
    upper_triangular_mask = torch.zeros_like(E_sample)
    indices = torch.triu_indices(n_max, n_max, offset=1, device=E_sample.device)
    upper_triangular_mask[:, indices[0], indices[1], :] = 1

    E_sample_upper = E_sample * upper_triangular_mask
    E_sample_sym = E_sample_upper + E_sample_upper.transpose(1, 2)
    E_sample_sym[:, torch.arange(n_max), torch.arange(n_max), 0] = 0

    X_sample = X_sample.to(node_mask.device)
    E_sample_sym = E_sample_sym.to(node_mask.device)
    return PlaceHolder(X=X_sample, E=E_sample_sym, y=None).mask(node_mask)


