<filename>pgfmutate.py
"""Author: <NAME>

The purpose of this library is to provide a way to mutate PGF files from
the HEKA PatchMaster program in simple ways.

Specifically, this library provides a way to use a "template" PGF file
to generate another PGF file with slightly different properties. The user
can add records to any of the 3 mutable (non-root) levels of the tree
structure by copying and mutating existing records, as well as being able
to mutate existing records in-place.

NOTE: It should be noted that this library _can_ be used to generate
PGF files from scratch. However, this can be quite difficult due to the
complexity of the records at each of the tree levels, and therefore this
will not be the primary intent of the library and convenience functions
relating specifically to this purpose only will not be created.

Specifications for the PGF file format are taken from the official format
specification as described in the file accessed by
ftp://server.hekahome.de/pub/FileFormat/Patchmasterv9/FileFormat_v9.zip
      and
http://www.heka.com/downloads/software/manual/m_patchmaster.pdf
      as of January 2017.

This library was developed and tested in Python 2.7. It is highly likely to
work in other python versions as well (including possibly 3.x), but I can
make no promises.

This library does not include much in the way of error checking. If something
goes wrong (e.g. the file format is invalid or the file cannot be read),
library utilities will fail fast and hard, bubbling up all exceptions
generated by whatever python standard library call fails."""

# --------    PGF TREE; USE THIS TO DO STUFF!    --------
import struct
import copy
class PGFTree:
    """Represent the tree structure of a PGF file."""

    def __init__(self, path=None):
        """Initialize structure.

        If a path is given, tries to load from file."""
        self.tree = []

        if path:
            self.load(path)

    def load(self, path):
        """Loads structure from file."""
        inpf = INPFAutoEndian(path)

        # figure out file bit order (endian)
        magic = _unpackone('i', inpf.read(4))
        if magic == 0x054726565:
            inpf.must_reverse = False
        elif magic == 0x65657254:
            inpf.must_reverse = True
        else:
            raise RuntimeError('Not a PGF file!')

        # do checks
        nlvls = _unpackone('i', str(inpf.read(4)))
        if nlvls != 4:
            raise RuntimeError('Not a PGF file! (Too many levels.)')
        for lvl_i in range(4):
            lvlsz = _unpackone('i', str(inpf.read(4)))
            if lvlsz != (584, 248, 400, 80)[lvl_i]:
                raise RuntimeError('Unsupported PGF version! (Wrong level sizes.)')

        # read records
        # my oh my, this is a mess...
        lvl_progression_classes = \
            [RootRecord, StimulationRecord, ChannelRecord, StimSegmentRecord]
        self.tree.append([RootRecord(inpf), []])
        last_recrecs = [self.tree[0]]
        curr_recrecs = []
        for lvl in range(1, 4):
            for recrec in last_recrecs:
                for x in range(recrec[0].num_children):
                    recrec[1].append([lvl_progression_classes[lvl](inpf), []])
                    curr_recrecs.append(recrec[1][-1])
            last_recrecs = copy.copy(curr_recrecs)
            curr_recrecs = []

        inpf.close()

    def save(self, path):
        """Saves structure to file."""
        fout = open(path, 'wb')

        # write initial stuff
        fout.write(struct.pack('i', 0x054726565))
        fout.write(struct.pack('i', 4))
        fout.write(struct.pack('iiii', 584, 248, 400, 80))
        # write records in BFS order
        bfs_order = collections.deque()
        bfs_order.append(self.tree[0])
        while bfs_order:
            nxt_recrec = bfs_order.popleft()
            bfs_order.extend(nxt_recrec[1])
            fout.write(str(nxt_recrec[0]))
            fout.write(struct.pack('i', len(nxt_recrec[1])))

        fout.close()

    def get_nav(self):
        """Get a PGFTNavigator that is linked to this PGFTree."""
        return PGFTNavigator(self)

    def copy_record(self, where):
        """Make a copy of a record.

        Pass nav.here for the 'where' argument, where nav is a PGFTNavigator.
        Return nav location (for nav.here) of new record.

        NOTE: This does NOT copy any child records!"""
        this = self.__navloc_get(where)
        last = self.__navloc_get(where[:-1])
        last[1].append([copy.deepcopy(this[0]), []])
        last[0].num_children += 1
        # don't worry about the child count on last[0] for now; fix it
        # later in the output phase.
        where_cpy = copy.copy(where)
        where_cpy[-1] = len(last[1])-1
        return where_cpy

    def del_record(self, where):
        """Delete a record.

        NOTE: This deletes all child records!"""
        parent = self.__navloc_get(where[:-1])
        del parent[1][where[-1]]
        parent[0].num_children -= 1

    def transplant_record(self, src, dst):
        """Move a record between parents.

        'src' should lead to the record that should be moved.
        'dst' should lead to the new parent (of the record that's being moved)."""
        dst_rec = self.__navloc_get(dst)
        dst_rec[1].append(self.__navloc_get(src))
        dst_rec[0].num_children += 1
        self.del_record(src)

    def get_record(self, navloc):
        """Get a record."""
        return self.__navloc_get(navloc)[0]

    def __navloc_get(self, navloc):
        """Get a record-record."""
        this_recrec = self.tree[0]
        for nxt in navloc:
            this_recrec = this_recrec[1][nxt]
        return this_recrec

class PGFTNavigator:
    """Class to aid navigation through a PGFTree structure."""

    class Clipboard:
        def __init__(self, v):
            self.val = v

    def __init__(self, pgft):
        self.tree = pgft
        self.here = []
        self.clipboard = self.Clipboard(None)

    def ascend(self):
        """Go up one level."""
        if self.here:
            self.here.pop()
        return self
    a = ascend

    def descend(self, where):
        """Go down one level."""
        self.here.append(where)
        return self
    d = descend

    def get(self):
        """Get the record pointed to by this navigator."""
        return self.tree.get_record(self.here)
    g = get

    def delete(self):
        """Deletes the record pointed to by this navigator, and moves this
        navigator one level up (to the parent of the deleted record)."""
        self.tree.del_record(self.here)
        self.ascend()
        return self
    dl = delete

    def copy(self):
        """'Copy' the record pointed to by this navigator to this navigator's
        'clipboard', similar to the ctrl-C behavior on modern operating systems.

        Only one item may be on the clipboard at any one time; therefore, if
        the clipboard is already be occupied, this function will replace its
        contents."""
        self.clipboard.val = copy.copy(self.here)
        return self
    c = copy

    def paste(self):
        """'Paste' the contents of this navigator's clipboard to a location that
        is a child of the record pointed to by this navigator.

        Does NOT clear the clipboard.
        If the clipboard is empty, does nothing."""
        if not self.clipboard.val:
            return self
        self.tree.transplant_record(self.tree.copy_record(self.clipboard.val), self.here)
        return self
    p = paste

    def branch(self):
        """Create a 'branch' navigator; copies this navigator.

        The copy shares the same clipboard as this one!"""
        return copy.copy(self)
    b = branch

    def __copy__(self):
        other = PGFTNavigator(self.tree)
        other.here = copy.copy(self.here)
        other.clipboard = self.clipboard # just copy clipboard reference!
        return other

import collections
class Record:
    """Record classes represent the records found in PGF files at varying levels
    ("root", "group", "series", and "sweep") of the tree structure.

    All record parameters are stored as member variables in the record classes
    as instances of the most appropriate python types; e.g., parameters listed
    as INT16 and INT32 in the record spec will be stored as python ints, while
    parameters listed as REAL and LONGREAL will be stored as python floats.
    Member variable names will be _identical_ to those found in the PGF file
    format specification mentioned above.

    To mutate record parameters, simply set the corresponding member variable."""

    def __init__(self, inpf):
        """Initialize a PGF record from an input file."""
        if not inpf:
            return
        sz_read = 0 # tally for bytes read so that we can check it later.
        for param in self.params:
            sz_read += self.params[param].size
            setattr(self, param, self.params[param].load(inpf))
        if sz_read != self.rec_size:
            raise RuntimeError(
                'Oops, the pgfmutate developer messed up! (%s)'%self.__class__)

        # num children is essentially part of all records, but is not
        #     pedantically part of the record itself. Therefore, we input it
        #     separately here instead of letting the classes put it in their
        #     params dictionary.
        self.num_children = TY_INT32.load(inpf)

    def __str__(self):
        """Get the binary representation of this record."""
        rep = ''
        for param in self.params:
            rep += self.params[param].to_raw(getattr(self, param))
        return rep

    def __deepcopy__(self, memo):
        """Produce a deep copy of this record."""
        other = self.__class__(None)
        for attr in vars(self):
            setattr(other, attr, copy.copy(getattr(self, attr)))
        return other

# --------    PGF SPECIFICATION VARIABLE HANDLERS    --------
# The sizes (in bytes) of different variable types as described by the PGF
# specification.
class PGFTY(object):
    """A PGF Type Handler."""
    @classmethod
    def load(cls, inpf):
        """Load this value from file."""
        return _unpackone(cls.conv, inpf.read(cls.size))
    @classmethod
    def to_raw(cls, val):
        """Convert value to binary representation."""
        return struct.pack(cls.conv, val)
class TY_BYTE(PGFTY): size, conv = 1, 'B'
class TY_CHAR(PGFTY): size, conv = 1, 'c'
class TY_ENUM(PGFTY): size, conv = 1, 'B'
class TY_BOOLEAN(PGFTY): size, conv = 1, '?'
class TY_INT16(PGFTY): size, conv = 2, 'h'
class TY_CARD16(PGFTY): size, conv = 2, '2s'
class TY_SET16(PGFTY): size, conv = 2, '2s'
class TY_INT32(PGFTY): size, conv = 4, 'i'
class TY_CARD32(PGFTY): size, conv = 4, '4s'
class TY_REAL(PGFTY): size, conv = 4, 'f'
class TY_LONGREAL(PGFTY): size, conv = 8, 'd'
def TY_ARRAY(ty, sz):
    template = \
    '''class TY_ARRAY_{ty}_{sz}(PGFTY):
        size = {totsz}
        innercls = {ty}
        len = {sz}
        @classmethod
        def load(cls, inpf):
            l = []
            for x in range(cls.len):
                l.append(cls.innercls.load(inpf))
            return l
        @classmethod
        def to_raw(cls, l):
            s = ''
            for val in l:
                s += cls.innercls.to_raw(val)
            return s
        @classmethod
        def __len__(cls):
            return cls.len
    '''.format(ty=ty.__name__, sz=str(sz), totsz=str(sz*ty.size))
    exec template in globals()
    return globals()['TY_ARRAY_%s_%s'%(ty.__name__, str(sz))]
def DEF_TY_STRING(sz):
    template = \
    '''class TY_STRING{sz}(TY_ARRAY(TY_CHAR, {sz})):
        @classmethod
        def load(cls, inpf):
            return ''.join(super(cls, cls).load(inpf))
    '''.format(sz=str(sz))
    exec template in globals()
    return globals()['TY_STRING%s'%str(sz)]
DEF_TY_STRING(8)
DEF_TY_STRING(32)
DEF_TY_STRING(128)

class RootRecord(Record):
    """Represent a PGF root record (level 0/root).

    The root record is the only record that cannot be duplicated, as the file
    specification does not support the use of multiple root nodes."""
    params = collections.OrderedDict([
        ('roVersion', TY_INT32),
        ('roMark', TY_INT32),
        ('roVersionName', TY_STRING32),
        ('roMaxSamples', TY_INT32),
        ('roFiller1', TY_INT32),
        ('roParams', TY_ARRAY(TY_LONGREAL, 10)),
        ('roParamText', TY_ARRAY(TY_STRING32, 10)),
        ('roReserved', TY_STRING128),
        ('roFiller2', TY_INT32),
        ('roCRC', TY_CARD32)])
    rec_size = 584

    def __deepcopy__(self):
        # disable copying of root nodes
        raise NotImplementedError

class StimulationRecord(Record):
    """Represent a PGF stimulation record (level 1/group)."""
    params = collections.OrderedDict([
        ('stMark', TY_INT32),
        ('stEntryName', TY_STRING32),
        ('stFileName', TY_STRING32),
        ('stAnalName', TY_STRING32),
        ('stDataStartSegment', TY_INT32),
        ('stDataStartTime', TY_LONGREAL),
        ('stSampleInterval', TY_LONGREAL),
        ('stSweepInterval', TY_LONGREAL),
        ('stLeakDelay', TY_LONGREAL),
        ('stFilterFactor', TY_LONGREAL),
        ('stNumberSweeps', TY_INT32),
        ('stNumberLeaks', TY_INT32),
        ('stNumberAverages', TY_INT32),
        ('stActualAdcChannels', TY_INT32),
        ('stActualDacChannels', TY_INT32),
        ('stExtTrigger', TY_BYTE),
        ('stNoStartWait', TY_BOOLEAN),
        ('stUseScanRates', TY_BOOLEAN),
        ('stNoContAq', TY_BOOLEAN),
        ('stHasLockIn', TY_BOOLEAN),
        ('stOldStartMacKind', TY_CHAR),
        ('stOldEndMacKind', TY_BOOLEAN),
        ('stAutoRange', TY_BYTE),
        ('stBreakNext', TY_BOOLEAN),
        ('stIsExpanded', TY_BOOLEAN),
        ('stLeakCompMode', TY_BOOLEAN),
        ('stHasChirp', TY_BOOLEAN),
        ('stOldStartMacro', TY_STRING32),
        ('stOldEndMacro', TY_STRING32),
        ('sIsGapFree', TY_BOOLEAN),
        ('sHandledExternally', TY_BOOLEAN),
        ('stFiller1', TY_BOOLEAN),
        ('stFiller2', TY_BOOLEAN),
        ('stCRC', TY_CARD32)])
    rec_size = 248

class ChannelRecord(Record):
    """Represent a PGF channel record (level 2/series)."""
    params = collections.OrderedDict([
        ('chMark', TY_INT32),
        ('chLinkedChannel', TY_INT32),
        ('chCompressionFactor', TY_INT32),
        ('chYUnit', TY_STRING8),
        ('chAdcChannel', TY_INT16),
        ('chAdcMode', TY_BYTE),
        ('chDoWrite', TY_BOOLEAN),
        ('stLeakStore', TY_BYTE),
        ('chAmplMode', TY_BYTE),
        ('chOwnSegTime', TY_BOOLEAN),
        ('chSetLastSegVmemb', TY_BOOLEAN),
        ('chDacChannel', TY_INT16),
        ('chDacMode', TY_BYTE),
        ('chHasLockInSquare', TY_BYTE),
        ('chRelevantXSegment', TY_INT32),
        ('chRelevantYSegment', TY_INT32),
        ('chDacUnit', TY_STRING8),
        ('chHolding', TY_LONGREAL),
        ('chLeakHolding', TY_LONGREAL),
        ('chLeakSize', TY_LONGREAL),
        ('chLeakHoldMode', TY_BYTE),
        ('chLeakAlternate', TY_BOOLEAN),
        ('chAltLeakAveraging', TY_BOOLEAN),
        ('chLeakPulseOn', TY_BOOLEAN),
        ('chStimToDacID', TY_SET16),
        ('chCompressionMode', TY_SET16),
        ('chCompressionSkip', TY_INT32),
        ('chDacBit', TY_INT16),
        ('chHasLockInSine', TY_BOOLEAN),
        ('chBreakMode', TY_BYTE),
        ('chZeroSeg', TY_INT32),
        ('chStimSweep', TY_INT32),
        ('chSine_Cycle', TY_LONGREAL),
        ('chSine_Amplitude', TY_LONGREAL),
        ('chLockIn_VReversal', TY_LONGREAL),
        ('chChirp_StartFreq', TY_LONGREAL),
        ('chChirp_EndFreq', TY_LONGREAL),
        ('chChirp_MinPoints', TY_LONGREAL),
        ('chSquare_NegAmpl', TY_LONGREAL),
        ('chSquare_DurFactor', TY_LONGREAL),
        ('chLockIn_Skip', TY_INT32),
        ('chPhoto_MaxCycles', TY_INT32),
        ('chPhoto_SegmentNo', TY_INT32),
        ('chLockIn_AvgCycles', TY_INT32),
        ('chImaging_RoiNo', TY_INT32),
        ('chChirp_Skip', TY_INT32),
        ('chChirp_Amplitude', TY_LONGREAL),
        ('chPhoto_Adapt', TY_BYTE),
        ('chSine_Kind', TY_BYTE),
        ('chChirp_PreChirp', TY_BYTE),
        ('chSine_Source', TY_BYTE),
        ('chSquare_NegSource', TY_BYTE),
        ('chSquare_PosSource', TY_BYTE),
        ('chChirp_Kind', TY_BYTE),
        ('chChirp_Source', TY_BYTE),
        ('chDacOffset', TY_LONGREAL),
        ('chAdcOffset', TY_LONGREAL),
        ('chTraceMathFormat', TY_BYTE),
        ('chHasChirp', TY_BOOLEAN),
        ('chSquare_Kind', TY_BYTE),
        ('chFiller1', TY_ARRAY(TY_CHAR, 5)),
        ('chSquare_BaseIncr', TY_LONGREAL),
        ('chSquare_Cycle', TY_LONGREAL),
        ('chSquare_PosAmpl', TY_LONGREAL),
        ('chCompressionOffset', TY_INT32),
        ('chPhotoMode', TY_INT32),
        ('chBreakLevel', TY_LONGREAL),
        ('chTraceMath', TY_STRING128),
        ('chFiller2', TY_INT32),
        ('chCRC', TY_CARD32)])
    rec_size = 400

class StimSegmentRecord(Record):
    """Represent a PGF stimulation segment record."""
    params = collections.OrderedDict([
        ('seMark', TY_INT32),
        ('seClass', TY_BYTE),
        ('seStoreKind', TY_BYTE),
        ('seVoltageIncMode', TY_BYTE),
        ('seDurationIncMode', TY_BYTE),
        ('seVoltage', TY_LONGREAL),
        ('seVoltageSource', TY_INT32),
        ('seDeltaVFactor', TY_LONGREAL),
        ('seDeltaVIncrement', TY_LONGREAL),
        ('seDuration', TY_LONGREAL),
        ('seDurationSource', TY_INT32),
        ('seDeltaTFactor', TY_LONGREAL),
        ('seDeltaTIncrement', TY_LONGREAL),
        ('seFiller1', TY_INT32),
        ('seCRC', TY_CARD32),
        ('seScanRate', TY_LONGREAL)])
    rec_size = 80

# --------    INPUT FILE WRAPPER    --------
# used to automatically reverse endian if needbe
# also returns bytearrays instead of strings
class INPFAutoEndian:
    """File wrapper that reverses bits on read if the must_reverse attribute is
    set."""
    def __init__(self, path):
        self.fh = open(path, 'rb')
        self.must_reverse = False
    def __del__(self):
        self.close()
    def close(self):
        self.fh.close()
    def read(self, nbytes):
        data =  bytearray(self.fh.read(nbytes))
        return data[::-1] if self.must_reverse else data

def _unpackone(dtype, data):
    """Convenience function for when only one value is being grabbed with
    struct.unpack()."""
    return struct.unpack(dtype, data)[0]
