import keras
from keras.datasets import mnist
from keras.models import Sequential
from keras.layers import Dense, Dropout, Flatten, Activation, Reshape
from keras.layers import Input, Conv2D, MaxPooling2D, UpSampling2D
import numpy as np
import pandas as pd
from collections import OrderedDict
import itertools
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import random
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import OneHotEncoder
import math
from math import sqrt
from keras.models import Model, load_model

def create_blank_image(size):
    data = np.ndarray(shape=(size, size))
    for i in range(0, size):
        for j in range(0, size):
            data[[i], [j]] = 0
    #print(data)

    return data

def circle_randomizer():
    number_of_circles = random.randint(4,10)
    intensity = np.ndarray(shape=(128, 128))
    #print(number_of_circles)
    radius_list = []


    for i in range(number_of_circles):
        radius_list.append(random.uniform(8, 10))
    #print(radius_list)

    center_coords = np.zeros((2,1))
    center_coords[[0],[0]] = random.uniform(0,size)
    center_coords[[1],[0]] = random.uniform(0,size)

    for i in range(number_of_circles):
      #temp_array = np.ndarray(shape=(2,1))
      #temp_array[[0],[0]] = random.uniform(0,size)
      #temp_array[[1],[0]] = random.uniform(0,size)

      if i > 0:
          j = 0
          #print(i,j)
          while j in range(i):
              #print(i,j)
              #print(center_coords)
              temp_array = np.ndarray(shape=(2,1))
              temp_array[[0],[0]] = random.uniform(0,size)
              temp_array[[1],[0]] = random.uniform(0,size)
              #while sqrt((center_coords[[0],[i]] - center_coords[[0],[j]])**2 + (center_coords[[1],[i]] - center_coords[[1],[j]])**2) < radius_list[i] + radius_list[j]:
              while sqrt((temp_array[[0],[0]] - center_coords[[0],[j]])**2 + (temp_array[[1],[0]] - center_coords[[1],[j]])**2) < radius_list[i] + radius_list[j]:
                  temp_array[[0],[0]] = random.uniform(0,size)
                  temp_array[[1],[0]] = random.uniform(0,size)
                  j = 0
              center_coords = np.concatenate((center_coords,temp_array), axis = 1)
              j = j + 1
              #print('loop ran ' + str(j) + ' times')

    return radius_list, center_coords

def image_creator(centers, radii, img_data, size):
    x = np.arange(1, size, 1)
    y = np.arange(1, size, 1)

    for c in range(len(centers)):
        x0 = centers[[c],[0]]
        y0 = centers[[c],[1]]
        radius = radii[c]
        for i in range(0, size-1):
            for j in range(0, size-1):
                height2 = radius**2 - (x[i]-x0)**2 - (y[j]-y0)**2
                if height2 >= 0:
                    img_data[[i], [j]] = sqrt(radius**2 - (x[i]-x0)**2 - (y[j]-y0)**2)

    return img_data

def make_ellipses(size, radii, center_coords):
    # idea: use a random number generator to create a random rotation of the x,y axes for the ellipse

    # size is the length of a side of the square
    # length is the length of the ellipse
    # defined as equal to the radius of the circle later

    my_label = np.ndarray(shape=(size, size))
    x = np.arange(1, size, 1)
    y = np.arange(1, size, 1)

    # inefficiently zero the array
    for i in range(0, size):
        for j in range(0, size):
            my_label[[i], [j]] = 0
            # print(my_label)
    for c in range(len(center_coords)):
        x0 = center_coords[[c],[0]]
        y0 = center_coords[[c],[1]]
        #theta = random.uniform(0, 6.28318)
        theta = 0.775

        for i in range(0, size - 1):
            for j in range(0, size - 1):
                xprime = (x[i] - x0) * math.cos(theta) + (y[j] - y0) * math.sin(theta)
                yprime = -(x[i] - x0) * math.sin(theta) + (y[j] - y0) * math.cos(theta)
                height2 = (0.5 * radii[c]) ** 2 - 0.25 * xprime ** 2 - yprime ** 2
                if height2 >= 0:
                    my_label[[i], [j]] = sqrt((0.5 * radii[c]) ** 2 - 0.25 * xprime ** 2 - yprime ** 2)

    return my_label

size = 128
radii, centers = circle_randomizer()
#print(radii)
#print(centers)

#Make labels and samples consistent with rest of code
N = 100
circle_images = []
ellipse_images = []
coords = []
for sample in range(0, N):
    blank_image = create_blank_image(size)
    radii, centers = circle_randomizer()
    temp_image = image_creator(centers, radii, blank_image, size)
    circle_images.append(temp_image)
    temp_output = make_ellipses(size, radii, centers)
    ellipse_images.append(temp_output)
    coords.append(centers)
#print(labels)
#print(samples[0][40])

filenames = []
for i in range(0,N):
  np.save('ellipses_' + str(i) + '.npy', ellipse_images[i])
  filenames.append('ellipses_' + str(i) + '.npy')
  np.save('circles_' + str(i) + '.npy', circle_images[i])
circles_stack = np.stack(circle_images,axis=0)
ellipses_stack = np.stack(ellipse_images,axis=0)
np.save('ellipses_stack.npy', ellipses_stack)
np.save('circles_stack.npy', circles_stack)

# load training images and corresponding "labels"
# training samples
training_images_path = 'circles_stack.npy'
labels_path = 'ellipses_stack.npy'

X = np.load(training_images_path,'r')/20.
y = np.load(labels_path,'r')/20.

# Preprocessing for training images
def preprocessing_X(image_data, image_size):
    image_data = image_data.reshape(image_data.shape[0], image_size[0], image_size[1], 1)
    image_data = image_data.astype('float32')
    image_data = (image_data - np.amin(image_data))/(np.amax(image_data) - np.amin(image_data))
    return image_data


​
# preprocessing for "labels" (ground truth)
def preprocessing_Y(image_data, image_size):
    n_images = 0
    label = np.array([])
    for idx in range(image_data.shape[0]):
        img = image_data[idx,:,:]
        n, m = img.shape
        img = np.array(OneHotEncoder(n_values=nb_classes).fit_transform(img.reshape(-1,1)).todense())
        img = img.reshape(n, m, nb_classes)
        label = np.append(label, img)
        n_images += 1
    label_4D = label.reshape(n_images, image_size[0], image_size[1], nb_classes)
    return label_4D

# Split into train/test and make the shapes of tensors compatible with tensorflow format
nb_classes = 10
target_size = (128, 128)

#Below line randomizes which images are picked for train/test sets. ~20% will go to test.
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size = 0.2, random_state = 42)
X_train = preprocessing_X(X_train, target_size)
X_test = preprocessing_X(X_test, target_size)
y_train = preprocessing_Y(y_train, target_size)
y_test = preprocessing_Y(y_test, target_size)