"""
# Reminders
- ## Run with GPU accelerator or you'll hate yourself (need to verify phone # with kaggle)
- ## Turn on Internet in order to pip install shit

# So far
- Tried with a generic CNN on the dataset as is, got around 50% accuracy
- Digits are demarcated in each individual image, using openCV

# TO DO
- Get a model trained to identify INDIVIDUAL digits, then find a way to pass in images consisting of MULTIPLE digits (i.e. the original data) to that model, and then concatenate the outputs.
    - from a friend: "the output of our neural network is just 1 label, and then we combine multiple labels together"
    - my interpretation: each image passes thru the opencv pipeline, each digit is individually run on a prediction model, and then the labels are put together
- note from a friend: "don't use cv2 to resize the cropped image, need to resize by a scale"

"""
import numpy as np 
import pandas as pd
import h5py
import matplotlib.pyplot as plt
from sklearn.model_selection import RepeatedKFold
from keras.models import Sequential
from keras.layers import Dense
from sklearn.metrics import accuracy_score
from keras.utils.np_utils import to_categorical,normalize
!pip install imutils
# Input data files are available in the read-only "../input/" directory
# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using "Save & Run All" 
# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session

import os
for dirname, _, filenames in os.walk('/kaggle/input'):
    for filename in filenames:
        print(os.path.join(dirname, filename))

f = h5py.File('/kaggle/input/mnistsynthetic/MNIST_synthetic.h5', 'r')
list(f.keys())
# Load in data from h5py format
X_train = f['train_dataset'].value
X_test = f['test_dataset'].value
y_train = f['train_labels'].value

# Reduce certain channels (see below image); idk if this does anything
X_train = np.where(X_train<100, 0, X_train)
X_test = np.where(X_test<100, 0, X_test)


print("Train shape:", X_train.shape)
print("Test shape:", X_test.shape)
print("Train labels shape:", y_train.shape)
"""
# Preprocessing

Use OpenCV to apply a threshold (to reduce signal-to-noise), extract contours, and draw rectangles around digits. Our goal is to make the number of digits extracted from each image AS CLOSE AS POSSIBLE to the actual number (which we must calculate). After, we can throw this all into the CNN.

From here: https://github.com/alexanderhale/COMP551/blob/master/Project_3/submission-notebooks/digit_localization.ipynb
"""
import cv2
import imutils
# print(cv2.__version__) # 4.4.0

##### DIGIT EXTRACTION PARAMETERS #####
# We can tweak these (using grid search or something)

threshold = 50    # binary threshold; pixels above this value are set to white, else set to black
min_digit_width = 5  # number of pixels wide a contour box must be to be detected as a digit
min_digit_height = 5  # number of pixels tall a contour box must be to be detected as a digit
max_digit_width = 30
max_digit_height = 30

allDigitWindows = []   # holds the windows containing digits (in black & white)

for i, img in enumerate(X_train):
        # threshold the image to only black (non-digit) or white (digit)
    _, t_img = cv2.threshold(img.copy(), threshold, 255, cv2.THRESH_BINARY) # pixels above `threshold` are set to 255 (white), otherwise set to black
    # ^ the first return is not needed, unless we use Otsu's binarization (for bimodal images)
    digitBoxes = t_img.copy()

    # find contours in the thresholded image
    contours, _ = cv2.findContours(t_img.astype(np.uint8).copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) # also try cv2.RETR_TREE for second arg

    # loop over candidates and extract real digits
    digitWindows = []
    for c in contours:
        # compute the bounding box of the contour
        (x, y, w, h) = cv2.boundingRect(c)

        # if the contour is sufficiently large, decide that it is a digit
        if w >= min_digit_width and w <= max_digit_width and h >= min_digit_height and h <= max_digit_height:
            digitWindows.append(np.asarray(t_img[y:y+h, x:x+w]))
            cv2.rectangle(digitBoxes, (x, y), (x + w, y + h), (130, 130, 255), 1)
        elif w >= min_digit_width or h >= min_digit_height:
            digitWindows.append(np.asarray(t_img[y:y+h, x:x+w]))

    # add digitContours to array
    allDigitWindows.append(digitWindows)
    
    #### visualization ####
    if(len(X_train) % 100*i == 0): # print some of the examples
        cv.drawContours(t_img, contours, -1, (155,115,150), 3)
        plt.figure(figsize = (20,20))
        plt.subplot(131), plt.imshow(img[:,:,0], vmin=0, vmax=255) # original image
        plt.subplot(132), plt.imshow(t_img, vmin=0, vmax=255) # thresholded image
        plt.subplot(133), plt.imshow(digitBoxes, vmin=0, vmax=255) # digits identified

        plt.show()
    ################



"""
# Model

Just a generic CNN model, from here: https://github.com/alexanderhale/COMP551/blob/master/Project_3/submission-notebooks/kaggle_cnn.ipynb
"""
from keras.layers import Dense, Dropout, Flatten, Conv2D, MaxPooling2D, LeakyReLU, ZeroPadding2D,Convolution2D

def cnn(): 
    # create model
    model = Sequential()
    model.add(ZeroPadding2D((1,1),input_shape=(64,64,1)))
    model.add(Conv2D(32, kernel_size=(3, 3), activation='relu'))
    model.add(ZeroPadding2D((1,1)))
    model.add(Conv2D(32, kernel_size=(3, 3), activation='relu'))
    model.add(MaxPooling2D(pool_size=(2, 2)))

    model.add(ZeroPadding2D((1,1)))
    model.add(Conv2D(64, kernel_size=(3, 3), activation='relu'))
    model.add(ZeroPadding2D((1,1)))
    model.add(Conv2D(64, kernel_size=(3, 3), activation='relu'))
    model.add(MaxPooling2D(pool_size=(2, 2)))
    model.add(Dropout(0.5))

    model.add(ZeroPadding2D((1,1)))
    model.add(Conv2D(128, kernel_size=(3, 3), activation='relu'))
    model.add(ZeroPadding2D((1,1)))
    model.add(Conv2D(128, kernel_size=(3, 3), activation='relu'))
    model.add(MaxPooling2D(pool_size=(2, 2)))
    model.add(Dropout(0.5))

    model.add(Flatten())
    model.add(Dense(512, activation='relu'))
    model.add(Dropout(0.5))
    model.add(Dense(5, activation='softmax'))
    # Compile
    model.compile(loss='binary_crossentropy', optimizer='adam', metrics=['accuracy'])
    return model
model = cnn()
history = model.fit(X_train, y_train, batch_size=100, epochs=10, shuffle=False)
predictions = np.argmax(model.predict(X_test), axis=1)