################################
# - Combine Umbrella Samples - #
################################

# This script combines MCMC samples from different umbrellas into a single chain
# See Gilbert 2022 for details

import os
import sys
import glob
import shutil
import warnings
from datetime import datetime
from timeit import default_timer as timer

print("")
print("+" * shutil.get_terminal_size().columns)
print("ALDERAAN Posterior chain inspection")
print("Initialized {0}".format(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
import argparse
import matplotlib as mpl

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(
        "--planet_no",
        default=None,
        type=int,
        required=True,
        help="ALDERAAN zero-indexed planet identifier (i.e. *NOT* KOI_ID)",
    )
    parser.add_argument(
        "--root_dir",
        default=None,
        type=str,
        required=True,
        help="Root directory for system",
    )
    parser.add_argument(
        "--project_dir",
        default=None,
        type=str,
        required=True,
        help="Project directory for accessing lightcurve data and saving outputs",
    )
    parser.add_argument(
        "--interactive",
        default=False,
        type=bool,
        required=False,
        help="'True' to enable interactive plotting; by default matplotlib backend will be set to 'Agg'",
    )

    args = parser.parse_args()
    MISSION = args.mission
    TARGET = args.target
    PLANET_NO = args.planet_no
    ROOT_DIR = args.root_dir
    PROJECT_DIR = ROOT_DIR + args.project_dir

    # set plotting backend
    if args.interactive == False:
        mpl.use("agg")

except:
    pass

# Set environment variables
sys.path.append(PROJECT_DIR)

# Build directory structure
FIGURE_DIR = PROJECT_DIR + "Figures/" + TARGET + "/"
TRACE_DIR = PROJECT_DIR + "Traces/" + TARGET + "/"
QUICK_TTV_DIR = PROJECT_DIR + "QuickTTVs/" + TARGET + "/"
DLC_DIR = PROJECT_DIR + "Detrended_lightcurves/" + TARGET + "/"
NOISE_DIR = PROJECT_DIR + "Noise_models/" + TARGET + "/"

# check if all the output directories exist and if not, create them
if os.path.exists(FIGURE_DIR) == False:
    os.mkdir(FIGURE_DIR)

if os.path.exists(TRACE_DIR) == False:
    os.mkdir(TRACE_DIR)

if os.path.exists(QUICK_TTV_DIR) == False:
    os.mkdir(QUICK_TTV_DIR)

if os.path.exists(DLC_DIR) == False:
    os.mkdir(DLC_DIR)

if os.path.exists(NOISE_DIR) == False:
    os.mkdir(NOISE_DIR)


# Import packages
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import arviz as az
import astropy
from astropy.io import fits
import corner

from alderaan.constants import *
from alderaan.utils import *
from alderaan.sampling import *
import alderaan.umbrella as umb

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

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

# check for interactive matplotlib backends
if np.any(np.array(["agg", "png", "svg", "pdf", "ps"]) == mpl.get_backend()):
    iplot = False
else:
    iplot = True


# Load MCMC posterior chains
# These can be generated by running "fit_transit_shape_independently.py"
files = glob.glob(TRACE_DIR + TARGET + "*{0}*".format(str(PLANET_NO).zfill(2)))
raw_data = {}

for i, f in enumerate(files):
    if ("N_indep" in f) or ("T_indep" in f) or ("G_indep" in f):
        with fits.open(f) as HDUList:
            ext = HDUList[0].header["UMBRELLA"]
            raw_data[ext] = pd.DataFrame(HDUList[1].data)

            # FITS stores big-endian data; we want little-endian
            for k in raw_data[ext].keys():
                if raw_data[ext][k].dtype == ">f8":
                    raw_data[ext][k] = raw_data[ext][k].astype("float")
                elif raw_data[ext][k].dtype == ">i8":
                    raw_data[ext][k] = raw_data[ext][k].astype("int")
                else:
                    warnings.warn(
                        "Expected all dtypes to be '>f8' or '>i8', but found '{0}'".format(
                            raw_data[ext][k].dtype
                        )
                    )

            # calculate convenience quantities
            for k in raw_data.keys():
                raw_data[k]["ROR"] = np.exp(raw_data[k]["LN_ROR"])
                raw_data[k]["DUR14"] = np.exp(raw_data[k]["LN_DUR14"])


# Visually inspect chains
fig, ax = plt.subplots(1, 1, figsize=(8, 5))
bins = np.linspace(0, 2, 21)

g = raw_data["T"].GAMMA

y, x, _ = ax.hist(
    g, bins=bins, density=True, histtype="step", color="C{0}".format(PLANET_NO), lw=2
)

ax.axvline(1, ls=":", color="k")
ax.text(
    0.95, y.max() * 1.13, r"$\leftarrow$ grazing", va="top", ha="right", fontsize=14
)
ax.text(
    1.05, y.max() * 1.13, r"non-grazing $\rightarrow$", va="top", ha="left", fontsize=14
)
ax.set_xlabel(r"$\gamma$", fontsize=20)
ax.set_ylim(0, y.max() * 1.15)
ax.set_yticks([])
ax.set_title(TARGET + "-{0}".format(PLANET_NO), fontsize=20)
plt.savefig(
    FIGURE_DIR + TARGET + "_transition_gamma_{0}.png".format(PLANET_NO),
    bbox_inches="tight",
)

if iplot:
    plt.show()
else:
    plt.close()

print(
    "{0} out of {1} transition samples are consistent with a grazing transit".format(
        np.sum(g < 1), len(g)
    )
)


# Calculate umbrella weights
gN = raw_data["N"].GAMMA
gT = raw_data["T"].GAMMA
gG = raw_data["G"].GAMMA

# Case 1: nongrazing transit near stellar center
if (np.sum(gT < 1) == 0) & (np.sum(gN < 2) == 0):
    trajectory = "N_only"
    psi_fxns = [umb.psi_N]
    psi_names = ["N"]

# Case 2: nongrazing transit near stellar limb
elif (np.sum(gT < 1) == 0) & (np.sum(gN < 2) != 0):
    trajectory = "N_or_T"
    psi_fxns = [umb.psi_N, umb.psi_T]
    psi_names = ["N", "T"]

# Case 3: ambiguous
elif (np.sum(gT < 1) != 0) & (np.sum(gT > 1) != 0):
    trajectory = "any"
    psi_fxns = [umb.psi_N, umb.psi_T, umb.psi_G]
    psi_names = ["N", "T", "G"]

# Case 4: weakly grazing
elif (np.sum(gT < 1) != 0) & (np.sum(gG > 0) != 0):
    trajectory = "T_or_G"
    psi_fxns = [umb.psi_T, umb.psi_G]
    psi_names = ["T", "G"]

# Case 5: strongly grazing
elif (np.sum(gT < 1) != 0) & (np.sum(gG > 0) == 0):
    trajectory = "G_only"
    psi_fxns = [umb.psi_G]
    psi_names = ["G"]

if len(psi_names) != len(psi_fxns):
    raise ValueError("mismatch between specified number of umbrellas")

plt.figure()
x = np.linspace(-1, 5, 1000)
for i, psi in enumerate(psi_fxns):
    plt.plot(x, psi(x) / psi(x).max(), lw=2.5, label=psi_names[i])
plt.xlabel("$\gamma$", fontsize=20)
plt.ylabel(r"$\psi(\gamma)$", fontsize=20)
plt.legend(fontsize=12)
if iplot:
    plt.show()
else:
    plt.close()


# Post-processed chains will only store core physical parameters
keys = "PERIOD T0 ROR IMPACT DUR14 LD_U1 LD_U2".split()

if len(psi_fxns) == 1:
    pname = psi_names[0]
    uweights = 1 / psi_fxns[0](raw_data[pname].GAMMA)

    samples = raw_data[pname][keys]
    weights = uweights / np.sum(uweights)

elif len(psi_fxns) > 1:
    coords = []

    for i, pname in enumerate(psi_names):
        coords.append(np.array(raw_data[pname].GAMMA))

    # starting estimate of z using single-bin WHAM
    z = np.zeros(len(psi_fxns))

    for i, psi in enumerate(psi_fxns):
        z[i] = np.mean(psi(coords[i]))

    z /= np.sum(z)

    print("Starting estimate for z:  ", z.round(4))

    # first iteration of EMUS
    F = umb.F_iter(z, psi_fxns, coords)
    z = umb.z_iter(F)

    print("First EMUS iteration:  ", z.round(4))
    print(z)

    # iterate until convergence
    for n in range(100):
        # for numerical stability
        z_old = np.copy(z)
        z_old[z_old < 1e-100] = 1e-100

        # Eq. 4 & 7 again
        F = umb.F_iter(z, psi_fxns, coords)
        z = umb.z_iter(F)

        # check if we have converged
        if np.max(np.abs(z - z_old) / z_old) < 1e-10:
            break

    z_emus = np.copy(z)

    print("Final converged z:  ", z.round(4))

    # compute weights as in eq 4 in Matthews et al.
    uweights = []
    for i in range(3):
        wd = 0.0

        for j in range(3):
            wd += psi_fxns[j](coords[i]) / z[j]

        uweights.append(1.0 / wd)

    # concatenate samples and track weights for ease of use
    samples = pd.concat(
        [raw_data["N"], raw_data["T"], raw_data["G"]], ignore_index=True
    )[keys]
    weights = np.hstack(uweights) / np.sum(np.hstack(uweights))

# Estimate KDE bandwidths
bandwidths = pd.DataFrame()

for k in samples.keys():
    print(k)

    bandwidths[k] = [1.0 * get_bw(samples[k], weights)]

# Generate synthetic samples with PDF-Oversampling (Gao+ 2014)
synth_samps = generate_synthetic_samples(
    samples.values, bandwidths.values, n_up=len(samples), weights=weights
)
synth_samps = pd.DataFrame(synth_samps, columns=samples.keys())

# Flag 10-sigma outliers
out = np.zeros(len(synth_samps.PERIOD), dtype="bool")

for k in synth_samps.keys():
    out += astropy.stats.sigma_clip(
        synth_samps[k], cenfunc="median", stdfunc="mad_std", sigma=10
    ).mask

# Make corner plot
fig = corner.corner(synth_samps[~out], color="C{0}".format(PLANET_NO))
plt.savefig(
    FIGURE_DIR + TARGET + "_corner_{0}.png".format(PLANET_NO), bbox_inches="tight"
)

if iplot:
    plt.show()
else:
    plt.close()


# Write combined MCMC samples to .fits file
print("\nWriting MCMC samples to .fits\n")

# set up header info
data_info = [
    # ['VARIABLE', 'UNITS',        'DESCRIPTION',                           'PRIOR'],
    ["PERIOD", "days", "Orbital period", "Normal"],
    ["T0", "days", "Mid-point of first transit", "Normal"],
    ["ROR", "dimensionless", "Planet-to-star radius ratio", "Log-Uniform"],
    ["IMPACT", "dimensionless", "Transit impact parameter", "Uniform"],
    ["DUR14", "days", "Transit duration: 1st-4th contact", "Log-Uniform"],
    [
        "LD_U1",
        "dimensionless",
        "1st quadratic stellar limbdark coeff",
        "Kipping (2013)",
    ],
    [
        "LD_U2",
        "dimensionless",
        "2nd quadratic stellar limbdark coeff",
        "Kipping (2013)",
    ],
]

# build primary HDU and header
hdu = fits.PrimaryHDU()
hdu.header["TARGET"] = (TARGET + "-" + str(PLANET_NO).zfill(2), "Target name")
hdu.header["PHOTOSRC"] = (MISSION, "Source of photometry")
hdu.header["DATETIME"] = (
    datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
    "Date and time of analysis",
)
hdu.header[
    "COMMENT"
] = "ALDERAAN single planet transit fit -- Combined Synthetic Samples"

f_name = PROJECT_DIR + "Traces/{0}/{0}-{1}_CSS_indep.fits".format(
    TARGET, str(PLANET_NO).zfill(2)
)
hdu.writeto(f_name, overwrite=True)

# Store data from dataframes into a single FITS file
with fits.open(f_name) as hduL:
    # Store trace data and summary data to HDUs for each model
    rec_data = synth_samps.to_records(index=False)
    hdu_data = fits.BinTableHDU(data=rec_data, name="SAMPLES")

    hduL.append(hdu_data)

    data_head = hduL["SAMPLES"].header
    for i, dinf in enumerate(data_info):
        data_head[dinf[0]] = (dinf[1], dinf[2])

    hduL.writeto(f_name, overwrite=True)


# Exit program
print("")
print("+" * shutil.get_terminal_size().columns)
print(
    "Umbrella sample processing complete {0}".format(
        datetime.now().strftime("%d-%b-%Y at %H:%M:%S")
    )
)
print("Total runtime = %.1f min" % ((timer() - global_start_time) / 60))
print("+" * shutil.get_terminal_size().columns)
