# Path: pysaliency/models.py
# def sample_from_logdensity(log_density, count=None, rst=None):
#     if count is None:
#         real_count = 1
#     else:
#         real_count = count
# 
#     height, width = log_density.shape
#     flat_log_density = log_density.flatten(order='C')
#     samples = sample_from_logprobabilities(flat_log_density, size=real_count, rst=rst)
#     sample_xs = samples % width
#     sample_ys = samples // width
# 
#     if count is None:
#         return sample_xs[0], sample_ys[0]
#     else:
#         return np.asarray(sample_xs), np.asarray(sample_ys)
# 
# Path: pysaliency/torch_utils.py
# def gaussian_filter(tensor, dim, sigma, truncate=4, kernel_size=None, padding_mode='replicate', padding_value=0.0):
#     if isinstance(dim, int):
#         dim = [dim]
# 
#     for k in dim:
#         tensor = gaussian_filter_1d(
#             tensor,
#             dim=k,
#             sigma=sigma,
#             truncate=truncate,
#             kernel_size=kernel_size,
#             padding_mode=padding_mode,
#             padding_value=padding_value,
#         )
# 
#     return tensor


import numpy as np
import torch
import torch.nn as nn
from scipy.ndimage import gaussian_filter as sp_gaussian_filter
from scipy.special import logsumexp
from torch.optim.optimizer import required
from tqdm import tqdm
from .models import sample_from_logdensity
from .torch_utils import gaussian_filter
                # out of the constraint
                projected_grad2 = projected_grad1 - torch.sum(projected_grad1 * constraint_grad) * normed_constraint_grad

                p.add_(projected_grad2, alpha=-group['lr'])

        return loss


def build_fixation_maps(Ns, Ys, Xs, batch_size, height, width, dtype=torch.float32):
    indices = torch.stack((Ns, Ys, Xs), axis=1).T
    src = torch.ones(indices.shape[1], dtype=dtype, device=indices.device)
    fixation_maps = torch.sparse_coo_tensor(indices, src, size=(batch_size, height, width)).to_dense()

    return fixation_maps


def torch_similarity(saliency_map, empirical_saliency_maps):
    normalized_empirical_saliency_maps = empirical_saliency_maps / torch.sum(empirical_saliency_maps, dim=[1, 2], keepdim=True)
    normalized_saliency_map = saliency_map / torch.sum(saliency_map)
    minimums = torch.min(normalized_empirical_saliency_maps, normalized_saliency_map[None, :, :])

    similarities = torch.sum(minimums, dim=[1, 2])

    return similarities


def compute_similarity(saliency_map, ns, ys, xs, batch_size, kernel_size, truncate_gaussian, dtype=torch.float32):
    height, width = saliency_map.shape
    fixation_maps = build_fixation_maps(ns, ys, xs, batch_size, height, width, dtype=dtype)


# Path: pysaliency/models.py
# def sample_from_logdensity(log_density, count=None, rst=None):
#     if count is None:
#         real_count = 1
#     else:
#         real_count = count
# 
#     height, width = log_density.shape
#     flat_log_density = log_density.flatten(order='C')
#     samples = sample_from_logprobabilities(flat_log_density, size=real_count, rst=rst)
#     sample_xs = samples % width
#     sample_ys = samples // width
# 
#     if count is None:
#         return sample_xs[0], sample_ys[0]
#     else:
#         return np.asarray(sample_xs), np.asarray(sample_ys)
# 
# Path: pysaliency/torch_utils.py
# def gaussian_filter(tensor, dim, sigma, truncate=4, kernel_size=None, padding_mode='replicate', padding_value=0.0):
#     if isinstance(dim, int):
#         dim = [dim]
# 
#     for k in dim:
#         tensor = gaussian_filter_1d(
#             tensor,
#             dim=k,
#             sigma=sigma,
#             truncate=truncate,
#             kernel_size=kernel_size,
#             padding_mode=padding_mode,
#             padding_value=padding_value,
#         )
# 
#     return tensor


# Path: pysaliency/metric_optimization_torch.py
import numpy as np
import torch
import torch.nn as nn
from scipy.ndimage import gaussian_filter as sp_gaussian_filter
from scipy.special import logsumexp
from torch.optim.optimizer import required
from tqdm import tqdm
from .models import sample_from_logdensity
from .torch_utils import gaussian_filter

                # out of the constraint
                projected_grad2 = projected_grad1 - torch.sum(projected_grad1 * constraint_grad) * normed_constraint_grad

                p.add_(projected_grad2, alpha=-group['lr'])

        return loss


def build_fixation_maps(Ns, Ys, Xs, batch_size, height, width, dtype=torch.float32):
    indices = torch.stack((Ns, Ys, Xs), axis=1).T
    src = torch.ones(indices.shape[1], dtype=dtype, device=indices.device)
    fixation_maps = torch.sparse_coo_tensor(indices, src, size=(batch_size, height, width)).to_dense()

    return fixation_maps


def torch_similarity(saliency_map, empirical_saliency_maps):
    normalized_empirical_saliency_maps = empirical_saliency_maps / torch.sum(empirical_saliency_maps, dim=[1, 2], keepdim=True)
    normalized_saliency_map = saliency_map / torch.sum(saliency_map)
    minimums = torch.min(normalized_empirical_saliency_maps, normalized_saliency_map[None, :, :])

    similarities = torch.sum(minimums, dim=[1, 2])

    return similarities


def compute_similarity(saliency_map, ns, ys, xs, batch_size, kernel_size, truncate_gaussian, dtype=torch.float32):
    height, width = saliency_map.shape
    fixation_maps = build_fixation_maps(ns, ys, xs, batch_size, height, width, dtype=dtype)


    empirical_saliency_maps = gaussian_filter(