import math
import time
import datetime
import os
from os import remove
from PIL import Image

#starting message
print("-"*79)
print("Starting...")

#path definition
diskPath = "C:/Users/Pavel/Downloads/scripts/"#"E:/_BRIX/_BRIX-repository/scripts/"
inputFolder = diskPath + "input/"
outputFolder = diskPath + "output/"

#printing path just to check
print("inputFolder is " + inputFolder)
print("outputFolder is " + outputFolder)

#started time
tt = time.time()
startedTime = datetime.datetime.fromtimestamp(tt).strftime('%H:%M:%S')



#open palette image
palette = Image.open(inputFolder + "openttd-palette-dos.png")
print("Opening: " + "openttd-palette-dos.png")  

p=[]

for b in range(0,palette.height):
  for a in range(0, palette.width):
    p.append(palette.getpixel((a,b)))
    
#print(p)

def rgb2palette(inputImage, frameNumber):
  #open input image
  i = Image.open(inputFolder + inputImage + frameNumber + ".png")
  print("Opening: " + inputImage + frameNumber + ".png")  
  
  #create new empty image for output
  imageOutput = Image.new("RGBA", (i.width,i.height), color=(0,0,0,0))
    
  for y in range (0, i.height):
    #timeStamp
    ts = time.time()
    timeStamp = datetime.datetime.fromtimestamp(ts).strftime('%H:%M:%S')
    print(timeStamp + " - " + inputImage + " row {}".format(y))
    for x in range (0, i.width):
      #defining winner variables
      winnerDistance = 100000000
      winnerID = 0
      
      #loading pixel from image and separating RGBA
      pixelNumber = x + (y * i.width)
      pix = i.getpixel((x,y))
      pixRed = pix[0]
      pixGreen = pix[1]
      pixBlue = pix[2]
      pixAlpha = pix[3]
            
      #check Alpha in pixel, and output alpha/color offset
      if pixAlpha < 128:
        finalAlpha = 0
        colorOffset = 0
      if pixAlpha >= 128 and pixAlpha < 178:
        finalAlpha = 255
        colorOffset = 1
      if pixAlpha >= 178 and pixAlpha < 230:
        finalAlpha = 255
        colorOffset = 2
      if pixAlpha >= 230:
        finalAlpha = 255
        colorOffset = 0
      
      #if alpha above 50%, do colour comparing to palette
      if pixAlpha >= 128:
        for z, (cr, cg, cb, ca) in enumerate(p):
          dr = pixRed - cr
          dg = pixGreen - cg
          db = pixBlue - cb
          distance = dr*dr + dg*dg + db*db

          if distance < winnerDistance:
            winnerDistance = distance
            winnerID = z
    
      #compare input RGB channels and output highest value
      if pixRed >= pixGreen and pixRed >= pixBlue:
        highestValue = pixRed
      if pixGreen >= pixRed and pixGreen >= pixBlue:
        highestValue = pixGreen
      if pixBlue >= pixRed and pixBlue >= pixGreen:
        highestValue = pixBlue     
      # set color offset +/- based on colour value
      if highestValue < 128:
        negation = -1
        colorOffset = colorOffset * negation
      #print("colorOffset is ... " + str(colorOffset) )
      
      #final color changed by colorOffset
      finalColor = p[winnerID - colorOffset]
      
      finalR = finalColor[0]
      finalG = finalColor[1]
      finalB = finalColor[2]
      #finalAlpha taken from the if output above palette colour comparing
      
      #put the final pixel into the output picture
      imageOutput.putpixel((x,y),(finalR,finalG,finalB,finalAlpha))
      #print("Pixel " + str(pixelNumber) + ": R= " + str(finalR) + ", G= " + str(finalG) + ", B= " + str(finalB) + ", A= "  + str(finalAlpha)   )
          
      
  os.makedirs(outputFolder, exist_ok = True)
  imageOutput.save(outputFolder + inputImage + "_8bpp.png")
  
  #finished time
  tx = time.time()
  finishedTime = datetime.datetime.fromtimestamp(tx).strftime('%H:%M:%S')
  
  print("Started:  " + startedTime)
  print("Finished: " + finishedTime)
  
  
      
      
def run():
  rgb2palette("TRACKS_EDIT", "")
  rgb2palette("BRIDGES_", "0000")
  rgb2palette("LAND_OUTPUT_", "0000")
  rgb2palette("ROADS_OUTPUT_", "0000")
  rgb2palette("TREES_OUTPUT_", "0000")
  rgb2palette("SIGNALS-01_", "0000")
  rgb2palette("SIGNALS-02_", "0000")



import traceback

try:

  run()

except Exception as e:

  traceback.print_exc()

input("Press enter to continue...")