# coding: utf-8

""" Class for reading data from SCF simulations """

from __future__ import division, print_function

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

# Standard library
import os, sys
import logging
import re

# Third-party
import numpy as np
import astropy.units as u
from astropy.constants import G
from astropy.table import Table

# Project
from .core import NBodyReader

# Create logger
logger = logging.getLogger(__name__)

__all__ = ["SCFReader"]

class SCFReader(NBodyReader):

    # TODO: way to list snapshot files?

    def _read_units(self):
        """ Read and parse the SCFPAR file containing simulation parameters
            and initial conditions. Right now, only parse out the simulation
            units.
        """
        pars = dict()

        pattr = re.compile("^([0-9\.e\+]+)\s+satellite\s+((scale)|(mass))\s\(([A-Za-z]+)\)")
        parfile = os.path.join(self.path, "SCFPAR")
        with open(parfile) as f:
            for line in f:
                spl = line.split()

                # find what G is set to
                try:
                    if spl[1].strip() == 'G':
                        pars['G'] = float(spl[0])
                except IndexError:
                    pass

                # parse out length scale and mass scale
                m = pattr.search(line)
                try:
                    grp = m.groups()
                    if grp[1] == 'scale':
                        pars['length'] = float(grp[0])
                    elif grp[1] == 'mass':
                        pars['mass'] = float(grp[0])
                except:
                    pass

                # exit if have length, mass, and G
                if pars.has_key('G') and pars.has_key('length')\
                    and pars.has_key('mass'):
                    break

        _G = G.decompose(bases=[u.kpc,u.M_sun,u.Myr]).value
        X = (_G / pars['length']**3 * pars['mass'])**-0.5

        length_unit = u.Unit("{0} kpc".format(pars['length']))
        mass_unit = u.Unit("{0} M_sun".format(pars['mass']))
        time_unit = u.Unit("{:08f} Myr".format(X))

        units = dict(length=length_unit,
                     mass=mass_unit,
                     time=time_unit,
                     speed=length_unit/time_unit,
                     dimensionless=u.dimensionless_unscaled)

        return units

    def read_snap(self, filename, units=None):
        """ Given a SNAP filename, read and return the data. By default,
            returns data in simulation units, but this can be changed with
            the `units` kwarg.

            Parameters
            ----------
            filename : str
                The name of the SNAP file to read.
            units : dict (optional)
                A unit system to transform the data to. If None, will return
                the data in simulation units.
        """

        # column names for SNAP file, in simulation units
        colnames = "m x y z vx vy vz s1 s2 tub".split()
        coltypes = "mass length length length speed speed speed dimensionless dimensionless time".split()
        colunits = [self.sim_units[x] for x in coltypes]

        # read the first line to get the numer of particles and timestep
        fullpath = os.path.join(self.path, filename)
        with open(fullpath) as f:
            firstline = f.readline()
            try:
                nparticles,time = firstline.split()
            except:
                raise ValueError("Invalid header line. Expected 'nparticles,time', "
                                 "got:\n\t\t{}".format(firstline))
        time = float(time)*self.sim_units['time']

        data = np.genfromtxt(os.path.join(self.path,filename),
                             skiprows=1, names=colnames)
        if units is not None:
            new_colunits = []
            for colname,colunit in zip(colnames,colunits):
                newdata = (data[colname]*colunit).decompose(units)
                data[colname] = newdata.value
                new_colunits.append(newdata.unit)

            time = time.decompose(units)
            colunits = new_colunits

        tbl = Table(data, meta=dict(time=time.value))
        for colname,colunit in zip(colnames,colunits):
            tbl[colname].unit = colunit

        return tbl
@@ -2884,286 +2884,8 @@
 %22%22%0A%0A
-        # column names for SNAP file, in simulation units%0A        colnames = %22m x y z vx vy vz s1 s2 tub%22.split()%0A        coltypes = %22mass length length length speed speed speed dimensionless dimensionless time%22.split()%0A        colunits = %5Bself.sim_units%5Bx%5D for x in coltypes%5D%0A%0A
     
@@ -2953,16 +2953,16 @@
 imestep%0A
+
         
@@ -3381,60 +3381,1125 @@
     
-data = np.genfromtxt(os.path.join(self.path,
+full_filename = os.path.join(self.path,filename)%0A%0A        with open(full_filename) as f:%0A            l0 = f.readline()%0A            l1 = f.readline()%0A            numcols = len(l1.split())%0A%0A        if numcols == 8:%0A            # not self gravitating%0A            logger.debug(%22Not a self-gravitating run: only 8 columns%22)%0A%0A            # column names for SNAP file, in simulation units%0A            colnames = %22m x y z vx vy vz s1%22.split()%0A            coltypes = %22mass length length length speed speed speed dimensionless%22.split()%0A            colunits = %5Bself.sim_units%5Bx%5D for x in coltypes%5D%0A%0A        elif numcols == 10:%0A            # not self gravitating%0A            logger.debug(%22A self-gravitating run: 10 columns%22)%0A%0A            # column names for SNAP file, in simulation units%0A            colnames = %22m x y z vx vy vz s1 s2 tub%22.split()%0A            coltypes = %22mass length length length speed speed speed dimensionless dimensionless time%22.split()%0A            colunits = %5Bself.sim_units%5Bx%5D for x in coltypes%5D%0A%0A        else:%0A            raise ValueError(%22Invalid SNAP file.%22)%0A%0A        data = np.genfromtxt(full_
 filename
 ),%0A 
@@ -4494,17 +4494,16 @@
 filename
-)
 ,%0A      

add support for different snap files