"""
photomosaic.py

Tworzenie fotomozaiki z danym obrazem docelowym i folderem obrazów wejściowych.

Autor: Mahesh Venkitachalam
"""

import sys, os, random, argparse
from PIL import Image
import imghdr
import numpy as np

def getAverageRGBOld(image):
  """
  Dla danego obiektu Image biblioteki PIL zwraca średnią wartość koloru jako (r, g, b)
  """
  # liczba pikseli w obrazie
  npixels = image.size[0]*image.size[1]
  # uzyskanie kolorów jako [(cnt1, (r1, g1, b1)), ...]
  cols = image.getcolors(npixels)
  # uzyskanie [(c1*r1, c1*g1, c1*g2),...]
  sumRGB = [(x[0]*x[1][0], x[0]*x[1][1], x[0]*x[1][2]) for x in cols] 
  # obliczanie (sum(ci*ri)/np, sum(ci*gi)/np, sum(ci*bi)/np)
  # zip daje nam [(c1*r1, c2*r2, ..), (c1*g1, c1*g2,...)...]
  avg = tuple([int(sum(x)/npixels) for x in zip(*sumRGB)])
  return avg

def getAverageRGB(image):
  """
  Dla każdego obrazu wejściowego zwraca średnią wartość koloru jako (r, g, b)
  """
  # uzyskanie każdego obrazu kafelkowego jako tablicy numpy
  im = np.array(image)
  # uzyskanie kształtu każdego obrazu wejściowego
  w,h,d = im.shape
  # uzyskanie średniej wartości RGB
  return tuple(np.average(im.reshape(w*h, d), axis=0))

def splitImage(image, size):
  """
  Dla danego obiektu Image i wymiarów (rows, cols) zwraca listę m*n obiektów Image 
  """
  W, H = image.size[0], image.size[1]
  m, n = size
  w, h = int(W/n), int(H/m)
  # lista obrazów
  imgs = []
  # generowanie listy wymiarów
  for j in range(m):
    for i in range(n):
      # załączanie przyciętego obrazu
      imgs.append(image.crop((i*w, j*h, (i+1)*w, (j+1)*h)))
  return imgs

def getImages(imageDir):
  """
  Dla danego folderu obrazów zwraca listę obiektów Image
  """
  files = os.listdir(imageDir)
  images = []
  for file in files:
    filePath = os.path.abspath(os.path.join(imageDir, file))
    try:
      # bezpośrednie ładowanie, abyśmy nie na napotkali niedoboru zasobów
      fp = open(filePath, "rb")
      im = Image.open(fp)
      images.append(im)
      # wymuszenie ładowania danych obrazu z pliku
      im.load() 
      # zamknięcie pliku
      fp.close() 
    except:
      # pominięcie
      print("Nieprawidłowy obraz: %s" % (filePath,))
  return images

def getImageFilenames(imageDir):
  """
  Dla danego folderu obrazów, zwraca listę nazw plików obrazów
  """
  files = os.listdir(imageDir)
  filenames = []
  for file in files:
    filePath = os.path.abspath(os.path.join(imageDir, file))
    try:
      imgType = imghdr.what(filePath) 
      if imgType:
        filenames.append(filePath)
    except:
      # pominięcie
      print("Nieprawidłowy obraz: %s" % (filePath,))
  return filenames

def getBestMatchIndex(input_avg, avgs):
  """
  Zwraca indeks najlepszego dopasowanie obiektu Image na podstawie odległości wartości RGB
  """

  # średnia obrazu wejściowego
  avg = input_avg
  
  # uzyskanie najbliższej wartości RGB dla danych wejściowych na podstawie odległości x/y/z
  index = 0
  min_index = 0
  min_dist = float("inf")
  for val in avgs:
    dist = ((val[0] - avg[0])*(val[0] - avg[0]) +
            (val[1] - avg[1])*(val[1] - avg[1]) +
            (val[2] - avg[2])*(val[2] - avg[2]))
    if dist < min_dist:
      min_dist = dist
      min_index = index
    index += 1

  return min_index


def createImageGrid(images, dims):
  """
  Dla danej listy obrazów i rozmiaru siatki (m, n) utworzenie siatki obrazów.  
  """
  m, n = dims

  # kontrola poprawności
  assert m*n == len(images)

  # uzyskanie maksymalnej wysokości i szerokości obrazów
  # nie zakładamy, że wszystkie są równe
  width = max([img.size[0] for img in images])
  height = max([img.size[1] for img in images])

  # tworzenie obrazu wyjściowego
  grid_img = Image.new('RGB', (n*width, m*height))
  
  # wklejanie obrazów kafelkowych do siatki
  for index in range(len(images)):
    row = int(index/n)
    col = index - n*row
    grid_img.paste(images[index], (col*width, row*height))
    
  return grid_img


def createPhotomosaic(target_image, input_images, grid_size,
                      reuse_images=True):
  """
  Tworzenie fotomozaiki dla danych obrazu docelowego i obrazów wejściowych.
  """

  print('dzielenie obrazu docelowego...')
  # dzieli obraz docelowy na kafelki 
  target_images = splitImage(target_image, grid_size)

  print('wyszukiwanie dopasowań dla obrazu...')
  # dla każdego kafelka wybiera jeden pasujący obraz wejściowy
  output_images = []
  # dla informacji zwrotnej dla użytkownika
  count = 0
  batch_size = int(len(target_images)/10)

  # obliczanie średnich obrazów wejściowych
  avgs = []
  for img in input_images:
    avgs.append(getAverageRGB(img))

  for img in target_images:
    # obliczanie średniej wartości RGB obrazu docelowego
    avg = getAverageRGB(img)
    # wyszukiwanie indeksu dopasowania najbliższej wartości RGB z listy
    match_index = getBestMatchIndex(avg, avgs)
    output_images.append(input_images[match_index])
    # informacja zwrotna dla użytkownika
    if count > 0 and batch_size > 10 and count % batch_size is 0:
      print('przetworzono %d z %d...' %(count, len(target_images)))
    count += 1
    # jeśli ustawiona jest flaga, usuwanie wybranego obrazu z danych wejściowych
    if not reuse_images:
      input_images.remove(match)

  print('tworzenie mozaiki...')
  # tworzenie obrazu fotomozaiki z kafelków
  mosaic_image = createImageGrid(output_images, grid_size)

  # zwracanie mozaiki
  return mosaic_image

# Zebranie kodu w funkcji main()
def main():
  # Argumentami wiersza poleceń są sys.argv[1], sys.argv[2] ..
  # sys.argv[0] to nazwa samego skryptu, która może być ignorowana

  # parsowanie argumentów
  parser = argparse.ArgumentParser(description='Tworzy fotomozaikę z obrazów wejściowych')
  # dodawanie argumentów
  parser.add_argument('--target-image', dest='target_image', required=True)
  parser.add_argument('--input-folder', dest='input_folder', required=True)
  parser.add_argument('--grid-size', nargs=2, dest='grid_size', required=True)
  parser.add_argument('--output-file', dest='outfile', required=False)

  args = parser.parse_args()

  ###### DANE WEJŚCIOWE ######

  # obraz docelowy
  target_image = Image.open(args.target_image)

  # obrazy wejściowe
  print('wczytywanie folderu wejściowego...')
  input_images = getImages(args.input_folder)

  # sprawdzanie, czy znalezione zostały jakieś prawidłowe obrazy wejściowe  
  if input_images == []:
      print('Nie znaleziono żadnych obrazów wejściowych w %s. Zamykanie programu.' % (args.input_folder, ))
      exit()

  # tasowanie listy w celu uzyskania bardziej urozmaiconych danych wyjściowych?
  random.shuffle(input_images)

  # rozmiar siatki
  grid_size = (int(args.grid_size[0]), int(args.grid_size[1]))

  # dane wyjściowe
  output_filename = 'mosaic.png'
  if args.outfile:
    output_filename = args.outfile
  
  # ponowne wykorzystywanie dowolnego obrazu z danych wejściowych
  reuse_images = True

  # zmiana rozmiaru danych wejściowych, aby pasowały do rozmiaru pierwotnego obrazu?
  resize_input = True

  ##### KONIEC DANYCH WEJŚCIOWYCH #####

  print('rozpoczęcie tworzenia fotomozaiki...')
  
  # jeśli obrazy nie mogą być ponownie używane, należy upewnić się, że m*n <= num_of_images 
  if not reuse_images:
    if grid_size[0]*grid_size[1] > len(input_images):
      print('rozmiar siatki jest mniejszy niż liczba obrazów')
      exit()
  
  # zmiana rozmiaru danych wejściowych
  if resize_input:
    print('zmiana rozmiaru obrazów...')
    # dla danego rozmiaru siatki obliczanie maksymalnej szerokości i wysokości kafelków
    dims = (int(target_image.size[0]/grid_size[1]), 
            int(target_image.size[1]/grid_size[0])) 
    print("maksymalne wymiary kafelka: %s" % (dims,))
    # zmiana rozmiaru
    for img in input_images:
      img.thumbnail(dims)

  # tworzenie fotomozaiki
  mosaic_image = createPhotomosaic(target_image, input_images, grid_size,
                                   reuse_images)

  # zapisywanie mozaiki
  mosaic_image.save(output_filename, 'PNG')

  print("dane wyjściowe zostały zapisane w %s" % (output_filename,))
  print('zrobione.')

# Standardowy kod do wywoływania funkcji main(),
# aby rozpocząć program.
if __name__ == '__main__':
  main()
