""" Main implementation of parcel model.
"""
from __future__ import print_function
from __future__ import absolute_import
from builtins import zip
from builtins import range
from builtins import object

import os

from scipy.optimize import bisect

# Parcel model imports
from . import constants as c
from . import output
from . aerosol import AerosolSpecies
# from . integrator import Integrator
from . thermo import *

__all__ = ['ParcelModel', ]

class ParcelModelError(Exception):
    """ Custom exception to throw during parcel model execution.
    """
    def __init__(self, error_str):
        self.error_str = error_str

    def __str__(self):
        return repr(self.error_str)

class ParcelModel(object):
    """ Wrapper class for instantiating and running the parcel model.

    The parcel model has been implemented in an object-oriented format to facilitate
    easy extensibility to different aerosol and meteorological conditions. A
    typical use case would involve specifying the initial conditions such as:

    >>> import pyrcel as pm
    >>> P0 = 80000.
    >>> T0 = 283.15
    >>> S0 = 0.0
    >>> V = 1.0
    >>> aerosol1 = pm.AerosolSpecies('sulfate',
    ...                              Lognorm(mu=0.025, sigma=1.3, N=2000.),
    ...                              bins=200, kappa=0.54)
    >>> initial_aerosols = [aerosol1, ]
    >>> z_top = 50.
    >>> dt = 0.01

    which initializes the model with typical conditions at the top of the boundary
    layer (800 hPa, 283.15 K, 100% Relative Humidity, 1 m/s updraft), and a simple
    sulfate aerosol distribution which will be discretized into 200 size bins to
    track. Furthermore the model was specified to simulate the updraft for 50
    meters (`z_top`) and use a time-discretization of 0.01 seconds. This
    timestep is used in the model output -- the actual ODE solver will generally
    calculate the trace of the model at many more times.

    Running the model and saving the output can be accomplished by invoking::

    >>> model = pm.ParcelModel(initial_aerosols, V, T0, S0, P0)
    >>> par_out, aer_out = pm.run(z_top, dt)

    This will yield `par_out`, a  :class:`pandas.DataFrame` containing the meteorological
    conditions in the parcel, and `aerosols`, a dictionary of :class:`DataFrame` objects
    for each species in `initial_aerosols` with the appropriately tracked size
    bins and their evolution over time.

    Attributes
    ----------
    V, T0, S0, P0, aerosols : floats
        Initial parcel settings (see **Parameters**).
    _r0s : array_like of floats
        Initial equilibrium droplet sizes.
    _r_drys : array_like of floats
        Dry radii of aerosol population.
    _kappas : array_like of floats
        Hygroscopicity of each aerosol size.
    _Nis : array_like of floats
        Number concentration of each aerosol size.
    _nr : int
        Number of aerosol sizes tracked in model.
    _model_set : boolean
        Flag indicating whether or not at any given time the model
        initialization/equilibration routine has been run with the current
        model settings.
    _y0 : array_like
        Initial state vector.

    Methods
    -------
    run(t_end, dt, max_steps=1000, solver="odeint", output_fmt="dataframes",\
        terminate=False, solver_args={})
        Execute model simulation.
    set_initial_conditions(V=None, T0=None, S0=None, P0=None, aerosols=None)
        Re-initialize a model simulation in order to run it.

    See Also
    --------
    _setup_run : companion routine which computes equilibrium droplet sizes
                 and sets the model's state vectors.

    """

    def __init__(self, aerosols, V, T0, S0, P0, console=False, accom=c.ac,
                 truncate_aerosols=False):
        """ Initialize the parcel model.

        Parameters
        ----------
        aerosols : array_like sequence of :class:`AerosolSpecies`
            The aerosols contained in the parcel.
        V, T0, S0, P0 : floats
            The updraft speed and initial temperature (K), pressure (Pa),
            supersaturation (percent, with 0.0 = 100% RH).
        console : boolean, optional
            Enable some basic debugging output to print to the terminal.
        accom : float, optional (default=:const:`constants.ac`)
            Condensation coefficient
        truncate_aerosols : boolean, optional (default=**False**)
            Eliminate extremely small aerosol which will cause numerical problems

        """
        self._model_set = False
        self.console = console
        self.trunc = truncate_aerosols

        self.V = V
        self.T0 = T0
        self.S0 = S0
        self.P0 = P0
        if self.trunc:
            self.aerosols = self._replace_aerosol(aerosols)
        else:
            self.aerosols = aerosols
        self.accom = accom

        # To be set by call to "self._setup_run()"
        self._r0s = None
        self._r_drys = None
        self._kappas = None
        self._Nis = None
        self._nr = None

        self._setup_run()

    def set_initial_conditions(self, V=None, T0=None, S0=None, P0=None, aerosols=None):
        """ Set the initial conditions and parameters for a new parcel
        model run without having to create a new :class:`ParcelModel` instance.

        Based on the aerosol population which has been stored in the model, this
        method will finish initializing the model. This has three major parts:

        1. concatenate the aerosol population information (their dry radii,
           hygroscopicities, etc) into single arrays which can be placed into the
           state vector for forward integration.
        2. Given the initial ambient water vapor concentration (computed from the
           temperature, pressure, and supersaturation), determine how much water
           must already be coated on the aerosol particles in order for their
           size to be in equilibrium.
        3. Set-up the state vector with these initial conditions.

        Once the state vector has been set up, the setup routine will record
        attributes in the parent instance of the :class:`ParcelModel`.

        Parameters
        ----------
        V, T0, S0, P0 : floats
            The updraft speed and initial temperature (K), pressure (Pa),
            supersaturation (percent, with 0.0 = 100% RH).
        aerosols : array_like sequence of :class:`AerosolSpecies`
            The aerosols contained in the parcel.

        Raises
        ------
        ParcelModelError
            If an equilibrium droplet size distribution could not be calculated.

        Notes
        -----
        The actual setup occurs in the private method `_setup_run()`; this
        method is simply an interface that can be used to modify an existing
        :class:`ParcelModel`.

        """

        if V:
            self.V = V
        if T0:
            self.T0 = T0
        if P0:
            self.P0 = P0
        if S0:
            self.S0 = S0
        if aerosols:
            if self.trunc:
                self.aerosols = self._replace_aerosol(aerosols)
            else:
                self.aerosols = aerosols

        if T0 or P0 or S0 or aerosols:
            self._setup_run()

    def _replace_aerosol(self, aerosols, smallest_rdry=1e-10):
        """ Truncates the bins with the smallest particles so that none
        have a dry radius of less than 0.1 nm.

        Parameters
        ----------
        aerosols : array_like sequence of :class:`AerosolSpecies`
            The aerosols contained in the parcel.
        smallest_rdry: float
            Smallest allowable dry radius bin.

        Returns
        -------
        fixed_aerosols : array_like sequence of :class:`AerosolSpecies`
            The modified sequence of aerosols.

        """
        fixed_aerosols = []

        for aer in aerosols:
            N_old = np.sum(aer.Nis)
            bin_count = np.count_nonzero(aer.r_drys[aer.r_drys < smallest_rdry])
            if bin_count > 0:
                aer_new = AerosolSpecies(aer.species,
                                         aer.distribution,
                                         kappa=aer.kappa,
                                         bins=aer.bins,
                                         r_min=smallest_rdry*1e6)
                N_new = np.sum(aer_new.Nis)
                if self.console:
                    print("%s: Removed %03d bins, N change: %5.1f -> %5.1f (%2.1f%%)" %
                          (aer.species, bin_count, N_old, N_new, (N_new-N_old)/N_new))
                fixed_aerosols.append(aer_new)
            else:
                fixed_aerosols.append(aer)

        return fixed_aerosols

    def _setup_run(self):
        """ Perform equilibration and initialization calculations for the
        parcel model simulation.

        .. note:: See `set_initial_conditions` for full details.

        """
        T0, S0, P0 = self.T0, self.S0, self.P0
        z0 = 0.0
        out = dict()

        if self.console:
            print("Setting up parcel model initial conditions...")

        # 1) Setup aerosols
        # a) grab all the initial aerosol size/concentrations
        species = []
        r_drys, Nis, kappas = [], [], []
        for aerosol in self.aerosols:
            r_drys.extend(aerosol.r_drys)
            kappas.extend([aerosol.kappa]*aerosol.nr)
            Nis.extend(aerosol.Nis)
            species.extend([aerosol.species]*aerosol.nr)

        r_drys = np.array(r_drys)
        kappas = np.array(kappas)
        Nis = np.array(Nis)

        out['r_drys'] = r_drys
        out['kappas'] = kappas
        out['Nis'] = Nis

        if self.console:
            # TODO: Fix print table formatting
            print("AEROSOL DISTRIBUTION")
            print("%8s %6s" % ("r", "N"))
            for sp, r, N in zip(species, r_drys, Nis):
                print("%10s %2.2e %4.1f" % (sp, r, N))
            print("\n"+"-"*44)

        # 2) Setup parcel initial conditions
        # a) water vapor
        #        RH    * (Rd/Rv = epsilon) * ( es / P - es )
        wv0 = (S0 + 1.)*(c.epsilon*es(T0-273.15)/(P0-es(T0-273.15)))  # Water Vapor mixing ratio, kg/kg

        # b) find equilibrium wet particle radius
        # wrapper function for quickly computing deviation from chosen
        # equilibrium supersaturation given current size and dry size
        f = lambda r, r_dry, kappa: Seq(r, r_dry, T0, kappa) - S0
        # Compute the equilibrium wet particle radii
        r0s = []
        # r_b, _ = kohler_crit(T0, r_drys[-1], kappas[-1])
        # for r_dry , kappa in zip(r_drys, kappas)[::-1]:
        for r_dry, kappa in zip(reversed(r_drys), reversed(kappas)):
        # Locate the critical value (f(r_crit) > 0), and use bisection from it and
        # r_dry (f(r_dry) < 0) to find equilibrium wet particle radius for a given S0
            r_b, _ = kohler_crit(T0, r_dry, kappa)
            r_a = r_dry

            r0 = bisect(f, r_a, r_b, args=(r_dry, kappa), xtol=1e-30, maxiter=500)
            r0s.append(r0)

        r0s = np.array(r0s[::-1])

        # Console logging output, if requested, of the equilibrium calcuations. Useful for
        # checking if the computations worked
        raised = False
        for (r, r_dry, sp, kappa) in zip(r0s, r_drys, species, kappas):
            ss = Seq(r, r_dry, T0, kappa)
            # rc, _ = kohler_crit(T0, r_dry, kappa)
            if r < 0:
                if self.console:
                    print("Found bad r", r, r_dry, sp)
                raised = True
            if np.abs(ss-S0) > 1e-4:
                if self.console:
                    print("Found S discrepancy", ss, S0, r_dry)
                raised = True
        if raised:
            raise ParcelModelError("Couldn't calculate initial aerosol population wet sizes.")
        out['r0s'] = r0s

        # c) compute equilibrium droplet water content
        water_vol = lambda r0, r_dry, Ni: (4.*np.pi/3.)*rho_w*Ni*(r0**3 - r_dry**3)
        wc0 = np.sum([water_vol(r0, r_dry, Ni) for r0, r_dry, Ni in zip(r0s, r_drys, Nis)])
        wc0 /= rho_air(T0, P0, 0.)

        # d) compute initial ice water content
        wi0 = 0.0

        # e) concatenate into initial conditions arrays
        y0 = [z0, P0, T0, wv0, wc0, wi0, S0]
        if self.console:
            print("PARCEL INITIAL CONDITIONS")
            print(("    " + "{:>9} "*6).format("P (hPa)", "T (K)", "wv (g/kg)",
                                               "wc (g/kg)", "wi (g/kg)", "S"))
            print("    " + "{:9.1f} {:9.2f} {:9.1e} {:9.1e} {:9.1e} {:9.3f}".format(
                           P0/100., T0, wv0*1e3, wc0*1e3, wi0*1e3, S0))
        y0.extend(r0s)
        y0 = np.array(y0)
        out['y0'] = y0
        self.y0 = y0

        # Store the model configuration
        self._r0s = r0s
        self._r_drys = r_drys
        self._kappas = kappas
        self._Nis = Nis
        self._nr = len(r_drys)

        self._model_set = True

        if self.console:
            print("Initial conditions set successfully.")

    def run(self, t_end,
            output_dt=1., solver_dt=None,
            max_steps=1000, solver="odeint", output_fmt="dataframes",
            terminate=False, terminate_depth=100., **solver_args):
        """ Run the parcel model simulation.

        Once the model has been instantiated, a simulation can immediately be
        performed by invoking this method. The numerical details underlying the
        simulation and the times over which to integrate can be flexibly set
        here.

        **Time** -- The user must specify two timesteps: `output_dt`, which is the
        timestep between output snapshots of the state of the parcel model, and
        `solver_dt`, which is the the interval of time before the ODE integrator
        is paused and re-started. It's usually okay to use a very large `solver_dt`,
        as `output_dt` can be interpolated from the simulation. In some cases though
        a small `solver_dt` could be useful to force the solver to use smaller
        internal timesteps.

        **Numerical Solver** -- By default, the model will use the `odeint` wrapper
        of LSODA shipped by default with scipy. Some fine-tuning of the solver tolerances
        is afforded here through the `max_steps`. For other solvers, a set of optional
        arguments `solver_args` can be passed.

        **Solution Output** -- Several different output formats are available by default.
        Additionally, the output arrays are saved with the `ParcelModel` instance so they
        can be used later.

        Parameters
        ----------
        t_end : float
            Total time over interval over which the model should be integrated
        output_dt : float
            Timestep intervals to report model output.
        solver_dt : float
            Timestep interval for calling solver integration routine.
        max_steps : int
            Maximum number of steps allowed by solver to satisfy error tolerances
            per timestep.
        solver : {'odeint', 'lsoda', 'lsode', 'vode', cvode'}
            Choose which numerical solver to use:
            * `'odeint'`: LSODA implementation from ODEPACK via
                SciPy's integrate module
            * `'lsoda'`: LSODA implementation from ODEPACK via odespy
            * `'lsode'`: LSODE implementation from ODEPACK via odespy
            * `'vode'` : VODE implementation from ODEPACK via odespy
            * `'cvode'` : CVODE implementation from Sundials via Assimulo
            * `'lsodar'` : LSODAR implementation from Sundials via Assimulo
        output_fmt : str, one of {'dataframes', 'arrays', 'smax'}
            Choose format of solution output.
        terminate : boolean
            End simulation at or shortly after a maximum supersaturation has been achieved
        terminate_depth : float, optional (default=100.)
            Additional depth (in meters) to integrate after termination criterion
            eached.

        Returns
        -------
        DataFrames, array, or float
            Depending on what was passed to the *output* argument, different
            types of data might be returned:

            - `dataframes': (default) will process the output into
               two pandas DataFrames - the first one containing profiles
               of the meteorological quantities tracked in the model,
               and the second a dictionary of DataFrames with one for
               each AerosolSpecies, tracking the growth in each bin
               for those species.
            - 'arrays': will return the raw output from the solver
               used internally by the parcel model - the state vector
               `y` and the evaluated timesteps converted into height
               coordinates.
            - 'smax': will only return the maximum supersaturation
               value achieved in the simulation.

        Raises
        ------
        ParcelModelError
            The parcel model failed to complete successfully or failed to initialize.

        See Also
        --------
        der : right-hand side derivative evaluated during model integration.

        """
        from . integrator import Integrator

        if output_fmt not in ["dataframes", "arrays", "smax"]:
            raise ParcelModelError("Invalid value ('%s') specified for output format." % output)

        if solver_dt is None:
            solver_dt = 10.*output_dt

        if not self._model_set:
            self._setup_run()

        y0 = self.y0
        r_drys = self._r_drys
        kappas = self._kappas
        Nis = self._Nis
        nr = self._nr

        # Setup/run integrator
        try:
            from .parcel_aux import der as der_fcn
        except ImportError:
            print("Could not load Cython derivative; using Python version.")
            from .parcel import der as der_fcn
        # Hack in Python derivative function
        # der_fcn = der

        # Is the updraft speed a function of time?
        v_is_func = hasattr(self.V, '__call__')
        if v_is_func:  # Re-wrap the function to correctly figure out V
            orig_der_fcn = der_fcn

            def der_fcn(y, t, *args):
                V_t = self.V(t)
                args[3] = V_t
                return orig_der_fcn(y, t, *args)

        # Will the simulation terminate early?
        if not terminate:
            terminate_depth = 0.
        else:
            if terminate_depth <= 0.:
                raise ParcelModelError("`terminate_depth` must be greater than 0!")

        if self.console:
            print()
            print("Integration control")
            print("----------------------------")
            print("        output dt: ", output_dt)
            print("    max solver dt: ", solver_dt)
            print(" solver int steps: ", int(solver_dt/output_dt))
            print("      termination: %r (%5dm)" % (terminate, terminate_depth))

        args = [nr, r_drys, Nis, self.V, kappas, self.accom]
        integrator_type = Integrator.solver(solver)
        integrator = integrator_type(der_fcn, output_dt, solver_dt, y0, args,
                                     terminate=terminate, terminate_depth=terminate_depth,
                                     console=self.console,
                                     **solver_args)
        success = False
        try:
            # Pack args as tuple for solvers
            args = tuple(args)

            if self.console:
                print("\nBEGIN INTEGRATION ->\n")
            x, t, success = integrator.integrate(t_end)
        except ValueError as e:
            raise ParcelModelError("Something failed during model integration: %r" % e)
        finally:
            if not success:
                raise ParcelModelError("Something failed during model integration.")

            # Success if reached this point!
            if self.console:
                print("\nEND INTEGRATION <-\n")

            self.x = x
            self.heights = self.x[:, c.STATE_VAR_MAP['z']]
            self.time = t

            if output_fmt == "dataframes":
                return output.parcel_to_dataframes(self)
            elif output_fmt == "arrays":
                return self.x, self.heights
            elif output_fmt == "smax":
                S = self.x[:, c.STATE_VAR_MAP['S']]
                return S.max()

    def write_summary(self, parcel_data, aerosol_data, out_filename):
        """ Write a quick and dirty summary of given parcel model output to the
        terminal.
        """
        from .activation import lognormal_activation
        # Check if parent dir of out_filename exists, and if not,
        # create it
        out_dir = os.path.dirname(out_filename)
        if not os.path.exists(out_dir):
            os.makedirs(out_dir)

        # Open a file to write to
        with open(out_filename, 'w') as out_file:
            # Write header
            out_file.write("PARCEL MODEL\n")
            out_file.write("--------------------------------------\n")

            # Write initial conditions
            out_file.write("V = %f\n" % self.V)
            out_file.write("P0 = %f\n" % self.P0)
            out_file.write("T0 = %f\n" % self.T0)
            out_file.write("S0 = %f\n" % self.S0)
            out_file.write("--------------------------------------\n")

            # Write aerosol details
            for aerosol in self.aerosols:
                out_file.write(aerosol.species+" = "+aerosol.summary_str()+"\n")
            out_file.write("--------------------------------------\n")

            # Write simulation summary results
            # 1) Maximum supersaturation in parcel
            S_max = parcel_data['S'].max()
            S_max_idx = np.argmax(parcel_data.S)
            out_file.write("S_max = %f\n" % S_max)

            # 2) Activated fraction of each species
            T_at_S_max = parcel_data['T'].ix[S_max_idx]
            total_number = 0.0
            total_activated = 0.0
            for aerosol in self.aerosols:
                act_frac = lognormal_activation(S_max, aerosol.mu*1e-6, aerosol.sigma,
                                                aerosol.N, aerosol.kappa, T=T_at_S_max)
                act_num = act_frac*aerosol.N
                out_file.write("%s - eq_act_frac = %f (%3.2f/%3.2f)\n" %
                               (aerosol.species, act_frac, act_num, aerosol.N))

                total_number += aerosol.N
                total_activated += act_num
            total_act_frac = total_activated/total_number
            out_file.write("Total activated fraction = %f (%3.2f/%3.2f)\n" %
                           (total_act_frac, total_activated, total_number))

    def save(self, filename=None, format="nc", other_dfs=None):
        output.write_parcel_output(filename=filename, format=format, parcel=self,
                                   other_dfs=other_dfs)

    @staticmethod
    def write_csv(parcel_data, aerosol_data, output_dir=None):
        """Write output to CSV files.

        Utilize pandas fast output procedures to write the model run output to a
        set of CSV files. Written as a static method so that any prior saved
        parcel model output can be saved to disk in batch, even after its
        associated model has been destroyed.

        **Args:
            * *parcel_data* -- Pandas DataFrame of the parcel thermodynamic profile
            * *aerosol_data* -- dictionary of pandas DataFrames with the aerosol radii \
                at each model step
            * *output_dir* -- String to location where output should be saved; if not \
                provided then the model will save to the current path.

        """
        if not output_dir:
            output_dir = os.getcwd()

        # Write parcel data
        parcel_data.to_csv(os.path.join(output_dir, "parcel.csv"))

        # Write aerosol data
        for species, data in list(aerosol_data.items()):
            data.to_csv(os.path.join(output_dir, "%s.csv" % species))

def der(y, t, nr, r_drys, Nis, V, kappas, accom=c.ac):
    """ Calculates the instantaneous time-derivate of the parcel model system.

    Given a current state vector `y` of the parcel model, computes the tendency
    of each term including thermodynamic (pressure, temperature, etc) and aerosol
    terms. The basic aerosol properties used in the model must be passed along
    with the state vector (i.e. if being used as the callback function in an ODE
    solver).

    This function is implemented in NumPy and Python, and is likely *very* slow
    compared to the available Cython version.

    Parameters
    ----------
    y : array_like
        Current state of the parcel model system,
            * y[0] = altitude, m
            * y[1] = Pressure, Pa
            * y[2] = temperature, K
            * y[3] = water vapor mass mixing ratio, kg/kg
            * y[4] = cloud liquid water mass mixing ratio, kg/kg
            * y[5] = cloud ice water mass mixing ratio, kg/kg
            * y[6] = parcel supersaturation
            * y[7:] = aerosol bin sizes (radii), m
    t : float
        Current simulation time, in seconds.
    nr : Integer
        Number of aerosol radii being tracked.
    r_drys : array_like
        Array recording original aerosol dry radii, m.
    Nis : array_like
        Array recording aerosol number concentrations, 1/(m**3).
    V : float
        Updraft velocity, m/s.
    kappas : array_like
        Array recording aerosol hygroscopicities.
    accom : float, optional (default=:const:`constants.ac`)
        Condensation coefficient.

    Returns
    -------
    x : array_like
        Array of shape (``nr``+7, ) containing the evaluated parcel model
        instaneous derivative.

    Notes
    -----
    This Python sketch of the derivative function shouldn't really be used for
    any computational purposes. Instead, see the cythonized version in the auxiliary
    file, **parcel_aux.pyx**. In the default configuration, once the code has been
    built, you can set the environmental variable **OMP_NUM_THREADS** to control
    the parallel for loop which calculates the condensational growth rate for each
    bin.

    """
    z = y[0]
    P = y[1]
    T = y[2]
    wv = y[3]
    wc = y[4]
    wi = y[5]
    S = y[6]
    rs = np.asarray(y[c.N_STATE_VARS:])

    T_c = T - 273.15   # convert temperature to Celsius
    pv_sat = es(T-T_c) # saturation vapor pressure
    wv_sat = wv/(S+1.)  # saturation mixing ratio
    Tv = (1.+0.61*wv)*T  # virtual temperature given parcel humidity
    e = (1. + S)*pv_sat  # water vapor pressure
    rho_air = P/(Rd*Tv)  # current air density accounting for humidity
    rho_air_dry = (P-e)/Rd/T  # dry air density

    # 1) Pressure
    dP_dt = -1.*rho_air*g*V

    # 2/3) Wet particle growth rates and droplet liquid water
    drs_dt = np.zeros(shape=(nr, ))
    dwc_dt = 0.
    for i in range(nr):
        r = rs[i]
        r_dry = r_drys[i]
        kappa = kappas[i]

        dv_r = dv(T, r, P, accom)
        ka_r = ka(T, rho_air, r)

        G_a = (rho_w*R*T)/(pv_sat*dv_r*Mw)
        G_b = (L*rho_w*((L*Mw/(R*T))-1.))/(ka_r*T)
        G = 1./(G_a + G_b)

        delta_S = S - Seq(r, r_dry, T, kappa)
        dr_dt = (G/r)*delta_S
        Ni = Nis[i]
        dwc_dt += Ni*(r*r)*dr_dt
        drs_dt[i] = dr_dt

        # if i == nr-1:
        #    print 1, r, r_dry, Ni
        #    print 2, dv(T, r, P, accom), ka(T, rho_air, r)
        #    print 3, G_a, G_b
        #    print 4, Seq(r, r_dry, T, kappa, 1.0)
        #    print 5, dr_dt

    dwc_dt *= 4.*np.pi*rho_w/rho_air_dry

    # 4) ice water content
    dwi_dt = 0.0

    # 5) Water vapor content
    dwv_dt = -1.*(dwc_dt + dwi_dt)

    # 6) Temperature
    dT_dt = -g*V/Cp - L*dwv_dt/Cp

    # 7) Supersaturation
    alpha = (g*Mw*L)/(Cp*R*(T**2)) - (g*Ma)/(R*T)
    gamma = (P*Ma)/(Mw*es(T-273.15)) + (Mw*L*L)/(Cp*R*T*T)
    dS_dt = alpha*V - gamma*dwc_dt

    dz_dt = V

    # Repackage tendencies for feedback to numerical solver
    x = np.zeros(shape=(nr+c.N_STATE_VARS, ))
    x[0] = dz_dt
    x[1] = dP_dt
    x[2] = dT_dt
    x[3] = dwv_dt
    x[4] = dwc_dt
    x[5] = dwi_dt
    x[6] = dS_dt
    x[c.N_STATE_VARS:] = drs_dt[:]

    return x


#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
策略值网络实现
依赖TensorFlow 1.10.0版本

@author: CoderJie
@email: coderjie@outlook.com
"""

import numpy as np
import tensorflow as tf

class PolicyValueNet:
    """
    策略值网络
    """
    def __init__(self, board_width, board_height, model_file=None):
        """
        构造函数
        :param board_width: 棋盘宽度
        :param board_height: 棋盘高度
        """

        self.__board_width = board_width
        self.__board_height = board_height

        # 定义模型输入占位符, 张量形状为[?, 4, height, width]
        # None表示张量的第一维为任意长度, 第一维即代表样本数量, 就是表示模型接收任意的样本数,
        # 并且每个样本包含4个 board_height*board_width 的表格
        self.__input_data = tf.placeholder(tf.float32,
                                           shape=[None, 4, board_height, board_width])

        # 定义训练所用的状态值占位符
        self.__train_value = tf.placeholder(tf.float32,
                                            shape=[None, 1])

        # 定义训练所用的动作概率值占位符
        self.__train_prob = tf.placeholder(tf.float32,
                                           shape=[None, board_height * board_width])

        # 定义学习速度占位符
        self.__learning_rate = tf.placeholder(tf.float32)

        # 组建神经网络
        self.__build_net()

        # 创建一个会话, 并进行初始化动作
        self.__session = tf.Session()
        self.__session.run(tf.global_variables_initializer())

        # 保存器, 用于保存或者恢复模型
        self.__saver = tf.train.Saver()

        # 从文件中恢复模型
        if model_file is not None:
            self.restore_model(model_file)

    def __del__(self):
        """
        析构函数
        """
        self.__session.close()

    def train(self, batch_states, batch_probs, batch_values, lr):
        """
        训练模型
        :param batch_states: 批量状态(形状为[?, 4, height, width])
        :param batch_probs: 批量概率(形状为[?, height*width])
        :param batch_values: 批量状态值(形状为[?, 1])
        :param lr: 学习速度(浮点数)
        :return: 损失值, 熵
        """
        loss, entropy, _ = self.__session.run(
                [self.__total_loss, self.__entropy, self.__optimizer],
                feed_dict={self.__input_data: batch_states,
                           self.__train_prob: batch_probs,
                           self.__train_value: batch_values,
                           self.__learning_rate: lr})
        return loss, entropy

    def policy_value(self, board):
        """
        获取策略(动作概率), 状态值
        :param board: 棋盘对象
        :return: 策略(动作概率), 状态值
        """
        state = board.state()
        input_data = state.reshape(-1, 4, self.__board_height, self.__board_width)
        acts_prob, state_value = self.__policy_value(input_data)

        # 只获取有效动作的概率
        avl_acts = board.avl_actions
        acts_prob = list(zip(avl_acts, acts_prob[0][avl_acts]))
        return acts_prob, state_value[0][0]

    def save_model(self, path):
        """
        保存模型
        :param path: 模型路径
        """
        self.__saver.save(self.__session, path)

    def restore_model(self, path):
        """
        恢复模型
        :param path: 模型路径
        """
        self.__saver.restore(self.__session, path)

    def __build_net(self):
        """
        组建卷积神经网络
        :return:
        """
        # 转置输入数据, 输出张量形状[?, height, width, 4]
        # [0, 2, 3, 1] 0 代表原始数据第一维, 1 代表原始数据第二维, 2 代表原始数据第三维, 3 代表原始数据第四维
        # 转置后可以把棋盘数据看作是通道数为4, 大小为height*width的图片
        input_trans = tf.transpose(self.__input_data, [0, 2, 3, 1])

        # 第一层卷积层, 32个卷积核, 输出张量形状[?, height, width, 32]
        # 卷积核数量32, 卷积核大小 3*3, 卷积核的通道数就是输入数据的通道数4, 卷积步长默认为1
        # 填充方式为"same" 表示边界填充0
        # channels_last表示通道为最后一维
        # 默认加上偏置, 偏置初始值默认为0
        # 激活函数选择斜坡函数
        conv1 = tf.layers.conv2d(inputs=input_trans,
                                 filters=32,
                                 kernel_size=[3, 3],
                                 padding="same",
                                 data_format="channels_last",
                                 activation=tf.nn.relu)

        # 第二层卷积层, 64个卷积核, 输出张量形状[?, height, width, 64]
        conv2 = tf.layers.conv2d(inputs=conv1,
                                 filters=64,
                                 kernel_size=[3, 3],
                                 padding="same",
                                 data_format="channels_last",
                                 activation=tf.nn.relu)

        # 第三层卷积层, 128个卷积核, 输出张量形状[?, height, width, 128]
        conv3 = tf.layers.conv2d(inputs=conv2,
                                 filters=128,
                                 kernel_size=[3, 3],
                                 padding="same",
                                 data_format="channels_last",
                                 activation=tf.nn.relu)

        # 策略端降维, 输出张量形状[?, height, width, 4]
        policy_conv = tf.layers.conv2d(inputs=conv3,
                                       filters=4,
                                       kernel_size=[1, 1],
                                       padding="same",
                                       data_format="channels_last",
                                       activation=tf.nn.relu)

        # 策略端修改数据维度, 输出张量形状[?, 4 * height * width]
        policy_flat = tf.reshape(policy_conv,
                                 [-1, 4 * self.__board_height * self.__board_width])

        # 策略端全连接层, 输出张量形状[?, height * width]
        # units代表输出大小
        # 默认加上偏置, 偏置初始值默认为0
        # 使用log_softmax激活函数
        self.__policy_fc = tf.layers.dense(inputs=policy_flat,
                                           units=self.__board_height * self.__board_width,
                                           activation=tf.nn.log_softmax)

        # 价值端降维, 输出张量形状[?, height, width, 2]
        value_conv = tf.layers.conv2d(inputs=conv3,
                                      filters=2,
                                      kernel_size=[1, 1],
                                      padding="same",
                                      data_format="channels_last",
                                      activation=tf.nn.relu)

        # 价值端修改数据维度, 输出张量形状[?, 2 * height * width]
        value_flat = tf.reshape(value_conv,
                                [-1, 2 * self.__board_height * self.__board_width])

        # 价值端64个神经元的全连接层, 输出张量形状[?, 64]
        # 使用斜坡激活函数
        value_fc1 = tf.layers.dense(inputs=value_flat,
                                    units=64,
                                    activation=tf.nn.relu)

        # 价值端1个神经元的全连接层, 输出张量形状[?, 1]
        # 使用双曲正切激活函数, 双曲正切函数的值域是(-1, 1)
        self.__value_fc2 = tf.layers.dense(inputs=value_fc1,
                                           units=1,
                                           activation=tf.nn.tanh)

        # 定义状态值损失函数(平方差)
        value_loss = tf.losses.mean_squared_error(self.__train_value, self.__value_fc2)

        # 定义策略损失函数
        policy_loss = tf.negative(tf.reduce_mean(tf.reduce_sum(
            tf.multiply(self.__train_prob, self.__policy_fc), 1)))

        # l2正则化系数
        l2_beta = 1e-4
        # 获取所有可训练变量
        vs = tf.trainable_variables()
        # 定义正则项
        l2_penalty = l2_beta * tf.add_n(
            [tf.nn.l2_loss(v) for v in vs if 'bias' not in v.name.lower()])

        # 定义总损失函数
        self.__total_loss = value_loss + policy_loss + l2_penalty

        # 定义优化器
        self.__optimizer = tf.train.AdamOptimizer(
            learning_rate=self.__learning_rate).minimize(self.__total_loss)

        # 计算策略的熵值, 用于监控网络
        self.__entropy = tf.negative(tf.reduce_mean(
                tf.reduce_sum(tf.exp(self.__policy_fc) * self.__policy_fc, 1)))

    def __policy_value(self, states):
        """
        获取策略(动作概率), 状态值
        :param states: 批量状态二值数据, 数据形状应该为[?, 4, height, width]
        :return: 策略(形状: [?, height*width]), 状态值(形状: [?, 1])
        """
        log_probs, values = self.__session.run([self.__policy_fc, self.__value_fc2],
                                               feed_dict={self.__input_data: states})
        acts_probs = np.exp(log_probs)
        return acts_probs, values


"""
Code for parsing Spice Double-Array Files (DAF). This is a superset of CK and SPK files. Note that while this
parses Spice files, it only makes minimal use of the spice libraries (date conversion only)
"""

import struct
from spiceypy import etcal  as cspice_etcal
from spiceypy import scdecd as cspice_scdecd
from spiceypy import sct2e  as cspice_sct2e

class daf_summary:
    @staticmethod
    def make(sr,daf,buf,name):
        if daf.subtype=="SPK":
            return daf_SPKsummary(sr,daf,buf,name)
        elif daf.subtype=="CK":
            return daf_CKsummary(sr,daf,buf,name)
        else:
            raise ValueError("Unrecognized subtype %s"%daf.subtype)
    def __init__(self,sr,daf,buf,name):
        """

        :param sr: Summary record which this summary is part of
        :param daf: Double-array file which holds the header (ND, NI) and can read the file
        :param buf: Data for this summary
        """
        self.sr=sr
        self.daf=daf
        self.D=[0.0]*self.daf.ND
        self.I=[0]*self.daf.NI
        self.name=name
        for i in range(self.daf.ND):
            self.D[i]=struct.unpack(self.daf.dformat,buf[i*8:(i+1)*8])[0]
        for i in range(self.daf.NI):
            self.I[i]=struct.unpack(self.daf.iformat,buf[self.daf.ND*8+i*4:self.daf.ND*8+(i+1)*4])[0]
    def __str__(self):
        result="%s\nD:\n" % self.name
        for i,D in enumerate(self.D):
            if i!=0:
                result+=","
            result+="%20.15e"%D
        result+="\nI:\n"
        for i,I in enumerate(self.I):
            if i!=0:
                result+=","
            result+="%d"%I
        return result

class daf_SPKsummary(daf_summary):
    def __init__(self,sr,daf,buf,name):
        super().__init__(sr,daf,buf,name)
    def getet0(self): return self.D[0]
    def getet1(self): return self.D[1]
    def gettarget(self): return self.I[0]
    def getcenter(self): return self.I[1]
    def getframe (self): return self.I[2]
    def gettype  (self): return self.I[3]
    def getaddr0 (self): return self.I[4]
    def getaddr1 (self): return self.I[5]
    et0   =property(fget=getet0   ,doc="Ephemeris time of beginning of segment")
    et1   =property(fget=getet1   ,doc="Ephemeris time of end of segment")
    target=property(fget=gettarget,doc="Target NAIF code")
    center=property(fget=getcenter,doc="Center NAIF code")
    frame =property(fget=getframe ,doc="Reference frame NAIF code")
    type  =property(fget=gettype  ,doc="SPK segment type")
    addr0 =property(fget=getaddr0 ,doc="Address of first element of segment data")
    addr1 =property(fget=getaddr1 ,doc="Address of last element of segment data")
    def __str__(self):
        result=(("%s\n"+
                 "ET0:    %30.14f (%s)\n"+
                 "ET1:    %30.14f (%s)\n"+
                 "Target: %d\n"+
                 "Center: %d\n"+
                 "Frame:  %d\n"+
                 "Type:   %d\n"+
                 "Addr0:  %d\n"+
                 "Addr1:  %d") %
                (self.name,self.et0,cspice_etcal(self.et0),self.et1,cspice_etcal(self.et1),self.target,self.center,self.frame,self.type,self.addr0,self.addr1))
        return result
    def segment(self):
        return daf_SPKSegment(self,self.daf)

class daf_CKsummary(daf_summary):
    def __init__(self,sr,daf,buf,name):
        super().__init__(sr,daf,buf,name)
    def sclk(self):return self.target // 1000 + (0 if self.target % 1000 == 0 else 1)
    def getsclk0(self): return self.D[0]
    def getsclk1(self): return self.D[1]
    def getet0(self): return cspice_sct2e(self.sclk(),self.sclk0)
    def getet1(self): return cspice_sct2e(self.sclk(),self.sclk1)
    def gettarget(self): return self.I[0]
    def getframe (self): return self.I[1]
    def gettype  (self): return self.I[2]
    def getrates (self): return self.I[3]
    def getaddr0 (self): return self.I[4]
    def getaddr1 (self): return self.I[5]
    sclk0 =property(fget=getsclk0 ,doc="DP encoded sclk of beginning of segment")
    sclk1 =property(fget=getsclk1 ,doc="DP encoded sclk of end of segment")
    et  =property(fget=getet0 ,doc="ephemeris time of beginning of segment")
    et0 =property(fget=getet0 ,doc="ephemeris time of beginning of segment")
    et1 =property(fget=getet1 ,doc="ephemeris time of end of segment")
    target=property(fget=gettarget,doc="Target NAIF code")
    frame =property(fget=getframe ,doc="Reference frame NAIF code")
    type  =property(fget=gettype  ,doc="CK segment type")
    rates =property(fget=getrates ,doc="Angular Rates Flag")
    addr0 =property(fget=getaddr0 ,doc="Address of first element of segment data")
    addr1 =property(fget=getaddr1 ,doc="Address of last element of segment data")
    def __str__(self):
        result ="%s\n"%self.name
        result+="SCLK0:    %20.15e" %self.sclk0
        try:
            et=cspice_sct2e(self.sclk(),self.sclk0)
            scdecd=cspice_scdecd(self.sclk(),self.sclk0)
            result+=" (%s ET=%16.6f %s)"%(scdecd,et,cspice_etcal(et))
        except:
            pass #If there is an error, it means that the correct kernels aren't available.
        result+="\n"
        result+="SCLK1:    %20.15e" %self.sclk1
        try:
            et=cspice_sct2e(self.sclk(),self.sclk1)
            scdecd=cspice_scdecd(self.sclk(),self.sclk1)
            result+=" (%s ET=%16.6f %s)"%(scdecd,et,cspice_etcal(et))
        except:
            pass #If there is an error, it means that the correct kernels aren't available.
        result+="\n"
        result+=(("Inst:  %7d\n"+
                  "Frame: %7d\n"+
                  "Type:  %7d\n"+
                  "Rates? %7d\n"+
                  "Addr0: %7d\n"+
                  "Addr1: %7d") %
                (self.target,self.frame,self.type,self.rates,self.addr0,self.addr1))
        return result
    def segment(self):
        return daf_CKSegment(self,self.daf)

class daf_SPKSegment:
    def __init__(self,summary,daf):
        self.summary=summary
        self.daf=daf
        daf.inf.seek(8*(summary.addr0-1))
        segment_size=summary.addr1-summary.addr0+1
        format="%s%d%s"%(daf.dformat[0],segment_size,daf.dformat[1])
        self.buf=struct.unpack(format,daf.inf.read(segment_size*8))
        self.N=int(self.buf[-1]) if self.summary.type not in (15,17) else 1
        result=self.line(0)
        self.csvheader = result.csvheader
    def line(self,i):
        if self.summary.type==1:
            return daf_SPK01line(self.N,i,self.buf)
        elif self.summary.type==2:
            return daf_SPK02line(self.N,i,self.buf)
        elif self.summary.type==3:
            return daf_SPK03line(self.N,i,self.buf)
        elif self.summary.type==9:
            return daf_SPK09line(self.N,i,self.buf)
        elif self.summary.type == 13:
            return daf_SPK13line(self.N, i, self.buf)
        elif self.summary.type == 15:
            return daf_SPK15line(self.N, i, self.buf)
        elif self.summary.type == 17:
            return daf_SPK17line(self.N, i, self.buf)
        else:
            raise ValueError("Unhandled SPK type %d"%self.summary.type)
    def __iter__(self):
        for i in range(self.N):
            result = self.line(i)
            yield result
    def __str__(self):
        result=(("N: %d") %
                (self.N))
        return result

class daf_CKSegment:
    def __init__(self,summary,daf):
        self.summary=summary
        self.daf=daf
        daf.inf.seek(8*(summary.addr0-1))
        segment_size=summary.addr1-summary.addr0+1
        format="%s%d%s"%(daf.dformat[0],segment_size,daf.dformat[1])
        self.buf=struct.unpack(format,daf.inf.read(segment_size*8))
        result = self.line(0)
        self.csvheader = result.csvheader
    def getN(self):
        if self.summary.type==1:
            return int(self.buf[-1])
        elif self.summary.type==2:
            return int(1+len(self.buf)/10.01)
        elif self.summary.type==3:
            return int(self.buf[-1])
        else:
            raise ValueError("Unknown CK type %d"%self.summary.type)
    N=property(fget=getN ,doc="Number of pointing intervals in file")
    def line(self,i):
        sclk=self.summary.target//1000+(0 if self.summary.target%1000==0 else 1)
        if self.summary.type==1:
            return daf_CK01line(sclk,self.N,i,self.buf,self.summary.rates)
        elif self.summary.type==2:
            return daf_CK02line(sclk,self.N,i,self.buf)
        elif self.summary.type == 3:
            return daf_CK03line(sclk,self.N,i,self.buf,self.summary.rates)
        else:
            raise ValueError("Unknown CK type %d" % self.summary.type)
    def __iter__(self):
        for i in range(self.N):
            result = self.line(i)
            yield result
    def __str__(self):
        result=(("N: %d") %
                (self.N))
        return result

class daf_SPKStateline:
    def __init__(self,N,i,buf):
        self.x   =buf[i*6+0]
        self.y   =buf[i*6+1]
        self.z   =buf[i*6+2]
        self.dxdt=buf[i*6+3]
        self.dydt=buf[i*6+4]
        self.dzdt=buf[i*6+5]
        self.et=buf[N*6+i]
        self.csvheader="ET,ETCAL*,x,y,z,dxdt,dydt,dzdt"
    def __str__(self):
        result="%30.14f,%s,%25.15e,%25.15e,%25.15e,%25.15e,%25.15e,%25.15e"%(self.et,cspice_etcal(self.et),self.x,self.y,self.z,self.dxdt,self.dydt,self.dzdt)
        return result

class daf_SPK09line(daf_SPKStateline):
    def __init__(self,N,i,buf):
        super().__init__(N,i,buf)

class daf_SPK13line(daf_SPKStateline):
    def __init__(self,N,i,buf):
        super().__init__(N,i,buf)

class daf_SPK02line:
    def __init__(self,N,i,buf):
        self.RSIZE=int(buf[-2])
        self.asize=(self.RSIZE-2)//3 #Number of elements in each component array
        self.INTLEN=buf[-3]
        self.INIT=buf[-4]
        self.mid    =buf[i*self.RSIZE+0]
        self.radius=buf[i*self.RSIZE+1]
        self.et=self.mid-self.radius
        self.x     =buf[i*self.RSIZE+2+self.asize*0:i*self.RSIZE+2+self.asize*1]
        self.y     =buf[i*self.RSIZE+2+self.asize*1:i*self.RSIZE+2+self.asize*2]
        self.z     =buf[i*self.RSIZE+2+self.asize*2:i*self.RSIZE+2+self.asize*3]
        self.csvheader="ET*,ETCAL*,ETEND*,ETCALEND*,MID,MIDCAL*,RADIUS"
        for i in range(self.asize):
            self.csvheader+=",x[%02d]"%i
        for i in range(self.asize):
            self.csvheader+=",y[%02d]" % i
        for i in range(self.asize):
            self.csvheader+=",z[%02d]" % i
    def __str__(self):
        result="%30.14f,%s,%30.14f,%s,%30.14f,%s,%30.14f"%(self.et,cspice_etcal(self.et),self.mid+self.radius,cspice_etcal(self.mid+self.radius),self.mid,cspice_etcal(self.mid),self.radius)
        for x in self.x:
            result+=",%25.15e"%x
        for y in self.y:
            result+=",%25.15e"%y
        for z in self.z:
            result+=",%25.15e"%z
        return result
    def eval(self,ET):
        pass

class daf_SPK03line:
    def __init__(self,N,i,buf):
        self.RSIZE=int(buf[-2])
        self.asize=(self.RSIZE-2)//6 #Number of elements in each component array
        self.INTLEN=buf[-3]
        self.INIT=buf[-4]
        self.mid    =buf[i*self.RSIZE+0]
        self.radius=buf[i*self.RSIZE+1]
        self.et=self.mid-self.radius
        self.x      =buf[i*self.RSIZE+2+self.asize*0:i*self.RSIZE+2+self.asize*1]
        self.y      =buf[i*self.RSIZE+2+self.asize*1:i*self.RSIZE+2+self.asize*2]
        self.z      =buf[i*self.RSIZE+2+self.asize*2:i*self.RSIZE+2+self.asize*3]
        self.dx     =buf[i*self.RSIZE+2+self.asize*3:i*self.RSIZE+2+self.asize*4]
        self.dy     =buf[i*self.RSIZE+2+self.asize*4:i*self.RSIZE+2+self.asize*5]
        self.dz     =buf[i*self.RSIZE+2+self.asize*5:i*self.RSIZE+2+self.asize*6]
        self.csvheader="ET*,ETCAL*,ETEND*,ETCALEND*,MID,MIDCAL*,RADIUS"
        for i in range(self.asize):
            self.csvheader+=",x[%02d]"%i
        for i in range(self.asize):
            self.csvheader+=",y[%02d]" % i
        for i in range(self.asize):
            self.csvheader+=",z[%02d]" % i
        for i in range(self.asize):
            self.csvheader+=",dx[%02d]"%i
        for i in range(self.asize):
            self.csvheader+=",dy[%02d]" % i
        for i in range(self.asize):
            self.csvheader+=",dz[%02d]" % i
    def __str__(self):
        result="%30.14f,%s,%30.14f,%s,%30.14f,%s,%30.14f"%(self.et,cspice_etcal(self.et),self.mid+self.radius,cspice_etcal(self.mid+self.radius),self.mid,cspice_etcal(self.mid),self.radius)
        for x in self.x:
            result+=",%25.15e"%x
        for y in self.y:
            result+=",%25.15e"%y
        for z in self.z:
            result+=",%25.15e"%z
        for dx in self.dx:
            result+=",%25.15e"%dx
        for dy in self.dy:
            result+=",%25.15e"%dy
        for dz in self.dz:
            result+=",%25.15e"%dz
        return result
    def eval(self,ET):
        pass

class daf_SPK15line:
    def __init__(self,N,i,buf):
        self.et   =buf[0]
        self.polex    =buf[1]
        self.poley    =buf[2]
        self.polez    =buf[3]
        self.perix    =buf[4]
        self.periy    =buf[5]
        self.periz    =buf[6]
        self.l=buf[7]
        self.e =buf[8]
        self.j2flag=buf[9]
        self.centx=buf[10]
        self.centy=buf[11]
        self.centz=buf[12]
        self.centgm=buf[13]
        self.centj2=buf[14]
        self.centre=buf[15]
        self.csvheader="ETPERI,ETPERICAL*,polex,poley,polez,perix,periy,periz,l,e,j2flag,centx,centy,centz,centgm,centj2,centre"
    def __str__(self):
        result="%30.14f,%s,%.15f,%.15f,%.15f,%.15f,%.15f,%.15f,%.15e,%.15e,%1.0f,%17.14f,%17.14f,%17.14f,%.15e,%.15e,%.15e"%(
            self.et,cspice_etcal(self.et),
            self.polex,self.poley,self.polez,self.perix,self.periy,self.periz,
            self.l,self.e,self.j2flag,
            self.centx, self.centy, self.centz,
            self.centgm,self.centj2,self.centre
        )
        return result
    def eval(self,ET):
        pass

class daf_SPK17line:
    def __init__(self,N,i,buf):
        self.et   =buf[0]
        self.a    =buf[1]
        self.h    =buf[2]
        self.k    =buf[3]
        self.L    =buf[4]
        self.p    =buf[5]
        self.q    =buf[6]
        self.omegabardot=buf[7]
        self.Ldot =buf[8]
        self.Omegadot=buf[9]
        self.polera=buf[10]
        self.poledec=buf[11]
        self.csvheader="ETPERI,ETPERICAL*,a,h,k,L,p,q,omegabardot,Ldot,Omegadot,polera,poledec"
    def __str__(self):
        result="%30.14f,%s,%30.14f,%30.14f,%30.14f,%30.14f,%30.14f,%30.14f,%30.14f,%30.14f,%30.14f,%30.14f,%30.14f"%(
            self.et,cspice_etcal(self.et),
            self.a,self.h,self.k,self.L,self.p,self.q,
            self.omegabardot,self.Ldot,self.Omegadot,
            self.polera,self.poledec
        )
        return result
    def eval(self,ET):
        pass

class daf_SPK01line:
    n_elements=71
    def __init__(self,N,i,buf):
        self.e   =buf[i*self.n_elements:(i+1)*self.n_elements]
        self.et=buf[N*self.n_elements+i]
        import numpy as np
        self.FD=np.zeros(15)*float('nan')

        # initialize arrays for spke01
        self.G = np.zeros(15)
        self.REFPOS = np.zeros(3)
        self.REFVEL = np.zeros(3)
        self.KQ = np.array([0, 0, 0])
        self.FC = np.zeros(15)
        self.FC[0] = 1.0
        self.WC = np.zeros(13)
        self.W = np.zeros(17)

        self.spke01(self.e[0],self.e)
        self.csvheader= "ET,ETCAL*,TL"
        for i in range(15):
            self.csvheader += ",G[%02d]" % (i)
        self.csvheader += ",X,DXDT,Y,DYDT,Z,DZDT"
        for i in range(3):
            for j in range(15):
                self.csvheader += ",DT[%02d%02d]" % (i, j)
                self.csvheader += ",KQMAX1,KQ[0],KQ[1],KQ[2]"

    #        self.eval(self.TL)
    def __str__(self):
        result="%30.14f,%s" %(self.et,cspice_etcal(self.et))
        for i in range(self.n_elements):
            result+=",%25.15e"%self.e[i]
        return result
    """
    def getTL(self): return self.e[0]
    def getG (self): return self.e[1:15+1]
    def getREFPOS(self): return self.e[16:22:2]
    def getREFVEL(self): return self.e[17:23:2]
    def getDT    (self): return self.e[22:22+45]
    def getKQMAX1(self): return int(self.e[67])
    def getKQ    (self): return self.e[68:71]
    TL    =property(fget=getTL    ,doc="Ephemeris time of end of segment")
    G     =property(fget=getG     ,doc="Stepsize function vector")
    REFPOS=property(fget=getREFPOS,doc="Reference position vector")
    REFVEL=property(fget=getREFVEL,doc="Reference velocity vector")
    DT    =property(fget=getDT    ,doc="Modified divided difference arrays")
    KQMAX1=property(fget=getKQMAX1,doc="Maximum integration order plus 1")
    KQ    =property(fget=getKQ    ,doc="Integration order array")
    """
    def getFC(self,J):
        print("Accessing element %d, value %f"%(J,self.FD[J]))
        return self.FD[J]
    def setFC(self,J,val):
        print("Writing element %d, was value %f, will be value %f"%(J,self.FD[J],val))
        self.FD[J]=val
    def eval(self,ET):
        import numpy as np
        STATE=np.zeros(6)
        DELTA = ET - self.TL
        TP = DELTA
        MQ2 = self.KQMAX1 - 2
        KS = self.KQMAX1 - 1
        self.setFC(0,1.0)
        WC = np.zeros(13)
        W  = np.zeros(17)

        #     This is clearly collecting some kind of coefficients.
        #     The problem is that we have no idea what they are...
        #
        #     The G coefficients are supposed to be some kind of step size
        #     vector.
        #
        #     TP starts out as the delta t between the request time
        #     and the time for which we last had a state in the MDL file.
        #     We then change it from DELTA  by the components of the stepsize
        #     vector G.
        for J in range(1, MQ2 + 1):
            self.setFC(J,TP / self.G[J-1])
            WC[J-1] = DELTA / self.G[J-1]
            TP = DELTA + self.G[J-1]
        #     Collect KQMAX1 reciprocals.
        for J in range(1, self.KQMAX1 + 1):
            W[J-1] = 1.0 / float(J)

        #     Compute the W(K) terms needed for the position interpolation
        #     (Note,  it is assumed throughout this routine that KS, which
        #     starts out as KQMAX1-1 (the ``maximum integration'')
        #     is at least 2.
        JX = 0
        KS1 = KS - 1
        while KS >= 2:
            JX = JX + 1
            for J in range(1, JX + 1):
                W[J + KS - 1] = self.getFC(J) * W[J + KS1 - 1] - WC[J - 1] * W[J + KS - 1]
            KS = KS1
            KS1 = KS1 - 1

        #     Perform position interpolation: (Note that KS = 1 right now.
        #     We don't know much more than that.)
        for I in range(3):
            KQQ = int(self.KQ[I])
            SUM = 0.0
            for J in range(KQQ, 0, -1):
                SUM = SUM + self.DT[J - 1+I*15] * W[J + KS - 1]
            STATE[I] = self.REFPOS[I] + DELTA * (self.REFVEL[I] + DELTA * SUM)

        #     Again we need to compute the W(K) coefficients that are
        #     going to be used in the velocity interpolation.
        #     (Note, at this point, KS = 1, KS1 = 0.)
        for J in range(1, JX + 1):
            W[J + KS - 1] = self.getFC(J) * W[J + KS1 - 1] - WC[J - 1] * W[J + KS - 1]
        KS = KS - 1

        #     Perform velocity interpolation:
        for I in range(1, 3 + 1):
            KQQ = int(self.KQ[I])
            SUM = 0.0

            for J in range(KQQ, 0, -1):
                SUM = SUM + self.DT[J - 1+I*15] * W[J + KS - 1]

            STATE[I + 3] = self.REFVEL[I] + DELTA * SUM

        #     That's all folks.  We don't know why we did anything, but
        #     at least we can tell structurally what we did.
        return STATE
    def test(self,et0,N=100):
        """
        Compare the calculation of the state in pure python to that provided by spiceypy (which is an interface
        to the cspice library). This should be near bit-perfect.
        :param ET0: Time to begin at. Since the record only contains the end time, we need a beginning time. The
                    beginning time *might* be encoded in the DeltaT's, but this isn't known for sure yet
        :param N: Number of points to evaluate at. If you pass 100, you will get the point at the beginning,
                  99 points in between, and the point at the end, so a total of 101 points.
        :return: Sum of squares of differences of position components. Ideally should be zero, should be only a couple
                 of ulps.
        """
        from spiceypy import spkezr as cspice_spkezr
        import numpy as np
        et1=self.et
        result=0.0
        for i in range(N):
            et=et0+i*(et1-et0)/N
            #et=et1
            ref=cspice_spkezr("-189",et,"J2000","NONE","399")[0][0:3]
            test=np.array(self.spke01(et,self.e)[0:3])
            result+=np.sum((ref-test)**2)
        return result
    def spke01(self, ET, RECORD):
        import numpy as np
        """Compute position and velocity from a Modified Difference Array record

        Inputs:
            ET: Epoch time to evaluate position and velocity (seconds since J2000)
            RECORD: A record of Modified Difference Array
        Returns: STATE
            STATE: A numpy array which contains position and velocity
        """

        # This method has been translated from SPKE01 of SPICE Toolkit and
        # modified by Shushi Uetsuki.
        #
        # SPICE Toolkit for FORTRAN : http://naif.jpl.nasa.gov/naif/toolkit_FORTRAN.html
        # SPK Required Reading : http://naif.jpl.nasa.gov/pub/naif/toolkit_docs/C/req/spk.html
        #
        # Original FORTRAN code uses 'SAVE' directive, and it means all variable
        # should be saved for next call.  So i decided to make almost all
        # variables to be instance variable.  Some of them are initialized in
        # __init__ method.

        STATE = np.zeros(6)

        #     Variable  I/O  Description
        #     --------  ---  --------------------------------------------------
        #     ET         I   Target epoch.
        #     RECORD     I   Data record.
        #     STATE      O   State (position and velocity).
        #
        # $ Detailed_Input
        #
        #     ET          is a target epoch, at which a state vector is to
        #                 be computed.
        #
        #     RECORD      is a data record which, when evaluated at epoch ET,
        #                 will give the state (position and velocity) of some
        #                 body, relative to some center, in some inertial
        #                 reference frame.
        #
        # $ Detailed_Output
        #
        #     STATE       is the state. Units are km and km/sec.
        #
        # $ Parameters
        #
        #     None.
        #
        # $ Exceptions
        #
        #     None.
        #
        # $ Files
        #
        #     None.
        #
        # $ Particulars
        #
        #     The exact format and structure of type 1 (difference lines)
        #     segments are described in the SPK Required Reading file.
        #
        #     Difference lines (DL's) are generated by JPL navigation
        #     system programs P and PV. Each data record is equivalent
        #     to the (slightly rearranged) 'P' portion of a NAVIO PV file
        #     data record.
        #
        #     SPKE01 is a specialized version of Fred Krogh's subroutine DAINT.
        #     Only the calling sequence has been changed.
        #
        #     Because the original version was undocumented, only Fred
        #     knows how this really works.
        #
        # $ Examples
        #
        #     None.
        #
        # $ Restrictions
        #
        #     Unknown.
        #
        # $ Literature_References
        #
        #     NAIF Document 168.0, "S- and P- Kernel (SPK) Specification and
        #     User's Guide"
        #
        # $ Author_and_Institution
        #
        #     F.T. Krogh      (JPL)
        #     I.M. Underwood  (JPL)
        #
        # $ Version
        #
        # -    SPICELIB Version 1.1.0, 14-FEB-1997 (WLT)
        #
        #        The goto's were removed and loop and if structures
        #        revealed.  We still don't know exactly what's going
        #        on, but at least the bones of this routine have been
        #        cleaned off and are ready for assembly. (WLT)
        #
        # -    SPICELIB Version 1.0.4, 30-OCT-1996 (WLT)
        #
        #        Removed redundant SAVE statements from the declaration
        #        section.  Thanks to Steve Schlaifer for finding this
        #        error.
        #
        # -    SPICELIB Version 1.0.3, 10-MAR-1992 (WLT)
        #
        #        Comment section for permuted index source lines was added
        #        following the header.
        #
        # -    SPICELIB Version 1.0.2, 23-AUG-1991 (HAN)
        #
        #        SPK01 was removed from the Required_Reading section of the
        #        header. The information in the SPK01 Required Reading file
        #        is now part of the SPK Required Reading file.
        #
        # -    SPICELIB Version 1.0.1, 22-MAR-1990 (HAN)
        #
        #        Literature references added to the header.
        #
        # -    SPICELIB Version 1.0.0, 31-JAN-1990 (IMU) (FTK)
        #
        # -&
        #
        # $ Index_Entries
        #
        #     evaluate type_1 spk segment
        #
        # -&

        #
        #     Unpack the contents of the MDA array.
        #
        #        Name    Dimension  Description
        #        ------  ---------  -------------------------------
        #        TL              1  Final epoch of record
        #        G              15  Stepsize function vector
        #        REFPOS          3  Reference position vector
        #        REFVEL          3  Reference velocity vector
        #        DT         15,NTE  Modified divided difference arrays
        #        KQMAX1          1  Maximum integration order plus 1
        #        KQ            NTE  Integration order array
        #
        #     For our purposes, NTE is always 3.
        #
        self.TL = RECORD[0]
        self.G = RECORD[1:1 + 15]
        #
        #     Collect the reference position and velocity.
        #
        self.REFPOS[0] = RECORD[16]
        self.REFVEL[0] = RECORD[17]

        self.REFPOS[1] = RECORD[18]
        self.REFVEL[1] = RECORD[19]

        self.REFPOS[2] = RECORD[20]
        self.REFVEL[2] = RECORD[21]

        self.DT = np.reshape(RECORD[22:22 + 45], (15, 3), order='F')

        self.KQMAX1 = int(RECORD[67])
        self.KQ[0] = int(RECORD[68])
        self.KQ[1] = int(RECORD[69])
        self.KQ[2] = int(RECORD[70])
        #
        #     Next we set up for the computation of the various differences
        #
        self.DELTA = ET - self.TL
        self.TP = self.DELTA
        self.MQ2 = self.KQMAX1 - 2
        self.KS = self.KQMAX1 - 1
        #
        #     This is clearly collecting some kind of coefficients.
        #     The problem is that we have no idea what they are...
        #
        #     The G coefficients are supposed to be some kind of step size
        #     vector.
        #
        #     TP starts out as the delta t between the request time
        #     and the time for which we last had a state in the MDL file.
        #     We then change it from DELTA  by the components of the stepsize
        #     vector G.
        #
        for J in range(1, self.MQ2 + 1):
            self.FC[J] = self.TP / self.G[J - 1]
            self.WC[J - 1] = self.DELTA / self.G[J - 1]
            self.TP = self.DELTA + self.G[J - 1]
        #
        #     Collect KQMAX1 reciprocals.
        #
        for J in range(1, self.KQMAX1 + 1):
            self.W[J - 1] = 1.0 / float(J)
        #
        #     Compute the W(K) terms needed for the position interpolation
        #     (Note,  it is assumed throughout this routine that KS, which
        #     starts out as KQMAX1-1 (the ``maximum integration'')
        #     is at least 2.
        #
        self.JX = 0
        self.KS1 = self.KS - 1

        while self.KS >= 2:

            self.JX = self.JX + 1

            for J in range(1, self.JX + 1):
                self.W[J + self.KS - 1] = self.FC[J] * self.W[J + self.KS1 - 1] - self.WC[J - 1] * self.W[
                    J + self.KS - 1]

            self.KS = self.KS1
            self.KS1 = self.KS1 - 1
        #
        #     Perform position interpolation: (Note that KS = 1 right now.
        #     We don't know much more than that.)
        #
        for I in range(1, 3 + 1):

            self.KQQ = self.KQ[I - 1]
            self.SUM = 0.0

            for J in range(self.KQQ, 0, -1):
                self.SUM = self.SUM + self.DT[J - 1, I - 1] * self.W[J + self.KS - 1]

            #x=x0+delta*dx0+delta**2*sum
            STATE[I - 1] = self.REFPOS[I - 1] + self.DELTA * (self.REFVEL[I - 1] + self.DELTA * self.SUM)
        #
        #     Again we need to compute the W(K) coefficients that are
        #     going to be used in the velocity interpolation.
        #     (Note, at this point, KS = 1, KS1 = 0.)
        #
        for J in range(1, self.JX + 1):
            self.W[J + self.KS - 1] = self.FC[J] * self.W[J + self.KS1 - 1] - self.WC[J - 1] * self.W[J + self.KS - 1]

        self.KS = self.KS - 1

        #
        #     Perform velocity interpolation:
        #
        for I in range(1, 3 + 1):
            self.KQQ = self.KQ[I - 1]
            self.SUM = 0.0

            for J in range(self.KQQ, 0, -1):
                self.SUM = self.SUM + self.DT[J - 1, I - 1] * self.W[J + self.KS - 1]

            STATE[I + 3 - 1] = self.REFVEL[I - 1] + self.DELTA * self.SUM

        #
        #     That's all folks.  We don't know why we did anything, but
        #     at least we can tell structurally what we did.
        #

        return STATE

class daf_CK01line:
    def __init__(self,clknum,N,i,buf,rates):
        self.clknum=clknum
        self.rates=rates
        self.csvheader="SCLK,SCLK*,ET*,ETCAL*,q0,q1,q2,q3"
        if self.rates:
            self.n_elements = 7
            self.w = buf[i * self.n_elements+4:(i + 1) * self.n_elements]
            self.csvheader+=",wx,wy,wz"
        else:
            self.n_elements=4
        self.q = buf[i * self.n_elements:i * self.n_elements+4]

        self.sclk=buf[N*self.n_elements+i]
        self.et =cspice_sct2e(self.clknum,self.sclk)
    def __str__(self):
        result="%25.15e"%self.sclk
        try:
            et=cspice_sct2e(self.clknum, self.sclk)
            result+=",%s,%16.6f,%s" %(cspice_scdecd(self.clknum,self.sclk),et,cspice_etcal(et))
        except:
            result+=",,,"
        for i in range(4):
            result+=",%25.15e"%self.q[i]
        if self.rates:
            for i in range(3):
                result += ",%25.15e" % self.w[i]
        else:
            result+=",,,"
        return result

class daf_CK02line:
    def __init__(self,clknum,N,i,buf):
        self.clknum=clknum
        self.n_elements = 8
        self.w = buf[i * self.n_elements+4:i * self.n_elements+7]
        self.q = buf[i * self.n_elements:i * self.n_elements+4]
        self.rate=buf[i*self.n_elements+7]
        self.sclk=buf[N*self.n_elements+i]
        self.sclk1=buf[N*self.n_elements+N+i]
        self.et =cspice_sct2e(self.clknum,self.sclk)
        self.et1 =cspice_sct2e(self.clknum,self.sclk1)
        self.csvheader="SCLK,SCLK*,SCLK1,SCLK1*,ET*,ETCAL*,ET1*,ETCAL1*,q0,q1,q2,q3,wx,wy,wz,rate"
    def __str__(self):
        result="%25.15e"%self.sclk
        result+=",%s,%25.15e,%s,%16.6f,%s,%16.6f,%s" %(cspice_scdecd(self.clknum,self.sclk),self.sclk1,cspice_scdecd(self.clknum,self.sclk1),self.et,cspice_etcal(self.et),self.et1,cspice_etcal(self.et1))
        for i in range(4):
            result+=",%25.15e"%self.q[i]
        for i in range(3):
            result += ",%25.15e" % self.w[i]
        result += ",%25.15e" % self.rate
        return result

class daf_CK03line:
    def __init__(self,clknum,N,i,buf,rates):
        self.clknum=clknum
        self.rates=rates
        self.csvheader="SCLK,SCLK*,ET*,ETCAL*,q0,q1,q2,q3"
        if self.rates:
            self.n_elements = 7
            self.w = buf[i * self.n_elements+4:(i + 1) * self.n_elements]
            self.csvheader+=",wx,wy,wz"
        else:
            self.n_elements=4
        self.q = buf[i * self.n_elements:i * self.n_elements+4]

        self.sclk=buf[N*self.n_elements+i]
        self.et =cspice_sct2e(self.clknum,self.sclk)
    def __str__(self):
        result="%25.15e"%self.sclk
        try:
            et=cspice_sct2e(self.clknum, self.sclk)
            result+=",%s,%16.6f,%s" %(cspice_scdecd(self.clknum,self.sclk),et,cspice_etcal(et))
        except:
            result+=",,,"
        for i in range(4):
            result+=",%25.15e"%self.q[i]
        if self.rates:
            for i in range(3):
                result += ",%25.15e" % self.w[i]
        else:
            result+=",,,"
        return result

class daf_summary_record:
    def __init__(self,parent):
        """
        Read a record from a file-like object opened in binary mode
        :param parent: object of class daf
        """
        self.parent=parent
        self.dformat=self.parent.dformat
        self.buf=self.read(1024) #Read the whole block first
        self.namebuf=self.read(1024) #Read the name record next
        self.NEXT=int(struct.unpack(self.dformat,self.buf[ 0: 0+8])[0]) #Address of next summary record in file (zero if last)
        self.PREV=int(struct.unpack(self.dformat,self.buf[ 8: 8+8])[0]) #Address of prev summary record in file (zero if first)
        self.NSUM=int(struct.unpack(self.dformat,self.buf[16:16+8])[0]) #Number of summaries actually in this summary record
        self.SS=self.parent.SS
    def read(self,len):
        return self.parent.inf.read(len)
    def tell(self):
        return self.parent.inf.tell()
    def __iter__(self):
        for i in range(self.NSUM):
            result = daf_summary.make(self,self.parent,self.buf[8*(3+i*self.SS):8*(3+(i+1)*self.SS)],str(self.namebuf[i*self.parent.NC:(i+1)*self.parent.NC],encoding="cp437").strip(" \t\n\r\0"))
            yield result
    def __str__(self):
        result=(("NEXT: %d\n" +
                 "PREV: %s\n" +
                 "NSUM: %d") %
                (self.NEXT,self.PREV,self.NSUM))
        return result

class double_array_file:
    def __init__(self,inf):
        """
        Read a double-array file. This class reads the DAF header itself, and delegates reading the
        records, comments, and data. It therefore has fields representing the header, and lists
        of other class objects representing the data.
        :param inf: string filename or file-like object
        """
        try:
            self.inf=open(inf,"rb")
            self.needToClose=True
        except TypeError:
            self.inf=inf
            self.needToClose=False
        buf=self.inf.read(1024) #Read the whole header record first
        self.LOCIDW=str(buf[0:0+8],encoding='ASCII')
        if self.LOCIDW[0:4]=="NAIF":
            print("Old kernel type - guessing ")
            self.subtype="CK"
        elif self.LOCIDW[0:4]!="DAF/":
            raise ValueError("Not a DAF: Header is %s"%self.LOCIDW)
        else:
            self.subtype=self.LOCIDW[4:8].strip(" \t\n\r\0")
        self.LOCFMT=str(buf[88:88+8],encoding='ASCII') #IEEE format string
        if self.LOCFMT=="LTL-IEEE":
            self.endian="<"
        elif self.LOCFMT=="BIG-IEEE":
            self.endian=">"
        else:
            raise ValueError("Unknown double-precision format %s"%self.LOCFMT)
        self.iformat=self.endian+"i"
        self.dformat=self.endian+"d"
        self.ND=struct.unpack(self.iformat,buf[ 8: 8+4])[0] #Address  8, number of components in each array summary
        self.NI=struct.unpack(self.iformat,buf[12:12+4])[0] #Address 12, number of integer components in each array summary
        self.LOCIFN=str(buf[16:16+60],encoding='ASCII')     #Address 16, internal name or description of array file
        self.FWARD=struct.unpack(self.iformat,buf[76:76+4])[0] #Address 76, record number of initial summary in the file
        self.BWARD=struct.unpack(self.iformat,buf[80:80+4])[0] #Address 80, record number of final summary in the file
        self.FREE=struct.unpack(self.iformat,buf[84:84+4])[0]  #Address 84, first free address in file
        self.FTPSTR=str(buf[700:700+28],encoding='cp437') #Address 700, FTP validation string
        self.NC = 8 * (self.ND + (self.NI + 1) // 2)
        self.SS = self.ND + (self.NI + 1) // 2  # Size of a summary
        self.NS = 125 // self.SS  # Maximum number of summaries per summary record
    def __enter__(self):
        return self
    def __exit__(self,type,value,traceback):
        if self.needToClose:
            self.inf.close()
    def __str__(self):
        result=(("LOCIDW: %s\n"+
                 "ND:     %d\n"+
                 "NI:     %d\n" +
                 "NC*:    %d\n"+
                 "SS*:    %d\n" +
                 "NS*:    %d\n"+
                 "LOCIFN: %s\n" +
                 "FWARD:  %d\n" +
                 "BWARD:  %d\n" +
                 "FREE:   %d\n" +
                 "LOCFMT: %s") %
                (self.LOCIDW,self.ND,self.NI,self.NC,self.SS,self.NS,self.LOCIFN,self.FWARD,self.BWARD,self.FREE,self.LOCFMT))
        return result
    def __iter__(self):
        #Seek to the first summary record
        self.inf.seek((self.FWARD-1)*1024) #One for one-based, one for header record at beginning of file
        done=False
        while not done:
            result=daf_summary_record(self)
            yield result
            if result.NEXT==0:
                done=True
            else:
                self.inf.seek((result.NEXT-1)*1024)
    def dump(self,dump_lines=True,dump_comments=True):
        print(str(self))
        if dump_comments:
            for line in self.comments():
                print('"'+line+'"')
        for i,sr in enumerate(self):
            print("Summary record %d"%i)
            print(str(sr))
            for j,sum in enumerate(sr):
                print("Summary %d" % j)
                print(str(sum))
                print(str(sum.segment()))
                if dump_lines:
                    print("i*,dt," + sum.segment().csvheader)
                    et0=sum.et0
                    for k,line in enumerate(sum.segment()):
                        dt=line.et-et0
                        print(k,",%15.6f,"%dt,str(line))
#                        tval=line.test(et0)
#                        if True: #(tval>1e-10):
#                            print(line.et-et0,",",tval)
                        et0=line.et
    def comments(self):
        block=1
        done=False
        line=""
        while not done:
            self.inf.seek(block*1024)
            bb=self.inf.read(1000)
            for b in bb:
                if b==0:
                    yield line
                    line=""
                elif b==4:
                    yield line
                    return
                else:
                    line+=chr(b)
            block+=1

def spkdump():
    import argparse
    parser=argparse.ArgumentParser(description="Translate a Spice Double-Array File (DAF) into plain text")
    parser.add_argument('-x','--records',action='store_true',help="Dump all data records. May result in a very large amount of output!")
    parser.add_argument('-r','--comment',action='store_true',help="Dump file comments.")
    parser.add_argument('-a','--all',action='store_true',help="Dump everything that can be dumped")
    parser.add_argument('inf',nargs=1,help="Kernel to translate. Must be a bc or bsp kernel.")
    parser.add_argument('kernel',nargs='*',help="Kernels to furnish. Typically these are sclk and/or lsk files, needed to interpret bc timestamps")
    result=parser.parse_args()
    dump_lines=result.records or result.all
    dump_comments=result.comment or result.all
    inf=result.inf[0]

    from spiceypy import furnsh as cspice_furnsh
    for furnsh in result.kernel:
        cspice_furnsh(furnsh)
    with double_array_file(inf) as in_daf:
        cspice_furnsh(inf) #Just to run the test
        in_daf.dump(dump_lines=dump_lines,dump_comments=dump_comments)
#    with double_array_file("../../Data/spice/insight/insight_nom_2016e09o_edl_v1.bsp") as in_daf:
#        in_daf.dump(dump_lines)
#    with double_array_file("../../Data/spice/phoenix/phx_edl_rec_traj.bsp") as in_daf:
#        in_daf.dump(dump_lines)

if __name__=="__main__":
    spkdump()


#!/usr/bin/env python

import os
import sys
import time
import pdb
import argparse
import numpy
import crowdsource.psf as psfmod
from astropy.io import fits
import crowdsource
from unwise_psf import unwise_psf
# implicit dependency for WISE runs only
# https://github.com/legacysurvey/unwise_psf
import crowdsource.unwise_primary as unwise_primary
from astropy import wcs
from collections import OrderedDict
from pkg_resources import resource_filename

extrabits = {'crowdsat': 2**25,
             'nebulosity': 2**26,
             'w1brightoffedge': 2**7,
             'w2brightoffedge': 2**8,
             'hyperleda': 2**9}

nodeblend_bits = extrabits['hyperleda']
sharp_bits = (extrabits['w1brightoffedge'] | extrabits['w2brightoffedge'])

def wise_filename(basedir, coadd_id, band, _type, uncompressed=False,
                  drop_first_dir=False, epoch=-1):
    # type should be one of:
    # 'img-u', 'img-m', 'invvar-u', 'invvar-m', 'std-u', 'std-m'
    # 'n-u', 'n-m', 'frames', 'msk'

    # -msk is special because the info for both W1/W2 is in same file

    fname = 'unwise-' + coadd_id
    if _type != 'msk':
        fname += '-w' + str(band)

    fname += ('-' + _type + '.fits')

    path = [basedir, coadd_id[0:3], coadd_id, fname]
    if drop_first_dir:
        del path[1]
    if epoch >= 0:
        epochstr = 'e%03d' % epoch if _type != 'msk' else 'fulldepth'
        path = path[0:1] + [epochstr] + path[1:]
    fname = os.path.join(*path)

    if not uncompressed or _type == 'msk':
        if (_type != 'img-u') and (_type != 'img-m') and (_type != 'frames'):
            fname += '.gz'

    return fname

def read_blist(brightstars, raim, decim, hdr, maxsep):
    from astropy.coordinates.angle_utilities import angular_separation
    sep = angular_separation(numpy.radians(brightstars['ra']),
                             numpy.radians(brightstars['dec']),
                             numpy.radians(raim),
                             numpy.radians(decim))
    sep = numpy.degrees(sep)
    m = (sep < 3) & (brightstars['k_m'] < 5)
    brightstars = brightstars[m]
    wcs0 = wcs.WCS(hdr)
    yy, xx = wcs0.all_world2pix(brightstars['ra'], brightstars['dec'], 0)
    m = (xx > 0) & (xx < hdr['NAXIS1']) & (yy > 0) & (yy < hdr['NAXIS2'])
    xx, yy = xx[m], yy[m]
    mag = brightstars['k_m'][m]
    if not numpy.any(m):
        return None
    else:
        return [xx, yy, mag]

def massage_isig_and_dim(isig, im, flag, band, nm, nu, fac=None):
    """Construct a WISE inverse sigma image and add saturation to flag.

    unWISE provides nice inverse variance maps.  These however have no
    contribution from Poisson noise from sources, and so underestimate
    the uncertainties dramatically in bright regions.  This can pull the
    whole fit awry in bright areas, since the sky model means that every
    pixel feels every other pixel.

    It's not clear what the best solution is.  We make a goofy inverse
    sigma image from the original image and the inverse variance image.  It
    is intended to be sqrt(ivar) for the low count regime and grow like
    sqrt(1/im) for the high count regime.  The constant of proportionality
    should in principle be worked out; here I set it to 0.15, which worked
    once, and it doesn't seem like this should depend much on which
    WISE exposure the image came from?  It's ultimately something like the gain
    or zero point...
    """

    if fac is None:
        bandfacs = {1: 0.15, 2: 0.3}
        bandfloors = {1: 0.5, 2: 2}
        fac = bandfacs[band]
        floor = bandfloors[band]

    satbit = 16 if band == 1 else 32
    satlimit = 85000  # if band == 1 else 130000
    msat = ((flag & satbit) != 0) | (im > satlimit) | ((nm == 0) & (nu > 1))
    from scipy.ndimage import morphology
    # dilate = morphology.iterate_structure(
    #     morphology.generate_binary_structure(2, 1), 3)
    xx, yy = numpy.mgrid[-3:3+1, -3:3+1]
    dilate = xx**2+yy**2 <= 3**2
    msat = morphology.binary_dilation(msat, dilate)
    isig[msat] = 0
    flag = flag.astype('i8')
    # zero out these bits; we claim them for our own purposes.
    massagebits = (extrabits['crowdsat'] | crowdsource.nodeblend_maskbit |
                   crowdsource.sharp_maskbit | extrabits['nebulosity'])
    flag &= ~massagebits
    flag[msat] |= extrabits['crowdsat']
    flag[(flag & nodeblend_bits) != 0] |= crowdsource.nodeblend_maskbit
    flag[(flag & sharp_bits) != 0] |= crowdsource.sharp_maskbit

    sigma = numpy.sqrt(1./(isig + (isig == 0))**2 + floor**2 +
                       fac**2*numpy.clip(im, 0, numpy.inf))
    sigma[msat] = numpy.inf
    sigma[isig == 0] = numpy.inf
    return (1./sigma).astype('f4'), flag

def wise_psf_stamp(band, nosmooth=False):
    # psf noise: ~roughly 0.1 count in outskirts of W1 and W2
    if band >= 3:
        raise ValueError('Need to stare at W3+ PSF more!')
    psfnoise = 0.1
    stampfn = resource_filename('unwise_psf',
                                'data/psf_model_w'+str(band)+'.fits')
    stamp = fits.getdata(stampfn)
    edges = numpy.concatenate([stamp[0, 1:-1], stamp[-1, 1:-1],
                               stamp[1:-1, 0], stamp[1:-1, -1]])
    medval = numpy.median(edges[edges != 0]) / 2
    stamp[stamp == 0] = medval
    stamp -= medval
    from scipy import signal
    stamp[stamp < 0] = 0.
    # suppress spurious warnings in signal.wiener
    olderr = numpy.seterr(invalid='ignore', divide='ignore')
    # update to scipy.signal means that Wiener filter uses an FFT
    # to perform the various convolutions, which causes bad errors
    # here unless we cast to f8.  It's not that hard to do something
    # a bit better than scipy.signal.wiener---morally we really want to do
    # something like smooth in log space on radial lines---but I don't
    # want to go further down that rabbit hole today.
    stamp = signal.wiener(stamp.astype('f8'),  11, psfnoise)
    stamp = stamp.astype('f4')
    numpy.seterr(**olderr)
    # taper linearly over outer 60 pixels?
    stampszo2 = stamp.shape[0] // 2
    xx, yy = numpy.mgrid[-stampszo2:stampszo2+1, -stampszo2:stampszo2+1]
    edgedist = numpy.clip(stampszo2-numpy.abs(xx), 0,
                          stampszo2-numpy.abs(yy))
    stamp = stamp * numpy.clip(edgedist / 60., stamp < 10, 1)
    import psf
    stamp = psf.center_psf(stamp, censize=19)
    stamp = stamp / numpy.sum(stamp)
    return stamp

def wise_psf(band, coadd_id):
    stamp = wise_psf_stamp(band)
    stamp = unwise_psf.rotate_using_rd(stamp, coadd_id)
    psf = psfmod.SimplePSF(stamp)
    from functools import partial
    psf.fitfun = partial(psfmod.wise_psf_fit, psfstamp=stamp)
    return psf

def wise_psf_grid(band, coadd_id, basedir, uncompressed=False,
                  drop_first_dir=False, ngrid=None, epoch=-1):
    imagefn = wise_filename(basedir, coadd_id, band, 'img-m',
                            uncompressed=uncompressed,
                            drop_first_dir=drop_first_dir, epoch=epoch)
    hdr = fits.getheader(imagefn)
    if ngrid is None:
        rr, dd = hdr['CRVAL1'], hdr['CRVAL2']
        from astropy.coordinates import SkyCoord
        from astropy import units as u
        coord = SkyCoord(ra=rr*u.deg, dec=dd*u.deg, frame='icrs')
        coord = coord.geocentrictrueecliptic
        lam, bet = coord.lon.deg, coord.lat.deg
        dlam = 1.4/(numpy.abs(numpy.cos(numpy.radians(bet)))+1e-6)
        ngrid = numpy.floor(numpy.clip(dlam / 1, 4, 16)).astype('i4')
    x = numpy.linspace(0, 2047, ngrid)
    y = numpy.linspace(0, 2047, ngrid)
    wcs0 = wcs.WCS(hdr)
    stamp = wise_psf_stamp(band).astype('f4')
    stamps = numpy.zeros((len(x), len(y))+stamp.shape, dtype=stamp.dtype)
    unwise_psf.rotate_using_convolution.cache = None  # clear cache
    for i in range(len(x)):
        for j in range(len(y)):
            rr, dd = wcs0.all_pix2world(y[j], x[i], 0)
            stamps[i, j, ...] = unwise_psf.rotate_using_rd(
                stamp, coadd_id, ra=rr, dec=dd, cache=True)
    psf = psfmod.GridInterpPSF(stamps, x, y)
    from functools import partial
    psf.fitfun = partial(psfmod.wise_psf_fit, psfstamp=(stamps, x, y),
                         grid=True)
    return psf

def read_wise(coadd_id, band, basedir, uncompressed=False,
              drop_first_dir=False, epoch=-1):
    assert((band == 1) or (band == 2))
    assert(len(coadd_id) == 8)

    imagefn = wise_filename(basedir, coadd_id, band, 'img-m',
                            uncompressed=uncompressed,
                            drop_first_dir=drop_first_dir, epoch=epoch)
    ivarfn = wise_filename(basedir, coadd_id, band, 'invvar-m',
                           uncompressed=uncompressed,
                           drop_first_dir=drop_first_dir, epoch=epoch)
    flagfn = wise_filename(basedir, coadd_id, band, 'msk',
                           uncompressed=uncompressed,
                           drop_first_dir=drop_first_dir, epoch=epoch)
    nmfn = wise_filename(basedir, coadd_id, band, 'n-m',
                         uncompressed=uncompressed,
                         drop_first_dir=drop_first_dir, epoch=epoch)
    nufn = wise_filename(basedir, coadd_id, band, 'n-u',
                         uncompressed=uncompressed,
                         drop_first_dir=drop_first_dir, epoch=epoch)

    im, hdr = fits.getdata(imagefn, header=True)
    sqivar = numpy.sqrt(fits.getdata(ivarfn))
    flag = fits.getdata(flagfn)
    nm = fits.getdata(nmfn)
    nu = fits.getdata(nufn)
    sqivar, flag = massage_isig_and_dim(sqivar, im, flag, band, nm, nu)
    return im, sqivar, flag, hdr

def ivarmap(isig, psfstamp):
    from scipy.signal import fftconvolve
    ivarim = fftconvolve(isig**2., psfstamp[::-1, ::-1]**2., mode='same')
    return ivarim

def brightlist(brightstars, coadd_id, band, basedir, uncompressed=False,
               drop_first_dir=False, epoch=-1):
    imagefn = wise_filename(basedir, coadd_id, band, 'img-m',
                            uncompressed=uncompressed,
                            drop_first_dir=drop_first_dir, epoch=epoch)
    hdr = fits.getheader(imagefn)
    blist = read_blist(brightstars, hdr['CRVAL1'], hdr['CRVAL2'], hdr, 3)
    return blist

def collapse_unwise_bitmask(bitmask, band):
    # 2^0 = bright star core and wings
    # 2^1 = PSF-based diffraction spike
    # 2^2 = optical ghost
    # 2^3 = first latent
    # 2^4 = second latent
    # 2^5 = AllWISE-like circular halo
    # 2^6 = bright star saturation
    # 2^7 = geometric diffraction spike

    assert((band == 1) or (band == 2))

    bits_w1 = OrderedDict([('core_wings', 2**0 + 2**1),
                           ('psf_spike', 2**27),
                           ('ghost', 2**25 + 2**26),
                           ('first_latent', 2**13 + 2**14),
                           ('second_latent', 2**17 + 2**18),
                           ('circular_halo', 2**23),
                           ('saturation', 2**4),
                           ('geom_spike', 2**29)])

    bits_w2 = OrderedDict([('core_wings', 2**2 + 2**3),
                           ('psf_spike', 2**28),
                           ('ghost', 2**11 + 2**12),
                           ('first_latent', 2**15 + 2**16),
                           ('second_latent', 2**19 + 2**20),
                           ('circular_halo', 2**24),
                           ('saturation', 2**5),
                           ('geom_spike', 2**30)])

    bits = (bits_w1 if (band == 1) else bits_w2)

    # hack to handle both scalar and array inputs
    result = 0*bitmask

    for i, feat in enumerate(bits.keys()):
        result += (2**i)*(numpy.bitwise_and(bitmask, bits[feat]) != 0)

    # int8 would be fine here, but astropy.io.fits seems to read this
    # as a boolean... so we waste the extra 8 bits.
    return result.astype('i2')

def collapse_extraflags(bitmask, band):
    bits_w1 = OrderedDict([('bright_off_edge', 2**7),
                           ('resolved_galaxy', 2**9),
                           ('big_object', 2**10),
                           ('possible_bright_star_centroid', 2**21),
                           ('crowdsat', extrabits['crowdsat']),
                           ('nebulosity', extrabits['nebulosity']),
                           ('nodeblend', crowdsource.nodeblend_maskbit),
                           ('sharp', crowdsource.sharp_maskbit)])

    bits_w2 = OrderedDict([('bright_off_edge', 2**8),
                           ('resolved_galaxy', 2**9),
                           ('big_object', 2**10),
                           ('possible_bright_star_centroid', 2**22),
                           ('crowdsat', extrabits['crowdsat']),
                           ('nebulosity', extrabits['nebulosity']),
                           ('nodeblend', crowdsource.nodeblend_maskbit),
                           ('sharp', crowdsource.sharp_maskbit)])

    bits = (bits_w1 if (band == 1) else bits_w2)

    # hack to handle both scalar and array inputs
    result = 0*bitmask

    for i, feat in enumerate(bits.keys()):
        result += (2**i)*(numpy.bitwise_and(bitmask, bits[feat]) != 0)

    # could fit in a byte, but astropy.io.fits reads these as booleans,
    # so we waste the byte...
    return result.astype('i2')

if __name__ == "__main__":
    try:
        print('Running on host: ' + str(os.environ.get('HOSTNAME')))
    except Exception:
        print("Couldn't retrieve hostname!")

    parser = argparse.ArgumentParser(description='Run crowdsource on unWISE coadd image')
    parser.add_argument('coadd_id', type=str, nargs=1)
    parser.add_argument('band', type=int, nargs=1)
    parser.add_argument('outfn', type=str, nargs=1)

    parser.add_argument('basedir', type=str, nargs='?', default='/global/projecta/projectdirs/cosmo/work/wise/outputs/merge/neo4/fulldepth')
    parser.add_argument('--refit-psf', '-r', default=False, action='store_true')
    parser.add_argument('--verbose', '-v', default=False, action='store_true')
    parser.add_argument('--uncompressed', '-u', default=False, action='store_true')
    parser.add_argument('--brightcat', '-b',
                        default=os.environ.get('TMASS_BRIGHT', ''), type=str)
    parser.add_argument('--modelfn', '-m', default='', type=str,
                        help='file name for model image, if desired')
    parser.add_argument('--infoimfn', '-i', default='', type=str,
                        help='file name for info image, if desired')
    parser.add_argument('--masknebulosity', '-n', action='store_true')
    parser.add_argument('--forcecat', type=str, default='')
    parser.add_argument('--startsky', type=str, default='')
    parser.add_argument('--startpsf', type=str, default='')
    parser.add_argument('--noskyfit', default=False, action='store_true')
    parser.add_argument('--threshold', default=5, type=float,
                        help='find sources down to threshold*sigma')
    parser.add_argument('--epoch', type=int, default=-1,
                        help='epoch number of time-resolved WISE coadd')
    parser.add_argument('--release', type=str, default='')

    args = parser.parse_args()

    coadd_id = args.coadd_id[0]
    band = args.band[0]
    basedir = args.basedir

    im, sqivar, flag, hdr = read_wise(coadd_id, band, basedir,
                                      uncompressed=args.uncompressed,
                                      epoch=args.epoch)
    if len(args.startsky) > 0:
        startsky = fits.getdata(args.startsky, 'SKY')
    else:
        startsky = numpy.nan
    flag_orig = fits.getdata(wise_filename(basedir, coadd_id, band, 'msk',
                                           uncompressed=args.uncompressed,
                                           epoch=args.epoch))

    if args.masknebulosity:
        import nebulosity_mask
        nebfn = os.path.join(os.environ['WISE_DIR'], 'dat', 'nebnet',
                             'weights1', '1st_try')
        nebmod = nebulosity_mask.load_model(nebfn)
        nebmask = nebulosity_mask.gen_mask_wise(nebmod, im) == 2
        if numpy.any(nebmask):
            flag |= nebmask * extrabits['nebulosity']
            flag |= nebmask * crowdsource.sharp_maskbit
            print('Masking nebulosity, %5.2f' % (
                numpy.sum(nebmask)/1./numpy.sum(numpy.isfinite(nebmask))))

    psf = wise_psf_grid(band, coadd_id, basedir, epoch=args.epoch)
    if len(args.startpsf) > 0:
        startpsf = fits.getdata(args.startpsf, 'PSF').astype('f4')
        # there can be some endianness issues; astype('f4') converts to native
        modpsf = psf(1024, 1024, stampsz=psf.stamp.shape[-1])
        resid = startpsf-modpsf
        # need not sum to zero.
        newstamps = psf.stamp / psf.normstamp[:, :, None, None]
        newstamps += resid
        psf = psfmod.GridInterpPSF(newstamps, psf.x, psf.y)
        from functools import partial
        psf.fitfun = partial(psfmod.wise_psf_fit,
                             psfstamp=(newstamps, psf.x, psf.y), grid=True)

    if len(args.brightcat) > 0:
        brightstars = fits.getdata(args.brightcat)
        blist = brightlist(brightstars, coadd_id, band, basedir,
                           uncompressed=args.uncompressed, epoch=args.epoch)
    else:
        print('No bright star catalog, not marking bright stars.')
        blist = None

    if args.verbose:
        t0 = time.time()
        print('Starting %s, band %d, at %s' % (coadd_id, band, time.ctime()))
        sys.stdout.flush()

    if len(args.forcecat) == 0:
        res = crowdsource.fit_im(
            im, psf, weight=sqivar, dq=flag, refit_psf=args.refit_psf,
            verbose=args.verbose, ntilex=4, ntiley=4, derivcentroids=True,
            maxstars=30000*16, fewstars=50*16, blist=blist,
            threshold=args.threshold, psfvalsharpcutfac=0.5,
            psfsharpsat=0.8)
    else:
        forcecat = fits.getdata(args.forcecat, 1)
        x, y = forcecat['x'], forcecat['y']
        res = crowdsource.fit_im_force(
            im, x, y, psf, weight=sqivar, dq=flag, refit_psf=args.refit_psf,
            blist=blist, refit_sky=(not args.noskyfit),
            startsky=startsky, psfderiv=False, psfvalsharpcutfac=0.5,
            psfsharpsat=0.8)
    cat, model, sky, psf = res
    print('Finishing %s, band %d; %d sec elapsed.' %
          (coadd_id, band, time.time()-t0))

    outfn = args.outfn[0]

    x = cat['x']
    y = cat['y']

    wcs0 = wcs.WCS(hdr)
    ra, dec = wcs0.all_pix2world(y, x, 0)
    coadd_ids = numpy.zeros(len(ra), dtype='a8')
    bands = numpy.zeros(len(ra), dtype='i4')
    ids = numpy.zeros(len(ra), dtype='U20')
    coadd_ids[:] = coadd_id
    bands[:] = band
    if len(args.release) == 0:
        ids = ['%sw%1do%07d' % (coadd_id, band, num) for num in range(len(ra))]
    else:
        ids = ['%sw%1do%07dr%s' % (coadd_id, band, num, args.release)
               for num in range(len(ra))]

    nmfn = wise_filename(basedir, coadd_id, band, 'n-m',
                         uncompressed=args.uncompressed, epoch=args.epoch)
    nmim = fits.getdata(nmfn)
    nms = crowdsource.extract_im(cat['x'], cat['y'], nmim)
    flags_unwise = crowdsource.extract_im(
        cat['x'], cat['y'], collapse_unwise_bitmask(flag_orig, band))
    flags_infoim = collapse_extraflags(flag, band)
    flags_info = crowdsource.extract_im(cat['x'], cat['y'], flags_infoim)
    # cast to i2; astropy.io.fits seems to fail for bools?
    primary = unwise_primary.is_primary(coadd_id, ra, dec).astype('i2')

    import numpy.lib.recfunctions as rfn
    cat = rfn.drop_fields(cat, ['flags'])
    cat = rfn.append_fields(
        cat, ['ra', 'dec', 'coadd_id', 'band', 'unwise_detid', 'nm',
              'primary', 'flags_unwise', 'flags_info'],
        [ra, dec, coadd_ids, bands, ids, nms, primary, flags_unwise,
         flags_info])

    hdr['EXTNAME'] = 'PRIMARY'
    fits.writeto(outfn, None, hdr)
    fits.append(outfn, cat)
    if len(args.modelfn) > 0:
        hdulist = fits.open(args.modelfn, mode='append')
        compkw = {'compression_type': 'GZIP_1',
                  'quantize_method': 2, 'quantize_level': -0.5,
                  'tile_size': model.shape}
        hdr['EXTNAME'] = 'model'
        hdulist.append(fits.CompImageHDU(model, hdr, **compkw))
        hdr['EXTNAME'] = 'sky'
        hdulist.append(fits.CompImageHDU(sky, hdr, **compkw))
        hdulist.close(closed=True)

    if len(args.infoimfn) > 0:
        psffluxivar = ivarmap(sqivar, psf(1024, 1024, stampsz=59)).astype('f4')
        psfstamp = psf(1024, 1024, stampsz=325)
        hdulist = fits.open(args.infoimfn, mode='append')
        compkw = {'compression_type': 'GZIP_1',
                  'quantize_method': 2,
                  'tile_size': psffluxivar.shape}
        hdr['EXTNAME'] = 'psffluxivar'
        hdulist.append(fits.CompImageHDU(psffluxivar, hdr, **compkw))
        hdr['EXTNAME'] = 'infoflags'
        compkw = {'compression_type': 'GZIP_1',
                  'tile_size': flags_infoim.shape}
        # must recast flags_infoim as a u1; unsigned isn't supported
        # in tables, but signed int8 isn't supported in CompImageHDU.
        # ugh.
        hdulist.append(fits.CompImageHDU(flags_infoim.astype('u1'),
                                         hdr, **compkw))
        hdulist.append(fits.ImageHDU(psfstamp, None, name='psf'))
        hdulist.close(closed=True)