#!/usr/bin/env python
# Code to write out Tempo2-format PAR files based on Jodrell Crab monthly ephemeris parameters
# Update the text file with the monthly ephemeris:
# % curl -O https://www.jb.man.ac.uk/pulsar/crab/crab2.txt
import numpy as np
import astropy.units as u
import argparse
from astropy.table import Table
from astropy import log
import sys
import os.path as path

desc = "Write out Tempo2 par files based on Jodrell montly ephemeris"
parser = argparse.ArgumentParser(description=desc)
parser.add_argument("MJD", help="MJD to extract ephemeris for", type=float)

## Parse arguments
args = parser.parse_args()

if not path.exists("crab2.txt"):
    log.error(
        'crab2.txt does not exist!  Download with "curl -O https://www.jb.man.ac.uk/pulsar/crab/crab2.txt"'
    )
    sys.exit(2)

validmonths = [
    "JAN",
    "FEB",
    "MAR",
    "APR",
    "MAY",
    "JUN",
    "JUL",
    "JLY",
    "AUG",
    "SEP",
    "OCT",
    "NOV",
    "DEC",
]

day = []
mon = []
yr = []
mjd = []
tJPL = []
tacc = []
f0 = []
f1 = []
dm = []
dm1 = []
for line in open("crab2.txt"):
    line = line.strip()
    if len(line) == 0:
        continue
    line = (line.replace("(", " ")).replace(")", " ")
    cols = line.split()
    # Make sure second column is a valid month, or skip line. This skips headers, etc.
    if not cols[1] in validmonths:
        continue
    # Skip entries before 55910 when DM1 column was added.
    if float(cols[3]) < 55910:
        continue
    # print(line)
    day.append(int(cols[0]))
    mon.append(cols[1])
    yr.append(int(cols[2]))
    mjd.append(float(cols[3]))
    tJPL.append(float(cols[4]))
    tacc.append(float(cols[5]))
    f0.append(float(cols[6]))
    f1.append(float(cols[8]) * 1.0e-15)
    dm.append(float(cols[10]))
    dm1.append(float(cols[11]))

jtab = Table(
    [day, mon, yr, mjd, tJPL, tacc, f0, f1, dm, dm1],
    names=("day", "month", "year", "mjd", "tJPL", "tacc", "f0", "f1", "dm", "dm1"),
    dtype=("i", "S3", "i4", "f8", "f8", "f8", "f8", "f8", "f8", "f8"),
)

diff = np.abs(jtab["mjd"] - args.MJD)
idx = diff.argmin()

if diff[idx] > 18:
    log.warning(
        "Too large a difference ({} days) to closest ephemeris entry!".format(diff[idx])
    )
    # sys.exit(1)

f0 = jtab[idx]["f0"]
f1 = jtab[idx]["f1"]
p0 = 1.0 / f0
p1 = -f1 / (f0 * f0)
# This calculation of F2 comes from the C code at the end of the
# explanatory notes for the Jodrell ephemeris
f2 = 2.0 * p1 * p1 / (p0 * p0 * p0)

skel = """
PSR B0531+21
RAJ 05:34:31.97232
DECJ +22:00:52.069
EPHEM DE200
F0 {0:.15f}
F1 {1:.15g}
C F2 computed as F2=2.0*P1*P1/(P0*P0*P0) according to JBO ephemeris
F2 {2:.15g}
PEPOCH {3:.15f}
TZRSITE @
TZRFRQ 0.0
TZRMJD {4:.15f}
START {5:.4f}
FINISH {6:.4f}
UNITS TDB
CLOCK TT(TAI)
DM {7:.5f}
DM1 {8:.5g}
DMEPOCH {9:.15f}
TRES {10:.3f}
"""

# Note 5: The observed barycentric frequency and its first derivative
# are quoted at the arrival time, using the DE200 ephemeris
# So we use this value for both the period epoch and for TZRMJD
# This is in TDB units
mjd = jtab[idx]["mjd"]
pepoch = mjd + jtab[idx]["tJPL"] / 86400.0

print(
    skel.format(
        f0,
        f1,
        f2,
        pepoch,
        pepoch,
        mjd - 16,
        mjd + 16,
        jtab[idx]["dm"],
        jtab[idx]["dm1"],
        pepoch,
        jtab[idx]["tacc"],
    )
)
