#!/usr/bin/python2

#raster.py - a Bippy module - 17 Aug. 2004 - 0.0.2
#an original script written by Bill Allen for Python 2+
#this script has no price, limitations, or warranty - use at your own risk

"""
This module facilitates handling raster image data regardless of the input
  and output file formats, and includes functions to output to the utilitarian
  PGM and RAW file formats.
Please report problems to the author:
  <ballen@3dartist.com> SUBJECT: Python
  but do NOT attach files to messages without checking first.

file write functions index
--------------------------
pgmWrite - write a binary portable graphics grayscale file
rawWrite - write RAW file as-is at current bit depth & endianness

modification functions index
----------------------------
make8bit - equalize/convert any image raster type to 8-bit unsigned
stitch - put left & right frames into a stereo pair

class index
-----------
rasterDummyClass - empty raster class for passing only file header info
rasterObjectClass - basic class for all raster objects

history
0.0.1 - first version 6 April 2004
0.0.2 - 17 Aug. 2004: added rasterDummyClass
      - added class-type checking to write functions & make8bit
"""

#--- modules & constants ---

import array, os, sys
from string import split

thisDir = os.path.split(sys.argv[0])[0]
if not thisDir in sys.path: sys.path.append(thisDir)
from common import arrayTypeCode, fixPath

#--- file write functions ---

def pgmWrite(rasObj,pn,fn='',rem=None):
    """
    This function expects a rasterObjectClass with unsigned 8-bit image data,
      but, if the data is other than unsigned 8-bit, will on its own make a
      suitable copy of the data.
    pn - path name, required, path name, can include file name
    fn - file name if not part of path name
    rem - remark, optional, can be a string or a list of strings
    """
    errPfx = 'bippy:raster:pgmWrite: '
    if not rasObj: return 0,errPfx+'no raster object received'
    try:                                #make sure we have correct object type
        L = rasObj.lineage
        if L != 'bippy:raster:rasterObjectClass':
            return 0,errPfx+'wrong class, '+L
    except: return 0,errPfx+'object not recognized'
    if not pn: return 0,errPfx+'need path name'
    if fn: pn = fixPath(pn,fn)
    if rasObj.typecode() != 'B':
        return 0,errPfx+'need '+arrayTypeCode['B']          \
                 +' raster object, got '+rasObj.typecode()  \
                 +': '+arrayTypeCode[rasObj.typecode()]
    h = rasObj.height()
    max = 0
    w = rasObj.width()
    for y in range(h):
        for x in range(w):
            if rasObj.data[y][x] > max: max = rasObj.data[y][x]
    try: f = open(pn,'wb')
    except: return 0,errPfx+'couldn\'t open file '+pn
    f.write('P5\n')
    if not rem:
        rem = []
        for s in rasObj.history:
            if s:
                if s[0] == '>': rem.append(s)
    if rem:
        if type(rem) == type('string'): f.write('#'+rem+'\n')
        else:
            for s in rem: f.write('#'+s+'\n')
    f.write(str(w)+' '+str(h)+'\n'+str(max)+'\n')
    for y in range(h): rasObj.data[y].tofile(f)
    f.flush()
    f.close()
    return 1,'>'+errPfx+pn
    #end def pgmWrite

def rawWrite(rasObj,pn,fn='',embed=1):
    """
    rasObj is required, expected to be a rasterObjectClass object
    pn - path name, required, can include file name
    fn - file name, used only if file name is not part of path name
    embed - optional, put bit dimensions into output file name
    """
    errPfx = 'bippy:raster:rawWrite: '
    if not rasObj: return 0
    try:                                #make sure we have correct object type
        L = rasObj.lineage
        if L != 'bippy:raster:rasterObjectClass': return 0
    except: return 0
    if fn: pn = fixPath(pn,fn)
    if embed:                           #embed dimensions in file name
        code = rasObj.typecode()
        if code == 'B': s = 'x8.raw'
        elif code == 'H': s = 'x16.raw'
        else:
            print errPfx + 'code = ' + code + '?'
            return 0
        pn = pn[:-4] + '-' + str(rasObj.width()) + 'x'  \
             + str(rasObj.height()) + s
    try:
        f = open(pn,'wb')
        for y in range(rasObj.height()): rasObj.data[y].tofile(f)
        f.flush()
        f.close()
    except: return 0
    rasObj.history.append('saveRaw to '+pn)
    return 1
    #end def rawWrite

#--- modification functions ---

def make8bit(rasObj=None,max=None,min=None):
    """
    This function presumes unsigned data type. For signed, need to change
      the data, or to make sure that the algorithm will work with the data.
    Values max and min are high/low threshold levels to help control range
      (an untested feature).
    """
    errPfx = 'bippy:raster:make8bit: '
    if not rasObj: return None,errPfx+'no raster object received'
    try:                                #make sure we have correct object type
        L = rasObj.lineage
        if L != 'bippy:raster:rasterObjectClass':
            return None,errPfx+'wrong class, '+L
    except: return None,errPfx+'object not recognized'
    if rasObj.typecode() == 'B':
        return rasObj.copy(),errPfx+'is already an unsigned 8-bit raster object'
    high = 0
    low = sys.maxint                    #maximum integer value on this system
    w = rasObj.width()
    for y in range(rasObj.height()):
        for x in range(w):
            if rasObj.data[y][x] < low: low = rasObj.data[y][x]
            elif rasObj.data[y][x] > high: high = rasObj.data[y][x]
    history = ['>'+errPfx[:-2],
               ' type in '+rasObj.typecode()+' = '
                 +arrayTypeCode[rasObj.typecode()]
                 +' low='+str(low)+' high='+str(high)]
    r = 255.0 / float(high-low)         #ratio to apply to normalized value
    floor = low
    high = 0
    low = 255                           #maximum 8-bit value
    image = []                          #new values
    for y in range(rasObj.height()):
        a = array.array('B')
        for x in range(w):
            v = int(r*(rasObj.data[y][x]-floor))  #ratio applied to normalized
            if v < low: low = v
            if v > high: high = v
            if v > 255:                 #for safety & debugging
                print errPfx,'v = ',v
                v = 255
            a.append(v)
        image.append(a)
    newObj = rasterObjectClass(rasObj.parent,image,rasObj.history,
                               caller=errPfx[:-2])
    history.append(' type out B = '+arrayTypeCode['B']
                   +' low='+str(low)+' high='+str(high))
    for s in history: newObj.history.append(s)
    return newObj,'>'+errPfx
    #end make8bit

def stitch(rasObjL,rasObjR,verbose=0):
    """
    Stitches two rasterClassObjects into two, used for creating stereo pairs.
    Note: will probably fail if used with objects of dissimilar dimensions,
      and isn't set up to adjust for disparate tonal ranges (need to first
      do a group equalize range calculation).
    """
    errPfx = 'raster:stitch'
    for rasObj in [rasObjL,rasObjR]:
        if rasObj.typecode() != 'B':
            rasObj,msg = make8bit(rasObj)   #new 8-bit, equalized object
    image = []
    for y in range(rasObjR.height()):
        a = array.array('B')
        a = rasObjL.data[y]
        if y%2: a.append(0)  #***DEVnote: This is supposed to make a B&W dotted
        else: a.append(255)  #   vertical line, but creates black only - why?
        a += rasObjR.data[y]
        image.append(a)
    pn = rasObjL.filePath()
    pn = pn[:-4] + '-r' + pn[-4:]
    history = ['>'+errPfx,'>path '+pn]
    history.append('-left image---')
    history += rasObjL.history
    history.append('-right image---')
    history += rasObjR.history
    rasObj = rasterObjectClass(rasObjR.parent,image,history,errPfx)
    return rasObj,''
    #end def stitch

#--- classes ---

class rasterDummyClass:
    """
    This class is used when only the header is wanted and no raster data.
    """
    def __init__(self,parent=None,history=[],caller='?'):
        self.lineage = 'bippy:raster:rasterDummyClass'
        self.parent = parent            #from which this object is referenced
        self.data = []                  #empty list
        self.history = history          #actions performed on this data
        self.history.append('>'+self.lineage+' called by '+caller)
    #--- inspection functions ---
    def fileName(self):                 #tell original file name
        pn = self.filePath()
        if pn: return os.path.split(pn)[1]
        else: return ''
    def filePath(self):                 #tell original file path
        for s in self.history:
            t = split(s)
            if len(t) > 0:
                if t[0] == '>path': return t[1]
        return ''
    #class rasterDummyClass        

class rasterObjectClass:
    """
    This class holds image data and history in a file format-neutral manner
      and is central to all image handling in the Bippy scripts.
    """
    def __init__(self,parent=None,data=[],history=[],caller='?'):
        self.lineage = 'bippy:raster:rasterObjectClass'
        self.parent = parent            #from which this object is referenced
        self.data = data                #list of arrays
        self.history = history          #actions performed on this data
        self.history.append('>'+self.lineage+' called by '+caller)
    #--- inspection functions ---
    def fileName(self):                 #tell original file name
        pn = self.filePath()
        if pn: return os.path.split(pn)[1]
        else: return ''
    def filePath(self):                 #tell original file path
        for s in self.history:
            t = split(s)
            if len(t) > 0:
                if t[0] == '>path': return t[1]
        return ''
    def height(self):                   #tell row count
        return len(self.data)
    def itemsize(self):
        if self.len(): return self.data[0].itemsize
        else: return None
    def len(self):                      #tell element count
        if len(self.data) == 0: return 0
        if len(self.data[0]) == 0: return 0
        return len(self.data) * len(self.data[0])
    def length(self):                   #tell byte count
        i = self.len()
        if i: return i * self.data[0].itemsize
        else: return 0
    def typecode(self):                 #tell Python array type code
        if self.len(): return self.data[0].typecode
        else: return None
    def width(self):                    #column count
        if len(self.data) == 0: return 0
        return len(self.data[0])
    #--- action functions ---
    def array(self):                    #return a 1-dimensional array
        code = self.typecode()
        if not code: return None
        a = array.array(code)
        for y in range(self.height()): a += self.data[y]
        return a
    def copy(self):                     #return a clean, separated copy of self
        errPfx = self.lineage + ':copy: '
        data = []
        history = []
        code = self.typecode()
        w = self.width()
        for y in range(len(self.data)):
            a = array.array(code)
            for x in range(w): a.append(self.data[y][x])
            data.append(a)
        for s in self.history: history.append(s)
        history.append('self.copy')
        return rasterObject(self.parent,data,history,caller=errPfx)
    def flip(self,way='horizontal'):
        if not self.len(): return
        if way in ['both','vertical']: self.data.reverse()
        if way in ['both','horizontal']:
            for i in range(self.len()): self.data[i].reverse()
        self.history.append('flip '+way)
    #end class rasterObjectClass

#--- testing ---

if __name__ == '__main__':
    print """
This module (raster.py) is not a standalone program. It is a module
that contains functions and classes used by other scripts.
"""