#!/usr/bin/env python
"""
Example 3PC .fits reader. 

If you have the ftools installed, here's an easy way to see the file structure and column names:
ftlist 3PC_Catalog+SEDs_20230627.fits HC

WARNING! Some of the FITS columns that contain numerical data are actually stored as ASCII strings,
because they have upper limits preceded with "<" and missing values indicated by "NULL".

David.Smith@u-bordeaux.fr, 21 September 2022. Updated 27 June 2023.
Updated 2023-07-20 by PSR

"""

import astropy.io.fits as pyfits
import sys
import matplotlib.pyplot as plt
import numpy as np
import re


def col_to_arrays(col, unc_col=None):
    "Parse ASCII column including handling upper limits and bidirectional error bars"
    if type(col) == np.chararray:
        # Convert ASCII column to floats, with np.nan for NULL values
        # and set boolean array where the values are  upper limits
        dat = [float(xx.replace("<", "")) if xx.strip() != "" else np.nan for xx in col]
        uplims = [("<" in s) for s in col]
    else:
        dat = col
        uplims = np.zeros(len(dat), dtype=bool)

    if unc_col is None:
        unc = np.zeros(len(dat))
    else:
        if type(unc_col) == np.chararray:
            upper_errors = np.zeros(len(dat))
            lower_errors = np.zeros(len(dat))
            i = 0
            for vv in unc_col:
                if vv == "":
                    # Leave error at 0.0
                    pass
                elif vv.startswith("("):
                    try:
                        v1, v2 = re.findall("[*0-9.eEg+-]+", vv)
                    except ValueError:
                        print(f"vv = [{vv}]")
                        raise
                    if v1 != "*" and v2 != "*":
                        # Leave error at 0 for (*,*)
                        lower_errors[i] = float(v1)
                        upper_errors[i] = float(v2)
                else:
                    print(f"Unknown errorbar type [{vv}]")
                i += 1
            unc = [lower_errors, upper_errors]
        else:
            # Uncertainty is a numeric type. May have np.nan for NULL values
            unc = unc_col

    return (dat, unc, uplims)


if __name__ == "__main__":
    # This makes a simple scatter plot of EDOT and LUMG,
    # with upper limits and uncertainties
    # It should be easy to extend for other columns

    hdulist = pyfits.open("3PC_Catalog+SEDs_20230727.fits")
    # "tbf" means The Big File
    tbf = hdulist["PULSARS_BIGFILE"].data

    xdat, xunc, xuplims = col_to_arrays(tbf["EDOT"])

    ydat, yunc, yuplims = col_to_arrays(tbf["LUMG"], tbf["e_LUMG_stat"])

    fig, ax = plt.subplots()
    ax.errorbar(
        xdat,
        ydat,
        yerr=yunc,
        uplims=yuplims,
        marker=".",
        linestyle="none",
    )
    ax.set_xscale("log")
    ax.set_yscale("log")
    plt.show()

    sys.exit(0)

    # Below here is just trivial code to read and print data from both FITS extensions.
    for psr, code, edot, lumg in zip(tbf["PSRJ"], tbf["PSR_Code"], tbf["EDOT"]):
        print(f"{psr} {code:5s} {edot:.2e}")

    sedinfo = hdulist["LAT_Point_Source_Catalog"].data
    for psr, ts in zip(sedinfo["Source_Name"], sedinfo["TSSpec"]):
        print(psr, ts)