#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thurs Jul 21 18:07 2022

@author: MCR

Custom JWST DMS pipeline steps for Stage 4 (lightcurve fitting).
"""

from datetime import datetime
from exotic_ld import StellarLimbDarkening
import juliet
import numpy as np
import os
import pandas as pd
import ray
from tqdm import tqdm

from jwst import datamodels
from supreme_spoon.utils import fancyprint


def bin_at_pixel(flux, error, wave, npix):
    """Similar to bin_at_resolution, but will bin in widths of a set number of
    pixels instead of at a fixed resolution.

    Parameters
    ----------
    flux : array-like[float]
        Flux values.
    error : array-like[float]
        Flux error values.
    wave : array-like[float]
        Wavelength values.
    npix : int
        Number of pixels per bin.

    Returns
    -------
    wave_bin : np.ndarray[float]
        Central bin wavelength.
    wave_low : np.ndarray[float]
        Lower edge of wavelength bin.
    wave_up : np.ndarray[float]
        Upper edge of wavelength bin.
    dout : np.ndarray[float]
        Binned depth.
    derrout : np.ndarray[float]
        Error on binned depth.
    """

    # Calculate number of bins given wavelength grid and npix value.
    nint, nwave = np.shape(flux)
    # If the number of pixels does not bin evenly, trim from beginning and end.
    if nwave % npix != 0:
        cut = nwave % npix
        cut_s = int(np.floor(cut/2))
        cut_e = -1*(cut - cut_s)
        flux = flux[:, cut_s:cut_e]
        error = error[:, cut_s:cut_e]
        wave = wave[:, cut_s:cut_e]
        nint, nwave = np.shape(flux)
    nbin = int(nwave / npix)

    # Sum flux in bins and calculate resulting errors.
    flux_bin = np.nansum(np.reshape(flux, (nint, nbin, npix)), axis=2)
    err_bin = np.sqrt(np.nansum(np.reshape(error, (nint, nbin, npix))**2,
                                axis=2))
    # Calculate mean wavelength per bin.
    wave_bin = np.nanmean(np.reshape(wave, (nint, nbin, npix)), axis=2)
    # Get bin wavelength limits.
    up = np.concatenate([wave_bin[0][:, None],
                         np.roll(wave_bin[0][:, None], 1)], axis=1)
    lw = np.concatenate([wave_bin[0][:, None],
                         np.roll(wave_bin[0][:, None], -1)], axis=1)
    wave_up = (np.mean(up, axis=1) - wave_bin[0])[:-1]
    wave_up = np.insert(wave_up, -1, wave_up[-1])
    wave_up = np.repeat(wave_up, nint).reshape(nbin, nint).transpose(1, 0)
    wave_up = wave_bin + wave_up
    wave_low = (wave_bin[0] - np.mean(lw, axis=1))[1:]
    wave_low = np.insert(wave_low, 0, wave_low[0])
    wave_low = np.repeat(wave_low, nint).reshape(nbin, nint).transpose(1, 0)
    wave_low = wave_bin - wave_low

    return wave_bin, wave_low, wave_up, flux_bin, err_bin


def bin_at_resolution(inwave_low, inwave_up, flux, flux_err, res,
                      method='sum'):
    """Function that bins input wavelengths and transit depths (or any other
    observable, like flux) to a given resolution "res". Can handle 1D or 2D
    flux arrays.

    Parameters
    ----------
    inwave_low : array-like[float]
        Lower edge of wavelength bin. Must be 1D.
    inwave_up : array-like[float]
        Upper edge of wavelength bin. Must be 1D.
    flux : array-like[float]
        Flux values at each wavelength. Can be 1D or 2D. If 2D, the first axis
        must be the one corresponding to wavelength.
    flux_err : array-like[float]
        Errors corresponding to each flux measurement. Must be the same shape
        as flux.
    res : int
        Target resolution at which to bin.
    method : str
        Method to bin depths. Either "sum" or "average".

    Returns
    -------
    binned_waves : array-like[float]
        Wavelength of the given bin at the desired resolution.
    binned_werr : array-like[float]
        Half-width of the wavelength bin.
    binned_flux : array-like[float]
        Binned flux.
    binned_ferr : array-like[float]
        Error on binned flux.
    """

    def nextstep(w, r):
        return w * (2 * r + 1) / (2 * r - 1)

    # Sort quantities in order of increasing wavelength.
    waves = np.nanmean([inwave_low, inwave_up], axis=0)
    if np.ndim(waves) > 1:
        raise ValueError('Input wavelength array must be 1D.')
    ii = np.argsort(waves)
    waves, flux, flux_err = waves[ii], flux[ii], flux_err[ii]
    inwave_low, inwave_up = inwave_low[ii], inwave_up[ii]
    # Calculate the input resolution and check that we are not trying to bin
    # to a higher R.
    average_input_res = np.mean(waves[1:] / np.diff(waves))
    if res > average_input_res:
        raise ValueError('You are trying to bin at a higher resolution than '
                         'the input.')
    else:
        fancyprint('Binning from an average resolution of '
                   'R={:.0f} to R={}'.format(average_input_res, res))

    # Make the binned wavelength grid.
    outwave_low = []
    w_i, w_ip1 = waves[0], waves[0]
    while w_ip1 < waves[-1]:
        outwave_low.append(w_ip1)
        w_ip1 = nextstep(w_i, res)
        w_i = w_ip1
    outwave_low = np.array(outwave_low)
    outwave_up = np.append(outwave_low[1:], waves[-1])
    binned_waves = np.mean([outwave_low, outwave_up], axis=0)
    binned_werr = (outwave_up - outwave_low) / 2

    # Loop over all wavelengths in the input and bin flux and error into the
    # new wavelength grid.
    ii = 0
    for wl, wu in zip(outwave_low, outwave_up):
        first_time, count = True, 0
        current_flux = np.ones_like(flux[ii]) * np.nan
        current_ferr = np.ones_like(flux_err[ii]) * np.nan
        weight = []
        for i in range(ii, len(waves)):
            # If the wavelength is fully within the bin, append the flux and
            # error to the current bin info.
            if inwave_low[i] >= wl and inwave_up[i] < wu:
                if np.ndim(flux) == 1:
                    current_flux = np.hstack([flux[i], current_flux])
                    current_ferr = np.hstack([flux_err[i], current_ferr])
                else:
                    current_flux = np.vstack([flux[i], current_flux])
                    current_ferr = np.vstack([flux_err[i], current_ferr])
                count += 1
                weight.append(1)
            # For edge cases where one of the input bins falls on the edge of
            # the binned wavelength grid, linearly interpolate the flux into
            # the new bins.
            # Upper edge split.
            elif inwave_low[i] < wu <= inwave_up[i]:
                inbin_width = inwave_up[i] - inwave_low[i]
                in_frac = (inwave_up[i] - wu) / inbin_width
                weight.append(in_frac)
                if np.ndim(flux) == 1:
                    current_flux = np.hstack([flux[i], current_flux])
                    current_ferr = np.hstack([flux_err[i], current_ferr])
                else:
                    current_flux = np.vstack([flux[i], current_flux])
                    current_ferr = np.vstack([flux_err[i], current_ferr])
                count += 1
            # Lower edge split.
            elif inwave_low[i] < wl <= inwave_up[i]:
                inbin_width = inwave_up[i] - inwave_low[i]
                in_frac = (wl - inwave_low[i]) / inbin_width
                weight.append(in_frac)
                if np.ndim(flux) == 1:
                    current_flux = np.hstack([flux[i], current_flux])
                    current_ferr = np.hstack([flux_err[i], current_ferr])
                else:
                    current_flux = np.vstack([flux[i], current_flux])
                    current_ferr = np.vstack([flux_err[i], current_ferr])
                count += 1
            # Since wavelengths are in increasing order, once we exit the bin
            # completely we're done.
            if inwave_low[i] >= wu or i == len(waves)-1:
                if count != 0:
                    # If something was put into this bin, bin it using the
                    # requested method.
                    weight.append(0)
                    weight = np.array(weight)
                    if method == 'sum':
                        if np.ndim(current_flux) != 1:
                            thisflux = np.nansum(current_flux * weight[:, None],
                                                 axis=0)
                        else:
                            thisflux = np.nansum(current_flux * weight, axis=0)
                        thisferr = np.sqrt(np.nansum(current_ferr**2, axis=0))
                    elif method == 'average':
                        if np.ndim(current_flux) != 1:
                            thisflux = np.nansum(current_flux * weight[:, None],
                                                 axis=0)
                            thisflux /= np.nansum(weight)
                        else:
                            thisflux = np.nansum(current_flux * weight, axis=0)
                            thisflux /= np.nansum(weight)
                        thisferr = np.sqrt(np.nansum(current_ferr**2, axis=0))
                        thisferr /= np.nansum(weight)
                    else:
                        raise ValueError('Unknown method.')
                else:
                    # If nothing is in the bin (can happen if the output
                    # reslution is higher than the local input resolution),
                    # append NaNs
                    if np.ndim(flux) == 1:
                        thisflux, thisferr = np.nan, np.nan
                    else:
                        thisflux = np.ones_like(flux[0]) * np.nan
                        thisferr = np.ones_like(flux[0]) * np.nan
                # Store the binned quantities.
                if ii == 0:
                    binned_flux = thisflux
                    binned_ferr = thisferr
                else:
                    binned_flux = np.vstack([binned_flux, thisflux])
                    binned_ferr = np.vstack([binned_ferr, thisferr])
                # Move to the next bin.
                ii = i-1
                break

    # If the input was 1D, reformat to match.
    if np.ndim(flux) == 1:
        binned_flux = binned_flux[:, 0]
        binned_ferr = binned_ferr[:, 0]

    return binned_waves, binned_werr, binned_flux, binned_ferr


@ray.remote
def fit_data(data_dictionary, priors, output_dir, bin_no, num_bins):
    """Functional wrapper around run_juliet to make it compatible for
    multiprocessing with ray.
    """

    fancyprint('Fitting bin {} / {}'.format(bin_no, num_bins))

    # Get key names.
    all_keys = list(data_dictionary.keys())

    # Unpack fitting arrays
    t = {'SOSS': data_dictionary['times']}
    flux = {'SOSS': data_dictionary['flux']}
    flux_err = {'SOSS': data_dictionary['error']}

    # Initialize GP and linear model regressors.
    gp_regressors = None
    linear_regressors = None
    if 'GP_parameters' in all_keys:
        gp_regressors = {'SOSS': data_dictionary['GP_parameters']}
    if 'lm_parameters' in all_keys:
        linear_regressors = {'SOSS': data_dictionary['lm_parameters']}

    fit_results = run_juliet(priors, t, flux, flux_err, output_dir,
                             gp_regressors, linear_regressors)

    return fit_results


def fit_lightcurves(data_dict, prior_dict, order, output_dir, fit_suffix,
                    nthreads=4):
    """Wrapper about both the juliet and ray libraries to parallelize juliet's
    lightcurve fitting functionality.

    Parameters
    ----------
    data_dict : dict
        Dictionary of fitting data: time, flux, and flu error.
    prior_dict : dict
        Dictionary of fitting priors in juliet format.
    order : int
        SOSS diffraction order.
    output_dir : str
        Path to directory to which to save results.
    fit_suffix : str
        String to label the results of this fit.
    nthreads : int
        Number of cores to use for multiprocessing.

    Returns
    -------
    results : juliet.fit object
        The results of the juliet fit.
    """

    # Initialize results dictionary and keynames.
    results = dict.fromkeys(data_dict.keys(), [])
    keynames = list(data_dict.keys())

    # Format output directory
    if output_dir[-1] != '/':
        output_dir += '/'

    # Initialize ray with specified number of threads.
    ray.shutdown()
    ray.init(num_cpus=nthreads)

    # Set juliet fits as remotes to run parallel with ray.
    all_fits = []
    num_bins = np.arange(len(keynames))+1
    for i, keyname in enumerate(keynames):
        outdir = output_dir + 'speclightcurve{2}/order{0}_{1}'.format(order, keyname, fit_suffix)
        all_fits.append(fit_data.remote(data_dict[keyname],
                                        prior_dict[keyname],
                                        output_dir=outdir,
                                        bin_no=num_bins[i],
                                        num_bins=len(num_bins)))
    # Run the fits.
    ray_results = ray.get(all_fits)

    # Reorder the results based on the key name.
    for i in range(len(keynames)):
        keyname = keynames[i]
        results[keyname] = ray_results[i]

    return results


def gen_ld_coefs(spectrace_ref, wavebin_low, wavebin_up, order, m_h, logg, teff,
                 ld_data_path, model_type='stagger'):
    """Generate estimates of quadratic limb-darkening coefficients using the
    ExoTiC-LD package.

    Parameters
    ----------
    spectrace_ref : str
        Path to spectrace reference file.
    wavebin_low : array-like[float]
        Lower edge of wavelength bins.
    wavebin_up: array-like[float]
        Upper edge of wavelength bins.
    order : int
        SOSS diffraction order.
    m_h : float
        Stellar metallicity as [M/H].
    logg : float
        Stellar log gravity.
    teff : float
        Stellar effective temperature in K.
    ld_data_path : str
        Path to ExoTiC-LD model data.
    model_type : str
        Identifier for type of stellar model to use. See
        https://exotic-ld.readthedocs.io/en/latest/views/supported_stellar_grids.html
        for supported grids.

    Returns
    -------
    c1s : array-like[float]
        c1 parameter for the quadratic limb-darkening law.
    c2s : array-like[float]
        c2 parameter for the quadratic limb-darkening law.
    """

    # Set up the stellar model parameters - with specified model type.
    sld = StellarLimbDarkening(m_h, teff, logg, model_type, ld_data_path)

    # Load the most up to date throughput info for SOSS
    spec_trace = datamodels.SpecTraceModel(spectrace_ref)
    wavelengths = spec_trace.trace[order-1].data['WAVELENGTH'] * 10000
    throughputs = spec_trace.trace[order-1].data['THROUGHPUT']
    # Note that custom throughputs are used.
    mode = 'custom'

    # Compute the LD coefficients over the given wavelength bins.
    c1s, c2s = [], []
    for wl, wu in tqdm(zip(wavebin_low * 10000, wavebin_up * 10000),
                       total=len(wavebin_low)):
        wr = [wl, wu]
        try:
            c1, c2 = sld.compute_quadratic_ld_coeffs(wr, mode, wavelengths,
                                                     throughputs)
        except ValueError:
            c1, c2 = np.nan, np.nan
        c1s.append(c1)
        c2s.append(c2)
    c1s = np.array(c1s)
    c2s = np.array(c2s)

    return c1s, c2s


def read_ld_coefs(filename, wavebin_low, wavebin_up, ld_model='quadratic'):
    """Unpack limb darkening coefficients and interpolate to the wavelength
    grid of data being fit. File must be comma separated with three columns:
    wavelength c1 and c2.

    Parameters
    ----------
    filename : str
        Path to file containing model limb darkening coefficients.
    wavebin_low : array-like[float]
        Lower edge of wavelength bins being fit.
    wavebin_up : array-like[float]
        Upper edge of wavelength bins being fit.
    ld_model : str
        Limb darkening model.

    Returns
    -------
    prior_q1 : array-like[float]
        Model estimates for q1 parameter.
    prior_q2 : array-like[float]
        Model estimates for q2 parameter.
    """

    # Open the LD model file and convert c1 and c2 parameters to q1 and q2 of
    # the Kipping (2013) parameterization.
    ld = pd.read_csv(filename, comment='#', sep=',')
    q1s, q2s = juliet.reverse_q_coeffs(ld_model, ld['c1'].values,
                                       ld['c2'].values)

    # Get model wavelengths and sort in increasing order.
    waves = ld['wave'].values
    ii = np.argsort(waves)
    waves = waves[ii]
    q1s, q2s = q1s[ii], q2s[ii]

    # Check that coeffs in file span the wavelength range of the observations
    # with at least the same resolution.
    if np.min(waves) < np.min(wavebin_low) or np.max(waves) > np.max(wavebin_up):
        msg = 'LD coefficient file does not span the full wavelength range ' \
              'of the observations.'
        raise ValueError(msg)
    if len(waves) < len(wavebin_low):
        msg = 'LD coefficient file has a coarser wavelength grid than ' \
              'the observations.'
        raise ValueError(msg)

    prior_q1, prior_q2 = [], []
    # Loop over all fitting bins. Calculate mean of model LD coefs within that
    # range.
    for wl, wu in zip(wavebin_low, wavebin_up):
        current_q1, current_q2 = [], []
        for w, q1, q2 in zip(waves, q1s, q2s):
            if wl < w <= wu:
                current_q1.append(q1)
                current_q2.append(q2)
            # Since LD model wavelengths are sorted in increasing order, once
            # we are above the upper edge of the bin, we can stop.
            elif w > wu:
                prior_q1.append(np.nanmean(current_q1))
                prior_q2.append(np.nanmean(current_q2))
                break

    return prior_q1, prior_q2


def run_juliet(priors, t_lc, y_lc, yerr_lc, out_folder,
               gp_regressors_lc, linear_regressors_lc):
    """Wrapper around the lightcurve fitting functionality of the juliet
    package.

    Parameters
    ----------
    priors : dict
        Dictionary of fitting priors.
    t_lc : dict
        Time axis.
    y_lc : dict
        Normalized lightcurve flux values.
    yerr_lc : dict
        Errors associated with each flux value.
    out_folder : str
        Path to folder to which to save results.
    gp_regressors_lc : dict
        GP regressors to fit, if any.
    linear_regressors_lc : dict
        Linear model regressors, if any.

    Returns
    -------
    res : juliet.fit object
        Results of juliet fit.
    """

    if np.isfinite(y_lc['SOSS']).all():
        # Load in all priors and data to be fit.
        dataset = juliet.load(priors=priors, t_lc=t_lc, y_lc=y_lc,
                              yerr_lc=yerr_lc, out_folder=out_folder,
                              GP_regressors_lc=gp_regressors_lc,
                              linear_regressors_lc=linear_regressors_lc)
        kwargs = {'print_progress': False}

        # Run the fit.
        try:
            res = dataset.fit(sampler='dynesty', **kwargs)
        except KeyboardInterrupt as err:
            raise err
        except:
            fancyprint('Exception encountered.', msg_type='WARNING')
            fancyprint('Skipping bin.', msg_type='WARNING')
            res = None
    else:
        fancyprint('NaN bin encountered.', msg_type='WARNING')
        fancyprint('Skipping bin.', msg_type='WARNING')
        res = None

    return res


def save_transmission_spectrum(wave, wave_err, dppm, dppm_err, order, outdir,
                               filename, target, extraction_type, resolution,
                               fit_meta='', occultation_type='transit'):
    """Write a transmission/emission spectrum to file.

    Parameters
    ----------
    wave : array-like[float]
        Wavelength values.
    wave_err : array-like[float]
        Bin half-widths for each wavelength bin.
    dppm : array-like[float]
        Transit/eclipse depth in each bin.
    dppm_err : array-like[float]
        Error on the transit/eclipse depth in each bin.
    order : array-like[int]
        SOSS order corresponding to each bin.
    outdir : str
        Firectory to whch to save outputs.
    filename : str
        Name of the file to which to save spectra.
    target : str
        Target name.
    extraction_type : str
        Type of extraction: either box or atoca.
    resolution: int, str
        Spectral resolution of spectrum.
    fit_meta: str
        Fitting metadata.
    occultation_type : str
        Type of occultation; either 'transit' or 'eclipse'.
    """

    # Pack the quantities into a dictionary.
    dd = {'wave': wave,
          'wave_err': wave_err,
          'dppm': dppm,
          'dppm_err': dppm_err,
          'order': order}
    # Save the dictionary as a csv.
    df = pd.DataFrame(data=dd)
    if os.path.exists(outdir + filename):
        os.remove(outdir + filename)

    # Re-open the csv and append some critical info the header.
    f = open(outdir + filename, 'w')
    f.write('# Target: {}\n'.format(target))
    f.write('# Instrument: NIRISS/SOSS\n')
    f.write('# Pipeline: supreme-SPOON\n')
    f.write('# 1D Extraction: {}\n'.format(extraction_type))
    f.write('# Spectral Resoluton: {}\n'.format(resolution))
    f.write('# Author: {}\n'.format(os.environ.get('USER')))
    f.write('# Date: {}\n'.format(datetime.utcnow().replace(microsecond=0).isoformat()))
    f.write(fit_meta)
    f.write('# Column wave: Central wavelength of bin (micron)\n')
    f.write('# Column wave_err: Wavelength bin halfwidth (micron)\n')
    if occultation_type == 'transit':
        f.write('# Column dppm: (Rp/R*)^2 (ppm)\n')
        f.write('# Column dppm_err: Error in (Rp/R*)^2 (ppm)\n')
    else:
        f.write('# Column dppm: (Fp/F*) (ppm)\n')
        f.write('# Column dppm_err: Error in (Fp/F*) (ppm)\n')
    f.write('# Column order: SOSS diffraction order\n')
    f.write('#\n')
    df.to_csv(f, index=False)
    f.close()
