"""AUREOLE v1: small, explicit linear-Gaussian reference operations.

This is a CPU research witness, not a neural renderer or DLSS implementation.
All matrices use float64. Measurements are conditionally independent unless
an explicit joint noise covariance is provided. Never reuse correlated rays
as independent observations.
"""
from __future__ import annotations
import numpy as np


def sym(a):
    return (a + a.T) * 0.5


def sqrt_psd(a):
    vals, vecs = np.linalg.eigh(sym(a))
    if vals.min() < -1e-10:
        raise ValueError("Matrix is not positive semidefinite")
    return (vecs * np.sqrt(np.maximum(vals, 0))) @ vecs.T


def observe(mean, covariance, h, value, noise_variance):
    """One scalar Gaussian update, using Joseph form for numerical stability."""
    h = np.asarray(h, dtype=float)
    if noise_variance <= 0:
        raise ValueError("Noise variance must be positive")
    ph = covariance @ h
    innovation_var = float(noise_variance + h @ ph)
    k = ph / innovation_var
    m = mean + k * (value - h @ mean)
    residual = np.eye(len(mean)) - np.outer(k, h)
    p = residual @ covariance @ residual.T + noise_variance * np.outer(k, k)
    return m, sym(p)


def query_value(covariance, task_metric, h, noise_variance):
    """Exact one-query quadratic Bayes risk reduction; future task map is fixed."""
    ph = covariance @ h
    return float(ph @ task_metric @ ph / (noise_variance + h @ ph))


def batch_covariance(p, h, noise):
    """Exact joint update, including correlated measurement noise."""
    v = h @ p @ h.T + noise
    return sym(p - p @ h.T @ np.linalg.solve(v, h @ p))


def future_metric(dynamics, task_maps, weights):
    """W=sum w_tau Phi_tau.T C_tau.T C_tau Phi_tau; tau starts at zero."""
    p = np.eye(dynamics.shape[0])
    w = np.zeros_like(dynamics, dtype=float)
    for c, weight in zip(task_maps, weights):
        w += weight * p.T @ c.T @ c @ p
        p = dynamics @ p
    return sym(w)


def closed_observation_basis(dynamics_list, channel_rows, tol=1e-10):
    """Smallest common A.T-invariant column span containing all C.T and H.T.

    Returns orthonormal U with quotient z=U.T x. This exact finite-dimensional
    construction is expensive for a full scene; intended for tests/local blocks.
    """
    def orth(x):
        u, s, _ = np.linalg.svd(x, full_matrices=False)
        return u[:, s > tol]
    basis = orth(np.asarray(channel_rows, float).T)
    while True:
        enlarged = orth(np.concatenate([basis] + [a.T @ basis for a in dynamics_list], axis=1))
        if enlarged.shape[1] == basis.shape[1]:
            return enlarged
        basis = enlarged


def risk(p, w):
    return float(np.trace(w @ p))


def transform_coding(source_covariance, task_metric, rank):
    """Optimal exact rank-r linear encoding of a Gaussian source innovation.

    This source covariance is NOT automatically the renderer's posterior P.
    The encoder must actually possess the source being compressed.
    """
    root = sqrt_psd(source_covariance)
    vals, vectors = np.linalg.eigh(sym(root @ task_metric @ root))
    order = np.argsort(vals)[::-1]
    vals, vectors = vals[order], vectors[:, order]
    kept = vectors[:, :rank]
    error_cov = root @ (np.eye(len(vals)) - kept @ kept.T) @ root
    return kept, np.maximum(vals, 0), sym(error_cov)


def area_mixture(operators, weights):
    """Convex area mixing under the manuscript's independent patch assumptions."""
    weights = np.asarray(weights, float)
    if np.min(weights) < 0 or not np.isclose(weights.sum(), 1):
        raise ValueError("Mixture weights must form a simplex")
    return np.einsum("k,kij->ij", weights, operators)