#!/usr/bin/env python3
"""Chronology and joint longevity/maturation checks fixed in analysis-plan.md."""

from __future__ import annotations

from pathlib import Path

import numpy as np
import pandas as pd
from scipy.optimize import minimize_scalar


ROOT = Path(__file__).resolve().parent
RESULTS = ROOT / "results"

NAMES = ["Shem", "Arphaxad", "Cainan", "Shelah", "Eber", "Peleg", "Reu", "Serug", "Nahor", "Terah"]
FATHER_AGES = np.array([100, 135, 130, 130, 134, 130, 132, 130, 79, 130])
LIFESPANS = np.array([600, 565, 460, 533, 504, 339, 339, 330, 208, 205])


def chronology():
    # Shem is time zero. Each next birth occurs at the preceding named man's
    # recorded fathering age.
    births = [0]
    for g in range(1, len(NAMES)):
        births.append(births[-1] + int(FATHER_AGES[g - 1]))
    deaths = np.asarray(births) + LIFESPANS
    peleg_death = deaths[NAMES.index("Peleg")]
    rows = []
    for g, name in enumerate(NAMES):
        rows.append(
            {
                "generation": g,
                "name": name,
                "birth_relative_to_Shem": births[g],
                "death_relative_to_Shem": deaths[g],
                "birth_relative_to_Peleg_death": births[g] - peleg_death,
                "born_before_Peleg_death": births[g] < peleg_death,
            }
        )
    return pd.DataFrame(rows)


def combined_maturation():
    """Probability of the delayed-run / break / immediate rebound pattern.

    Arphaxad starts homozygous. Six transmissions produce Cainan through
    Serug. A full ambient outcross then makes Nahor heterozygous (and early)
    with probability one. Terah must receive the founder allele from his
    heterozygous father (1/2) and from his mother (q_post), at every required
    recessive factor.
    """
    rows = []
    qpres = [0.75, 0.80, 0.85, 0.885, 0.90, 0.95, 1.0]
    qposts = [0.25, 0.50, 0.75, 1.0]
    for m in range(1, 7):
        for qpre in qpres:
            for qpost in qposts:
                p_pre_run = qpre ** (6 * m)
                p_rebound = (qpost / 2) ** m
                rows.append(
                    {
                        "M_recessive_factors": m,
                        "q_pre_maternal_founder_allele": qpre,
                        "q_Terah_mother_transmits_founder_allele": qpost,
                        "P_seven_delayed_pre_Babel": p_pre_run,
                        "P_Nahor_early_given_full_outcross": 1.0,
                        "P_Terah_delayed_rebound": p_rebound,
                        "P_complete_pattern_conditional_on_Babel_outcross": p_pre_run * p_rebound,
                    }
                )
    return pd.DataFrame(rows)


def clustering_null():
    # The seven ages Arphaxad..Serug lie in [130,135]. Report both a window
    # fixed in advance and a range selected after seeing the values.
    n = 7
    lo, hi = 25, 300
    discrete_values = hi - lo + 1
    fixed_window_values = 6  # 130,131,...,135
    p_fixed = (fixed_window_values / discrete_values) ** n
    # Continuous approximation for the probability that the sample range is
    # at most five years anywhere in a 275-year interval.
    x = 5 / (hi - lo)
    p_selected = n * x ** (n - 1) - (n - 1) * x**n
    return pd.DataFrame(
        [
            {
                "null": "seven random named births uniformly from ages 25..300, fixed 130..135 window",
                "probability": p_fixed,
            },
            {
                "null": "seven random named births uniformly from ages 25..300, any five-year sample range",
                "probability": p_selected,
            },
        ]
    )


def developmental_clock_comparison():
    """Compare a constant late-onset model with uniform lifespan scaling."""
    # The seven-person cluster used in the stated argument: Arphaxad..Serug.
    ages = FATHER_AGES[1:8].astype(float)
    lives = LIFESPANS[1:8].astype(float)
    const_mean = ages.mean()
    pred_const = np.repeat(const_mean, len(ages))
    # Uniform slow-clock model through the biological origin: onset=c*lifespan.
    c = float(np.dot(lives, ages) / np.dot(lives, lives))
    pred_scaled = c * lives
    rows = []
    for model, pred, parameter in [
        ("constant maturation age", pred_const, const_mean),
        ("uniform clock: age proportional to lifespan", pred_scaled, c),
    ]:
        rss = float(np.sum((ages - pred) ** 2))
        # Both models have one fitted mean parameter, so their Gaussian AIC
        # comparison reduces to n*log(RSS/n); additive constants cancel.
        aic_relative = len(ages) * np.log(max(rss / len(ages), 1e-300)) + 2
        rows.append(
            {
                "model": model,
                "fitted_parameter": parameter,
                "RSS": rss,
                "relative_AIC_component": aic_relative,
                "RMSE": np.sqrt(rss / len(ages)),
                "predicted_Serug_fathering_age": pred[-1],
                "observed_Serug_fathering_age": ages[-1],
            }
        )
    out = pd.DataFrame(rows)
    out["delta_AIC_from_best"] = out.relative_AIC_component - out.relative_AIC_component.min()
    return out


def spouse_and_rebound_inference():
    """Additive phenotype-equivalent scores of the two post-Babel mothers."""
    serug, nahor, terah = 330.0, 208.0, 205.0
    # In an additive mid-parent model L_child=(L_father+L_mother)/2.
    nahor_mother_equivalent = 2 * nahor - serug
    terah_mother_equivalent = 2 * terah - nahor
    rows = []
    for ambient in [40, 60, 80, 100]:
        d = 600 - ambient
        q_from_terah_mother = (terah_mother_equivalent - ambient) / d
        for m in range(1, 7):
            rebound = (q_from_terah_mother / 2) ** m
            rows.append(
                {
                    "ambient": ambient,
                    "Nahor_mother_lifespan_equivalent": nahor_mother_equivalent,
                    "Terah_mother_lifespan_equivalent": terah_mother_equivalent,
                    "Terah_mother_founder_fraction_if_traits_track": q_from_terah_mother,
                    "M_recessive_factors": m,
                    "P_Terah_rebound_if_maturation_allele_frequency_tracks_fraction": rebound,
                    "P_full_pattern_with_qpre_0.885": (0.885 ** (6 * m)) * rebound,
                }
            )
    return pd.DataFrame(rows)


def castle_wright_pre_babel():
    """Opus-style effective-factor estimate using only Shem through Serug."""
    rows = []
    for label, lives in [
        ("Shelah_533_primary", LIFESPANS[:8].astype(float)),
        ("Shelah_460_sensitivity", np.array([600, 565, 460, 460, 504, 339, 339, 330], dtype=float)),
    ]:
        generations = np.arange(8, dtype=float)
        for ambient in [40, 60, 80, 100]:
            d = 600 - ambient
            objective = lambda r: np.sum((lives - (ambient + d * r**generations)) ** 2)
            z = minimize_scalar(objective, bounds=(0.5, 0.9999), method="bounded")
            p = z.x**generations
            residuals = lives - (ambient + d * p)
            observed_sd = np.std(residuals, ddof=1)
            binomial_constant = d * np.sqrt(np.mean(p * (1 - p)))
            n_eff = (binomial_constant / observed_sd) ** 2
            rows.append(
                {
                    "series": label,
                    "ambient": ambient,
                    "retention_pre_Babel": z.x,
                    "outloss_pre_Babel": 1 - z.x,
                    "observed_residual_SD": observed_sd,
                    "predicted_scatter_numerator": binomial_constant,
                    "N_effective_if_all_residual_scatter_is_segregation": n_eff,
                }
            )
    return pd.DataFrame(rows)


def environmental_variance_sensitivity(cw: pd.DataFrame):
    rows = []
    for row in cw.itertuples(index=False):
        for env_sd in [0, 10, 20, 30, 40, 45]:
            genetic_variance = row.observed_residual_SD**2 - env_sd**2
            n_eff = (
                row.predicted_scatter_numerator**2 / genetic_variance
                if genetic_variance > 0
                else np.inf
            )
            rows.append(
                {
                    "series": row.series,
                    "ambient": row.ambient,
                    "assumed_non_genetic_SD": env_sd,
                    "remaining_genetic_variance": max(0.0, genetic_variance),
                    "N_effective": n_eff,
                }
            )
    return pd.DataFrame(rows)


def cross_text_half_excess():
    traditions = {
        "LXX-derived supplied series": (
            ["Shem", "Arphaxad", "Cainan", "Shelah", "Eber", "Peleg", "Reu", "Serug", "Nahor", "Terah"],
            [600, 565, 460, 533, 504, 339, 339, 330, 208, 205],
        ),
        "Masoretic lifespan series": (
            ["Shem", "Arphaxad", "Shelah", "Eber", "Peleg", "Reu", "Serug", "Nahor", "Terah"],
            [600, 438, 433, 464, 239, 239, 230, 148, 205],
        ),
    }
    rows = []
    for tradition, (names, lives) in traditions.items():
        for i in range(1, len(lives)):
            implied = 2 * lives[i] - lives[i - 1]
            rows.append(
                {
                    "tradition": tradition,
                    "boundary": f"{names[i-1]}->{names[i]}",
                    "ambient_implied_by_full_additive_outcross": implied,
                    "inside_40_100": 40 <= implied <= 100,
                }
            )
    return pd.DataFrame(rows)


def main():
    c = chronology()
    m = combined_maturation()
    z = clustering_null()
    d = developmental_clock_comparison()
    s = spouse_and_rebound_inference()
    cw = castle_wright_pre_babel()
    env = environmental_variance_sensitivity(cw)
    xt = cross_text_half_excess()
    c.to_csv(RESULTS / "babel_chronology.csv", index=False)
    m.to_csv(RESULTS / "combined_maturation_pattern.csv", index=False)
    z.to_csv(RESULTS / "fathering_cluster_null.csv", index=False)
    d.to_csv(RESULTS / "developmental_clock_comparison.csv", index=False)
    s.to_csv(RESULTS / "spouse_equivalent_and_rebound.csv", index=False)
    cw.to_csv(RESULTS / "castle_wright_pre_babel.csv", index=False)
    env.to_csv(RESULTS / "castle_wright_environment_sensitivity.csv", index=False)
    xt.to_csv(RESULTS / "cross_text_half_excess.csv", index=False)
    print(c.to_string(index=False))
    print("\nMaturation examples at q_pre=.885:\n")
    print(
        m[(m.q_pre_maternal_founder_allele == 0.885) & (m.q_Terah_mother_transmits_founder_allele.isin([0.5, 1.0]))]
        .to_string(index=False)
    )
    print("\nClustering nulls:\n", z.to_string(index=False))
    print("\nDevelopmental clock comparison:\n", d.to_string(index=False))
    print("\nSpouse-equivalent rebound inference:\n", s.to_string(index=False))
    print("\nCastle-Wright-style pre-Babel estimate:\n", cw.to_string(index=False))
    print("\nEnvironmental variance sensitivity (ambient 80 only):\n", env[env.ambient == 80].to_string(index=False))
    print("\nCross-text half-excess hits in ambient range:\n", xt[xt.inside_40_100].to_string(index=False))


if __name__ == "__main__":
    main()
