"""Conditional cluster holdout experiment, not a per-star age posterior.

Train one temperature–period relation on the fixed Praesepe author subset;
compare two explicitly imposed age exponents on other clusters. Period errors
are absent in this combined table, so no precise per-star likelihood is invented.
"""
from pathlib import Path
import argparse
import json
import numpy as np
from inspect_cluster_samples import read_benchmarks, read_ruprecht


def fit_reference(stars, robust=True, degree=4):
    temperature = np.array([row['teff_k'] for row in stars])
    period = np.array([row['period_days'] for row in stars])
    coordinate = (temperature - 5100) / 900
    design = np.polynomial.polynomial.polyvander(coordinate, degree)
    weights = np.ones(len(stars))
    for _ in range(12 if robust else 1):
        coefficients = np.linalg.lstsq(design * np.sqrt(weights[:, None]), np.log(period) * np.sqrt(weights), rcond=None)[0]
        residual = np.log(period) - design @ coefficients
        scale = max(.005, 1.4826 * np.median(abs(residual - np.median(residual))))
        weights = np.minimum(1., 1.5 * scale / np.maximum(abs(residual), 1e-12))
    return coefficients


def reference_period(coefficients, temperature):
    return np.exp(np.polynomial.polynomial.polyval((np.asarray(temperature) - 5100) / 900, coefficients))


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--inputs', type=Path, default=Path('inputs/curtis2020-source'))
    parser.add_argument('--output', type=Path, default=Path('results/age-holdout-results.json'))
    arguments = parser.parse_args()
    stars = read_benchmarks(arguments.inputs / 'Benchmark_Clusters.txt')
    chosen = [row for row in stars if 4200 <= row['teff_k'] <= 6000]
    training = [row for row in chosen if row['cluster'] == 'Praesepe']
    assert len(training) > 50
    coefficients = fit_reference(training)
    ordinary = fit_reference(training, robust=False)
    random = np.random.default_rng(20261103)
    bootstrap = [fit_reference([training[index] for index in random.integers(0, len(training), len(training))]) for _ in range(300)]
    rows = []
    for cluster in ['NGC 6811', 'NGC 752', 'NGC 6819', 'Ruprecht 147']:
        cluster_stars = [row for row in chosen if row['cluster'] == cluster]
        for low, high in [(4200, 4600), (4600, 5000), (5000, 5500), (5500, 6000)]:
            selected = [row for row in cluster_stars if low <= row['teff_k'] < high]
            if not selected:
                continue
            temperatures = np.array([row['teff_k'] for row in selected]); periods = np.array([row['period_days'] for row in selected])
            for exponent in [.5, .62]:
                inferred = .67 * (periods / reference_period(coefficients, temperatures)) ** (1 / exponent)
                boot_medians = [float(np.median(.67 * (periods / reference_period(fit, temperatures)) ** (1 / exponent))) for fit in bootstrap]
                alternatives = [.67 * (periods / reference_period(ordinary, temperatures)) ** (1 / exponent)]
                for shift in [-50, 50]:
                    shifted = np.clip(temperatures + shift, 4200, 6000)
                    alternatives.append(.67 * (periods / reference_period(coefficients, shifted)) ** (1 / exponent))
                fixed_age = selected[0]['age_gyr']
                rows.append({'withheld_cluster': cluster, 'fixed_published_age_gyr': fixed_age, 'temperature_bin_k': [low, high], 'stars': len(selected),
                             'imposed_age_exponent': exponent, 'median_predicted_age_gyr': float(np.median(inferred)),
                             'median_age_ratio_to_fixed_cluster_age': float(np.median(inferred) / fixed_age),
                             'median_age_training_star_bootstrap16_84_gyr': np.percentile(boot_medians, [16, 84]).tolist(),
                             'median_age_ordinary_reference_fit_gyr': float(np.median(alternatives[0])),
                             'temperature_plusminus50k_clipped_to_domain_median_age_gyr': [float(np.median(value)) for value in alternatives[1:]],
                             'temperature_sensitivity_note': 'Shifted temperatures are clipped to the declared4200–6000K domain; this is a bounded sensitivity probe, not a temperature posterior.',
                             'qualification': 'Sparse bin, descriptive only' if len(selected) < 5 else 'Conditional temperature-limited author subset; no new membership vetting'})
    assert any(row['withheld_cluster'] == 'NGC 6811' and row['temperature_bin_k'] == [4200, 4600] and row['median_age_ratio_to_fixed_cluster_age'] < .9 for row in rows)
    report = {'status': 'COMPLETED conditional cluster holdout, no per-star likelihood/posterior or independent absolute age claim',
              'training': {'cluster': 'Praesepe', 'age_gyr_fixed': .67, 'temperature_domain_k': [4200, 6000], 'stars': len(training), 'degree_log_period_polynomial': 4, 'coefficients_in_scaled_temperature': coefficients.tolist(), 'robust_rule': 'Huber IRLS on log period, twelve iterations, threshold1.5 robust MAD scales; ordinary least squares comparison retained', 'star_bootstrap_replicates': 300, 'seed': 20261103},
              'qualification': 'The imposed separable laws are deliberately tested against stalling. Cluster ages and existing author membership/temperature scale are conditioned on. Bootstrap changes only training-star composition, not absolute cluster ages, temperature systematics, age spread, detection selection or per-star period errors. All periods and star rows remain local inputs.',
              'withheld_bins': rows}
    arguments.output.parent.mkdir(parents=True, exist_ok=True)
    arguments.output.write_text(json.dumps(report, indent=2) + '\n')
    print(json.dumps({'status': report['status'], 'training_stars': len(training), 'bins': [{key: row[key] for key in ['withheld_cluster', 'temperature_bin_k', 'stars', 'imposed_age_exponent', 'median_predicted_age_gyr']} for row in rows]}, indent=2))


if __name__ == '__main__': main()