# === Start Python 2/3 compatibility
from __future__ import absolute_import, division, print_function, unicode_literals
from future.builtins import *  # noqa  pylint: disable=W0401, W0614
from future.builtins.disabled import *  # noqa  pylint: disable=W0401, W0614

# === End Python 2/3 compatibility


import numpy as np
import tempfile
from subprocess import run, CalledProcessError
import os
from shutil import rmtree


class Timespec(object):
    """Emulate C `timespec` struct."""

    def __init__(self, tv_sec, tv_nsec=None):
        if tv_nsec is None:
            self.tv_sec = np.floor(tv_sec)
            self.tv_nsec = int((tv_sec - self.tv_sec) * 1e9)
        else:
            self.tv_sec = tv_sec
            self.tv_nsec = tv_nsec

    def __repr__(self):
        return "tv_sec={:.1f}, tv_nsec={:.1f}".format(self.tv_sec, self.tv_nsec)


class Polyco(object):
    """Hold a pulsar timing solution and provide time of arrival predictions."""

    def __init__(self, tmid, dm, phase_ref, f0, coeff, seg=None):
        self.tmid = tmid
        self.dm = dm
        self.phase_ref = phase_ref
        self.f0 = f0
        self.coeff = coeff
        self.seg = seg

    def mjd2phase(self, t):
        dt = (t - self.tmid) * 1440
        phase = self.phase_ref + dt * 60 * self.f0
        for i in range(len(self.coeff)):
            phase += self.coeff[i] * dt ** i

        return phase

    def unix2phase(self, t):
        t_mjd = unix2mjd(t)

        return self.mjd2phase(t_mjd)

    def next_toa(self, t, freq):
        # Adjust time for dispersion delay
        dm_delay = 4140.0 * self.dm * freq ** (-2)
        psr_t = add_nsec(t, -dm_delay * 1e9)

        phase = self.unix2phase(psr_t)

        # time until next pulse in s
        return (1.0 - (phase - np.floor(phase))) / self.f0


def add_nsec(t, nsec):
    nsec_sum = int(t.tv_nsec + nsec)
    return Timespec(t.tv_sec + nsec_sum // 1000000000, nsec_sum % 1000000000)


def unix2mjd(t):
    # number of days between UNIX epoch and MJD epoch is 40587
    return (t.tv_sec / 86400.0) + (t.tv_nsec / 86400.0 / 1e9) + 40587


def mjd2unix(t):
    t_unix = (t - 40587.0) * 86400
    return Timespec(int(t_unix), (t_unix - int(t_unix)) * 1e9)


class PolycoFile(object):
    """Read a polyco file generated by TEMPO2 in the `newpolyco.dat` format."""

    head_type = [
        ("name", str),
        ("date", str),
        ("UTC", str),
        ("TMID", float),
        ("DM", float),
        ("doppler", float),
        ("log10rms", float),
        ("RPHASE", float),
        ("F0", float),
        ("obs", str),
        ("span", float),
        ("num_coeff", int),
        ("freq", float),
        ("bin_phase", float),
    ]

    def __init__(self, fname):
        self.polyco_specs = []
        self.polycos = []
        self.dm = None
        self.width = None
        self.name = None
        with open(fname, "r") as fh:
            line = fh.readline()
            while line != "":
                if line[:6] == "TEMPO2":
                    # read header
                    new_header = {}
                    for i in range(len(self.head_type)):
                        line = fh.readline()
                        if self.head_type[i][0] == "bin_phase":
                            continue
                        new_header[self.head_type[i][0]] = self.head_type[i][1](
                            line.strip()
                        )
                    self.polyco_specs.append(new_header)
                    # read coeffecients
                    coeff = []
                    for i in range(new_header["num_coeff"]):
                        line = fh.readline()
                        coeff.append(float(line.strip()))
                    params = [new_header[k] for k in ("TMID", "DM", "RPHASE", "F0")] + [
                        coeff
                    ]
                    params += [new_header["span"]]
                    self.polycos.append(Polyco(*params))
                line = fh.readline()
        self.tmid = np.array([p["TMID"] for p in self.polyco_specs])
        for p in self.polyco_specs:
            if p["span"] != self.polyco_specs[0]["span"]:
                raise Exception(
                    "Polycos in file {} have segments of".format(fname)
                    + " different lengths."
                )
        self.spans = np.array(
            [
                np.array((-0.5, 0.5)) * self.polyco_specs[i]["span"] / 60.0 / 24.0
                + self.tmid[i]
                for i in range(len(self.tmid))
            ]
        )

    def get_closest(self, t):
        return self.polycos[np.argmin(np.abs(self.tmid - t))]

    def config_block(self, start_t, end_t=None):
        if end_t is not None:
            start_ind = np.argmin(np.abs(start_t - self.tmid))
            end_ind = np.argmin(np.abs(end_t - self.tmid))
            poly = self.polycos[start_ind : end_ind + 1]
        else:
            poly = [self.get_closest(start_t)]

        return {
            "coeff": [p.coeff for p in poly],
            "t_ref": [p.tmid for p in poly],
            "phase_ref": [p.phase_ref for p in poly],
            "rot_freq": poly[0].f0,
            "dm": poly[0].dm if self.dm is None else self.dm,
            "segment": poly[0].seg * 60,
            "pulsar_name": self.polyco_specs[0]["name"]
            if self.name is None
            else self.name,
            "enabled": True,
            "pulse_width": self.width,
        }

    @classmethod
    def generate(
        cls,
        start,
        end,
        parfile,
        dm=None,
        seg=300.0,
        ncoeff=12,
        max_ha=12.0,
        tempo_dir=None,
    ):
        if tempo_dir is not None:
            env = os.environ.copy()
            env["TEMPO2"] = tempo_dir
        else:
            env = None
        cmd = [
            "tempo2",
            "-f",
            parfile,
            "-polyco",
            '"{:f} {:f} {:d} {:d} {:f} chime inf"'.format(
                start, end, int(seg), ncoeff, max_ha
            ),
            "-tempo1",
        ]
        cmd = " ".join(cmd)

        tmp_dir = tempfile.mkdtemp()
        try:
            run(cmd, cwd=tmp_dir, shell=True, check=True, env=env)
        except CalledProcessError as e:
            print(
                "Command '{}' failed with return code {:d}.".format(e.cmd, e.returncode)
            )
            print(e.output)
            rmtree(tmp_dir)
            return None

        new_polyco = cls(tmp_dir + "/newpolyco.dat")
        if dm is not None:
            new_polyco.dm = dm

        rmtree(tmp_dir)

        return new_polyco
