from keras import losses
from keras import optimizers
from keras.models import Sequential
from keras.layers import Dense, Dropout, Activation, Flatten
from keras.layers import Convolution2D, MaxPooling2D, Dense

from keras.utils import np_utils

from sklearn.preprocessing import MinMaxScaler

from rect import Rect
from PIL import Image, ImageDraw
from tqdm import tqdm

import matplotlib.pyplot as plt
import numpy as np
import math, random
import utils, os

batch_size = 32
epochs = 100

grid_size = 13
cell_size = 32
rects_in_cell = 5
class_num = 43

cell_out_size = class_num + 5
cell_vec_size = cell_out_size * rects_in_cell

input_width = grid_size * cell_size
input_height = grid_size * cell_size

def loss_function(y_true, y_pred):
  return 0

def load_dataset(data_path):
  x = []
  y = []
  image_info = {}
  file_order = []

  for f in tqdm(utils.list_files(data_path, 'ppm')):
    img = Image.open(f)

    p, name = os.path.split(f)
    image_info[name] = [len(x), img.size]
    img = img.resize((input_width, input_height), Image.NEAREST)

    file_order.append(name)
    x.append(np.asarray(img))

  rects_map = {}
  rects = utils.read_description(os.path.join(data_path, 'gt.txt'))
  for r in rects:
    name = r[0]
    x1 = int(r[1])
    y1 = int(r[2])
    x2 = int(r[3])
    y2 = int(r[4])
    cl = int(r[5])

    assert x1 < x2
    assert y1 < y2

    info = image_info[name]
    w, h = info[1]
    x_mult = float(input_width) / float(w)
    y_mult = float(input_height) / float(h)

    x1 = round(x1 * x_mult)
    x2 = round(x2 * x_mult)
    y1 = round(y1 * y_mult)
    y2 = round(y2 * y_mult)

    assert x1 < x2
    assert y1 < y2

    v = [Rect(x1, y1, x2 - x1, y2 - y1), cl]
    if name in rects_map:
      rects_map[name].append(v)
    else:
      rects_map[name] = [ v ]

  for f in file_order:
    file_data = []
    for cy in range(grid_size):
      row = []
      for cx in range(grid_size):
        cell = [0.0] * cell_vec_size

        # get coordinates of cell rect
        cell_rect = Rect(cx * cell_size,
                        cy * cell_size,
                        cell_size,
                        cell_size)

        # find intersected rects
        save_rects = []
        if f in rects_map:
          for r in rects_map[f]:
            if cell_rect.intersection(r[0]) != None:
              save_rects.append(r)

          # write rects to output array
          #print('Save: {}, {}'.format(len(save_rects), f))
          for i in range(len(save_rects)):
            r = save_rects[i][0]
            c = save_rects[i][1]

            offset = i * cell_out_size
            cell[offset] = r.x()
            cell[offset + 1] = r.y()
            cell[offset + 2] = r.width()
            cell[offset + 3] = r.height()
            cell[offset + 4] = 1.0
            cell[offset + 5 + cl] = 1.0

        row.append(cell)

      file_data.append(row)

    y.append(file_data)

  return (x, y)

print ("Prepare train data...")
x_train, y_train = load_dataset('data/train')

print ("Prepare test data...")
x_test, y_test = load_dataset('data/test')

# TEST ANY INPUT IMAGE
# index = 78
# img = Image.fromarray(x_train[index], 'RGB')
# draw = ImageDraw.Draw(img)
# for r in y_train[index]:
#   draw.rectangle(r[:-1], outline=(0, 255, 0))
# img.show()
# TEST END

input_shape = (input_width, input_height, 3)

# Use simple architecture: http://machinethink.net/blog/object-detection-with-yolo/
model = Sequential()
model.add(Convolution2D(16, kernel_size=(3, 3),
                        activation='relu',
                        strides=1,
                        input_shape=input_shape))

model.add(MaxPooling2D(pool_size=(2, 2),
                       strides=2))

model.add(Convolution2D(32, kernel_size=(3, 3),
                        activation='relu',
                        strides=1))

model.add(MaxPooling2D(pool_size=(2, 2),
                       strides=2))

model.add(Convolution2D(64, kernel_size=(3, 3),
                        activation='relu',
                        strides=1))

model.add(MaxPooling2D(pool_size=(2, 2),
                       strides=2))

model.add(Convolution2D(128, kernel_size=(3, 3),
                        activation='relu',
                        strides=1))

model.add(MaxPooling2D(pool_size=(2, 2),
                       strides=2))

model.add(Convolution2D(256, kernel_size=(3, 3),
                        activation='relu',
                        strides=1))

model.add(MaxPooling2D(pool_size=(2, 2),
                       strides=2))

model.add(Convolution2D(512, kernel_size=(3, 3),
                        activation='relu',
                        strides=1))

model.add(MaxPooling2D(pool_size=(2, 2),
                       strides=1))

model.add(Convolution2D(1024, kernel_size=(3, 3),
                        activation='relu',
                        strides=1))

model.add(Convolution2D(1024, kernel_size=(3, 3),
                        activation='relu',
                        strides=1))

model.add(Convolution2D(cell_vec_size, kernel_size=(1, 1),
                        activation='relu',
                        strides=1))

# model.add(Dense(256))
# model.add(Dense(100))


model.compile(loss='mse',
              optimizer=optimizers.Adadelta(),
              metrics=['accuracy'])

model.fit(np.array(x_train), np.array(y_train),
          batch_size=batch_size,
          epochs=epochs,
          verbose=1,
          validation_data=(np.array(x_test), np.array(y_test)))