#!/usr/bin/env python

import sys

from optparse import OptionParser

import tables

types = {
    'name': 'name',
    'i32': 'tables.Int32Col()',
    'i64': 'tables.Int64Col()',
    'f32': 'tables.Float32Col()',
    'f64': 'tables.Float64Col()',
    }

converters = {
    'name': 'name',
    'i32': int,
    'i64': long,
    'f32': float,
    'f64': float,
    }


def main():
    op = OptionParser(usage="txt2pyt.py [opts] file1.txt [file2.txt [...]]")

    opts, args = op.parse_args()

    if len(args) == 0:
        op.print_usage()
        return 1
            
    for fname in args:
        if fname[-4:].lower() != '.txt':
            sys.stderr.write("Skipping non-text file `" + fname + "'...\n")
            sys.stderr.flush()
            continue

        infp = open(fname, 'r')

        header = infp.readline()

        if header[0] != '#':
            sys.stderr.write("File `" + fname + "' does not begin with `#', skipping...\n")
            sys.stderr.flush()
            continue

        header = header.lstrip('#').lstrip().strip()
        cols = header.split()

        ntname   = "nt"
        col_defs = dict()
        col_conv = dict()
        col_names = list()
        for col in cols:
            try:
                name, dtype = col.split(':')
            except IndexError, e:
                sys.stderr.write("Error parsing ntuple definition: " + e + '\n')
                sys.stderr.write("Aborting...")
                return 2
                
            if name == 'ntname':
                ntname = dtype
                continue

            col_names.append(name)
            
            try:
                col_defs[name] = types[dtype]
                col_conv[name] = converters[dtype]
            except KeyError, e:
                sys.stderr.write("Error parsing ntuple definition: ")
                sys.stderr.write("Unknown datatype `" + dtype + "'\n")
                sys.stderr.write("Choices are: " + str(types.keys()) + "\n")
                sys.stderr.write("Aborting...\n")
                sys.stderr.flush()
                return 3
        
#        class_str = "import tables\n\n"

        class_str = "class " + ntname + "(tables.IsDescription):\n"
        for colname, coltype in col_defs.iteritems():
            class_str += "    " + colname + " = " + coltype + "\n"
        class_str += '\n'

        classfp = open(fname[:-4] + '.py', 'w')
        classfp.write(class_str)
        classfp.close()

        try:
            exec(class_str)
        except SyntaxError, e:
            sys.stderr.write("Error parsing ntuple definition: ")
            sys.stderr.write(str(e) + "'\n")
            sys.stderr.write("Aborting...\n")
            sys.stderr.flush()
            return 4

        outfp = tables.openFile(fname[:-4] + '.h5', 'w', 
                                'generated by txt2pyt.py')
        tblstr = "outfp.createTable(outfp.root, '" + ntname.lower() + "', " + ntname + ", " + "'generated by txt2pyt.py')"

        tbl = eval(tblstr)

        for line in infp:
            if not line.strip():
                continue
            
            txt_cols = line.strip().split()

            if len(txt_cols) != len(col_names):
                sys.stderr.write("Error in ntuple data: ")
                sys.stderr.write("Number of data columns does not equal number of columns in definition.\n")
                sys.stderr.write("Aborting...\n")
                sys.stderr.flush()
                return 5
            
            for i in xrange(len(col_names)):
                tbl.row[col_names[i]] = col_conv[col_names[i]](txt_cols[i])

            tbl.row.append()

        infp.close()

        tbl.flush()

        outfp.close()


if __name__ == "__main__":
    sys.exit(main())