"""
ssd
F. Pierfederici (fpierfed@lsst.org)

Ephemerides generation module. It wraps custon fortran code developed using
JPL libraries.
"""
import os
import sys
import numpy
import _ssd as _ephem


# If we have a private copy of CAET_DATA, make sure that $CAET_DATA is defined
# and points to our private copy. Do not do anything otherwise.
this_dir = os.path.dirname(os.path.realpath(__file__))
our_copy = os.path.join(this_dir, '..', '..', 'caet_data')
CAET_DATA = os.environ.get('CAET_DATA', None)
if(os.path.exists(our_copy) and (not CAET_DATA or CAET_DATA != our_copy)):
    CAET_DATA = os.path.join(this_dir, '..', '..', 'caet_data')
    os.environ['CAET_DATA'] = CAET_DATA
# Now we can init the SSD module.
_ephem.initialize()


__all__ = ['ephemerides', 'coarseEphem', 'covar2src']


def ephemerides(orbitElements,
                orbitEpoch,
                times,
                obsCode,
                absMag,
                gMag=0.150,
                covariance=None):
    """ephemerides(orbitElements, orbitEpoch, absMag, gMag,
       times, obsCode, covariance)

    Given one orbit (described by its orbital elements), a list of times, and a
    position on Earth (specified by an observatory code), compute the predictied
    position (i.e. ephemerides) of that orbit at each of the input times. Return
    positions, predicted magnitudes and positional uncertainties.
    
    INPUT (all angles in degrees)
      orbitElements: numpy.array containing the orbital elements
        (dimension <= 6) in COMETARY format (q, e, i, node, peri, tp).
        All angles in degrees.
      orbitEpoch: epoch of the orbit (UTC MJD).
      obsCode: MPC observatory code.
      times: list of ephemerides times (UTC MJD).
      absMag: absolute H magnitude.
      gMag: magnitude g parameter (default 0.150).
      covariance: optional square root covariance matrix in diagonal form as
        21-element numpy.array.
    
    OUTPUT
      numpy.array of numpy.arrays of the form
        [RA, Dec, mag, t, RAErr, DecERR, err1, err2, PA].
      RA: Right Ascension (deg).
      Dec: Declination (deg).
      t: ephemerides time (UTC MJD).
      mag: apparent magnitude (if available, None otherwise).
      RAErr: RA 3-sigma uncertainty (arcsec).
      DecErr: Dec 3-sigma uncertainty (arcsec).
      err1: semi-major axis of the 3-sigma positional error ellipse (arcsec).
      err1: semi-major axis of the 3-sigma positional error ellipse (arcsec).
      PA: position angle of the 3-sigma positional error ellipse (deg).
    """
    # We redirect STDOUT since the JPL code writes a lot of crap there...
    # out_log = file('/dev/null', 'a+')
    # err_log = file('/dev/null', 'a+', 0)
    # dev_null = file('/dev/null', 'r')

    # os.dup2(out_log.fileno(), sys.stdout.fileno())
    # os.dup2(err_log.fileno(), sys.stderr.fileno())
    # os.dup2(dev_null.fileno(), sys.stdin.fileno())
    
    if(not gMag):
        gMag = .150

    # Sort the times in reverse order.
    times.sort()
    times = times[::-1]

    res = []
    if(covariance != None and len(covariance) == 21 and
       None not in list(covariance)):
        out, err = _ephem.ephemerides(t_epoch=orbitEpoch,
                                      elems=orbitElements,
                                      obscod=obsCode,
                                      t_eph_uts=times,
                                      h_v=absMag,
                                      g=gMag,
                                      has_covar=1,
                                      covar=covariance)
    else:
        out, err = _ephem.ephemerides(t_epoch=orbitEpoch,
                                      elems=orbitElements,
                                      obscod=obsCode,
                                      t_eph_uts=times,
                                      h_v=absMag,
                                      g=gMag,
                                      has_covar=0,
                                      covar=numpy.zeros(21))
    # <-- end if
    res += list(out)

    # Restore stdout and quit.
    # sys.stdout = sys.__stdout__
    # sys.stderr = sys.__stderr__
    return(numpy.array(res))


def geocentricEphemerides(orbitElements,
                          orbitEpoch,
                          times,
                          absMag,
                          gMag=0.150,
                          covariance=None):
    """
    Convenience function: it is equivalent to calling ephemerides() with an
    obsCode=500.
    """
    return(ephemerides(orbitElements,
                       orbitEpoch,
                       times,
                       '500',
                       absMag,
                       gMag,
                       covariance))


def coarseEphem(orbitParameterList, ephemMJD, obscode,
                deltaMJD=1.):
    """
    This program calculates MOPS coarse ephemerides for a supplied set of
    orbital parameters.  For each object, three positions (ephemerides)
    are calculated, separated by a specified amount of time.  The there
    positions are used by the calling script to estimate the proximity of
    a given object to a list of fields so that a reduced set of orbital
    parameters can be used to calculate accurate ephemerides.

    For each set of orbital parameters, three lines are generated, one for
    each ephemeris.  If the three ephemerides are all below the user-specified
    detection limit, no output is generated for the object.  If an error
    occurs during ephemeris calculation, no output is generated.

    This is a python translation of the C program with the same name.

    orbitParameters: [id,
                      numpy.array([q,
                                   e,
                                   i,
                                   node,
                                   argPeri,
                                   timePeri]),
                      epoch,
                      h_v,
                      src]
    ephemMJD: MJD of the desired ephemerides time.
    deltaMJD: the ephemerides will be computed at ephemMJD-deltaMJD, ephemMJD,
              ephemMJD+deltaMJD
    obscode: MPC observatory code of where we are on Earth.

    Return {id: numpy.array([[ra, dec, mag, ephemMJD,
                              errRA, errDec, smaa, smia, pa],])}
    """
    if(ephemMJD < 0):
        # Ummm... this is not right: we do not want to operate in the middle
        # ages ;-)
        sys.stderr.write('ephemMJD cannot be negative!\n')
        return

    # Precompute the ephem times.
    deltaMJD = abs(deltaMJD)
    times = numpy.array([ephemMJD+deltaMJD, ephemMJD, ephemMJD-deltaMJD])

    # Propagate each orbit.
    positions = {}
    for orbit in orbitParameterList:
        if(not _isValid(*(list(orbit[1]) + [orbit[2]]))):
            # print('Invalid orbit for coarseEphem.')
            continue
        # <-- end if

        # The orbit is valid compute ephems.
        orbitId = orbit[0]
        positions[orbitId] = ephemerides(orbitElements=orbit[1],
                                         orbitEpoch=orbit[2],
                                         times=times,
                                         obsCode=obscode,
                                         absMag=orbit[3],
                                         gMag=None,
                                         covariance=None)
    # <-- end for
    return(positions)


def _isValid(q, e, i, node, argPeri, timePeri, epoch):
    """
    Make sure that MOPS is happy with this orbit.
    """
    if(q < 0):
        sys.stderr.write('EINVALIDINPUT %d 1\n' %(objId))
        return(False)
    if(e < 0):
        sys.stderr.write('EINVALIDINPUT %d 2\n' %(objId))
        return(False)
    if(i > 180 or i < 0):
        sys.stderr.write('EINVALIDINPUT %d 3\n' %(objId))
        return(False)
    if(q < 0.05 and abs(timePeri - epoch) < 5.):
        sys.stderr.write('ESUNGRAZER JPLEPH-ERROR: SUN %d\n' %(objId))
        return(False)
    return(True)
    
    
    


def covar2src(covar):
    """
    Given a vector stored (21 element) covariance matrix, compute the
    corresponding SRC (Square Root Covariance matrix), also in vector stored
    form.
    """
    return(_ephem.covar2src(covar))
