#!/usr/bin/python -tt
#
# Werner Lustermann, Dominik Neise
# ETH Zurich, TU Dortmund
#
import os
import sys
import numpy as np
import ROOT
ROOT.gROOT.SetBatch(True)

#--------- BUILDING OF THE SHARED OBJECT FILES -------------------------------
if __name__ == '__main__' and len(sys.argv) > 1 and 'build' in sys.argv[1]:
    ROOT.gSystem.AddLinkedLibs("-lz")
    root_make_string = ROOT.gSystem.GetMakeSharedLib()
    if "-std=c++0x" not in root_make_string:
        make_string = root_make_string.replace('$Opt', '$Opt -std=c++0x -D HAVE_ZLIB')
    ROOT.gSystem.SetMakeSharedLib(make_string)
    ROOT.gROOT.ProcessLine(".L extern_Mars_mcore/izstream.h+O")
    ROOT.gROOT.ProcessLine(".L extern_Mars_mcore/fits.h+O")
    ROOT.gROOT.ProcessLine(".L extern_Mars_mcore/zfits.h+O")
    ROOT.gROOT.ProcessLine(".L extern_Mars_mcore/factfits.h+O")
    if "-std=c++0x" not in root_make_string:
        make_string = root_make_string.replace(
            '$Opt',
            "$Opt -std=c++0x "
            "-D HAVE_ZLIB "
            "-D'PACKAGE_NAME=\"PACKAGE_NAME\"' "
            "-D'PACKAGE_VERSION=\"PACKAGE_VERSION\"' "
            "-D'REVISION=\"REVISION\"' "
            )
    ROOT.gSystem.SetMakeSharedLib(make_string)
    ROOT.gROOT.ProcessLine(".L extern_Mars_mcore/DrsCalib.h+O")

    ROOT.gInterpreter.GenerateDictionary(
        "map<string,fits::Entry>",
        "map;string;extern_Mars_mcore/fits.h")
    ROOT.gInterpreter.GenerateDictionary(
        "pair<string,fits::Entry>",
        "map;string;extern_Mars_mcore/fits.h")
    ROOT.gInterpreter.GenerateDictionary(
        "map<string,fits::Table::Column>",
        "map;string;extern_Mars_mcore/fits.h")
    ROOT.gInterpreter.GenerateDictionary(
        "pair<string,fits::Table::Column>",
        "map;string;extern_Mars_mcore/fits.h")
    ROOT.gInterpreter.GenerateDictionary(
        "vector<DrsCalibrate::Step>",
        "vector;extern_Mars_mcore/DrsCalib.h")

#-------- USAGE -------------------------------------------------------------
if __name__ == '__main__' and len(sys.argv) < 3:
    print """ Usage:
    ----------------------------------------------------------------------
    To just build the shared object libs call:
    python pyfact.py build

    To build (if necessary) and open an example file
    python -i pyfact.py /path/to/data_file.zfits /path/to/calib_file.drs.fits.gz

    Any of the 3 file 'types': fits.gz    zfits    fits
    should be supported.

    To clean all of automatically built files, do something like:
    rm *.so *.d AutoDict_* extern_Mars_mcore/*.so extern_Mars_mcore/*.d
    ----------------------------------------------------------------------
    """
    sys.exit(1)

#--------  START OF PYFACT MODULE --------------------------------------------
path = os.path.dirname(os.path.realpath(__file__))
ROOT.gSystem.Load(path+'/AutoDict_map_string_fits__Entry__cxx.so')
ROOT.gSystem.Load(path+'/AutoDict_map_string_fits__Table__Column__cxx.so')
ROOT.gSystem.Load(path+'/AutoDict_pair_string_fits__Entry__cxx.so')
ROOT.gSystem.Load(path+'/AutoDict_pair_string_fits__Table__Column__cxx.so')
ROOT.gSystem.Load(path+'/AutoDict_vector_DrsCalibrate__Step__cxx.so')
ROOT.gSystem.Load(path+'/extern_Mars_mcore/fits_h.so')
ROOT.gSystem.Load(path+'/extern_Mars_mcore/izstream_h.so')
ROOT.gSystem.Load(path+'/extern_Mars_mcore/zfits_h.so')
ROOT.gSystem.Load(path+'/extern_Mars_mcore/factfits_h.so')
ROOT.gSystem.Load(path+'/extern_Mars_mcore/DrsCalib_h.so')
del path

env_name = 'FACT_DATA_BASE_DIR'
if env_name in os.environ:
    base_path = os.environ[env_name]
    raw_base = os.path.join(base_path, 'raw')
    aux_base = os.path.join(base_path, 'aux')
else:
    raw_base = os.path.join('fact', 'raw')
    aux_base = os.path.join('fact', 'aux')


def stub_to_auxfile(s, s2=None, basep=aux_base, basen="*"):
    """ FIXME :-)
        It's not yet clear how to generate aux file names from stubs, or iterables
    """
    if s2 is not None:
        s = "{0}_{1:03d}".format(s, s2)
    y = s[0:4]
    m = s[4:6]
    d = s[6:8]
    date = s[0:8]
    #r = s[9:12]
    return os.path.join(basep, y, m, d, date + '.' + basen + '.fits')


def make_path(s):
    """ make path to fits file

        s : string or iterable

        Fits files in pyfact do not need to be specified by full paths
        A so called stub path is sufficient. A valid stub path looks like:
        yyyymmdd_rrr
        or in case of a drs calibration file like
        yyyymmdd_rrr.drs

        In addition a file can be specified by a 2-element integer iterable,
        e.g. a list or a tuple, whose first element is the integer denoting the
        'night' : yyyymmdd
        the 2nd integer is the 'run' : rrr
        so valid specifiers are also:
        (yyyymmdd, rrr) or [yyyymmdd, rrr]

        the pyfact.Fits class can handle .fits, .fits.gz, and .fits.fz files.
        stubs dont's specify which one should be used, so we need to check
        which one exists, if more than one file exists fitting to the stub,
        the order of precendence is:
            * .fits    -- since access should be pretty fast
            * .fits.fz -- newer format should be mostly the case
            * .fits.gz -- still possible
    """
    import re
    import os
    if not isinstance(s, (str, unicode)):
        # s is (hopefully) some kind of iterable, so we read out two integers
        night = int(s[0])
        run = int(s[1])
        _s = "{0:08d}_{1:03d}".format(night, run)
    else:
        _s = s

    # so now s is some kind of string
    # maybe it is already a path:
    if os.path.isabs(_s):
        # s is already an absolute path.
        # maybe the file is not there, but we don't care about that
        return _s

    # it was not an absolute path, maybe it's a relative path, let's check if
    # it points to a file.
    elif os.path.exists(_s):
        # s is certainly a path, it even points to a file :-)
        return os.path.abspath(_s)

    # okay ... s was not a path, so let's generate a nice absolute path
    # from the stub, in case it is a stub ..
    elif re.match('\d{8}_\d{3}', _s):
        # s is a stub, so we can make a full path from it
        y = _s[0:4]
        m = _s[4:6]
        d = _s[6:8]
        #r = _s[9:12]
        start = os.path.join(raw_base, y, m, d, _s[:12] + '.fits')
        candidates = [start, start+'.fz', start+'.gz']
        for c in candidates:
            if os.path.exists(c):
                return c
        # if we reached this point, s was a stub, but we were not
        # able to find a file on the file system...
        # let's return our best guess, the '.fz' file path
        # the Fits.__init__() will complain, if it can't find the file.
        return candidates[1]
    else:
        raise IOError("{0} can't be used for path generation".format(s))


class Fits(object):
    """ General FITS file access

    Wrapper for factfits class from Mars/mcore/factfits.h
    (factfits might not be suited to read any FITS file out there, but certainly
    all FACT FITS files can be read using this class)
    """
    __module__ = 'pyfact'

    def __init__(self, path):
        """ Open fits file, parse header and setup table colums

        path : string : path to the fits file (.fits, .fits.gz, .zfits)
        """
        path = make_path(path)
        self._path = path
        self._current_row = None
        if not os.path.exists(path):
            raise IOError(path+' was not found')
        self.f = ROOT.factfits(path)
        self._make_header()
        self._setup_columns()

    def __repr__(self):
        return '{s.__class__.__name__}(path="{s._path}")'.format(s=self)

    def __str__(self):
        s = """Fits object:
        path = {s._path}
        row = {s._current_row}
        self.cols.keys = {col_names}
        """.format(s=self, col_names=self.cols.keys())
        return s

    def __len__(self):
        """ return the number of events, in case it's a FACT physics data file
        or simply the number of rows in case it's a slow data (so called aux) file.
        """
        return self.f.GetNumRows()

    def _make_header(self):
        """
        """
        str_to_bool = {'T': True, 'F': False}
        type_conversion = {'I': int, 'F': float, 'T': str, 'B': str_to_bool.__getitem__}

        self.header = {}
        self.header_comments = {}
        for key, entry in self.f.GetKeys():
            try:
                self.header[key] = type_conversion[entry.type](entry.value)
            except KeyError:
                raise IOError(
                    "Error: entry type unknown.\n "
                    "Is %s, but should be one of: [I,F,T,B]" % (entry.type))
            self.header_comments[key] = entry.comment

    def _setup_columns(self):
        """
        """
        col_type_to_np_type_map = {
            'L': 'b1',
            'A': 'a1',
            'B': 'i1',
            'I': 'i2',
            'J': 'i4',
            'K': 'i8',
            'E': 'f4',
            'D': 'f8'}
        self.cols = {}
        for key, col in self.f.GetColumns():
            self.cols[key] = np.zeros(col.num, col_type_to_np_type_map[col.type])
            if col.num != 0:
                self.f.SetPtrAddress(key, self.cols[key])

    def __iter__(self):
        return self

    def next(self, row=None):
        """
        """
        if row is None:
            if self._current_row is None:
                self._current_row = 0
            else:
                self._current_row += 1
            if self.f.GetNextRow() is False:
                self._current_row = None
                raise StopIteration
        else:
            row = int(row)
            self._current_row = row
            if self.f.GetRow(row) is False:
                self._current_row = None
                raise StopIteration
        return self


class AuxFile(Fits):
    """ easy(?) access to FACT aux files
    """
    __module__ = 'pyfact'

    def __init__(self, path, verbose=False):
        self._verbose = verbose
        super(AuxFile, self).__init__(path)

    def _setup_columns(self):
        col_type_to_np_type_map = {
            'L': 'b1',
            'A': 'a1',
            'B': 'i1',
            'I': 'i2',
            'J': 'i4',
            'K': 'i8',
            'E': 'f4',
            'D': 'f8'}
        self._cols = {}
        self.cols = {}

        N = len(self)
        for key, col in self.f.GetColumns():
            self._cols[key] = np.zeros(col.num, col_type_to_np_type_map[col.type])
            self.cols[key] = np.zeros((N, col.num), col_type_to_np_type_map[col.type])
            if col.num != 0:
                self.f.SetPtrAddress(key, self._cols[key])

        for i, row in enumerate(self):
            if self._verbose:
                try:
                    step = int(self._verbose)
                except:
                    step = 10
                if i % step == 0:
                    print "reading line", i, 'of', self.header['NAXIS2']

            for key in self._cols:
                self.cols[key][i, :] = self._cols[key]


class RawData(Fits):
    """ Special raw data FITS file access (with DRS4 calibration)

        During iteration the C++ method DrsCalibration::Apply is being called.
    """
    __module__ = 'pyfact'

    def __init__(self, data_path, calib_path):
        """ -constructor-
        *data_path*  : fits or fits.gz file of the data including the path
        *calib_path* : fits or fits.gz file containing DRS calibration data
        """
        super(RawData, self).__init__(data_path)
        self.cols['CalibData'] = np.zeros(self.cols['Data'].shape, np.float32)
        self.cols['CalibData2D'] = self.cols['CalibData'].reshape(self.header['NPIX'], -1)
        if not self.cols['CalibData2D'].base is self.cols['CalibData']:
            print "Error seomthing went wrong!"
        self.drs_calibration = ROOT.DrsCalibration()
        self.drs_calibration.ReadFitsImp(calib_path)

        self.drs_calibrate = ROOT.DrsCalibrate()
        self.list_of_previous_start_cells = []

    def __iter__(self):
        """ iterator """
        return self

    def next(self, row=None):
        """
        """
        super(RawData, self).next(row)

        self.drs_calibration.Apply(
            self.cols['CalibData'],
            self.cols['Data'],
            self.cols['StartCellData'],
            self.header['NROI']
            )

        for previous_start_cells in self.list_of_previous_start_cells:
            self.drs_calibrate.CorrectStep(
                self.cols['CalibData'],
                self.header['NPIX'],
                self.header['NROI'],
                previous_start_cells,
                self.cols['StartCellData'],
                self.header['NROI']+10)
            self.drs_calibrate.CorrectStep(
                self.cols['CalibData'],
                self.header['NPIX'],
                self.header['NROI'],
                previous_start_cells,
                self.cols['StartCellData'],
                3)
        self.list_of_previous_start_cells.append(self.cols['StartCellData'])
        if len(self.list_of_previous_start_cells) > 5:
            self.list_of_previous_start_cells.pop(0)

        for ch in range(self.header['NPIX']):
            self.drs_calibrate.RemoveSpikes3(self.cols['CalibData2D'][ch], self.header['NROI'])

        return self


if __name__ == '__main__':
    """ Example """
    print "Example for calibrated raw-file"
    f = RawData(sys.argv[1], sys.argv[2])

    print "number of events:", len(f)
    print "date of observation:", f.header['DATE']
    print "The files has these cols:", f.cols.keys()

    for counter, row in enumerate(f):
        print "Event Id:", row.cols['EventNum']
        print "shape of column 'StartCellData'", row.cols['StartCellData'].shape
        print "dtype of column 'Data'", row.cols['StartCellData'].dtype
        if counter > 3:
            break
    # get next row
    f.next()
    print "Event Id:", f.cols['EventNum']
    # get another row
    f.next(10)
    print "Event Id:", f.cols['EventNum']
    # Go back again
    f.next(3)
    print "Event Id:", f.cols['EventNum']