comprehensive_exomoon_search.py

#!/usr/bin/env python3
"""
Comprehensive Exomoon Detection Pipeline
=========================================
This script implements multiple exomoon detection methods with improvements:
- Skew/Asymmetry Detection (improved)
- Variability Detection (improved)
- Transit Timing Variation (TTV) Detection (new)
- Shoulder Detection (improved)

All methods include:
- Proper error handling and validation
- Statistical significance testing
- Data quality metrics
- Comprehensive result reporting
"""

import os
import sys
import numpy as np
import pandas as pd
import argparse
from functools import partial
from concurrent.futures import ProcessPoolExecutor, as_completed
from tqdm import tqdm
import warnings
warnings.filterwarnings('ignore')

from exomoon_core import (
    validate_parameters,
    ingress_egress_area_per_transit,
    compute_skew_series,
    compute_skew_distribution_score,
    compute_skew_distribution_multi_metrics,
    analyze_periodogram_structure,
    compute_multi_ttv_metrics,
    detect_periodic_signal,
    compute_shoulder_snr,
    compute_variability_metrics,
    extract_transit_midpoints_and_durations,
    compute_physical_constraints,
    optimize_period_search_range,
)
from astropy.timeseries import BoxLeastSquares
from pipeline.ephemeris import build_ephemeris_prior
from pipeline.transit_search import refine_transit_params_with_guard
from pipeline.lightcurve_io import load_tess_lightcurve

# ============================================================================
# Helper Functions
# ============================================================================


def extract_transit_times(t, flux, flux_err, P, T0, dur, window_factor=2.0):
    """Extract individual transit mid-times."""
    n_min = int(np.floor((t.min() - T0) / P))
    n_max = int(np.ceil((t.max() - T0) / P))
    
    transit_times = []
    transit_errors = []
    
    for n in range(n_min, n_max + 1):
        expected_time = T0 + n * P
        window = window_factor * dur
        
        mask = np.abs(t - expected_time) < window
        if mask.sum() < 5:
            continue
            
        t_transit = t[mask]
        flux_transit = flux[mask]
        err_transit = flux_err[mask] if flux_err is not None else np.ones_like(flux_transit) * np.nanstd(flux[~mask])
        
        try:
            bls = BoxLeastSquares(t_transit, flux_transit)
            periods = np.array([P])
            durations = np.array([dur])
            
            result = bls.power(periods, durations)
            if len(result.transit_time) > 0:
                transit_time = float(result.transit_time[0])
                model = bls.model(t_transit, P, dur, transit_time)
                residuals = flux_transit - model
                time_uncertainty = dur / (2 * np.sqrt(np.sum(1 / err_transit**2)))
                
                transit_times.append(transit_time)
                transit_errors.append(time_uncertainty)
        except:
            continue
    
    return np.array(transit_times), np.array(transit_errors)


def fit_linear_ephemeris(transit_times, transit_errors=None):
    """Fit linear ephemeris to transit times."""
    if len(transit_times) < 2:
        raise ValueError("Need at least 2 transits")
    
    if transit_errors is None:
        transit_errors = np.ones_like(transit_times) * np.nanstd(np.diff(transit_times))
    
    n = np.arange(len(transit_times))
    weights = 1.0 / transit_errors**2
    
    sum_w = np.sum(weights)
    sum_wn = np.sum(weights * n)
    sum_wt = np.sum(weights * transit_times)
    sum_wn2 = np.sum(weights * n**2)
    sum_wnt = np.sum(weights * n * transit_times)
    
    det = sum_w * sum_wn2 - sum_wn**2
    if det == 0:
        raise ValueError("Singular matrix in ephemeris fit")
    
    P = (sum_w * sum_wnt - sum_wn * sum_wt) / det
    T0 = (sum_wn2 * sum_wt - sum_wn * sum_wnt) / det
    
    residuals = transit_times - (T0 + n * P)
    chi2_val = np.sum(weights * residuals**2)
    dof = len(transit_times) - 2
    sigma = np.sqrt(chi2_val / dof) if dof > 0 else 1.0
    
    P_err = sigma * np.sqrt(sum_w / det)
    T0_err = sigma * np.sqrt(sum_wn2 / det)
    
    return P, T0, P_err, T0_err


def detect_ttv_period(transit_times, oc_residuals, oc_errors=None, 
                     min_period=None, max_period=None):
    """Search for periodic TTV signal."""
    if len(transit_times) < 3:
        return np.nan, np.nan, np.nan, np.nan, np.nan
    
    time_span = transit_times.max() - transit_times.min()
    if min_period is None:
        min_period = 0.1 * np.median(np.diff(transit_times))
    if max_period is None:
        max_period = min(10 * np.median(np.diff(transit_times)), time_span / 2)
    
    if min_period >= max_period:
        return np.nan, np.nan, np.nan, np.nan, np.nan
    
    try:
        if oc_errors is not None:
            ls = LombScargle(transit_times, oc_residuals, dy=oc_errors)
        else:
            ls = LombScargle(transit_times, oc_residuals)
        
        freqs = np.linspace(1/max_period, 1/min_period, 10000)
        power = ls.power(freqs)
        periods = 1.0 / freqs
        
        idx_max = np.argmax(power)
        best_period = periods[idx_max]
        best_power = float(power[idx_max])
        
        N = len(transit_times)
        fap = N * np.exp(-best_power)
        
        return periods, power, best_period, best_power, fap
    except:
        return np.nan, np.nan, np.nan, np.nan, np.nan




# ============================================================================
# Main Processing Function
# ============================================================================

def process_planet(row_dict, cadence_policy="short_then_any", cache_dir=None, cache_enabled=True):
    """
    Process a single planet with all detection methods.
    
    Parameters
    ----------
    row_dict : dict
        Dictionary with planet information (from CSV row)
        
    Returns
    -------
    result : dict
        Comprehensive results dictionary
    """
    planet_name = row_dict.get('planet_name', 'Unknown')
    tic_id = row_dict.get('tic_id', '')
    P_expected = float(row_dict.get('orbital_period_days', np.nan))
    M_star = row_dict.get('M_star', row_dict.get('m_star', np.nan))
    M_planet = row_dict.get('M_planet', row_dict.get('m_planet', np.nan))
    R_planet = row_dict.get('R_planet_jup', row_dict.get('r_planet_jup', np.nan))
    rho_planet = row_dict.get('rho_planet', np.nan)
    
    result = {
        'planet_name': planet_name,
        'tic_id': tic_id,
        'P_expected': P_expected,
        'Status': 'Error',
        'Reason': '',
        'num_transits_observed': 0,
        'expected_transits': np.nan,
        'transit_coverage': np.nan,
        'data_quality_score': 0.0,
    }
    
    # Initialize all method results
    for prefix in ['skew', 'ttv', 'shoulder', 'variability']:
        for key in ['P_refined', 'T0', 'dur', 'power', 'fap', 'amplitude', 'snr']:
            result[f'{prefix}_{key}'] = np.nan
    for key in ['duration_cv', 'duration_period', 'duration_power', 'duration_fap', 'bls_warning']:
        result[key] = np.nan
    for key in ['skew_dist_pvalue', 'skew_dist_score', 'skew_dist_inconsistency_score', 'skew_dist_n_neg',
                'skew_dist_n_zero', 'skew_dist_n_pos', 'skew_distribution_score']:
        result[key] = np.nan
    for key in [
        'single_skew_score', 'single_ttv_score', 'single_shoulder_score',
        'single_variability_score', 'single_duration_score',
        'single_skew_dist_consistency_score', 'single_data_quality_bonus',
        'single_transit_reliability', 'combined_single_moon_score',
        'multi_skew_structure_score', 'multi_ttv_structure_score',
        'multi_distribution_score', 'multi_shoulder_score',
        'multi_variability_score', 'multi_duration_score',
        'multi_data_quality_bonus', 'multi_transit_reliability',
        'combined_multi_moon_score',
        'multi_skew_dominant_period', 'multi_skew_dominant_power',
        'multi_skew_peak_ratio_2to1', 'multi_skew_spectral_entropy',
        'multi_skew_topk_periods', 'multi_skew_topk_powers',
        'multi_skew_distribution_score',
        'multi_ttv_rms', 'multi_ttv_dominant_period',
        'multi_ttv_dominant_power', 'multi_ttv_dominant_fap',
        'multi_ttv_fap_penalty', 'multi_ttv_orbit_gate', 'multi_ttv_rms_gate',
        'multi_ttv_score_effective',
        'multi_ttv_peak_ratio_2to1', 'multi_ttv_spectral_entropy',
        'multi_ttv_topk_periods', 'multi_ttv_topk_powers',
    ]:
        result[key] = np.nan
    for key in ['skew_score', 'ttv_score', 'shoulder_score', 'variability_score',
                'duration_score', 'data_quality_bonus']:
        result[key] = np.nan
    
    try:
        # Download light curve using TIC ID
        if not tic_id or pd.isna(tic_id):
            raise ValueError("ERR_NO_TIC")
        
        # Handle TIC ID format - extract the numeric part
        tic_str = str(tic_id).strip()
        
        # Extract numeric part (remove "TIC " prefix if present, handle any whitespace)
        tic_num_str = tic_str.replace('TIC', '').strip()
        
        # Try to convert to integer and use string format "TIC {number}"
        try:
            tic_num = int(tic_num_str)
            search_str = f"TIC {tic_num}"
        except ValueError:
            # If conversion fails, try using original string
            search_str = tic_str
        
        # Search for light curve with cadence policy (use_temp_download to avoid disk bloat)
        t, flux, flux_err, meta = load_tess_lightcurve(
            search_str,
            cadence_policy=cadence_policy,
            cache_dir=cache_dir,
            cache_enabled=cache_enabled,
            use_temp_download=True,
        )
        if t is None or flux is None:
            err_code = meta.get("error", "ERR_LC_NOT_FOUND") if meta else "ERR_LC_NOT_FOUND"
            raise ValueError(err_code)
        if len(t) < 50:
            raise ValueError(f'ERR_INSUFFICIENT_POINTS: {len(t)}')
        
        # Refine transit parameters with shared BLS + ephemeris prior
        P_expected_val = P_expected if np.isfinite(P_expected) and P_expected > 0 else None
        prior = build_ephemeris_prior(P_expected_val)
        bls_ref = refine_transit_params_with_guard(t, flux, P_expected=P_expected_val)
        if bls_ref.get("warning_flag") == "ERR_BLS_DEGENERATE":
            raise ValueError("ERR_BLS_DEGENERATE")
        if bls_ref.get("warning_flag"):
            result["Reason"] = bls_ref.get("warning_flag")
            result["bls_warning"] = bls_ref.get("warning_flag")
        P_ref = bls_ref["period"]
        dur = bls_ref["duration"]
        T0 = bls_ref["t0"]
        
        # Validate parameters
        validate_parameters(P_ref, dur, t, flux)
        
        # Create transit model and compute residuals
        bls = BoxLeastSquares(t, flux)
        model = bls.model(t, P_ref, dur, T0)
        resid = flux - model
        
        # Compute data quality metrics
        num_transits = int((t.max() - t.min()) / P_ref) if P_ref > 0 else 0
        expected_transits = num_transits
        # Estimate observed transits from extracted midpoints when possible
        try:
            t_times, _, _, _ = extract_transit_midpoints_and_durations(t, flux, P_ref, T0)
            observed_transits = len(t_times)
            if expected_transits > 0:
                result['transit_coverage'] = observed_transits / expected_transits
        except Exception:
            observed_transits = num_transits
        result['expected_transits'] = expected_transits
        result['num_transits_observed'] = observed_transits

        snr_per_transit = np.nanstd(model) / np.nanstd(resid) if np.nanstd(resid) > 0 else 0
        data_completeness = 1.0 - np.sum(np.isnan(flux)) / len(flux)
        
        result['data_quality_score'] = min(snr_per_transit * data_completeness / 10.0, 1.0)
        
        # ====================================================================
        # Optional: physical constraints for period search
        # ====================================================================
        period_search_params = None
        try:
            if all(np.isfinite(x) for x in [M_star, M_planet, R_planet, rho_planet]):
                phys = compute_physical_constraints(P_ref, float(M_star), float(M_planet),
                                                   float(R_planet), float(rho_planet))
                period_search_params = optimize_period_search_range(
                    t,
                    phys['beat_period_5th'],
                    phys['beat_period_95th'],
                )
        except Exception:
            period_search_params = None

        # Save series for dual single/multi scoring later.
        S_n = None
        ttv_times = None
        ttv_oc = None

        # ====================================================================
        # 1. SKEW DETECTION
        # ====================================================================
        try:
            Tn, Ai, Ae = ingress_egress_area_per_transit(t, resid, P_ref, T0, dur)
            Tn_ok, S_n = compute_skew_series(Tn, Ai, Ae)
            if period_search_params:
                P_skew, skew_power, skew_fap = detect_periodic_signal(
                    Tn_ok, S_n,
                    min_period=period_search_params['min_period'],
                    max_period=period_search_params['max_period'],
                )
            else:
                P_skew, skew_power, skew_fap = detect_periodic_signal(Tn_ok, S_n)
            
            result['skew_P_refined'] = P_ref
            result['skew_T0'] = T0
            result['skew_dur'] = dur
            result['skew_P_skew'] = P_skew
            result['skew_power'] = skew_power
            result['skew_fap'] = skew_fap
            result['skew_amplitude'] = float(np.nanstd(S_n))
            result['skew_num_transits'] = len(Tn_ok)
            result.update(compute_skew_distribution_score(S_n))
            multi_struct = analyze_periodogram_structure(
                Tn_ok, S_n,
                min_period=period_search_params['min_period'] if period_search_params else None,
                max_period=period_search_params['max_period'] if period_search_params else None,
                max_frequency=None,
                top_k=3,
            )
            result['multi_skew_dominant_period'] = multi_struct['dominant_period']
            result['multi_skew_dominant_power'] = multi_struct['dominant_power']
            result['multi_skew_peak_ratio_2to1'] = multi_struct['peak_power_ratio_2to1']
            result['multi_skew_spectral_entropy'] = multi_struct['spectral_entropy']
            result['multi_skew_topk_periods'] = multi_struct['topk_periods']
            result['multi_skew_topk_powers'] = multi_struct['topk_powers']
            result['multi_skew_structure_score'] = multi_struct['multi_peak_score']
            result.update(compute_skew_distribution_multi_metrics(S_n))
            result['multi_distribution_score'] = result.get('multi_skew_distribution_score', np.nan)
        except Exception as e:
            result['skew_Status'] = f'Error: {str(e)}'

        # ====================================================================
        # 1b. DURATION PERIODICITY
        # ====================================================================
        try:
            t_times, t_depths, t_durs, t_quals = extract_transit_midpoints_and_durations(
                t, flux, P_ref, T0
            )
            valid = np.isfinite(t_durs)
            if valid.sum() >= 3:
                duration_cv = float(np.nanstd(t_durs[valid]) / np.nanmean(t_durs[valid]))
                if period_search_params:
                    dur_period, dur_power, dur_fap = detect_periodic_signal(
                        t_times[valid], t_durs[valid],
                        min_period=period_search_params['min_period'],
                        max_period=period_search_params['max_period'],
                    )
                else:
                    dur_period, dur_power, dur_fap = detect_periodic_signal(
                        t_times[valid], t_durs[valid]
                    )
                result['duration_cv'] = duration_cv
                result['duration_period'] = dur_period
                result['duration_power'] = dur_power
                result['duration_fap'] = dur_fap
        except Exception:
            pass
        
        # ====================================================================
        # 2. TTV DETECTION
        # ====================================================================
        try:
            transit_times, transit_errors = extract_transit_times(
                t, flux, flux_err, P_ref, T0, dur
            )
            
            if len(transit_times) >= 3:
                P_linear, T0_linear, P_err, T0_err = fit_linear_ephemeris(
                    transit_times, transit_errors
                )
                
                # Compute O-C residuals
                n = np.round((transit_times - T0_linear) / P_linear).astype(int)
                expected_times = T0_linear + n * P_linear
                oc_residuals = transit_times - expected_times
                oc_errors = np.sqrt(transit_errors**2 + T0_err**2)
                ttv_times = transit_times
                ttv_oc = oc_residuals
                
                # Search for periodic TTV
                _, _, P_ttv, ttv_power, ttv_fap = detect_ttv_period(
                    transit_times, oc_residuals, oc_errors,
                    min_period=0.1*P_ref,
                    max_period=10*P_ref
                )
                
                ttv_amplitude = float(np.nanstd(oc_residuals))
                
                result['ttv_P_refined'] = P_ref
                result['ttv_T0'] = T0_linear
                result['ttv_dur'] = dur
                result['ttv_P_ttv'] = P_ttv
                result['ttv_power'] = ttv_power
                result['ttv_fap'] = ttv_fap
                result['ttv_amplitude'] = ttv_amplitude
                result['ttv_num_transits'] = len(transit_times)
                multi_ttv = compute_multi_ttv_metrics(
                    transit_times, oc_residuals,
                    min_period=0.1 * P_ref,
                    max_period=min(10 * P_ref, (transit_times.max() - transit_times.min()) / 2)
                    if len(transit_times) > 1 else None,
                    max_frequency=None,
                    top_k=3,
                    planet_period=P_ref,
                )
                result.update(multi_ttv)
                result['multi_ttv_structure_score'] = multi_ttv.get('multi_ttv_multi_peak_score', np.nan)
            else:
                result['ttv_Status'] = f'Insufficient transits: {len(transit_times)}'
        except Exception as e:
            result['ttv_Status'] = f'Error: {str(e)}'
        
        # ====================================================================
        # 3. SHOULDER DETECTION
        # ====================================================================
        try:
            S_lead, S_trail = compute_shoulder_snr(t, resid, P_ref, T0, dur)
            
            result['shoulder_P_refined'] = P_ref
            result['shoulder_T0'] = T0
            result['shoulder_dur'] = dur
            result['shoulder_S_lead'] = S_lead
            result['shoulder_S_trail'] = S_trail
            result['shoulder_snr'] = max(abs(S_lead) if not np.isnan(S_lead) else 0,
                                        abs(S_trail) if not np.isnan(S_trail) else 0)
        except Exception as e:
            result['shoulder_Status'] = f'Error: {str(e)}'
        
        # ====================================================================
        # 4. VARIABILITY DETECTION
        # ====================================================================
        try:
            sigma_ing, sigma_egr, sigma_base, V_ratio = compute_variability_metrics(
                t, resid, P_ref, T0, dur
            )
            
            result['variability_P_refined'] = P_ref
            result['variability_T0'] = T0
            result['variability_dur'] = dur
            result['variability_sigma_ing'] = sigma_ing
            result['variability_sigma_egr'] = sigma_egr
            result['variability_sigma_base'] = sigma_base
            result['variability_V_ratio'] = V_ratio
        except Exception as e:
            result['variability_Status'] = f'Error: {str(e)}'
        
        # ====================================================================
        # COMPUTE COMBINED SCORES (single-moon and multi-moon)
        # ====================================================================
        single_score = 0.0
        multi_score = 0.0
        
        def _fap_penalty(fap):
            if fap is None or np.isnan(fap):
                return 1.0
            fap = min(max(float(fap), 0.0), 1.0)
            penalty = -np.log10(fap + 1e-10) / 3.0
            return min(max(penalty, 0.0), 1.0)

        # Transit reliability modulation (discourage low-N saturation).
        observed_n = result.get('num_transits_observed', np.nan)
        if np.isnan(observed_n):
            observed_n = 0.0
        transit_reliability = float(np.clip(np.log10(observed_n + 1.0) / 2.0, 0.0, 1.0))
        result['single_transit_reliability'] = transit_reliability
        result['multi_transit_reliability'] = transit_reliability

        # --- single-moon score ---
        # Skew contribution (0-0.22)
        if not np.isnan(result.get('skew_power', np.nan)):
            skew_score = min(result['skew_power'] / 10.0, 1.0)
            if not np.isnan(result.get('skew_fap', np.nan)) and result['skew_fap'] > 0:
                skew_score *= _fap_penalty(result['skew_fap'])
            skew_score *= transit_reliability
            single_score += 0.22 * skew_score
            result['skew_score'] = 0.22 * skew_score
            result['single_skew_score'] = result['skew_score']

        # Single-moon consistency from skew distribution GOF:
        # high p-value => more consistent with single-moon occupancy.
        if not np.isnan(result.get('skew_dist_pvalue', np.nan)):
            consistency = float(np.clip(result['skew_dist_pvalue'], 0.0, 1.0))
            consistency *= transit_reliability
            single_score += 0.06 * consistency
            result['single_skew_dist_consistency_score'] = 0.06 * consistency
            result['skew_distribution_score'] = 0.06 * consistency
        
        # TTV contribution (0-0.38) - highest weight
        if not np.isnan(result.get('ttv_power', np.nan)):
            ttv_score = min(result['ttv_power'] / 10.0, 1.0)
            if not np.isnan(result.get('ttv_fap', np.nan)) and result['ttv_fap'] > 0:
                ttv_score *= _fap_penalty(result['ttv_fap'])
            ttv_score *= transit_reliability
            single_score += 0.38 * ttv_score
            result['ttv_score'] = 0.38 * ttv_score
            result['single_ttv_score'] = result['ttv_score']
        
        # Shoulder contribution (0-0.12)
        if not np.isnan(result.get('shoulder_snr', np.nan)):
            shoulder_score = min(result['shoulder_snr'] / 5.0, 1.0)
            single_score += 0.12 * shoulder_score
            result['shoulder_score'] = 0.12 * shoulder_score
            result['single_shoulder_score'] = result['shoulder_score']
        
        # Variability contribution (0-0.12)
        if not np.isnan(result.get('variability_V_ratio', np.nan)):
            var_score = min(result['variability_V_ratio'] / 3.0, 1.0)
            single_score += 0.12 * var_score
            result['variability_score'] = 0.12 * var_score
            result['single_variability_score'] = result['variability_score']

        # Duration periodicity contribution (0-0.11)
        if not np.isnan(result.get('duration_power', np.nan)):
            dur_score = min(result['duration_power'] / 10.0, 1.0)
            if not np.isnan(result.get('duration_fap', np.nan)) and result['duration_fap'] > 0:
                dur_score *= _fap_penalty(result['duration_fap'])
            dur_score *= transit_reliability
            single_score += 0.11 * dur_score
            result['duration_score'] = 0.11 * dur_score
            result['single_duration_score'] = result['duration_score']
        
        # Data quality bonus (0-0.05)
        single_score += 0.05 * result['data_quality_score']
        result['data_quality_bonus'] = 0.05 * result['data_quality_score']
        result['single_data_quality_bonus'] = result['data_quality_bonus']

        # --- multi-moon score ---
        if not np.isnan(result.get('multi_skew_structure_score', np.nan)):
            comp = 0.26 * float(np.clip(result['multi_skew_structure_score'], 0.0, 1.0)) * transit_reliability
            multi_score += comp
            result['multi_skew_structure_score'] = float(np.clip(result['multi_skew_structure_score'], 0.0, 1.0))
        if not np.isnan(result.get('multi_ttv_structure_score', np.nan)):
            ttv_struct = float(np.clip(result['multi_ttv_structure_score'], 0.0, 1.0))
            dom_fap = result.get('multi_ttv_dominant_fap', np.nan)
            # Additional guardrail so weak-significance structure cannot dominate.
            if np.isfinite(dom_fap) and dom_fap <= 1.0:
                ttv_weight = 1.0
            elif np.isfinite(dom_fap) and dom_fap <= 10.0:
                ttv_weight = 0.5
            else:
                ttv_weight = 0.0
            comp = 0.28 * ttv_struct * transit_reliability * ttv_weight
            multi_score += comp
            result['multi_ttv_structure_score'] = float(np.clip(result['multi_ttv_structure_score'], 0.0, 1.0))
            result['multi_ttv_score_effective'] = comp
        if not np.isnan(result.get('multi_distribution_score', np.nan)):
            comp = 0.10 * float(np.clip(result['multi_distribution_score'], 0.0, 1.0)) * transit_reliability
            multi_score += comp
        if not np.isnan(result.get('shoulder_snr', np.nan)):
            comp = 0.10 * min(result['shoulder_snr'] / 5.0, 1.0)
            multi_score += comp
            result['multi_shoulder_score'] = comp
        if not np.isnan(result.get('variability_V_ratio', np.nan)):
            comp = 0.08 * min(result['variability_V_ratio'] / 3.0, 1.0)
            multi_score += comp
            result['multi_variability_score'] = comp
        if not np.isnan(result.get('duration_power', np.nan)):
            # Multi-moon can show broad/structured TDV; penalize less aggressively.
            dur_multi = min(result['duration_power'] / 10.0, 1.0) * transit_reliability
            multi_score += 0.13 * dur_multi
            result['multi_duration_score'] = 0.13 * dur_multi
        multi_score += 0.05 * result['data_quality_score']
        result['multi_data_quality_bonus'] = 0.05 * result['data_quality_score']
        
        result['combined_single_moon_score'] = min(single_score, 1.0)
        result['combined_multi_moon_score'] = min(multi_score, 1.0)
        # Backward-compat alias (phase 1).
        result['combined_exomoon_score'] = result['combined_single_moon_score']
        result['Status'] = 'OK'
        
    except ValueError as e:
        result['Status'] = 'Error'
        msg = str(e)[:200]
        # Normalize to reason codes when possible
        if msg.startswith("ERR_") or msg.startswith("WARN_"):
            result['Reason'] = msg
        else:
            result['Reason'] = f"ERR_UNHANDLED: {msg}"
    except Exception as e:
        result['Status'] = 'Error'
        msg = str(e)[:200]
        result['Reason'] = f"ERR_UNHANDLED: {msg}"
    
    return result


# ============================================================================
# Main Execution
# ============================================================================

def main():
    parser = argparse.ArgumentParser(
        description='Comprehensive exomoon detection pipeline',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python comprehensive_exomoon_search.py
  python comprehensive_exomoon_search.py --input ranked_transiting_planets.csv --output results.csv
  python comprehensive_exomoon_search.py --workers 8 --start 0 --end 100
        """
    )
    parser.add_argument('--input', type=str, default='ranked_transiting_planets.csv',
                       help='Input CSV file with planet data')
    parser.add_argument('--output', type=str, default='comprehensive_exomoon_results.csv',
                       help='Output CSV file for results')
    parser.add_argument('--timestamped-output', action='store_true',
                       help='Append timestamp to output filename')
    parser.add_argument('--workers', type=int, default=4,
                       help='Number of parallel workers')
    parser.add_argument('--start', type=int, default=0,
                       help='Start index (for resuming)')
    parser.add_argument('--end', type=int, default=None,
                       help='End index (None = process all)')
    parser.add_argument('--shard-id', type=int, default=None,
                       help='Shard ID for parallel execution')
    parser.add_argument('--num-shards', type=int, default=1,
                       help='Total number of shards')
    parser.add_argument('--cadence', type=str, default='short_then_any',
                       choices=['short', 'long', 'short_then_any', 'any'],
                       help='Lightcurve cadence policy')
    parser.add_argument('--cache-dir', type=str, default='.cache/lightcurves',
                       help='Cache directory for preprocessed light curves')
    parser.add_argument('--cache', action='store_true',
                       help='Enable light curve caching (default: disabled to save disk)')
    
    args = parser.parse_args()
    
    # Read input CSV
    print(f"Reading planet list from {args.input}...")
    df = pd.read_csv(args.input)
    print(f"Found {len(df)} planets")
    
    # Timestamp output if requested
    if args.timestamped_output:
        base = args.output.replace(".csv", "")
        ts = pd.Timestamp.now().strftime("%Y%m%d_%H%M%S")
        args.output = f"{base}_{ts}.csv"

    # Apply sharding if specified
    if args.shard_id is not None:
        df = df[df.index % args.num_shards == args.shard_id].reset_index(drop=True)
        args.output = args.output.replace('.csv', f'_shard{args.shard_id}.csv')
        print(f"Processing shard {args.shard_id}/{args.num_shards}: {len(df)} planets")
    
    # Apply start/end limits
    if args.end is not None:
        df = df.iloc[args.start:args.end].reset_index(drop=True)
    else:
        df = df.iloc[args.start:].reset_index(drop=True)
    
    print(f"Processing {len(df)} planets with {args.workers} worker(s)...")
    
    # Check for existing results
    if os.path.exists(args.output) and os.path.getsize(args.output) > 0:
        try:
            # Try reading with header, if that fails, check if first row looks like data
            df_existing = pd.read_csv(args.output)
            
            # Check if 'planet_name' column exists (might be first column without header)
            if 'planet_name' in df_existing.columns:
                processed = set(df_existing['planet_name'].dropna().tolist())
            elif len(df_existing.columns) > 0:
                # Assume first column is planet_name if no header
                first_col = df_existing.columns[0]
                processed = set(df_existing[first_col].dropna().astype(str).tolist())
            else:
                processed = set()
            
            if len(processed) > 0:
                df = df[~df['planet_name'].isin(processed)].reset_index(drop=True)
                print(f"Resuming: {len(df)} planets remaining (skipping {len(processed)} already processed)")
            else:
                print("Existing output file found but no processed planets detected - starting fresh")
        except Exception as e:
            print(f"Warning: Could not read existing output file ({e}) - starting fresh")
            # Backup the old file and start fresh
            if os.path.exists(args.output):
                backup_name = args.output + '.backup'
                import shutil
                shutil.move(args.output, backup_name)
                print(f"Backed up existing file to {backup_name}")
    
    if len(df) == 0:
        print("All planets already processed!")
        return
    
    # Prepare output file - don't create empty file, let first write create it with proper header
    
    # Process planets
    results = []
    worker = partial(
        process_planet,
        cadence_policy=args.cadence,
        cache_dir=args.cache_dir,
        cache_enabled=args.cache,
    )
    if args.workers <= 1:
        with tqdm(total=len(df), desc='Processing') as pbar:
            for _, row in df.iterrows():
                try:
                    result = worker(row.to_dict())
                    results.append(result)
                    df_result = pd.DataFrame([result])
                    file_exists = os.path.exists(args.output) and os.path.getsize(args.output) > 0
                    df_result.to_csv(args.output, mode='a', header=not file_exists, index=False)
                    pbar.update(1)
                    pbar.set_postfix({
                        'Status': result['Status'],
                        'Score': f"{result.get('combined_exomoon_score', 0):.2f}"
                    })
                except Exception as e:
                    pbar.update(1)
                    print(f"\nError processing planet: {e}")
    else:
        with ProcessPoolExecutor(max_workers=args.workers) as executor:
            futures = {executor.submit(worker, row.to_dict()): row 
                      for _, row in df.iterrows()}
            
            with tqdm(total=len(futures), desc='Processing') as pbar:
                for future in as_completed(futures):
                    try:
                        result = future.result()
                        results.append(result)
                        
                        # Append to CSV incrementally
                        df_result = pd.DataFrame([result])
                        file_exists = os.path.exists(args.output) and os.path.getsize(args.output) > 0
                        df_result.to_csv(args.output, mode='a', header=not file_exists, index=False)
                        
                        pbar.update(1)
                        pbar.set_postfix({
                            'Status': result['Status'],
                            'Score': f"{result.get('combined_exomoon_score', 0):.2f}"
                        })
                    except Exception as e:
                        pbar.update(1)
                        print(f"\nError processing planet: {e}")
    
    print(f"\nProcessing complete! Results saved to {args.output}")
    print(f"Processed {len(results)} planets")
    
    # Print summary statistics
    if results:
        df_results = pd.DataFrame(results)
        print("\nSummary Statistics:")
        print(f"  Successful: {sum(r['Status'] == 'OK' for r in results)}")
        print(f"  Errors: {sum(r['Status'] == 'Error' for r in results)}")
        
        # Check if combined_exomoon_score exists before accessing
        if 'combined_exomoon_score' in df_results.columns:
            df_results['combined_exomoon_score'] = pd.to_numeric(
                df_results['combined_exomoon_score'], errors='coerce'
            )
            valid_scores = df_results['combined_exomoon_score'].dropna()
            if len(valid_scores) > 0:
                print(f"  Average score: {valid_scores.mean():.3f}")
                print(f"  Max score: {valid_scores.max():.3f}")
                
                # Top candidates
                try:
                    top_candidates = df_results.nlargest(min(10, len(df_results)), 'combined_exomoon_score')
                except TypeError:
                    top_candidates = df_results.sort_values(
                        'combined_exomoon_score', ascending=False
                    ).head(min(10, len(df_results)))
                print("\nTop 10 Candidates:")
                for idx, row in top_candidates.iterrows():
                    score = row.get('combined_exomoon_score', 0)
                    if not pd.isna(score):
                        print(f"  {row['planet_name']}: score={score:.3f}")
            else:
                print("  No successful detections with scores")
        else:
            print("  No scores computed (all planets failed)")


if __name__ == '__main__':
    main()