# Copyright (c) 2026 INNOVATIO SAS
# SPDX-License-Identifier: MIT
"""Reproduce the synthetic MSC-P-007 alpha/omega reference calculation."""

import numpy as np

SEED = 20260819
N = 300
B = 2000
LOADINGS = np.array([0.95, 0.90, 0.80, 0.55, 0.35, 0.20])
RESIDUAL_VARIANCES = 1 - LOADINGS**2


def alpha(data: np.ndarray) -> float:
    correlation = np.corrcoef(data, rowvar=False)
    k = correlation.shape[0]
    return k / (k - 1) * (1 - np.trace(correlation) / correlation.sum())


def omega_one_factor(data: np.ndarray) -> tuple[float, bool, int]:
    """Estimate total omega with a deterministic one-factor PAF iteration."""
    correlation = np.corrcoef(data, rowvar=False)
    communalities = np.clip(1 - 1 / np.diag(np.linalg.inv(correlation)), 1e-6, 0.999)
    converged = False
    iterations = 0
    for iterations in range(1, 201):
        reduced = correlation.copy()
        np.fill_diagonal(reduced, communalities)
        values, vectors = np.linalg.eigh(reduced)
        loading = vectors[:, -1] * np.sqrt(max(values[-1], 0))
        if loading.sum() < 0:
            loading = -loading
        updated = np.clip(loading**2, 1e-6, 0.999)
        if np.max(np.abs(updated - communalities)) < 1e-9:
            converged = True
            break
        communalities = updated
    omega = loading.sum() ** 2 / (loading.sum() ** 2 + np.sum(1 - loading**2))
    return omega, converged, iterations


def population_values() -> tuple[float, float]:
    matrix = np.outer(LOADINGS, LOADINGS)
    np.fill_diagonal(matrix, 1)
    k = len(LOADINGS)
    alpha_value = k / (k - 1) * (1 - np.trace(matrix) / matrix.sum())
    omega_value = LOADINGS.sum() ** 2 / (LOADINGS.sum() ** 2 + RESIDUAL_VARIANCES.sum())
    return alpha_value, omega_value


def main() -> None:
    rng = np.random.default_rng(SEED)
    estimates = np.empty((B, 3))
    non_convergences = 0
    max_iterations = 0
    for index in range(B):
        factor = rng.normal(size=(N, 1))
        errors = rng.normal(size=(N, 6)) * np.sqrt(RESIDUAL_VARIANCES)
        data = factor * LOADINGS + errors
        estimates[index, 0] = alpha(data)
        estimates[index, 1], converged, iterations = omega_one_factor(data)
        non_convergences += int(not converged)
        max_iterations = max(max_iterations, iterations)
        estimates[index, 2] = estimates[index, 1] - estimates[index, 0]

    alpha_population, omega_population = population_values()
    print("population", alpha_population, omega_population, omega_population - alpha_population)
    print("paf_diagnostics", "non_convergences", non_convergences, "max_iterations", max_iterations)
    for column, name in enumerate(("alpha", "omega", "omega_minus_alpha")):
        lower, upper = np.quantile(estimates[:, column], [0.025, 0.975])
        print(name, estimates[:, column].mean(), lower, upper)


if __name__ == "__main__":
    main()