import numpy as np
import pandas as pd
from scipy.stats import norm
import torch
import torch.nn as nn
import types
from sklearn.neighbors import KernelDensity

import torch
from torch.distributions import Normal, MultivariateNormal




def names2idxes(dict_n2i: dict, list_names: list):
    list_idxes = []
    for ith_name in list_names:
        list_idxes.append(dict_n2i[ith_name])
    #idxes = torch.Tensor(list_idxes).long()
    return list_idxes #idxes


def check_S_set(set_S: str or list):
    if isinstance(set_S, str):
        set_S = [set_S]
    return set_S




def mapping_c2z_all(model, arr_cov: torch.Tensor):
    list_z = []
    for ith_cov in model.covariate_names:
        cov_idx = model.dict_cov_idx[ith_cov]
        current_cov_val = arr_cov[..., [cov_idx]]
        current_z = model.mapping_c2z[ith_cov](current_cov_val)
        list_z.append(current_z)

    arr_z = torch.cat(list_z, dim=-1)
    return arr_z

def mapping_z2c_all(model, arr_z: torch.Tensor):
    list_cov = []
    for ith_cov in model.covariate_names:
        cov_idx = model.dict_cov_idx[ith_cov]
        current_z_val = arr_z[..., [cov_idx]]
        current_cov = model.mapping_z2c[ith_cov](current_z_val)
        list_cov.append(current_cov)

    arr_cov = torch.cat(list_cov, dim=-1)
    return arr_cov

def dict2arr(list_keys, dict_):
    list_ = []
    for ith_key in list_keys:
        list_.append(dict_[ith_key])
    arr_ = torch.cat(list_, dim=-1)
    return arr_



