# coding: utf-8

""" Code for helping to select stars from the nearby Sgr wraps. """

from __future__ import division, print_function

__author__ = "adrn <adrn@astro.columbia.edu>"

# Standard library
import os, sys
import gc

# Third-party
import h5py
import numpy as np
import numexpr
import astropy.units as u
from astropy.io import ascii
from astropy.table import vstack, Table, Column
import astropy.coordinates as coord

# Project
from ..coordinates.frame import heliocentric, galactocentric
from ..dynamics import Particle, ObservedParticle, Orbit

__all__ = ["read_table", "read_hdf5"]

def read_table(filename, expr=None, N=None):
    _table = np.genfromtxt(filename, names=True)

    if expr is not None:
        idx = numexpr.evaluate(str(expr), _table)
        _table = _table[idx]

    if N is not None and N > 0:
        np.random.shuffle(_table)
        _table = _table[:min(N,len(_table))]

    return _table

def read_hdf5(h5file):
    """ Read particles and satellite from a given HDF5 file. """

    return_dict = dict()
    with h5py.File(h5file, "r") as f:
        try:
            ptcl = f["particles"]
            satl = f["satellite"]
        except KeyError:
            raise ValueError("Invalid HDF5 file. Missing 'particles' or "
                             "'satellite' group.")

        if "error" in ptcl.keys():
            p = ObservedParticle(ptcl["data"].value.T, ptcl["error"].value.T,
                                 frame=heliocentric,
                                 units=[u.Unit(x) for x in ptcl["units"]])
            p.tub = ptcl["tub"].value
            return_dict["true_particles"] = Particle(ptcl["true_data"].value.T,
                                                     frame=heliocentric,
                                                     units=[u.Unit(x) for x in ptcl["units"]])
        else:
            p = Particle(ptcl["data"].value.T,
                         frame=heliocentric,
                         units=[u.Unit(x) for x in ptcl["units"]])
            p.tub = ptcl["tub"].value
        return_dict["particles"] = p

        if "error" in satl.keys():
            s = ObservedParticle(satl["data"].value.T, satl["error"].value.T,
                                 frame=heliocentric,
                                 units=[u.Unit(x) for x in satl["units"]])
            return_dict["true_satellite"] = Particle(satl["true_data"].value.T,
                                                     frame=heliocentric,
                                                     units=[u.Unit(x) for x in satl["units"]])

        else:
            s = Particle(satl["data"].value.T,
                         frame=heliocentric,
                         units=[u.Unit(x) for x in satl["units"]])
        s.m = satl["m"].value
        s.v_disp = satl["v_disp"].value
        return_dict["satellite"] = s

        if "simulation" in f.keys():
            return_dict["t1"] = float(f["simulation"]["t1"].value)
            return_dict["t2"] = float(f["simulation"]["t2"].value)

    return return_dict
@@ -1573,32 +1573,36 @@
 tcl%5B%22tub%22%5D.value
+*0.%0A
 %0A            ret
@@ -1602,37 +1602,14 @@
     
-return_dict%5B%22true_particles%22%5D
+true_p
  = P
@@ -1673,39 +1673,16 @@
         
-               
-        
 frame=he
@@ -1693,39 +1693,16 @@
 entric,%0A
-                       
         
@@ -1757,32 +1757,126 @@
 ptcl%5B%22units%22%5D%5D)%0A
+            true_p.tub = ptcl%5B%22tub%22%5D.value%0A            return_dict%5B%22true_particles%22%5D = true_p%0A
         else:%0A  

store true particles