import torch


def match_z_v(marginal_entropies, conditional_entropies, top_mi_discrete_normed, top_mi_discrete_idx):
    """
    Matches latent dimension z_j with generative factor v_k with the highest
    estimated mutual information I(z_j; v_k) using dSprites.
    
    Returns dictionary containing results of matching for both continuous and
    discrete MI approximations.
    """
    
    factor_entropies = [3, 6, 40, 32, 32]
    mi_info = lambda x: dict(latent_dim=int(top_mi_idx[x]), 
                              mi_normed=float(top_mi_normed[x]), 
                              mi_normed_gap=float(mi_sorted_normed_gap[x].item()),  # normalized gap b/w highest and second highest MI, dimwise
                              discrete_latent_dim=int(top_mi_discrete_idx[x]),
                              discrete_mi_normed=float(top_mi_discrete_normed[x]))

    mutual_infos = torch.unsqueeze(marginal_entropies, dim=0) - conditional_entropies
    mi_sorted, mi_sorted_idx = torch.sort(mutual_infos.clamp(min=0), dim=1, descending=True)
    mi_sorted_normed = mi_sorted / torch.Tensor(factor_entropies).log()[:, None]
    
    mi_sorted_normed_gap = mi_sorted_normed[:,0] - mi_sorted_normed[:,1]

    top_mi_normed = mi_sorted_normed[:,0].numpy()
    top_mi_idx = mi_sorted_idx[:,0].numpy()
    
    z_v_mapping = dict(shape=mi_info(0),
                       scale=mi_info(1),
                       rotation=mi_info(2),
                       posX=mi_info(3),
                       posY=mi_info(4))
    
    return z_v_mapping