"""Small, differentiable Fisher-observability utilities."""
import torch

def fisher_information(jacobians, noise_vars, prior_cov=None, eps=1e-8):
    """Accumulate J^T R^-1 J for a sequence of (observation_dim,state_dim) J."""
    if not torch.is_tensor(noise_vars):
        noise_vars = torch.as_tensor(noise_vars, dtype=jacobians[0].dtype, device=jacobians[0].device)
    d = jacobians[0].shape[-1]
    I = torch.zeros((d, d), dtype=jacobians[0].dtype, device=jacobians[0].device)
    inv = torch.diag(1.0 / noise_vars)
    for J in jacobians:
        I = I + J.T @ inv @ J
    if prior_cov is not None:
        I = I + torch.linalg.inv(prior_cov)
    return I

def observability_regularizer(I, alpha=1.0, beta=0.0, eps=1e-6):
    """The proposed -logdet plus optional smooth condition-number penalty."""
    eig = torch.linalg.eigvalsh(I)
    logdet = torch.logdet(I + eps * torch.eye(I.shape[-1], device=I.device, dtype=I.dtype))
    cond = eig[-1] / (eig[0] + eps)
    return alpha * (-logdet) + beta * cond

def bearing_brightness_jacobian(state, brightness=True, c=1.0):
    x, y = state[0], state[1]
    r2 = x*x + y*y
    Hb = torch.stack((-y/r2, x/r2, torch.zeros_like(x), torch.zeros_like(x))).reshape(1,4)
    if not brightness:
        return Hb
    r4 = r2*r2
    Hl = torch.stack((-2*c*x/r4, -2*c*y/r4, torch.zeros_like(x), torch.zeros_like(x))).reshape(1,4)
    return torch.cat((Hb, Hl), dim=0)