from ...imports import *
from .extract import *
from .visualize import *

# Details of PandExo files are [here](https://natashabatalha.github.io/PandExo/jwstdict.html?highlight=rawdata).


def read_pandexo(filename, extract=False):
    """
    Read a PandExo output file, including S/N estimates.

    When calling PandExo, please use the following settings:
    - set "Baseline" to "Fraction of time: in/out" = 1
    - set "Number of Transits" = 1
    - set "Constant Minimum Noise" = 0 (unless you have a good reason to believe otherwise)

    Parameters
    ----------
    filename : str
        The filepath to a `.p` pickle file generated by Pandexo.
    extract : bool
        Should we try to extract a S/N from a 2D image?
        (This currently likely works only for MIRI/LRS.
        It definitely won't work for NIRISS/SOSS.)

    Returns
    -------
    results : dict
        A dictionary containing noise estimates and intermediate ingredients.
        The typical keys are:
            `1D` = tabular results along the wavelength axis
            `2D` = image results along the wavelength axis and one spatial axis
            `3D` = cube results along the wavelength axis and two spatial axes
    """

    # read the pickle file
    f = open(filename, "rb")
    d = pickle.load(f)

    # get the overall pandexo depth and uncertainty results
    s = d["FinalSpectrum"]
    spectra = {}

    # all results are per-pixel
    spectra["pixel_number"] = np.arange(len(s["wave"]))
    spectra["wavelength"] = s["wave"]
    spectra["planet_model"] = s["spectrum"]
    spectra["planet_realization"] = s["spectrum_w_rand"]
    spectra["depth_uncertainty"] = s["error_w_floor"]
    spectra["snr_per_transit"] = 1 / spectra["depth_uncertainty"]

    # get some metadata about the timing.
    t = d["timing"]
    metadata = {}

    metadata["pandexo_input"] = d["input"]
    metadata["transit_duration"] = t["Transit Duration"]
    metadata["observation_duration"] = t["Transit+Baseline, no overhead (hrs)"]
    metadata["time_per_integration"] = t["Time/Integration incl reset (sec)"]
    metadata["observing_efficiency"] = 0.01 * t["Observing Efficiency (%)"]
    metadata["number_of_groups_per_integration"] = t["APT: Num Groups per Integration"]
    metadata["number_of_integrations_per_transit"] = t["Num Integrations In Transit"]
    metadata["number_of_groups_per_integration"] = t["APT: Num Groups per Integration"]
    metadata["time_per_group"] = t["Seconds per Frame"]

    # make sure that the number of transits is 1
    assert t["Number of Transits"] == 1

    # make sure the in/out transit ratio is about 1
    assert metadata["transit_duration"] / metadata["observation_duration"] > 0.4
    assert metadata["transit_duration"] / metadata["observation_duration"] < 0.6

    # collect some warnings as metadata
    w = d["PandeiaOutTrans"]["warnings"]
    for k in w:
        if "saturated" in k:
            metadata[k] = w[k]

    # collect the raw pandeia results, all per-pixel
    # Details [here](https://jwst-docs.stsci.edu/jwst-exposure-time-calculator-overview/jwst-etc-outputs-overview/jwst-etc-downloads)
    pandeia_results = d["PandeiaOutTrans"]["1d"]
    for k in pandeia_results:
        # skip inputs for calculations that are not necessarily on the pixel grid
        if k in ["wave_calc", "target", "fp", "bg", "bg_rate", "total_flux"]:
            continue
        # most inputs have a wavelength axis embedded with them
        if len(pandeia_results[k]) == 2:
            this_wave = pandeia_results[k][0]
            assert np.all(this_wave == spectra["wavelength"])
            N_w = len(this_wave)
            N_y = len(pandeia_results[k][1])
            if N_y == N_w:
                spectra[k] = pandeia_results[k][1]
                # print(k, len(spectra[k]))
            else:
                print(f"skipping {k} because its size {N_y} isn't {N_w} wavelengths")

    # estimate from pure photon noise (no read noise or 1/f!)
    t_integration = metadata["time_per_integration"] * metadata["observing_efficiency"]
    n_integrations = metadata["number_of_integrations_per_transit"]
    convert_integration_to_transit = np.sqrt(n_integrations) / np.sqrt(2)

    N_photons = spectra["extracted_flux"] * t_integration
    sigma_N_photons = np.sqrt(spectra["extracted_flux_plus_bg"] * t_integration)
    spectra["snr_per_transit_from_photons_only"] = (
        N_photons / sigma_N_photons * convert_integration_to_transit
    )

    # get the per-integration noise from the direct ETC result
    spectra["snr_per_integration_from_etc"] = (
        spectra["extracted_flux"]
        / spectra["extracted_noise"]
        / convert_integration_to_transit
    )
    spectra["snr_per_integration_from_photons_only"] = N_photons / sigma_N_photons

    # make sure the number of integrations per transit lines up
    estimated_integrations_per_transit = (
        metadata["transit_duration"] * 60 * 60 / metadata["time_per_integration"]
    )
    assert np.isclose(
        estimated_integrations_per_transit,
        metadata["number_of_integrations_per_transit"],
        atol=2,
    )

    # make sure the number of groups makes sense
    N_groups = metadata["time_per_integration"] / metadata["time_per_group"] - 1
    assert np.isclose(N_groups, metadata["number_of_groups_per_integration"])
    assert np.isclose(metadata["observing_efficiency"], (N_groups - 1) / (N_groups + 1))

    t = Table(spectra, meta=metadata)

    # get some images
    images = {}
    disperser = d["input"]["Disperser"]
    for k in ["detector", "snr", "saturation"]:
        images[k] = trim_image(d["PandeiaOutTrans"]["2d"][k], disperser=disperser)
    images["snr"] /= np.sqrt(metadata["number_of_integrations_per_transit"])

    if extract:
        t["snr_extracted_from_image"] = extract_sn_from_image(images)
    t["snr_per_integration_from_pandexo_depth"] = 1 / (
        t["depth_uncertainty"]
        * np.sqrt(metadata["number_of_integrations_per_transit"])
        / np.sqrt(2)
    )

    return {"1D": t, "2D": images}
