# SPDX-License-Identifier: MIT
"""Check analyse.weibull_fit on synthetic data with known answers.

Drives get Weibull lifetimes, join the study at a random age (left
truncation: only drives still alive at entry are kept) and are censored after
a random follow-up. The fit should recover beta and eta, the profile 95%
interval should cover the true beta about 95% of the time, and the
likelihood-ratio test should reject beta = 1 only about 5% of the time when
the truth is exponential.
"""
import numpy as np

from analyse import weibull_fit


def synth(rng, n, beta, eta):
    entry = np.where(rng.random(n) < 0.3, rng.uniform(0, 4, n), 0.0)
    life = eta * rng.weibull(beta, n)
    keep = life > entry
    entry, life = entry[keep], life[keep]
    follow = rng.uniform(0.5, 6, len(entry))
    exit_ = np.minimum(life, entry + follow)
    event = (life <= entry + follow).astype(int)
    return entry, exit_, event


def main():
    rng = np.random.default_rng(5)
    for beta, eta in ((0.7, 40.0), (1.0, 30.0), (1.6, 12.0), (3.0, 8.0)):
        covered, rejected, est = 0, 0, []
        reps = 200
        for _ in range(reps):
            e, x, d = synth(rng, 6000, beta, eta)
            f = weibull_fit(e, x, d)
            est.append(f["beta"])
            covered += f["beta_lo"] <= beta <= f["beta_hi"]
            rejected += f["p_exponential"] < 0.05
        print(f"beta {beta:3.1f}: mean estimate {np.mean(est):.3f}, CI coverage {covered / reps:.0%}, "
              f"rejects beta=1 in {rejected / reps:.0%}")


if __name__ == "__main__":
    main()