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

#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 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(input_image, x_start, x_end):
  #open input image
  i = Image.open(inputFolder + input_image + ".png")
  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)   )
          
      
  os.makedirs(outputFolder, exist_ok = True)
  imageOutput.save(outputFolder + input_image[-5:] + "_8bpp.png") #assumes _#### frame number format (removes last 5 characters)
  
  #finished time
  tx = time.time()
  finishedTime = datetime.datetime.fromtimestamp(tx).strftime('%H:%M:%S')
  
  print("Started:  " + startedTime)
  print("Finished: " + finishedTime)
  
  
      
      
def run():
  # ----------------------------------------------------------------------------------------------------------------
  #VARIABLES
  # ----------------------------------------------------------------------------------------------------------------
  job_list = [
      'test0'
      #'BRIDGES_0000',
      #'LAND_OUTPUT_0000'
      ]
  # ----------------------------------------------------------------------------------------------------------------
  thread_count = 4
  # ----------------------------------------------------------------------------------------------------------------
  # ----------------------------------------------------------------------------------------------------------------
  # ----------------------------------------------------------------------------------------------------------------

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

    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)

      q.append( Process(target = rgb2palette, args= (job,start,end)) )
    
    #for ch in range(0, chunk_modulo_size):
    #  job_chunks[ch] += 1
    all_jobs.append(job_chunks)

    printList_thread = []
    printList_start = []
    printList_end = []

    for a_job in all_jobs:
      print('-'*32)#
      thread_id = 0
      for b_thread in a_job:
        if thread_id == 0:
          printList_thread.append(b_thread[0])
          thread_id += 1
          print('Job: ' + str(b_thread[0]))#
          print(' '*10 + 'Start, ' + 'End')#
        printList_start.append(b_thread[1])
        printList_end.append(b_thread[2])
        print('Thread ' + str(thread_id) + ': ' + str(b_thread[1]) + ', ' + str(b_thread[2]))#
    
    print(all_jobs)
    
    # ----------------------------------------------------------------------------------------------------------------
    # PROCESS
    # ----------------------------------------------------------------------------------------------------------------
    
    for process in range(thread_count):
      print('Process count' + str(process))
      q[process].start()
    for process in range(thread_count):
      q[process].join()
      

  #for a_text in printList_thread:
  #  print('Job:' + str(a_text))
  #  for b_text in printList_start:
  #   #list_id = 0
  #   print('Thread:' + str(b_text))
      

import traceback

try:

  run()
    

except Exception as e:

  traceback.print_exc()

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