from __future__ import absolute_import, division, print_function

from os import environ, getcwd

from os.path import join

import shutil

import re

import os

import argparse

import keras

import numpy as np

import pandas as pd

import sklearn as skl

import tensorflow as tf

from keras.applications.vgg19 import VGG19

from keras.applications import DenseNet169, InceptionResNetV2, DenseNet201

from keras.applications import NASNetMobile

from keras.layers import Dense, GlobalAveragePooling2D

from keras.metrics import binary_accuracy, binary_crossentropy, kappa_error

from keras.models import Model

from keras.optimizers import Adam

from keras.preprocessing.image import ImageDataGenerator

from custom_layers import *

from mura import Mura

pd.set_option('display.max_rows', 20)

pd.set_option('precision', 4)

np.set_printoptions(precision=4)

environ['TF_CPP_MIN_LOG_LEVEL'] = '2'  # Shut up tensorflow!

print("tf : {}".format(tf.__version__))

print("keras : {}".format(keras.__version__))

print("numpy : {}".format(np.__version__))

print("pandas : {}".format(pd.__version__))

print("sklearn : {}".format(skl.__version__))

# Hyper-parameters / Globals

BATCH_SIZE = 4  # tweak to your GPUs capacity

IMG_HEIGHT = 420  # ResNetInceptionv2 & Xception like 299, ResNet50/VGG/Inception 224, NASM 331

IMG_WIDTH = IMG_HEIGHT

CHANNELS = 3

DIMS = (IMG_HEIGHT, IMG_WIDTH, CHANNELS)  # blame theano

MODEL_TO_EVAL1 = './models/DenseNet169_420_HUMERUS.hdf5'

MODEL_TO_EVAL2 = './models/DenseNet169_420_HAND.hdf5'

MODEL_TO_EVAL3 = './models/DenseNet169_420_FINGER.hdf5'

MODEL_TO_EVAL4 = './models/DenseNet169_420_FOREARM.hdf5'

MODEL_TO_EVAL5 = './models/DenseNet169_420_ELBOW.hdf5'

MODEL_TO_EVAL6 = './models/DenseNet169_420_SHOULDER.hdf5'

MODEL_TO_EVAL7 = './models/DenseNet169_420_WRIST.hdf5'

MODEL_TO_EVAL8 = './models/DenseNet169_420_NEW_HIST.hdf5'

DATA_DIR = 'data_val/'

EVAL_CSV = 'valid.csv'

EVAL_DIR = 'data/val/'

parser = argparse.ArgumentParser(description='Input Path')

parser.add_argument('input_filename',default='valid_image_paths.csv', type=str)

parser.add_argument('output_path', default='prediction.csv', type=str)

proc_data_dir = join(os.getcwd(), 'data/val/')

proc_train_dir = join(proc_data_dir, 'train')

proc_val_dir = join(proc_data_dir, 'val')

class ImageString(object):

    _patient_re = re.compile(r'patient(\d+)')

    _study_re = re.compile(r'study(\d+)')

    _image_re = re.compile(r'image(\d+)')

    _study_type_re = re.compile(r'XR_(\w+)')

    def __init__(self, img_filename):

        self.img_filename = img_filename

        self.patient = self._parse_patient()

        self.study = self._parse_study()

        self.image_num = self._parse_image()

        self.study_type = self._parse_study_type()

        self.image = self._parse_image()

        self.normal = self._parse_normal()

        self.valid = self._parse_valid()

    def flat_file_name(self):

        return "{}_{}_patient{}_study{}_{}_image{}.png".format(self.valid,  self.study_type, self.patient, self.study,

                                                            self.normal, self.image)

    def _parse_patient(self):

        return int(self._patient_re.search(self.img_filename).group(1))

    def _parse_study(self):

        return int(self._study_re.search(self.img_filename).group(1))

    def _parse_image(self):

        return int(self._image_re.search(self.img_filename).group(1))

    def _parse_study_type(self):

        return self._study_type_re.search(self.img_filename).group(1)

    def _parse_normal(self):

        return "normal" if ("negative" in self.img_filename) else "abnormal"

    def _parse_normal_label(self):

        return 1 if("negative" in self.img_filename) else 0

    def _parse_valid(self):

        return "valid" if ("valid" in self.img_filename) else "test"

def preprocess_img(img):

    # Histogram normalization in v channel

    hsv = color.rgb2hsv(img)

    hsv[:, :, 2] = exposure.equalize_hist(hsv[:, :, 2])

    img = color.hsv2rgb(hsv)

    # central square crop

    min_side = min(img.shape[:-1])

    centre = img.shape[0] // 2, img.shape[1] // 2

    img = img[centre[0] - min_side // 2:centre[0] + min_side // 2,

              centre[1] - min_side // 2:centre[1] + min_side // 2,

              :]

    # rescale to standard size

    img = transform.resize(img, (IMG_SIZE, IMG_SIZE))

    # roll color axis to axis 0

    img = np.rollaxis(img, -1)

    return img

def eval(args=None):

    args= parser.parse_args()

    # load up our csv with validation factors

    data_dir = join(getcwd(), DATA_DIR)

    eval_csv = join(data_dir, EVAL_CSV)

    true_labels=[]

    ###########################################

    df = pd.read_csv(args.input_filename, names=['img', 'label'], header=None)

    samples = [tuple(x) for x in df.values]

 #   for img, label in samples:

 #       #assert ("negative" in img) is (label is 0)

 #       enc = ImageString(img)

 #       true_labels.append(enc._parse_normal_label())

 #       cat_dir = join(proc_val_dir, enc.normal)

 #       if not os.path.exists(cat_dir):

 #           os.makedirs(cat_dir)

 #       shutil.copy2(enc.img_filename, join(cat_dir, enc.flat_file_name()))

    ###########################################

    eval_datagen = ImageDataGenerator(rescale=1./255

#                                    , histogram_equalization=True

                                      )

    eval_generator = eval_datagen.flow_from_directory(

         EVAL_DIR, class_mode='binary', shuffle=False,target_size=(IMG_HEIGHT, IMG_WIDTH), batch_size=BATCH_SIZE)

    n_samples = eval_generator.samples

    base_model = DenseNet169(input_shape=DIMS, weights='imagenet', include_top=False)  #weights='imagenet'

    x = base_model.output

    x = GlobalAveragePooling2D(name='avg_pool')(x)  # comment for RESNET

 #   x = WildcatPool2d()(x)

    x = Dense(1, activation='sigmoid', name='predictions')(x)

    model = Model(inputs=base_model.input, outputs=x)

    model.load_weights(MODEL_TO_EVAL8)

    model.compile(optimizer=Adam(lr=1e-3)

                  , loss=binary_crossentropy

#                  , loss=kappa_error

                  , metrics=['binary_accuracy'])

    score, acc = model.evaluate_generator(eval_generator, n_samples / BATCH_SIZE)

    print(model.metrics_names)

    print('==> Metrics with eval')

    print("loss :{:0.4f} \t Accuracy:{:0.4f}".format(score, acc))

    y_pred = model.predict_generator(eval_generator, n_samples / BATCH_SIZE)

#    print(y_pred)

#    df_filenames = pd.Series(np.array(eval_generator.filenames), name='filenames')

#    df_classes   = pd.Series(np.array(y_pred), name='classes')

#    prediction_data = pd.concat([df_filenames, df_classes,])

#    prediction_data.to_csv(args.output_path + "/prediction.csv")

    mura = Mura(eval_generator.filenames, y_true = eval_generator.classes, y_pred1=y_pred, y_pred2=y_pred, y_pred3=y_pred, y_pred4= y_pred, y_pred5= y_pred, output_path= args.output_path)

    print(mura.metrics_by_encounter())

if __name__ == '__main__':

    eval()