# standard_all_png.py

class pngIDAT:
    def __init__(self,width,height,IDATs):
        if type(IDATs) in (list,tuple):IDAT=''.join(IDATs)
        else:IDAT=IDATs
        del IDATs
        try:self.IDAT=IDAT.decode('zlib')
        except:self.IDAT=IDAT
        self.width,self.height=width,height
        self.size=lambda:(self.width,self.height)
        l=len(self.IDAT)
        if l%self.height!=0 or (l/self.height-1)%self.width!=0 or (l/self.height-1)/self.width not in (1,2,3,4):raise TypeError('It not a valid  8 bit IDAT')
        self.IDAT=[x[1:] for x in yields.yieldSliceLength(self.IDAT,len(self.IDAT)/self.height)]
        #self.IDAT=string2D(self.width,self.height,list(''.join([x[1:] for x in yieldSliceLength(self.IDAT,len(self.IDAT)/self.height)])))
        self.planes=(l/self.height-1)/self.width
        self.modes=('L','LA','RGB','RGBA')
        self.mode=self.modes[self.planes-1]
        self.R=lambda:self.getChannel(0)
        self.G=lambda:self.getChannel(1)
        self.B=lambda:self.getChannel(2)
        self.A=lambda:self.getChannel(3)
    def getChannel(self,n):
        if n==0:N=0
        elif n==4:
            if self.mode in ('L','RGB'):N=None
            elif self.mode=='LA':N=2
            elif self.mode=='RGBA':N=4
        elif n==3:
            if self.mode in ('L','RGB'):N=0
            elif self.mode=='RGBA':N=3
            elif self.mode=='LA':N=2
        elif n==2:
            if self.mode in ('L','LA'):N=0
            elif self.mode in ('RGB','RGBA'):N=2
        elif n==1:
            if self.mode in ('L','LA'):N=0
            elif self.mode in ('RGB','RGBA'):N=1
        if N==None:return ('\0'*self.width for x in xrange(self.height))
        elif self.planes==1:return (x for x in self.IDAT)
        else:
            return (x[N::self.planes] for x in self.IDAT)
    def toFile(self,path=None):
        if path==None:path=tempfile.mktemp('.png')
        pngWriteAll(path,genIHDR(self.width,self.height,self.planes),('\0'+'\0'.join(self.IDAT)).encode('zlib'))
        return path
    def toImage(self,path=None):
        path=self.toFile(path)
        return Image.open(path)
    def show(self,temp=None):
        i=self.toImage(temp)
        showImage(i)
        del i
    def IHDR(self):return genIHDR(self.width,self.height,self.planes)
    def dictIHDR(self):return IHDR_decode(self.IHDR())
    #
    #

    #


def pngAddAlpha(source,alpha_path,target=None):
    if type(source) is not Image:
        img=pngChunks(source,1,'IHDR').next()
        iBitdepth=img[1]['bitdepth']
        if iBitdepth!=8:source=Image.open(source)
        del img
    if type(alpha_path) is not Image:
        mask=pngChunks(alpha_path,1,'IHDR').next()
        mBitdepth=mask[1]['bitdepth']
        if mBitdepth!=8:alpha_path=Image.open(alpha_path)
        del mask
    if type(source) is Image:
        _source=tempfile.mktemp('.png')
        source.save(_source,bpp=24)
        del source
        source=_source
    if type(alpha_path) is Image:
        _alpha_path=tempfile.mktemp('.png')
        alpha_path.save(_alpha_path,bpp=8)
        del alpha_path
        alpha_path=_alpha_path
    img=pngChunksDict(source,1)
    mask=pngChunksDict(alpha_path,1)
    IHDR=img['IHDR'][0]
    iBitdepth=IHDR['bitdepth']
    mBitdepth=mask['IHDR'][0]['bitdepth']
    IHDR['color_type']=6
    IHDR['bitdepth']=8
    width,height=IHDR['width'],IHDR['height']
    IHDR['compression']=0
    IHDR['interlace']=0
    IHDR=IHDR_encode(IHDR)
    RGB=yieldSliceLength(''.join(img['IDAT']).decode('zlib'),width*3+1)
    A=iter(yieldSliceLength(''.join(mask['IDAT']).decode('zlib'),width+1))
    RGBA=list()
    for row in RGB:
        RGBs=row[1:]
        if len(RGBs)!=width*3:RGBs=rowToRGB(RGBs,width)
        As=A.next()[1:]
        if len(As)==width*3:As=As[::3]
        elif len(As)==width*2:As=As[1::2]
        elif len(As)==width*4:As=As[3::4]
        temp=[' ']*(len(RGBs)*4/3)
        temp[::4]=RGBs[::3]
        temp[1::4]=RGBs[1::3]
        temp[2::4]=RGBs[2::3]
        temp[3::4]=As
        RGBA.extend(['\0']+temp)
    RGBA=''.join(RGBA)
    RGBA=RGBA.encode('zlib')
    if target==None:target='%smasked%s'%os.path.split(source)
    pngWriteAll(target,IHDR,RGBA)
def rowToRGB(row,width):
    l=len(row)
    if l==width:case='L'
    elif l==width*3:return row
    elif l==width*2:case='LA'
    temp=[' ']*width*3
    if case=='LA':row=row[::2]
    temp[::3]=temp[1::3]=temp[2::3]=row
    rt=''.join(temp)
    del temp
    return rt
def genIHDR(width,height,planes):
    if planes==3:ty=2
    elif planes==4:ty=6
    else:ty=2
    return IHDR_encode({'width':width, 'height':height, 'bitdepth':8, 'color_type':ty,'compression':0, 'filter':0,'interlace':0})




def IHDR_decode(data):
    try:return dict(zip(
        ('width', 'height', 'bitdepth', 'color_type','compression', 'filter','interlace') ,struct.unpack("!2I5B", data)) )
    except:raise TypeError('It is not a valid IHDR data')
def IHDR_encode(data):
    d=[data[x] for x in ('width', 'height', 'bitdepth', 'color_type','compression', 'filter','interlace')]
    return struct.pack('!2I5B',*d)
def is_png(f):
    signature='\x89PNG\r\n\x1a\n'
    if type(f)  in (str,unicode):
        try:return open('f','wb').read(8)==signature
        except:raise TypeError
    else:
        try:
            f.seek(0)
            return f.read(8)==signature
        except:raise TypeError
def pngChunk(f,result=0,req=None):
    head=f.read(8)
    try:
        length,ty=struct.unpack('!I4s', head)
        if ty=='IHDR':
            data=IHDR_decode(f.read(length+4)[:-4])
        elif result==0:
            pos=f.tell()
            f.seek(pos+length+4)
            data=(pos,length)
        elif result in (1,2,3):
            data=f.read(length+4)[:-4]
            if ty=='IDAT' and result==2:
                pass
            elif ty=='IDAT' and result==3:data=len(data)
        if req==None or req==ty:return (ty,data)
        else:return pngChunk(f,result,req)
    except:raise ValueError('It is not a valid checksum')

def pngChunks(path,result=0,req=None):
    f=open(path,'rb')
    if not is_png(f):raise TypeError('It is not a png file')
    while True:
            try:yield pngChunk(f,result,req)
            except ValueError:
                f.close()
                del f
                break
def pngChunksDict(path,result=0):
    rt=dict()
    for x in iter(pngChunks(path,result)):
        ty,data=x
        if ty not in rt:rt[ty]=[]
        rt[ty].append(data)
    return rt

def pngWriteChunk(f, ty, data=''):
    f.write(struct.pack("!I", len(data)))
    f.write(ty)
    f.write(data)
    checksum=zlib.crc32(ty)
    checksum=zlib.crc32(data, checksum)
    checksum&=2**32-1
    f.write(struct.pack("!I", checksum))
def pngWriteAll(path,IHDR,IDAT):
    signature='\x89PNG\r\n\x1a\n'
    f=file(path,'wb')
    f.write(signature)
    pngWriteChunk(f,'IHDR',IHDR)
    pngWriteChunk(f,'IDAT',IDAT)
    pngWriteChunk(f,'IEND')
    f.close()

def pngCeil(path,target,n):
    img=pngChunksDict(path,1)
    IHDR=IHDR_encode(img['IHDR'][0])
    IDAT=''.join(img['IDAT']).decode('zlib').translate(ceilMask(n)).encode('zlib')
    pngWriteAll(target,IHDR,IDAT)