import tensorflow as tf
from google.colab.patches import cv2_imshow

tf.logging.set_verbosity(tf.logging.INFO)

model_pathed_name = "./circle_model"
training_steps = 2000
batch_size = 32
eval_set_size = 500

IMAGE_SIZE = 64
image_size = IMAGE_SIZE
max_circle_radius = image_size // 2
noise_level = 2

import numpy as np
import tensorflow as tf

IMAGE_SIZE = 64

def noisy_circle(size, radius, noise):
    img = np.zeros((size, size), dtype=np.float)

    # Circle
    row = np.random.randint(size)
    col = np.random.randint(size)
    rad = np.random.randint(10, max(10, radius))
    # draw_circle(img, row, col, rad)

    # Noise
    img += noise * np.random.rand(*img.shape)
    return (row, col, rad), img

def find_circle(img):
    model_path_name = "./circle_model"
    img_data = np.zeros((1, img.shape[0], img.shape[1]))
    img_data[0, :, :] = img
    predict_input_function = tf.estimator.inputs.numpy_input_fn(
        x={"x": img_data},
        num_epochs=1,
        shuffle=False)
    tf_circle_detector = tf.estimator.Estimator(model_fn=cnn_m, model_dir=model_path_name)
    prediction_result = tf_circle_detector.predict(input_fn=predict_input_function)
    prediction_results = list(prediction_result)
    prediction = cnn_trainer.IMAGE_SIZE * prediction_results[0]['location']

    return prediction[0], prediction[1], prediction[2]

def iou(params0, params1):
    row0, col0, rad0 = params0
    row1, col1, rad1 = params1

    shape0 = Point(row0, col0).buffer(rad0)
    shape1 = Point(row1, col1).buffer(rad1)

    return (
        shape0.intersection(shape1).area /
        shape0.union(shape1).area
    )

def draw_circle(img, row, col, rad):
    rr, cc, val = circle_perimeter_aa(row, col, rad)
    valid = (
        (rr >= 0) &
        (rr < img.shape[0]) &
        (cc >= 0) &
        (cc < img.shape[1])
    )
    img[rr[valid], cc[valid]] = val[valid]

import numpy as np
from shapely.geometry.point import Point
from skimage.draw import circle_perimeter_aa
import tensorflow as tf

def create_training_data(samples, image_size, max_radius, noise_level):
    training_images = np.zeros((samples, image_size, image_size))
    training_labels = np.zeros((samples, 3), dtype=np.float64)

    image = np.zeros((image_size, image_size), dtype=np.float)
    for i in range(samples):
        params, image = noisy_circle(image_size, max_radius, noise_level)
        training_images[i, :, :] = image
        training_labels[i] = params

    # Normalize to relative image coordinates: every image has dimensions 1.0 x 1.0, pixels are stored as floats
    training_labels /= image_size

    return training_images, training_labels

# Adapted from Tensorflow mnist digit recognition example:
# https://github.com/tensorflow/docs/blob/master/site/en/tutorials/estimators/cnn.ipynb
def model_function(features, labels, mode):
    """Model function for CNN."""
    # Input layer
    input_layer = tf.reshape(features["x"], [-1, IMAGE_SIZE, IMAGE_SIZE, 1])

    # Convolutional, pooling layer 1
    conv1 = tf.layers.conv2d(
        inputs=input_layer,
        filters=32,
        kernel_size=[5, 5],
        padding="same",
        activation=tf.nn.relu)

    pool1 = tf.layers.max_pooling2d(inputs=conv1, pool_size=[2, 2], strides=2)

    # Convolutional, pooling layer 2
    conv2 = tf.layers.conv2d(
        inputs=pool1,
        filters=64,
        kernel_size=[5, 5],
        padding="same",
        activation=tf.nn.relu)

    pool2 = tf.layers.max_pooling2d(inputs=conv2, pool_size=[2, 2], strides=2)

    # Dense layer, 1024 units is unchanged from example
    pool2_flat = tf.reshape(pool2, [-1, IMAGE_SIZE // 4 * IMAGE_SIZE // 4 * 64])
    dense = tf.layers.dense(inputs=pool2_flat, units=1024, activation=tf.nn.relu)
    dropout = tf.layers.dropout(inputs=dense, rate=0.4, training=mode == tf.estimator.ModeKeys.TRAIN)

    # Output layer, very simple
    output_layer = tf.layers.dense(inputs=dropout, units=3)
    predictions = {"location": output_layer}

    if mode == tf.estimator.ModeKeys.PREDICT:
        return tf.estimator.EstimatorSpec(mode=mode, predictions=predictions)

    # Calculate Loss (for both TRAIN and EVAL modes)
    # This has to be after the PREDICT part, since when in PREDICT mode the labels will be None.
    loss = tf.losses.mean_squared_error(labels, output_layer)

    if mode == tf.estimator.ModeKeys.TRAIN:
        optimizer = tf.train.GradientDescentOptimizer(learning_rate=.08)
        train_op = optimizer.minimize(loss=loss, global_step=tf.train.get_global_step())
        return tf.estimator.EstimatorSpec(mode=mode, loss=loss, train_op=train_op)

    if mode == tf.estimator.ModeKeys.EVAL:
        eval_metric_ops = {"MSE": tf.metrics.mean_squared_error(labels=labels, predictions=output_layer)}
        return tf.estimator.EstimatorSpec(mode=mode, loss=loss, eval_metric_ops=eval_metric_ops)

    tf.logging.set_verbosity(tf.logging.INFO)

    model_pathed_name = "./circle_detection_model"
    training_steps = 8000
    batch_size = 32
    eval_set_size = 500

    image_size = IMAGE_SIZE
    max_circle_radius = image_size // 2
    noise_level = .5

    train_data, train_labels = create_training_data(training_steps, image_size, max_circle_radius, noise_level)
    eval_data, eval_labels = create_training_data(eval_set_size, image_size, max_circle_radius, noise_level)

    tf_circle_detector = tf.estimator.Estimator(model_fn=model_function, model_dir=model_pathed_name)

    tensors_to_log = {}
    logging_hook = tf.train.LoggingTensorHook(tensors=tensors_to_log, every_n_iter=500)

    # Train the model
    train_input_fn = tf.estimator.inputs.numpy_input_fn(
        x={"x": train_data},
        y=train_labels,
        batch_size=batch_size,
        num_epochs=None,
        shuffle=True)

    tf_circle_detector.train(input_fn=train_input_fn, steps=training_steps, hooks=[logging_hook])

    eval_input_fn = tf.estimator.inputs.numpy_input_fn(
        x={"x": eval_data},
        y=eval_labels,
        num_epochs=1,
        shuffle=False)

    eval_results = tf_circle_detector.evaluate(input_fn=eval_input_fn)

import cv2 as cv

number = 1000
model_path_name = "./circle_detection_model"
image_size = IMAGE_SIZE
max_radius = 32
noise_level = .5

tf.logging.set_verbosity(tf.logging.ERROR)
inference_data, inference_labels = create_training_data(number, image_size, max_radius, noise_level)

predict_input_function = tf.estimator.inputs.numpy_input_fn(
    x={"x": inference_data},
    num_epochs=1,
    shuffle=False)
tf_circle_detector = tf.estimator.Estimator(model_fn=model_function, model_dir=model_path_name)
inference_result = tf_circle_detector.predict(input_fn=predict_input_function)

# Results will be returned from Tensorflow as a generator,
# wrapped in the Tensorflow verbiage defined in PREDICT mode
inference_results = list(inference_result)

# Batch statistics: average IOU. Doing this separately from the display loop below because I don't expect
# that to actually accumulate all results
ious = np.zeros(number)
for i, prediction in enumerate(inference_results):
    prediction = IMAGE_SIZE * prediction['location']
    true_x, true_y, true_r = IMAGE_SIZE * inference_labels[i]
    ious[i] = iou((true_x, true_y, max(true_r, 1)), (prediction[0], prediction[1], prediction[2]))
print()
print("{0} samples. Average IOU: {1:2.2f} Min IOU: {2:2.2f}".format(number, np.average(ious), np.min(ious)))
print("Samples with IOU > .5: {0}%".format(100 * len(np.where(ious > .5)[0]) / number))
print("------------------------")

for i, prediction in enumerate(inference_results):
    # Unwrap the Tensorflow output.
    prediction = IMAGE_SIZE * prediction['location']
    prediction_int = [int(round(prediction[0])), int(round(prediction[1])), max(int(round(prediction[2])), 1)]

    image = np.array(np.reshape(inference_data[i], (image_size, image_size)) * 256, dtype=np.uint8)

    overlay_image = cv.cvtColor(np.copy(image), cv.COLOR_GRAY2BGR)
    # Reverse the order of the coordinates to go from matrix indexing to image indexing
    cv.circle(overlay_image, (prediction_int[1], prediction_int[0]), prediction_int[2], color=(200, 0, 255), thickness=2)

    upscale_factor = 6
    upscaled_image = cv.resize(image, dsize=(0, 0), fx=upscale_factor, fy=upscale_factor, interpolation=cv.INTER_NEAREST)
    overlay_image = cv.resize(overlay_image, dsize=(0, 0), fx=upscale_factor, fy=upscale_factor, interpolation=cv.INTER_NEAREST)

    #cv.imwrite("upscaled.png", upscaled_image)
    #cv.imwrite("overlay.png", overlay_image)

    true_x, true_y, true_r = IMAGE_SIZE * inference_labels[i]
    print("True x, y, r: {0} {1} {2}".format(int(true_x), int(true_y), int(true_r)))
    print("Pred x, y, r: {0} {1} {2}".format(prediction_int[0], prediction_int[1], prediction_int[2]))
    print("IOU: {0:2.2f} \n".format(iou((true_x, true_y, true_r), (prediction[0], prediction[1], prediction[2]))))

    while True:
        k = cv.waitKey(1)
        # Press j or f for next image, q or escape to quit
        if k == 'q' or k == 27 or k == 'Q' or k == 1048603 or k == 1048689:
            sys.exit("Quitting")
        elif k == 'j' or k == 102 or k == 'f' or k == 106 or k == 65363:
            break

        cv2_imshow(upscaled_image)
        cv2_imshow(overlay_image)