from torch import nn
import torch
import numpy as np

def nv_loss_func_np(params, w_pred, z_real):
    h_list = params['h_list']
    b_list = params['b_list']
    dim_sol = params['dim_sol']
    b_size = np.shape(w_pred)[0]
    loss = 0
    for i in range(b_size):
        for l in range(dim_sol):
           # print("in np nv loss", np.maximum(w_pred[i, l] - z_real[i, l], 0),np.maximum(-w_pred[i, l] + z_real[i, l], 0))
            loss += np.maximum(w_pred[i, l] - z_real[i, l], 0)*h_list[l] + np.maximum(-w_pred[i, l] + z_real[i, l], 0)*b_list[l]
    return loss/b_size


class entropy_loss_func(nn.Module):
    def __init__(self, params):
        super().__init__()
        self.dim_sol = params['dim_sol']

    def forward(self, p_pred, z_real_idx):
        z_real_idx = z_real_idx.type(torch.LongTensor)
        # #print('in cross entropy loss, z_real_index', z_real_index)
        # #print('in cross entropy loss, p_pred', p_pred.size())
        EntropyLoss = torch.nn.CrossEntropyLoss()
        entropy_loss = EntropyLoss(p_pred, z_real_idx)
        return entropy_loss

