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

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

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

os.makedirs(outputFolder, exist_ok = True)

#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 and loading palette: " + "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('Palette loaded.')

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


def rgb2palette(args):
  input_image = args[0]
  x_start = args[1]
  x_end = args[2]
  
  #open input image
  i = Image.open(inputFolder + input_image + ".png")
  #i.show()
  print("Opening: " + input_image + ".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 + " - " + input_image + " row {}".format(y))
    for x in range (x_start, x_end):
      #print(timeStamp + " - " + input_image + " row {}".format(x))
      #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, (r1, g1, b1, ca) in enumerate(p):
          #dr = pixRed - cr
          #dg = pixGreen - cg
          #db = pixBlue - cb
          #distance = dr*dr + dg*dg + db*db
          l1 = (r1*299 + g1*587 + b1*114) / 255000
          l2 = (pixRed*299 + pixGreen*587 + pixBlue*114) / 255000
          dL = l1-l2
          dR = (r1-pixRed)/255
          dG = (g1-pixGreen)/255
          dB = (b1-pixBlue)/255
          distance = (dR*dR*0.299 + dG*dG*0.587 + dB*dB*0.114)*0.75 + dL*dL
          
          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)   )
      
    
  print('Saving picture to Output')
  print(str(finalColor))
  imageOutput.save(outputFolder + input_image[-5:] + "WTF_TESTING_8bpp.png") #assumes _#### frame number format (removes last 5 characters)

      

      
      
def run():
  # ----------------------------------------------------------------------------------------------------------------
  #VARIABLES
  # ----------------------------------------------------------------------------------------------------------------
  job_list = [
      #'test0'
      #'STATUE'
      'BRIDGES_0000'
      #'LAND_OUTPUT_0000'
      ]
  # ----------------------------------------------------------------------------------------------------------------
  thread_count = 16
  # ----------------------------------------------------------------------------------------------------------------
  # ----------------------------------------------------------------------------------------------------------------
  # ----------------------------------------------------------------------------------------------------------------

  for job in job_list:
    all_jobs = []
    job_chunks = []
    queue = []

    extra = 0
    chunk_image = Image.open(inputFolder + job + ".png")
    chunk_average_size = math.floor(chunk_image.width/thread_count)
    chunk_modulo_size = chunk_image.width % thread_count 
    for thread in range(0, thread_count):
      start = chunk_average_size * thread
      end = (chunk_average_size * (thread+1)) -1

      start += extra
      if thread < chunk_modulo_size:   
        extra += 1
      end += extra

      start_and_end = [job, start, end]
      
      job_chunks.append(start_and_end)

      queue.append( [job, start, end] )
      #print(str(queue))

    all_jobs.append(job_chunks)

    for a_job in all_jobs:
      #print('-'*32)
      thread_id = 0
      for b_thread in a_job:
        if thread_id == 0:
          thread_id += 1
          print('Job: ' + str(b_thread[0]))#
          print(' '*10 + 'Start, ' + 'End')#
          print('Thread ' + str(thread_id) + ': ' + str(b_thread[1]) + ', ' + str(b_thread[2]))#
        
        
    print(queue)
    
    # ----------------------------------------------------------------------------------------------------------------
    # PROCESS
    # ----------------------------------------------------------------------------------------------------------------
    
    pool = multiprocessing.Pool(processes = thread_count)
    pool.map(rgb2palette, queue)
    pool.close()
    pool.join()

    
    #imageOutput.save(outputFolder + input_image[-5:] + "_8bpp.png") #assumes _#### frame number format (removes last 5 characters)
    #print('SAVING IMAGE')

    
    #finished time
    tx = time.time()
    finishedTime = datetime.datetime.fromtimestamp(tx).strftime('%H:%M:%S')
    
    print("Started:  " + startedTime)
    print("Finished: " + finishedTime)
    

#run()

if __name__ == '__main__':
  #freeze_support()
  run()