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

# ACTION_SIZE = 10


class RewardLoss(torch.nn.Module):
    def __init__(self, hyper_params):
        super().__init__()
        self.hyper_params = hyper_params

    def forward(self, output, action, delta):
        return F.binary_cross_entropy(output[range(action.size(0)), action], delta)


class CustomLoss(torch.nn.Module):
    def __init__(self, hyper_params):
        super().__init__()
        self.hyper_params = hyper_params
        self.lamda = hyper_params["lamda"]

    def forward(self, output, action, delta, prop):
        risk = 1.0 - delta

        loss = (risk - self.lamda) * (output[range(action.size(0)), action] / prop)

        return torch.mean(loss)


class CustomLossRec(torch.nn.Module):
    def __init__(self, hyper_params):
        super().__init__()
        self.hyper_params = hyper_params
        self.lamda = hyper_params["lamda"]

    def forward(self, output, delta, prop):
        risk = 1.0 - delta

        loss = (risk - self.lamda) * (output / prop)

        return torch.mean(loss)


# def KLLoss(output, action, prop):
#     h_scores = output[range(action.size(0)), action]
#     return torch.sum(h_scores * torch.log(H_scores / prop + 1e-8))


def KLLoss(output, action, prop, action_size=10):
    h_scores = output[range(action.size(0)), action]
    all_values = h_scores * torch.log(h_scores / prop + 1e-8)
    index = action
    src = all_values
    out = torch.scatter_reduce(
        torch.zeros(action_size).to(output.device),
        dim=0,
        index=index,
        src=src,
        reduce="mean",
    )
    return torch.sum(out)


def KLLossRec(output, action, prop):
    action_size = torch.max(action) + 1
    h_scores = output
    all_values = h_scores * torch.log(h_scores / prop + 1e-8)
    index = action
    src = all_values
    out = torch.scatter_reduce(
        torch.zeros(action_size).to(output.device),
        dim=0,
        index=index,
        src=src,
        reduce="mean",
    )
    return torch.sum(out)


def KLLossRev(output, action, prop, action_size=10):
    h_scores = output[range(action.size(0)), action]
    all_values = prop * torch.log(prop / (h_scores + 1e-8) + 1e-8)
    index = action
    src = all_values
    out = torch.scatter_reduce(
        torch.zeros(action_size).to(output.device),
        dim=0,
        index=index,
        src=src,
        reduce="mean",
    )
    return torch.sum(out)


def KLLossRevRec(output, action, prop):
    action_size = torch.max(action) + 1
    h_scores = output
    all_values = prop * torch.log(prop / (h_scores + 1e-8) + 1e-8)
    index = action
    src = all_values
    out = torch.scatter_reduce(
        torch.zeros(action_size).to(output.device),
        dim=0,
        index=index,
        src=src,
        reduce="mean",
    )
    return torch.sum(out)


def SupKLLoss(output, action, delta, prop, eps, action_size=10):
    h_scores = output[range(action.size(0)), action]
    all_values = h_scores * torch.log(h_scores * (delta + eps) / prop + 1e-8)
    index = action
    src = all_values
    out = torch.scatter_reduce(
        torch.zeros(action_size).to(output.device),
        dim=0,
        index=index,
        src=src,
        reduce="mean",
    )
    return torch.sum(out)
