#! /usr/bin/env python
# -*- coding: latin-1 -*-

import os
from os.path import split, splitext, exists, join
from numpy import *
from struct import *


def GetNumDataType(typ):
    if   typ == "s": return 'S'
    elif typ == "b": return 'i1'
    elif typ == "h": return 'i2'
    elif typ == "i": return 'i4'
    elif typ == "B": return 'u1'
    elif typ == "H": return 'u2'
    elif typ == "I": return 'u4'
    elif typ == "f": return 'f4'
    elif typ == "d": return 'f8'
    else           : return None


def GetDataType(typ):
    if   typ == "C" : return 's'
    elif typ == "I1": return 'b'
    elif typ == "I2": return 'h'
    elif typ == "I4": return 'i'
    elif typ == "U1": return 'B'
    elif typ == "U2": return 'H'
    elif typ == "U4": return 'I'
    elif typ == "F4": return 'f'
    elif typ == "F8": return 'd'
    else            : return None


def GetRecordDesc(rec_pth):
    f = open(rec_pth, "r")
    lines = f.read().splitlines()
    rec_dsc = []
    for ll in lines:
        l = ll.strip()
        if len(l) == 0: continue
        if l[0] == "#": continue
        nam, typ, lth, dim = l.split()
        dims = map(int, dim.split(","))
        if typ != "C": typ += lth
        rec_dsc.append([nam, GetDataType(typ), int(lth), dims, prod(array(dims))])
    f.close()
    return rec_dsc


def GetRecord(f, dsc, bigend, raw=0):
    rec_flds = []
    for d in dsc:
        if bigend: fmt = ">"
        else: fmt = "<"
        if d[1] == 's':
            for i in range(int(d[4])): fmt += "%d%s" % (d[2], d[1])
        else:
            fmt += "%d%s" % (d[4], d[1])
        fld = unpack(fmt, f.read(d[2]*d[4]))
        if raw: rec_flds.append(fld)
        elif d[4] == 1: rec_flds.append(fld[0])
        else: rec_flds.append(reshape(array(fld), d[3]))
    return rec_flds


def PrintRecord(flds, dsc):
    for d in range(len(dsc)):
        print dsc[d],' : '
        print flds[d]


def GetDsetIndex(dsc, sdname):
    for i in range(len(dsc)):
        if dsc[i][0] == sdname: return i
    return -1


def GetRawData(file, hdr_files, rec_file, bigend):

    # Récupération des descripteurs de headers
    hdr_dscs = []
    if hdr_files is not None:
        if not isinstance(hdr_files, list): hdr_files = [hdr_files]
##         hdr_files = hdr_files.split()
        for h in hdr_files: hdr_dscs.append(GetRecordDesc(h))

    # Récupération du descripteur de record
    rec_dsc = GetRecordDesc(rec_file)

    # Ouverture du fichier
    f = open(file, "r")
    
    # Récupération des données des headers
    headers = []
    for hdr_dsc in hdr_dscs: headers.append(GetRecord(f, hdr_dsc, bigend))

    # Récupération des données des records
    records = []
    records.append(GetRecord(f, rec_dsc, bigend))
    try:
        while True: records.append(GetRecord(f, rec_dsc, bigend))
    except:
        f.close()

    # On retourne le contenu du fichier
    return (hdr_dscs, headers, rec_dsc, records)


def MakeDataArray(rec_dsc, records, isds):
    try:
        dims = [len(records)]
        if sum(rec_dsc[isds][3]) > 1: dims += rec_dsc[isds][3]
        rec_arr = zeros(dims, dtype(GetNumDataType(rec_dsc[isds][1])))
        for i in range(len(records)): rec_arr[i] = array(records[i][isds])
        return rec_arr
    except:
        return None


def ReadRawFile(file, hdr_files, rec_file, bigend, dset_names=None):

    hdr_dscs, headers, rec_dsc, records = GetRawData(file, hdr_files, rec_file, bigend)

    # Fabrication des données désirées
    hdrs, dsets = (dict(), dict())

    for idsc in range(len(hdr_dscs)):
        for h in range(len(hdr_dscs[idsc])): hdrs[hdr_dscs[idsc][h][0]] = headers[idsc][h]

    if dset_names is None: dset_names = [r[0] for r in rec_dsc]
    elif not isinstance(dset_names, list): dset_names = [dset_names]
    for nam in dset_names: dsets[nam] = MakeDataArray(rec_dsc, records, GetDsetIndex(rec_dsc, nam))
        
    # On retourne les données cherchées
    return hdrs, dsets