import tensorflow as tf
from tensorflow import keras
import numpy as np
import cv2
import os
import sklearn as skl
from sklearn import model_selection
import pandas as pd
from tensorflow.keras import backend as K
import math
import matplotlib.pyplot as plt
import scipy.stats


def weightedContr(a, w):
    average = np.average(a, weights=w)
    variance = np.average((a - average) ** 2, weights=w)
    return math.sqrt(variance)


def sigmoid(x):
    return 1 / (1 + 2 ** (-x))


def gammaToLinear(x):
    return x ** 2.2


def linearToGamma(x):
    return x ** (1 / 2.2)


def normalize(a, weights=None):
    array = a
    #	shape = array.shape
    #	array = scipy.stats.yeojohnson(array.flatten())[0].reshape(shape)
    array -= np.average(array, weights=weights)
    array /= weightedContr(array, weights)
    return array


def scaleAndCrop(openCVimg, length):
    height, width, channels = openCVimg.shape
    va = height / width
    ha = width / height
    if va > ha:
        dim = (length, int(va * length))
    if ha > va:
        dim = (int(ha * length), length)
    else:
        dim = (length, length)
    resized = cv2.resize(openCVimg, dim, cv2.INTER_AREA)
    height, width, channels = resized.shape
    y = int((height - length) / 2)
    x = int((width - length) / 2)
    return resized[y:y + length, x:x + length]


def adjustExposure(imgLinear):
    return imgLinear / np.mean(imgLinear) * 0.5


def preprocess(imgFiles, size):
    processed = []

    for i, f in enumerate(imgFiles):
        print((i + 1) / len(imgFiles)) if i % 10 == 0 else None

        img = scaleAndCrop(cv2.imread(f), size)
        YCrCb = cv2.cvtColor(img, cv2.COLOR_BGR2YCR_CB)
        Y, Cr, Cb = cv2.split(YCrCb)
        Y, Cr, Cb = [np.array(channel).astype(np.float64) for channel in (Y, Cr, Cb)]
        Y /= 255
        Cr /= 255
        Cb /= 255

        Y = adjustExposure(gammaToLinear(Y))
        Y = linearToGamma(Y)

        lumaWeights = np.copy(Y)
        chromaWeights = np.copy(Y)
        blackThreshold = 0.1

        for row, rowData in enumerate(Y):
            for dot, dotValue in enumerate(rowData):
                lumaWeights[row][dot] = sigmoid((dotValue - blackThreshold) / blackThreshold ** 2)
                chromaWeights[row][dot] = Y[row][dot]  # max(0, math.log(Y[row][dot] + (1 / 255), 255) + 1)

        Y = normalize(Y, lumaWeights)
        Cb = normalize(Cb, chromaWeights)
        Cr = normalize(Cr, chromaWeights)
        Y *= lumaWeights
        Cb *= chromaWeights
        Cr *= chromaWeights

        Y = sigmoid(Y)
        Cr = sigmoid(Cr)
        Cb = sigmoid(Cb)

        '''
        fig, axs = plt.subplots(nrows=2, ncols=2)
        axs[0][0].imshow(cv2.cvtColor((cv2.merge([Y, Cr, Cb]) * 255).astype(np.uint8), cv2.COLOR_YCR_CB2RGB), vmin=0, vmax=1)
        axs[1][0].imshow(Y, cmap='jet', vmin=0, vmax=1)
        axs[0][1].imshow(Cb, cmap='jet', vmin=0, vmax=1)
        axs[1][1].imshow(Cr, cmap='jet', vmin=0, vmax=1)
        plt.show()

        processed.append({
            'Y' : Y
        ,	'Cr' : Cr
        ,	'Cb' : Cb
        })
        '''

        merged = (cv2.merge([Y, Cr, Cb]) * 255).astype(np.uint8)
        processed.append(merged)

        '''
        cv2.imwrite('Processed Train Images/' + f.rsplit('/', 1)[-1], merged)
        print((i + 1) / len(imgFiles))
        '''


    return np.array(processed)


##########


trainCSV = pd.read_csv('../input/aptos2019-blindness-detection/train.csv')


def loads(folder, csv):
    images = []
    for id_code in csv['id_code']:
        path = os.path.join(folder, id_code + '.png')
        idata = cv2.imread(path)
        images.append(idata)
    print(images)
    return np.array(images)


allImgData = loads('../input/processed-train-images/processed train images/Processed Train Images', trainCSV)
print(allImgData.shape)

GROUP_COUNT = 5

#diagnosis = np.array([[int(y > i) for i in range(0, GROUP_COUNT - 1)] for y in trainCSV['diagnosis']]).astype(np.float64)
diagnosis = np.array([[int(y == i) for i in range(0, GROUP_COUNT)] for y in trainCSV['diagnosis']]).astype(np.float64)

X_train, X_valid, y_train, y_valid = skl.model_selection.train_test_split(allImgData, diagnosis, train_size=0.9)


# create model
model = keras.Sequential()

# add model layers
model.add(keras.layers.Conv2D(256, kernel_size=9, activation='selu', input_shape=(256, 256, 3), data_format='channels_last'))
model.add(keras.layers.MaxPooling2D(pool_size=(2, 2), strides=None, padding='same', data_format='channels_last'))
model.add(keras.layers.Conv2D(256, kernel_size=7, activation='selu', data_format='channels_last'))
model.add(keras.layers.MaxPooling2D(pool_size=(2, 2), strides=None, padding='same', data_format='channels_last'))
model.add(keras.layers.Conv2D(128, kernel_size=5, activation='selu', data_format='channels_last'))
model.add(keras.layers.MaxPooling2D(pool_size=(2, 2), strides=None, padding='same', data_format='channels_last'))
model.add(keras.layers.Conv2D(64, kernel_size=5, activation='selu', data_format='channels_last'))
model.add(keras.layers.Flatten())
model.add(keras.layers.Dense(256, activation='selu'))
model.add(keras.layers.Dense(128, activation='selu'))
model.add(keras.layers.Dense(64, activation='selu'))
model.add(keras.layers.Dense(5, activation='softmax'))


def loss(y_true, y_pred):
    weights = K.cast(K.abs(K.argmax(y_true, axis=1) - K.argmax(y_pred, axis=1))/(K.int_shape(y_pred)[1] - 1), dtype='float32')
    return (1.0 + weights) * keras.losses.categorical_crossentropy(y_true, y_pred)


#compile model using accuracy to measure model performance
model.compile(optimizer=keras.optimizers.Adam(lr=0.000001, beta_1=0.9, beta_2=0.99999, epsilon=K.epsilon(), decay=0.5)
			  , loss=loss, metrics=['accuracy'])

#train the model
model.fit(X_train, y_train, validation_data=(X_valid, y_valid), epochs=50)


def round04(x):
    return max(0, min(4, round(x)))


testCSV = pd.read_csv('../input/aptos2019-blindness-detection/test.csv')
transformedTestImages = preprocess([f'../input/aptos2019-blindness-detection/test_images/{im}.png' for im in testCSV['id_code']], 256)

print(transformedTestImages)
print(transformedTestImages.shape)

rawOutput = model.predict(transformedTestImages)
print(rawOutput)

diagResults = [np.argmax(y) for y in rawOutput]  # [round04(np.sum(y)) for y in rawOutput]
print(diagResults)

submission = pd.DataFrame({
    'id_code' : testCSV['id_code'].tolist(),
    'diagnosis' : diagResults
})

output = submission.to_csv('submission.csv', index=False, header=True)  # OUTPUT