#!/usr/bin/env python
# coding: utf-8

# # Fit transit shape with nested sampling


import os
import sys
import json
import shutil
import warnings
import multiprocessing as multipro
from datetime import datetime
from timeit import default_timer as timer

import argparse
import batman
import dynesty
from dynesty import plotting as dyplot
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy.polynomial.polynomial as poly
import numpy as np

from aesara_theano_fallback import aesara as theano
from celerite2 import terms as GPterms


# #### Flush buffer and silence extraneous warnings


# flush buffer to avoid mixed outputs from progressbar
sys.stdout.flush()

# turn off FutureWarnings
warnings.filterwarnings("ignore", category=FutureWarning)


# #### Initialize timer


print("")
print("+" * shutil.get_terminal_size().columns)
print("ALDERAAN Detrending and TTV Estimation")
print(f"Initialized {datetime.now().strftime('%d-%b-%Y at %H:%M:%S')}")
print("+" * shutil.get_terminal_size().columns)
print("")

# start program timer
global_start_time = timer()


# #### Parse inputs


try:
    parser = argparse.ArgumentParser(
        description="Inputs for ALDERAAN transit fiting pipeline"
    )
    parser.add_argument(
        "--mission",
        default=None,
        type=str,
        required=True,
        help="Mission name; can be 'Kepler' or 'Simulated'",
    )
    parser.add_argument(
        "--target",
        default=None,
        type=str,
        required=True,
        help="Target name; format should be K00000 or S00000",
    )
    parser.add_argument(
        "--run_id", default=None, type=str, required=True, help="run identifier"
    )
    parser.add_argument(
        "--project_dir",
        default=None,
        type=str,
        required=True,
        help="Project directory for accessing lightcurve data and saving outputs",
    )
    parser.add_argument(
        "--verbose",
        default=False,
        type=bool,
        required=False,
        help="'True' to enable verbose logging",
    )
    parser.add_argument(
        "--iplot",
        default=False,
        type=bool,
        required=False,
        help="'True' to enable interactive matplotlib backend; default 'agg'",
    )

    args = parser.parse_args()
    MISSION = args.mission
    TARGET = args.target
    RUN_ID = args.run_id
    PROJECT_DIR = args.project_dir
    VERBOSE = args.verbose
    IPLOT = args.iplot

except SystemExit:
    warnings.warn("No arguments were parsed from the command line")


print("")
print(f"   MISSION : {MISSION}")
print(f"   TARGET  : {TARGET}")
print(f"   RUN ID  : {RUN_ID}")
print("")
print(f"   Project directory : {PROJECT_DIR}")
print("")
print(f"   theano cache : {theano.config.compiledir}")
print("")


# #### Build directory structure


# directories in which to place pipeline outputs for this run
RESULTS_DIR = os.path.join(PROJECT_DIR, "Results", RUN_ID, TARGET)
FIGURE_DIR = os.path.join(PROJECT_DIR, "Figures", RUN_ID, TARGET)

os.makedirs(RESULTS_DIR, exist_ok=True)
os.makedirs(FIGURE_DIR, exist_ok=True)

sys.path.append(PROJECT_DIR)


# #### Import ALDERAAN routines


import alderaan.io as io
import alderaan.dynesty_helpers as dynhelp
from alderaan.astro import make_transit_mask, set_oversample_factor
from alderaan.constants import scit, lcit
from alderaan.Ephemeris import Ephemeris
from alderaan.plotting import plot_omc


# #### Set matplotlib backend


# check for interactive matplotlib backends
if not IPLOT:
    mpl.use("agg")

if np.any(np.array(["agg", "png", "svg", "pdf", "ps"]) == mpl.get_backend()):
    IPLOT = False

# flush buffer
sys.stdout.flush()

# MAIN SCRIPT BEGINS HERE
def main():
    # # ################
    # # ----- DATA I/O -----
    # # ################

    print("\nLoading data...\n")

    # ## Read in planet and stellar properties

    path = os.path.join(RESULTS_DIR, f"{TARGET}_transit_parameters.csv")
    catalog = io.parse_catalog(path, MISSION, TARGET)

    KOI_ID = catalog.koi_id.to_numpy()[0]
    KIC_ID = catalog.kic_id.to_numpy()[0]

    NPL = catalog.npl.to_numpy()[0]

    PERIODS = catalog.period.to_numpy()
    DEPTHS = catalog.depth.to_numpy() * 1e-6
    DURS = catalog.duration.to_numpy() / 24

    U1 = catalog.limbdark_1.to_numpy()[0]
    U2 = catalog.limbdark_2.to_numpy()[0]

    if np.any(np.diff(PERIODS) <= 0):
        raise ValueError("Planets should be ordered by ascending period")

    # ## Read in filtered lightcurves
    # #### These can be generated by running the script "analyze_autocorrelated_noise.py"

    if os.path.exists(os.path.join(RESULTS_DIR, f"{TARGET}_lc_filtered.fits")):
        lc = io.load_detrended_lightcurve(
            os.path.join(RESULTS_DIR, f"{TARGET}_lc_filtered.fits")
        )
        lc.season = lc.quarter % 4
    else:
        lc = None

    if os.path.exists(os.path.join(RESULTS_DIR, f"{TARGET}_sc_filtered.fits")):
        sc = io.load_detrended_lightcurve(
            os.path.join(RESULTS_DIR, f"{TARGET}_sc_filtered.fits")
        )
        sc.season = sc.quarter % 4
    else:
        sc = None

    # ## Read in quick transit times
    # #### These can be generated by running the script "detrend_and_estimate_ttvs.py"

    # transit times
    epochs = np.zeros(NPL)
    periods = np.zeros(NPL)
    linear_ephemeris = [None] * NPL

    transit_inds = []
    indep_transit_times = []
    regular_transit_times = []

    for n in range(NPL):
        fname_in = os.path.join(RESULTS_DIR, f"{TARGET}_{n:02d}_quick.ttvs")
        data_in = np.genfromtxt(fname_in)

        transit_inds.append(np.array(data_in[:, 0], dtype="int"))
        indep_transit_times.append(np.array(data_in[:, 1], dtype="float"))
        regular_transit_times.append(np.array(data_in[:, 2], dtype="float"))

        # do a quick fit to get a linear ephemeris
        pfit = poly.polyfit(transit_inds[n], regular_transit_times[n], 1)

        epochs[n] = pfit[1]
        periods[n] = pfit[0]
        linear_ephemeris[n] = poly.polyval(transit_inds[n], pfit)

    # calculate centered transit indexes
    centered_transit_inds = [None] * NPL
    for n in range(NPL):
        centered_transit_inds[n] = transit_inds[n] - transit_inds[n][-1] // 2

    # make OMC plot
    fig, ax = plot_omc(linear_ephemeris, indep_transit_times, regular_transit_times)
    fig.savefig(os.path.join(FIGURE_DIR, TARGET + f"_omc.png"), bbox_inches="tight")
    if IPLOT:
        plt.show()
    else:
        plt.close()

    # # ####################
    # # --- PRELIMINARIES ---
    # # ####################

    print("\nRunning preliminaries...\n")

    # ## Establish time baseline

    time_min = []
    time_max = []

    if sc is not None:
        time_min.append(sc.time.min())
        time_max.append(sc.time.max())

    if lc is not None:
        time_min.append(lc.time.min())
        time_max.append(lc.time.max())

    TIME_START = np.min(time_min)
    TIME_END = np.max(time_max)

    # put epochs in range (TIME_START, TIME_START + PERIOD)
    for n in range(NPL):
        if epochs[n] < TIME_START:
            adj = 1 + (TIME_START - epochs[n]) // periods[n]
            epochs[n] += adj * periods[n]

        if epochs[n] > (TIME_START + periods[n]):
            adj = (epochs[n] - TIME_START) // periods[n]
            epochs[n] -= adj * periods[n]

    # ## Make transit masks

    if sc is not None:
        sc_mask = np.zeros((NPL, len(sc.time)), dtype="bool")
        for n in range(NPL):
            sc_mask[n] = make_transit_mask(
                sc.time,
                regular_transit_times[n],
                masksize=np.max([1 / 24, 1.5 * DURS[n]]),
            )

    if lc is not None:
        lc_mask = np.zeros((NPL, len(lc.time)), dtype="bool")
        for n in range(NPL):
            lc_mask[n] = make_transit_mask(
                lc.time,
                regular_transit_times[n],
                masksize=np.max([1 / 24, 1.5 * DURS[n]]),
            )

    # ## Identify overlapping transits

    overlap = [None] * NPL

    for i in range(NPL):
        overlap[i] = np.zeros(len(linear_ephemeris[i]), dtype="bool")

        for j in range(NPL):
            if i != j:
                for tt in linear_ephemeris[j]:
                    overlap[i] += np.abs(linear_ephemeris[i] - tt) < (
                        DURS[i] + DURS[j] + lcit
                    )

    # ## Predict number of points expected in transit

    # number of photometric points expected in transit
    count_expect_lc = np.zeros(NPL, dtype="int")
    count_expect_sc = np.zeros(NPL, dtype="int")

    for n in range(NPL):
        count_expect_lc[n] = np.max([1, int(np.floor(DURS[n] / lcit))])
        count_expect_sc[n] = np.max([1, int(np.floor(DURS[n] / scit))])

    # ## Grab data near transits

    # go quarter-by-quarter
    all_time = [None] * 18
    all_flux = [None] * 18
    all_error = [None] * 18
    all_mask = [None] * 18
    all_dtype = ["none"] * 18

    for q in range(18):
        if sc is not None:
            if np.isin(q, sc.quarter):
                use = (sc_mask.sum(0) != 0) * (sc.quarter == q)

                if np.sum(use) > np.min(count_expect_sc):
                    all_time[q] = sc.time[use]
                    all_flux[q] = sc.flux[use]
                    all_error[q] = sc.error[use]
                    all_mask[q] = sc_mask[:, use]
                    all_dtype[q] = "short"

                else:
                    all_dtype[q] = "short_no_transits"

        if lc is not None:
            if np.isin(q, lc.quarter):
                use = (lc_mask.sum(0) != 0) * (lc.quarter == q)

                if np.sum(use) > np.min(count_expect_lc):
                    all_time[q] = lc.time[use]
                    all_flux[q] = lc.flux[use]
                    all_error[q] = lc.error[use]
                    all_mask[q] = lc_mask[:, use]
                    all_dtype[q] = "long"

                else:
                    all_dtype[q] = "long_no_transits"

    all_dtype = np.array(all_dtype)

    # track which quarters have coverage
    quarters = np.arange(18)[(all_dtype == "short") + (all_dtype == "long")]

    # ## Track which quarter each transit falls in

    # get list of threshold times between quarters
    qthresh = np.zeros(len(quarters) + 1)
    qthresh[0] = TIME_START - 0.5

    for j, q in enumerate(quarters):
        if lc is not None:
            if np.isin(q, np.unique(lc.quarter)):
                qthresh[j + 1] = lc.time[lc.quarter == q].max() + lcit

        if sc is not None:
            if np.isin(q, np.unique(sc.quarter)):
                qthresh[j + 1] = sc.time[sc.quarter == q].max() + scit

    qthresh[-1] = TIME_END + 0.5

    # track which quarter each transit falls in
    transit_quarter = [None] * NPL

    for n in range(NPL):
        tts = regular_transit_times[n]
        transit_quarter[n] = np.zeros(len(tts), dtype="int")

        for j, q in enumerate(quarters):
            transit_quarter[n][(tts >= qthresh[j]) * (tts < qthresh[j + 1])] = q

    # ## Track mean and variance of each quarter

    # track mean and variance of each quarter
    mean_by_quarter = np.ones(18) * np.nan
    var_by_quarter = np.ones(18) * np.nan

    for q in range(18):
        if sc is not None:
            if np.isin(q, sc.quarter):
                mean_by_quarter[q] = np.mean(sc.flux[sc.quarter == q])
                var_by_quarter[q] = np.var(sc.flux[sc.quarter == q])

        if lc is not None:
            if np.isin(q, lc.quarter):
                mean_by_quarter[q] = np.mean(lc.flux[lc.quarter == q])
                var_by_quarter[q] = np.var(lc.flux[lc.quarter == q])

    # ## Set oversampling factor and exposure times

    # calculate oversample factor
    oversample_lc = set_oversample_factor(
        PERIODS, DEPTHS, DURS, lc.flux, lc.error, lcit
    )
    oversample_sc = 1

    # broadcast oversample factor
    oversample = np.zeros(18, dtype="int")
    oversample[all_dtype == "short"] = oversample_sc
    oversample[all_dtype == "long"] = oversample_lc

    # broadcast exposure times
    texp = np.zeros(18)
    texp[all_dtype == "short"] = scit
    texp[all_dtype == "long"] = lcit

    # precompute exposure integration time offsets
    texp_offsets = [None] * 18
    for j, q in enumerate(quarters):
        if all_dtype[q] == "short":
            texp_offsets[q] = np.array([0.0])
        elif all_dtype[q] == "long":
            texp_offsets[q] = np.linspace(-texp[q] / 2.0, texp[q] / 2.0, oversample[q])

    # ## Set GP noise priors

    # Read in noise model GP priors from analyze_autocorrelated_noise.py
    gp_percs = []

    for z in range(4):
        try:
            fname_in = os.path.join(RESULTS_DIR, f"{TARGET}_noise_gp_priors_{z}.txt")

            with open(fname_in) as infile:
                gp_percs.append(json.load(infile))

        except Exception:
            gp_percs.append(None)

    # convert the percentile priors into Gaussians
    gp_priors = []

    for z in range(4):
        if gp_percs[z] is not None:
            gpz = {}

            for k in gp_percs[z].keys():
                if k != "percentiles":
                    perc = np.array(gp_percs[z]["percentiles"])

                    med = np.array(gp_percs[z][k])[perc == 50.0][0]
                    err1 = np.array(gp_percs[z][k])[perc == 84.135][0]
                    err2 = np.array(gp_percs[z][k])[perc == 15.865][0]

                    dev = np.sqrt((err1 - med) ** 2 / 2 + (err2 - med) ** 2 / 2)

                    gpz[k] = (med, dev)

            gp_priors.append(gpz)

        else:
            # these are dummy values that effectively create a zero-amplitude kernel
            gpz = {}
            gpz["logw0"] = [np.log(2 * np.pi / (7 * DURS.max()))]
            gpz["logSw4"] = [-100.0]
            gpz["logQ"] = [np.log(1 / np.sqrt(2))]

            gp_priors.append(gpz)

    # calculate a few convenience quantities
    for z in range(4):
        gpz = gp_priors[z]

        logS = gpz["logSw4"][0] - 4 * gpz["logw0"][0]

        if len(gpz["logSw4"]) == 1:
            print(1)
            gp_priors[z]["logS"] = [logS]

        if len(gpz["logSw4"]) == 2:
            logS_var = gpz["logSw4"][1] ** 2 + 16 * gpz["logw0"][1] ** 2
            gp_priors[z]["logS"] = np.array([logS, np.sqrt(logS_var)])

    # # ############################
    # # ----- LIGHTCURVE FITTING -----
    # # ############################

    print("\nModeling Lightcurve...\n")

    # identify which quarters and seasons have data
    no_transits = (np.array(all_dtype) == "long_no_transits") + (
        np.array(all_dtype) == "short_no_transits"
    )
    which_quarters = np.sort(np.unique(np.hstack(transit_quarter)))
    which_quarters = which_quarters[
        ~np.isin(which_quarters, np.squeeze(np.argwhere(no_transits)))
    ]
    which_seasons = np.sort(np.unique(which_quarters % 4))

    # mean flux and jitter for each quarter
    nq = len(which_quarters)
    mbq = mean_by_quarter[which_quarters]
    vbq = var_by_quarter[which_quarters]

    # initialize alderaan.Ephemeris objects
    ephem = [None] * NPL
    for npl in range(NPL):
        ephem[npl] = Ephemeris(transit_inds[npl], regular_transit_times[npl])

    # initialize TransitParams objects
    theta = [None] * NPL
    for npl in range(NPL):
        theta[npl] = batman.TransitParams()
        theta[npl].per = ephem[npl].period
        theta[npl].t0 = 0.0  # t0 must be set to zero b/c we are warping TTVs
        theta[npl].rp = np.sqrt(DEPTHS[npl])
        theta[npl].b = 0.5
        theta[npl].T14 = DURS[npl]
        theta[npl].u = [U1, U2]
        theta[npl].limb_dark = "quadratic"

    # grab data
    t_ = [None] * nq
    f_ = [None] * nq
    e_ = [None] * nq
    m_ = [None] * nq  # m_ is "mask" not "model"

    for j, q in enumerate(which_quarters):
        m_[j] = all_mask[q].sum(axis=0) > 0
        t_[j] = all_time[q][m_[j]]
        f_[j] = all_flux[q][m_[j]]
        e_[j] = all_error[q][m_[j]]

    # build the GP kernel using a different noise model for each season
    kernel = [None] * 4
    for z in which_seasons:
        kernel[z] = GPterms.SHOTerm(
            S0=np.exp(gp_priors[z]["logS"][0]),
            w0=np.exp(gp_priors[z]["logw0"][0]),
            Q=np.exp(gp_priors[z]["logQ"][0]),
        )

    warped_t = []
    warped_x = []

    for npl in range(NPL):
        warped_t.append([])
        warped_x.append([])

        for j, q in enumerate(which_quarters):
            _warp_time, _warp_index = ephem[npl]._warp_times(t_[j], return_inds=True)
            _warp_legx = (_warp_index - transit_inds[npl][-1] // 2) / (
                transit_inds[npl][-1] / 2
            )

            warped_t[npl].append(_warp_time)
            warped_x[npl].append(_warp_legx)

    transit_legx = []

    for npl in range(NPL):
        transit_legx.append(centered_transit_inds[npl] / (transit_inds[npl][-1] / 2))

    phot_args = {}

    phot_args["time"] = t_
    phot_args["flux"] = f_
    phot_args["error"] = e_
    phot_args["quarters"] = which_quarters
    phot_args["warped_t"] = warped_t
    phot_args["warped_x"] = warped_x
    phot_args["exptime"] = texp
    phot_args["oversample"] = oversample
    phot_args["texp_offsets"] = texp_offsets

    ephem_args = {}
    ephem_args["transit_inds"] = centered_transit_inds
    ephem_args["transit_times"] = regular_transit_times
    ephem_args["transit_legx"] = transit_legx

    ncores = multipro.cpu_count() - 2
    ndim = 5 * NPL + 2
    logl = dynhelp.lnlike
    ptform = dynhelp.prior_transform
    logl_args = (NPL, theta, ephem_args, phot_args, [U1, U2], kernel)
    ptform_args = (NPL, DURS)
    chk_file = os.path.join(RESULTS_DIR, f"{TARGET}-dynesty.checkpoint")

    USE_MULTIPRO = False

    if USE_MULTIPRO:
        with dynesty.pool.Pool(
            ncores, logl, ptform, logl_args=logl_args, ptform_args=ptform_args
        ) as pool:
            sampler = dynesty.DynamicNestedSampler(
                pool.loglike,
                pool.prior_transform,
                ndim,
                bound="multi",
                sample="rwalk",
                pool=pool,
            )
            sampler.run_nested(
                checkpoint_file=chk_file, checkpoint_every=60, print_progress=VERBOSE
            )
            results = sampler.results

    else:
        sampler = dynesty.DynamicNestedSampler(
            logl,
            ptform,
            ndim,
            bound="multi",
            sample="rwalk",
            logl_args=logl_args,
            ptform_args=ptform_args,
        )
        sampler.run_nested(
            checkpoint_file=chk_file, checkpoint_every=60, print_progress=VERBOSE
        )
        results = sampler.results

    # ## Make diagnostic plots

    labels = []
    for npl in range(NPL):
        labels = labels + "C0_{0} C1_{0} r_{0} b_{0} T14_{0}".format(npl).split()

    labels = labels + "q1 q2".split()

    # run plot
    fig, ax = dyplot.runplot(results, logplot=True)
    plt.savefig(
        os.path.join(FIGURE_DIR, f"{TARGET}_dynesty_runplot.png"), bbox_inches="tight"
    )

    # trace plot
    for n in range(NPL):
        fig, ax = dyplot.traceplot(
            results,
            labels=labels[5 * n : 5 * (n + 1)],
            dims=np.arange(5 * n, 5 * (n + 1)),
        )
        plt.savefig(
            os.path.join(FIGURE_DIR, f"{TARGET}_dynesty_trace_{n:02d}.png"),
            bbox_inches="tight",
        )

    fig, ax = dyplot.traceplot(results, labels=labels[-2:], dims=np.array([-2, -1]))
    plt.savefig(
        os.path.join(FIGURE_DIR, f"{TARGET}_dynesty_trace_q1q2.png"),
        bbox_inches="tight",
    )

    # corner plot
    for n in range(NPL):
        fig, ax = dyplot.cornerplot(
            results,
            labels=labels[5 * n : 5 * (n + 1)],
            dims=np.arange(5 * n, 5 * (n + 1)),
        )
        plt.savefig(
            os.path.join(FIGURE_DIR, f"{TARGET}_dynesty_corner_{n:02d}.png"),
            bbox_inches="tight",
        )

    fig, ax = dyplot.cornerplot(results, labels=labels[-2:], dims=np.array([-2, -1]))
    plt.savefig(
        os.path.join(FIGURE_DIR, f"{TARGET}_dynesty_corner_q1q2.png"),
        bbox_inches="tight",
    )

    # ## Save sampling results

    hdu_list = io.dynesty_results_to_fits(results, PROJECT_DIR, MISSION, TARGET, RUN_ID)

    path = os.path.join(PROJECT_DIR, "Results", RUN_ID, TARGET)
    os.makedirs(path, exist_ok=True)
    hdu_list.writeto(os.path.join(path, f"{TARGET}-results.fits"), overwrite=True)

    # ## Exit program

    print("")
    print("+" * shutil.get_terminal_size().columns)
    print(
        f"Exoplanet recovery complete {datetime.now().strftime('%d-%b-%Y at %H:%M:%S')}"
    )
    print(f"Total runtime = {(timer()-global_start_time)/60:.1f} min")
    print("+" * shutil.get_terminal_size().columns)


if __name__ == "__main__":
    main()
