#!/usr/bin/env python2.7

import pydicom, cv2, re
import os, fnmatch, sys
import numpy as np
from keras.preprocessing.image import ImageDataGenerator
from keras import backend as K


from fcn_model import fcn_model
from helpers import center_crop, lr_poly_decay, get_SAX_SERIES

seed = 1234
np.random.seed(seed)

SAX_SERIES = get_SAX_SERIES()
SUNNYBROOK_ROOT_PATH = '/media/heshengji/Ubuntu/CardiacData/MICCAI2009'
# SUNNYBROOK_ROOT_PATH = os.path.abspath(os.path.join(os.getcwd(),'..','CardiacData','MICCAI2009'))

TRAIN_CONTOUR_PATH = os.path.join(SUNNYBROOK_ROOT_PATH,
                            'Sunnybrook Cardiac MR Database ContoursPart3',
                            'TrainingDataContours')
TRAIN_IMG_PATH = os.path.join(SUNNYBROOK_ROOT_PATH,
                        'challenge_training')


def shrink_case(case):
    toks = case.split('-')
    def shrink_if_number(x):
        try:
            cvt = int(x)
            return str(cvt)
        except ValueError:
            return x
    return '-'.join([shrink_if_number(t) for t in toks])


class Contour(object):
    def __init__(self, ctr_path):
        self.ctr_path = ctr_path
        match = re.search(r'/([^/]*)/contours-manual/IRCCI-expert/IM-0001-(\d{4})-.*', ctr_path)
        self.case = shrink_case(match.group(1))
        self.img_no = int(match.group(2))

    def __str__(self):
        return '<Contour for case %s, image %d>' % (self.case, self.img_no)

    __repr__ = __str__


def read_contour(contour, data_path):
    filename = 'IM-%s-%04d.dcm' % (SAX_SERIES[contour.case], contour.img_no)
    full_path = os.path.join(data_path, contour.case, filename)
    f = pydicom.read_file(full_path)
    img = f.pixel_array.astype('int')
    mask = np.zeros_like(img, dtype='uint8')
    coords = np.loadtxt(contour.ctr_path, delimiter=' ').astype('int')
    cv2.fillPoly(mask, [coords], 1)
    if img.ndim < 3:
        img = img[..., np.newaxis]
        mask = mask[..., np.newaxis]

    return img, mask


def map_all_contours(contour_path, contour_type, shuffle=True):
    contours = [os.path.join(dirpath, f)
        for dirpath, dirnames, files in os.walk(contour_path)
        for f in fnmatch.filter(files,
                        'IM-0001-*-'+contour_type+'contour-manual.txt')]
    if shuffle:
        print('Shuffling data')
        np.random.shuffle(contours)
    print('Number of examples: {:d}'.format(len(contours)))
    contours = list(map(Contour, contours))

    return contours


def export_all_contours(contours, data_path, crop_size):
    print('\nProcessing {:d} images and labels ...\n'.format(len(contours)))
    images = np.zeros((len(contours), crop_size, crop_size, 1))
    masks = np.zeros((len(contours), crop_size, crop_size, 1))
    for idx, contour in enumerate(contours):
        img, mask = read_contour(contour, data_path)
        img = center_crop(img, crop_size=crop_size)
        mask = center_crop(mask, crop_size=crop_size)
        images[idx] = img
        masks[idx] = mask

    return images, masks


if __name__== '__main__':
    # if len(sys.argv) < 3:
    #     sys.exit('Usage: python %s <i/o> <gpu_id>' % sys.argv[0])
    # contour_type = sys.argv[1]
    # os.environ['CUDA_VISIBLE_DEVICES'] = sys.argv[2]

    if len(sys.argv) < 2:
        sys.exit('Usage: python %s <i/o> <gpu_id>' % sys.argv[0])
    elif len(sys.argv) < 3:
        contour_type = sys.argv[1]
    else:
        contour_type = sys.argv[1]
        os.environ['CUDA_VISIBLE_DEVICES'] = sys.argv[2]
    # contour_type = 'i'

    crop_size = 100

    print('Mapping ground truth '+contour_type+' contours to images in train...')
    train_ctrs = map_all_contours(TRAIN_CONTOUR_PATH, contour_type, shuffle=True)
    print('Done mapping training set')

    split = int(0.1*len(train_ctrs))
    dev_ctrs = train_ctrs[0:split]
    train_ctrs = train_ctrs[split:]

    print('\nBuilding Train dataset ...')
    img_train, mask_train = export_all_contours(train_ctrs,
                                                TRAIN_IMG_PATH,
                                                crop_size=crop_size)
    print('\nBuilding Dev dataset ...')
    img_dev, mask_dev = export_all_contours(dev_ctrs,
                                            TRAIN_IMG_PATH,
                                            crop_size=crop_size)

    input_shape = (crop_size, crop_size, 1)
    num_classes = 2

    if contour_type == 'i':
        weights = 'weights/sunnybrook_i.h5'
    elif contour_type == 'o':
        weights = 'weights/sunnybrook_o.h5'
    else:
        sys.exit('\ncontour type "%s" not recognized\n' % contour_type)


    # model = fcn_model(input_shape, num_classes, weights=weights)
    model = fcn_model(input_shape, num_classes, weights=None)

    kwargs = dict(
        rotation_range=180,
        zoom_range=0.0,
        width_shift_range=0.0,
        height_shift_range=0.0,
        horizontal_flip=True,
        vertical_flip=True,
    )
    image_datagen = ImageDataGenerator(**kwargs)
    mask_datagen = ImageDataGenerator(**kwargs)

    epochs = 1000
    # epochs = 500
    mini_batch_size = 5

    image_generator = image_datagen.flow(img_train, shuffle=False,
                                    batch_size=mini_batch_size, seed=seed)
    mask_generator = mask_datagen.flow(mask_train, shuffle=False,
                                    batch_size=mini_batch_size, seed=seed)
    train_generator = zip(image_generator, mask_generator)

    max_iter = (len(train_ctrs) // mini_batch_size) * epochs
    curr_iter = 0
    base_lr = K.eval(model.optimizer.lr)
    lrate = lr_poly_decay(model, base_lr, curr_iter, max_iter, power=0.5)


    import datetime

    op_acc=0.99
    for e in range(epochs):
        starttime = datetime.datetime.now()
        print('Main Epoch {:d}'.format(e+1))
        print('Learning rate: {:6f}'.format(lrate))
        train_result = []
        for iteration in range(len(img_train) // mini_batch_size):
            img, mask = next(train_generator)
            res = model.train_on_batch(img, mask)
            curr_iter += 1
            lrate = lr_poly_decay(model, base_lr, curr_iter,
                                  max_iter, power=0.5)
            train_result.append(res)
        train_result = np.asarray(train_result)
        # train_result = np.nan_to_num(train_result)
        # print(train_result)
        train_result = np.mean(train_result, axis=0).round(decimals=10)
        print('Train result {}: {}'.format(model.metrics_names, train_result))
        print('Evaluating dev set ...')
        result = model.evaluate(img_dev, mask_dev, batch_size=32)
        result = np.round(result, decimals=10)
        print('Dev set result {}: {}'.format(model.metrics_names, result))

        if not os.path.exists('model_logs'):
            os.makedirs('model_logs')

        if result[1]>op_acc or e==epochs-1:
            op_acc=round(result[1],4)

            save_file = '_'.join(['sunnybrook', contour_type,
                                  'epoch', str(e+1),str(op_acc)]) + '.h5'
            save_path = os.path.join('model_logs', save_file)
            print('Saving model weights to {}'.format(save_path))
            model.save_weights(save_path)


        endtime = datetime.datetime.now()
        print('time:{:d}s'.format((endtime - starttime).seconds))