#!/usr/bin/env python3
"""Conditional, finite-window transit fit to a frozen real TESS light curve."""
import argparse
import csv
import hashlib
from importlib.metadata import version
import json
from pathlib import Path
import warnings
with warnings.catch_warnings():
    warnings.filterwarnings('ignore', message='invalid escape sequence.*', category=DeprecationWarning)
    warnings.filterwarnings('ignore', message='invalid escape sequence.*', category=SyntaxWarning)
    import batman
import numpy as np
from scipy import __version__ as scipy_version
from scipy.optimize import least_squares

PERIOD_START = 3.2888
EPOCH_REFERENCE = 1327.52
LOWER = np.array([.025, 0., 3., -.035, 3.2838])
UPPER = np.array([.35, 1.3, 35., .035, 3.2938])


def load_observations(folder):
    raw = (folder / 'wasp126-tess-windows.csv').read_bytes()
    provenance = json.loads((folder / 'observed-provenance.json').read_text())
    expected = next(entry['sha256'] for entry in provenance['products'] if entry['file'].endswith('.csv'))
    if hashlib.sha256(raw).hexdigest() != expected:
        raise ValueError('Frozen observed CSV hash mismatch')
    rows = list(csv.DictReader(raw.decode().splitlines()))
    accepted, rejected = [], []
    for row in rows:
        values = {key:float(value) if value else None for key,value in row.items()}
        valid = (values['quality'] == 0 and all(values[key] is not None and np.isfinite(values[key]) for key in ['time','pdcsap_flux','pdcsap_flux_err']) and values['pdcsap_flux_err'] > 0)
        (accepted if valid else rejected).append(values)
    normalization = float(np.median([row['pdcsap_flux'] for row in accepted]))
    return {'rows':accepted,'rejected':rejected,'normalization':normalization,
            'time':np.array([row['time'] - EPOCH_REFERENCE for row in accepted]),
            'flux':np.array([row['pdcsap_flux']/normalization for row in accepted]),
            'error':np.array([row['pdcsap_flux_err']/normalization for row in accepted]),
            'window':np.array([int(row['window']) for row in accepted]),
            'input_sha256':hashlib.sha256(raw).hexdigest(),'provenance':provenance}


def model_flux(parameters, times, coefficients, contamination, exposure_order):
    radius, impact, axis, epoch, period = parameters
    params = batman.TransitParams()
    params.t0, params.per, params.rp, params.a = epoch, period, radius, axis
    params.inc, params.ecc, params.w = float(np.degrees(np.arccos(impact/axis))), 0., 90.
    params.limb_dark, params.u = 'quadratic', list(coefficients)
    if exposure_order:
        nodes, weights = np.polynomial.legendre.leggauss(exposure_order)
        sample_times = (times[:,None] + nodes[None,:] * 60 / 86400).ravel()
        model = batman.TransitModel(params, sample_times, nthreads=1)
        samples = model.light_curve(params).reshape(len(times), exposure_order)
        if not np.all(np.isfinite(samples)):
            raise ValueError('Independent transit implementation returned nonfinite flux')
        host = np.sum(samples * (weights[None,:]/2), axis=1)
    else:
        model = batman.TransitModel(params, times, nthreads=1)
        host = model.light_curve(params)
    return (host + contamination)/(1 + contamination)


def fit_data(data, coefficients=(.45,.2), contamination=0., exposure_order=9, fixed_impact=None, start=None):
    if start is None:
        start = np.array([.08, .3, 7.8, 0., PERIOD_START])
    free_indices = np.array([0,2,3,4] if fixed_impact is not None else [0,1,2,3,4])
    def assemble(free):
        parameters = np.array(start, copy=True)
        parameters[free_indices] = free
        if fixed_impact is not None:
            parameters[1] = fixed_impact
        return parameters
    def prediction(parameters):
        transit = model_flux(parameters, data['time'], coefficients, contamination, exposure_order)
        prediction = np.empty_like(transit)
        baseline_coefficients = []
        for window in sorted(set(data['window'])):
            selected = data['window'] == window
            relative_time = data['time'][selected] - window*PERIOD_START
            design = np.column_stack((transit[selected], transit[selected]*relative_time))
            coefficients_fit = np.linalg.lstsq(design/data['error'][selected,None], data['flux'][selected]/data['error'][selected], rcond=None)[0]
            prediction[selected] = np.sum(design * coefficients_fit[None,:], axis=1)
            baseline_coefficients.append({'window':int(window),'intercept':float(coefficients_fit[0]),'slope_per_day':float(coefficients_fit[1])})
        return prediction, baseline_coefficients
    def residual(free):
        return (prediction(assemble(free))[0] - data['flux'])/data['error']
    fit = least_squares(residual, start[free_indices], bounds=(LOWER[free_indices],UPPER[free_indices]),
                        x_scale=np.array([.08,.3,8.,.002,.001])[free_indices],
                        ftol=1e-10, xtol=1e-10, gtol=1e-8, max_nfev=250)
    parameters = assemble(fit.x)
    prediction_values, baseline = prediction(parameters)
    total_contact = max(0., (1 + parameters[0])**2 - parameters[1]**2)
    duration = parameters[4]/np.pi * np.arcsin(np.sqrt(total_contact/(parameters[2]**2-parameters[1]**2)))
    return {'parameters':parameters,'model':prediction_values,'baseline':baseline,'chi_squared':float(np.sum(fit.fun**2)),
            'duration_hours':float(duration*24),'success':bool(fit.success),'evaluations':int(fit.nfev),
            'conditional_circular_density_kg_m3':float(3*np.pi*parameters[2]**3/(6.67430e-11*(parameters[4]*86400)**2))}


def summarize_fit(fit):
    return {key:([float(value) for value in entry] if isinstance(entry,np.ndarray) else entry)
            for key,entry in fit.items() if key != 'model'}


def run(folder):
    data = load_observations(folder)
    starts = [np.array([.08,impact,7.8,0.,PERIOD_START]) for impact in [.02,.4,.8,1.02]]
    candidates = [fit_data(data,start=start) for start in starts]
    best = min(candidates,key=lambda fit:fit['chi_squared'])
    profiles = []
    for impact in np.linspace(0,1.2,25):
        flat_candidate = fit_data(data,fixed_impact=float(impact),start=best['parameters'])
        physical_start = best['parameters'].copy()
        physical_start[0] = min(.34, max(float(physical_start[0]), float(impact) - 1 + .06))
        fit = fit_data(data,fixed_impact=float(impact),start=physical_start)
        fit = min([flat_candidate, fit], key=lambda result: result['chi_squared'])
        profiles.append({'fixed_impact':float(impact),'delta_chi_squared':fit['chi_squared']-best['chi_squared'], 'initial_best_radius_start_chi_squared':flat_candidate['chi_squared'], 'initialization':'retain original start; also seed a physically overlapping radius above impact-1; choose lower objective',**summarize_fit(fit)})
    sensitivities = []
    for name, coefficients, contamination, order in [
        ('uniform_star',(0.,0.),0.,9),('stronger_limb_darkening',(.65,.1),0.,9),
        ('additional_contamination_0.3',(.45,.2),.3,9),('instantaneous_samples',(.45,.2),0.,0),
        ('refined_exposure_17',(.45,.2),0.,17)]:
        fit = fit_data(data,coefficients,contamination,order,start=best['parameters'])
        sensitivities.append({'name':name,'limb_coefficients':coefficients,'extra_contamination':contamination,
                              'exposure_order':order,**summarize_fit(fit)})
    residual = data['flux'] - best['model']
    adjacent_left, adjacent_right, means = [], [], []
    for window in sorted(set(data['window'])):
        indices = np.flatnonzero(data['window']==window)
        for first,second in zip(indices[:-1],indices[1:]):
            if data['rows'][second]['source_row'] == data['rows'][first]['source_row']+1:
                adjacent_left.append(residual[first]);adjacent_right.append(residual[second])
        for start in range(0,len(indices)-14,15):
            block = indices[start:start+15]
            if data['rows'][block[-1]]['source_row']-data['rows'][block[0]]['source_row']==14:
                means.append(float(np.mean(residual[block])))
    exposure_33 = model_flux(best['parameters'],data['time'],(.45,.2),0.,33)
    exposure_65 = model_flux(best['parameters'],data['time'],(.45,.2),0.,65)
    exposure_9 = model_flux(best['parameters'],data['time'],(.45,.2),0.,9)
    return {'schema':'ac17-observed-fit-v1','kind':'conditional fit to real frozen TESS SPOC PDCSAP observations',
            'versions':{'numpy':np.__version__,'scipy':scipy_version,'batman_distribution':version('batman-package')},
            'input_sha256':data['input_sha256'],'accepted_rows':len(data['rows']),'rejected_rows':len(data['rejected']),
            'quality_policy':'QUALITY=0, finite time and PDC flux/error, positive error; no flux clipping',
            'flux_normalization_electrons_per_second':data['normalization'],
            'parameter_order':['radius_ratio','impact','a_over_rstar','epoch_minus1327.52_BTJD_days','period_days'],
            'assumptions':{'orbit':'circular; negligible planet mass; density is rho_circ conditional on that assumption',
                           'limb_coefficients':[.45,.2],'contamination':'zero additional after PDC; original CROWDSAP/FLFRCSAP not applied again',
                           'exposure':'uniform120s window approximation; Gauss9 fit; raw readout/live-time metadata retained',
                           'noise':'supplied pointwise errors in conditional Gaussian objective; no claim of independent residuals or calibrated posterior',
                           'baseline':'independent multiplicative intercept and linear trend per retained time window, analytically profiled at every fit'},
            'best':summarize_fit(best),'multistart': [summarize_fit(fit) for fit in candidates],
            'impact_profiles':profiles,'sensitivities':sensitivities,
            'exposure_convergence_at_fixed_best':{'max_flux_difference_order9_vs65':float(np.max(abs(exposure_9-exposure_65))),
                                                 'max_flux_difference_order33_vs65':float(np.max(abs(exposure_33-exposure_65)))},
            'residual_diagnostics':{'rms':float(np.sqrt(np.mean(residual**2))),
                                    'adjacent_correlation':float(np.corrcoef(adjacent_left,adjacent_right)[0,1]),
                                    'complete_15_sample_blocks':len(means),
                                    'binned_std_ratio_to_white':float(np.std(means,ddof=1)*np.sqrt(15)/np.std(residual,ddof=1))},
            'samples':[{'source_row':int(row['source_row']),'window':int(row['window']),
                        'time_btjd':row['time'],'flux':float(flux),'error':float(error),
                        'model':float(model),'residual':float(delta)}
                       for row,flux,error,model,delta in zip(data['rows'],data['flux'],data['error'],best['model'],residual)]}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--data',type=Path,default=Path('.'))
    parser.add_argument('--out',type=Path,default=Path('transit-fit-rerun.json'))
    args = parser.parse_args()
    result = run(args.data)
    args.out.parent.mkdir(parents=True,exist_ok=True)
    args.out.write_text(json.dumps(result,indent=2,allow_nan=False)+'\n')
    print('Accepted/rejected:',result['accepted_rows'],result['rejected_rows'])
    print('Best p,b,a/R,dt0,P:',result['best']['parameters'],'T14 hours:',result['best']['duration_hours'])
    print('Exposure comparison:',result['exposure_convergence_at_fixed_best'])
    print('Residual diagnostics:',result['residual_diagnostics'])


if __name__ == '__main__':
    main()