
"""


Basic support for reading the AFNI BRIK/HEAD file format.
This is quite quick-and-dirty, so please use with caution.


Floris van Vugt, April 2017

"""


import re
import numpy as np
import os




def read_header(fname):
    """ 
    Reads AFNI header file. That is, reads the HEAD from a BRIK/HEAD file pair.
    
    Arguments
    fname : the filename to be read (if it doesn't end in .HEAD this will be appended)
    
    Returns
    dict of key-values representing the header contents
    """ 
    
    if fname.endswith('.'):
        fname+="HEAD"
    if not fname.endswith('.HEAD'):
        fname+=".HEAD"
    
    # Read the file contents
    headf = open(fname,'r').read()
    
    # Now parse the contents
    remainder = headf[:]

    # No, this is not an insult, it's a pattern that matches the beginning of a chunk,
    # i.e. type=something, name=something_else, count=a_number
    chunkhead = re.compile(r'type\s*=\s*(integer|float|string)-attribute\s*name\s*=\s*(\w+)\s*count\s*=\s*(\d+)')

    header = {}
    types = {}

    m = re.search(chunkhead,remainder)

    type_regexp = {
        "integer":"[+-]?[0-9]",
        "float":"[-+]?(\d+([.,]\d*)?|[.,]\d+)([eE][-+]?\d+)?" # source: http://stackoverflow.com/questions/4703390/how-to-extract-a-floating-number-from-a-string
    }

    # While there is another chunk to be found...
    while m:

        # Parse the type/name/count fields
        endpos = m.end()
        (tp,name,count) = m.groups()
        count = int(count)
        header[name]=count
        types[name]=tp
        #print(m.start(),tp,name,count)

        # The rest of the file...
        remainder = remainder[endpos:]

        # Now read the actual contents
        if tp=='integer' or tp=="float":
            # Set up a regexp that will capture the next "count" ints
            contents = re.match(r'(\s*(%s)+){%i}'%(type_regexp[tp],count),remainder)
            cast = int if tp=="integer" else float
            if contents:
                values = [ cast(i) for i in contents.group().split() ]
                header[name]=values if count>1 else values[0]
            else:
                raise ValueError("Failed to parse contents for %s"%name)
                
            remainder = remainder[contents.end():]

        elif tp=="string":
            contents = re.match(r'\s*\'(.{%i})~'%(count-1),remainder,re.DOTALL)
            if contents:
                header[name]=contents.group(1)
            else:
                raise ValueError("Failed to parse contents for %s"%name)
                
            remainder = remainder[contents.end():]
            
        else:
            raise ValueError("Unknown data type %s for %s"%(tp,name))

            
        # Set up for the next iteration of the loop
        m = re.search(chunkhead,remainder)
        
    return header #,types







def read_brik(fname,header):
    """ Reads BRIK file. Presupposes that we have parsed the .HEAD file and 
    supplied at least the relevant portion of it.
    
    Arguments
    fname : the filename to be read (.BRIK will be added if necessary)
    header : the header information, a set of key-values
    
    Returns
    array containing the data
    """
    
    if fname.endswith('.'):
        fname+="BRIK"
    if not fname.endswith('.BRIK'):