JungleBookScript

import Image
import os.path
import itertools
import binascii

#------------------------------------------------------------------------------
class PpuSprite:
   def __init__(self, im):
     self.image = im
     self.imageData = im.getdata()
     self.colorTable = im.getcolors()
     self.colorTable.sort()
     self.firstColorsPos = []
     self.calcFirstColorsPoses()
     self.indexedData = []
     self.prepareIndexedData()

   def calcFirstColorsPoses(self):
     for _,colorInPalette in self.colorTable:
       firstColorPos = -1
       for color in self.imageData:
         firstColorPos+=1
         if color == colorInPalette:
           break
       self.firstColorsPos.append((firstColorPos,colorInPalette))
     self.firstColorsPos.sort()

   def prepareIndexedData(self):
     def calcColorHash(c):
       return c[0]<<16 | c[1]<<8 | c[2]
     colorHashes = map (lambda colorRec : calcColorHash(colorRec[1]), self.firstColorsPos)
     for pixel in self.imageData:
       pixelHash = calcColorHash(pixel)
       colorIndex = -1
       for colorHash in colorHashes:
         colorIndex +=1
         if pixelHash == colorHash:
           break
       self.indexedData.append(colorIndex)

   def isEqual(self, otherPpuSprite):
     if len(self.firstColorsPos) != len(otherPpuSprite.firstColorsPos):
       return False
     for fcp1, fcp2 in zip(self.firstColorsPos, otherPpuSprite.firstColorsPos):
       if  fcp1[0]!= fcp2[0]:
         return False
     if len(self.colorTable) != len(otherPpuSprite.colorTable):
       return False
     for colorTableRec1, colorTableRec2 in zip(self.colorTable, otherPpuSprite.colorTable):
       if colorTableRec1[0] != colorTableRec2[0]:
         return False
     for pixel1, pixel2 in zip(self.indexedData, otherPpuSprite.indexedData):
       if pixel1!=pixel2:
         return False
     return True

#------------------------------------------------------------------------------

def findFirstPicIndex(sprite, ppuSpriteTable):
  index = -1
  for ppuSprite in ppuSpriteTable:
    index+=1
    if ppuSprite.isEqual(sprite):
      return index
  return -1

def findUniquePicIndex(sprite, ppuSpriteTable):
  index = -1
  firstIndex = -1
  for ppuSprite in ppuSpriteTable:
    index+=1
    if ppuSprite.isEqual(sprite):
      if firstIndex!=-1:
        return -1
      else:
        firstIndex = index
  return firstIndex


def makeLevelMap(mem, lev):
  levTable = []
  for l in lev:
    levTable.append(findFirstPicIndex(l, mem))
  return levTable

def drawLevelFromData(levelArray, ppuSpriteTable, size, vertical = False):
  w,h = size
  im = Image.new("RGBA", (w*8,h*8))
  for i in xrange(len(levelArray)):
    if vertical:
      x,y = i/h, i%h
    else:
      x,y = i%w, i/w
    ind = levelArray[i]
    if (ind!=-1):
      im.paste(ppuSpriteTable[ind].image, (x*8,y*8,x*8+8,y*8+8))
  return im

#------------------------------------------------------------------------------
def parseLevelData(fname):
  f = open(fname)
  lines = f.readlines()
  f.close()
  def parseLine(line):
    ind = line.find(".BYTE")
    num = line[ind+5:ind+9].lstrip(" $")
    return int(num,16)
  return map(parseLine, lines)

#------------------------------------------------------------------------------
# levelSetNames = [
    # "A139", "A739", "DDC8", "DDC8",
    # "DBF4", "DC59", "AA49", "AAC0",
    # "DCBE", "DD23", "AB37", "ABAE",
    # "D974", "D9F4", "A7C9", "A849",
    # "DA74", "DAF4", "A8C9", "A949",
    # "DB74", "A9C9", "BBF3", "A825",
    # "A795", "A765", "A764", "A766",
    # "A7F5", "A7C5"]

#------------------------------------------------------------------------------------------------
#helpers
def prepareBinaries(binaryDump, levelSetNames):
  f = open(binaryDump,"rb")
  byt = f.read()
  f.close()
  for name in levelSetNames:
    ptr = int(name,16)
    bb = byt[ptr:ptr+256]
    p = os.path.expanduser("~/desktop/jungleBook/set1_"+name+".bin")
    f =open(p,"wb")
    f.write(bb)
    f.close()
#prepareBinaries("os.path.expanduser("~/desktop/jungleBook/jungleBook2.bin", levelSetNames)

def prepareBinary(binaryDump, addrBegin, addrLen, outName):
  f = open(binaryDump,"rb")
  byt = f.read()
  f.close()
  bb = byt[addrBegin:addrBegin+addrLen]
  f =open(outName,"wb")
  f.write(bb)
  f.close()
#------------------------------------------------------------------------------------------------
def parseBinData256(fname):
  f = open(fname, "rb")
  b = f.read()
  f.close()
  return map(ord, b)

#a3a4 = parseBinData256("a3a4_1test.bin")

def loadLevelSetElement(levelSetElemName):
  return parseBinData256(os.path.expanduser("~/desktop/jungleBook/set1_"+levelSetElemName+".bin"))

##levelSet = map(loadLevelSetElement, levelSetNames)

###b8x_inds 4F4F5151 4F4F5151 57575959 57575959
####b8x_inds = [[4,4,5,5], [4,4,5,5], [8,8,9,9], [8,8,9,9]]
b8x_inds = [[5,5,4,4], [7,7,6,6], [9,9,8,8], [11,11,10,10]]
###b9x_inds = 615F615F 69676967 615F615F 69676967
###b9x_inds = [[13,12,13,12], [17,16,17,16], [13,12,13,12], [17,16,17,16]]
b9x_inds = [[12,13,12,13], [14,15,14,15], [16,17,16,17], [18,19,18,19]]

def getLevelSetForCycle(levelSet,indsConst,loop):
  inds = map (lambda v: v[loop], indsConst)
  return map (lambda i: levelSet[i], inds)

#a3a4 = parseBinData256(os.path.expanduser("~/desktop/jungleBook/set1x_A2A3.bin"))
def get4Lines(begin,end, levelSet, a3a4, step, vertCount):
  cycleBegin = begin
  cycleEnd   = end
  a3a4ind = begin/4
  line4 = []
  for x in xrange(cycleBegin,cycleEnd):
    if x%4==0:
      a3a4ind += 1
    b83,b85,b87,b89 = getLevelSetForCycle(levelSet, b8x_inds,(x+2)%4)
    b93,b95,b97,b99 = getLevelSetForCycle(levelSet, b9x_inds,(x+2)%4)
    line4.extend(makeLine(a3a4, a3a4ind, b83,b85,b87,b89, b93,b95,b97,b99, step, vertCount))
  return line4


def makeLine(a3a4, a3a4BaseInd, b83,b85,b87,b89, b93,b95,b97,b99, step, vertCount):
  a3a4ind = a3a4BaseInd
  lines = []
  for globalRepeat in xrange(vertCount): ##5(screen size)
    predInd = a3a4[a3a4ind]
    ind2a,ind2b = -1,-1
    if predInd<128:
      ind2a = b83[predInd]
      ind2b = b87[predInd]
    else:
      ind2a = b85[predInd]
      ind2b = b89[predInd]

    res = [-1]*4
    if ind2a<128:
      res[0],res[1] = b93[ind2a], b97[ind2a]
    else:
      res[0],res[1] = b95[ind2a], b99[ind2a]
    if ind2b<128:
      res[2],res[3] = b93[ind2b], b97[ind2b]
    else:
      res[2],res[3] = b95[ind2b], b99[ind2b]
    lines.extend(res)
    a3a4ind+=step
  return lines

#>>> prepareBinary(binPath, 0xA138, 0x2000, outPath) #level 1 start address
#>>> a3a4 = parseBinData256(outPath)
#>>> ll = get4Lines(0,96*4)
#>>> drawLevelFromData(ll, sprites,(96*4,16*4),True).show()

#---
#helpers
def txt2bin(inname,outname):
  f = open(inname,"rt")
  l = f.readlines()
  l = l[0]
  f.close()
  bb = binascii.a2b_hex(l)
  f = open(outname,"wb")
  f.write(bb)
  f.close()

def str2SetAddr(trr):
      res = []
      while trr!="":
        res.append(trr[:4])
        trr = trr[4:]
      res2 = map (lambda v: v[2:]+v[:2], res)
      return res2
#---
def colorReplace(im, cmap):
  d = im.getdata()
  nd = []
  for pix in d:
    for c,r in cmap:
      if pix == c:
        nd.append(r)
  imgb = Image.new("RGB", im.size)
  imgb.putdata(nd)
  return imgb
#---

def makeJob (binaryName, spritesName, addr, baseAddrSet, w,h, levelNo):
  tempName1 = binaryName+".a2a3."+str(levelNo)+".bin"

  imNameTable = Image.open(spritesName)
  nameSize = imNameTable.size
  nameTable = [imNameTable.crop((x,y,x+8,y+8)) for y,x in itertools.product(xrange(0,nameSize[1],8),xrange(0,nameSize[0],8))]
  sprites = map (lambda im: PpuSprite(im),nameTable)

  prepareBinary(binaryName, addr, 0x2000, tempName1)
  a3a4 = parseBinData256(tempName1)

  prepareBinaries(binaryName, baseAddrSet)
  levelSet = map(loadLevelSetElement, baseAddrSet)
  ##print baseAddrSet
  ##print a3a4[0:50]
  ll = get4Lines(0,w*4, levelSet, a3a4, w,h)
  drawLevelFromData(ll, sprites,(w*4,h*4),True).show()


#------------------------------------------------------------------------------
##def main():
imNameTable = Image.open(os.path.expanduser("~/desktop/jungleBook/jungleBookTable1.png"))
imLevelTable = Image.open(os.path.expanduser("~/desktop/jungleBook/jungleBookLevel1.png"))

nameSize = imNameTable.size
levelSize = imLevelTable.size
nameTable = [imNameTable.crop((x,y,x+8,y+8)) for y,x in itertools.product(xrange(0,nameSize[1],8),xrange(0,nameSize[0],8))]
levelTable = [imLevelTable.crop((x,y,x+8,y+8)) for y,x in itertools.product(xrange(0,levelSize[1],8),xrange(0,levelSize[0],8))]
ppuSpriteMemTable = map (lambda im: PpuSprite(im),nameTable)
ppuSpriteLvlTable = map (lambda im: PpuSprite(im),levelTable)
#mapLevel = makeLevelMap(ppuSpriteMemTable, ppuSpriteLvlTable)
#return mapLevel