# This file is a modified version of 'observable_utils.py' from the 
# Google DeepMind 'flows_for_atomic_solids' repository:
# https://github.com/google-deepmind/flows_for_atomic_solids/blob/main/utils/observable_utils.py
#
# Original authors: Google DeepMind
# Modifications: Maximilian Schebek (2026)


import numpy as np
import scipy

def _volume_sphere(radius: np.ndarray, dim: int) -> np.ndarray:
    """Computes the volume of a Euclidean ball in `dim` dimensions."""
    c = np.pi ** (dim / 2) / scipy.special.gamma(1 + dim / 2)  # Volume of unit sphere.
    return c * radius**dim

def pairwise_distance_pbc(coordinates, box_length):
    """
    Compute pairwise distances under periodic boundary conditions (PBC).

    Args:
        coordinates: array of shape (N, dim)
        box_length: array of shape (dim,)

    Returns:
        distances: array of shape (N, N)
    """
    diff = coordinates[:, None, :] - coordinates[None, :, :]
    # Apply minimum image convention
    diff -= box_length * np.round(diff / box_length)
    dist = np.linalg.norm(diff, axis=-1)
    return dist

def radial_distribution_function(coordinates, box_length, num_bins=300):
    """
    Computes the radial distribution function g(r).

    Args:
        coordinates: array with shape [..., num_particles, dim]
        box_length: array of shape [dim]
        num_bins: number of bins for the histogram

    Returns:
        gr: array with shape [num_bins, 2], first column is r, second is g(r)
    """
    coordinates = np.asarray(coordinates)
    box_length = np.asarray(box_length)
    
    num_particles, dim = coordinates.shape[-2:]
    if box_length.shape != (dim,):
        raise ValueError(f"`box_length` must be a vector of length {dim}, got {box_length.shape}.")

    min_box_length = np.min(box_length)
    box_volume = np.prod(box_length)

    coordinates = coordinates.reshape(-1, num_particles, dim)
    batch_size = coordinates.shape[0]

    # Compute pairwise distances for each configuration
    all_dr = []
    for b in range(batch_size):
        dr_matrix = pairwise_distance_pbc(coordinates[b], box_length)
        # Take only upper-triangular distances (i < j)
        iu = np.triu_indices(num_particles, k=1)
        all_dr.append(dr_matrix[iu])
    all_dr = np.concatenate(all_dr)

    # Histogram
    hist, bins = np.histogram(all_dr, bins=num_bins, range=(0, min_box_length / 2.0))

    # RDF normalization
    density = num_particles / box_volume
    volume_shell = _volume_sphere(bins[1:], dim) - _volume_sphere(bins[:-1], dim)
    normaliser = volume_shell * density * batch_size * (num_particles - 1) / 2

    gr = np.column_stack(((bins[:-1] + bins[1:]) / 2, hist / normaliser))
    return gr

import torch
import torch.nn.functional as F


def _compute_importance_weights(model_log_probs: torch.Tensor, target_log_probs: torch.Tensor) -> torch.Tensor:
    """
    Returns the normalized importance weights.
    
    Args:
        model_log_probs: tensor of shape (B,)
        target_log_probs: tensor of shape (B,)
    Returns:
        normalized importance weights of shape (B,)
    """
    assert model_log_probs.shape == target_log_probs.shape
    log_diff = target_log_probs - model_log_probs
    return F.softmax(log_diff)

def compute_ess(model_log_probs: torch.Tensor, target_log_probs: torch.Tensor) -> torch.Tensor:
    """
    Compute the effective sample size (ESS) as a percentage.
    
    Args:
        model_log_probs: tensor of shape (B,)
        target_log_probs: tensor of shape (B,)
    Returns:
        ESS as a scalar percentage (0-100)
    """
    weights = _compute_importance_weights(model_log_probs, target_log_probs)
    ess = 100.0 / (torch.sum(weights**2) * weights.numel())
    return ess

def compute_logz(model_log_probs: torch.Tensor, target_log_probs: torch.Tensor) -> torch.Tensor:
    """
    Estimate log of the normalizer ratio: log(Z_target / Z_model)
    
    Args:
        model_log_probs: tensor of shape (B,)
        target_log_probs: tensor of shape (B,)
    Returns:
        log-normalizer difference as a scalar tensor
    """
    assert model_log_probs.shape == target_log_probs.shape
    log_diff = target_log_probs - model_log_probs
    log_sum_exp = torch.logsumexp(log_diff, dim=0)
    return log_sum_exp - torch.log(torch.tensor(target_log_probs.numel(), dtype=log_diff.dtype, device=log_diff.device))