#!/usr/bin/env python

from __future__ import division, print_function

import os

import glob

import numpy as np

import matplotlib.pyplot as plt

import cv2

from rvseg import opts, patient, dataset, models

def load_patient_images(path, normalize=True):

    p = patient.PatientData(path)

    # reshape to account for channel dimension

    images = np.asarray(p.images, dtype='float64')[:,:,:,None]

    # maybe normalize images

    if normalize:

        dataset.normalize(images, axis=(1,2))

    return images, p.index, p.labeled, p.rotated

def get_contours(mask):

    mask_image = np.where(mask > 0.5, 255, 0).astype('uint8')

    im2, coords, hierarchy = cv2.findContours(mask_image, cv2.RETR_LIST, cv2.CHAIN_APPROX_NONE)

    if not coords:

        print("No contour detected.")

        coords = np.ones((1, 1, 1, 2), dtype='int')

    if len(coords) > 1:

        print("Multiple contours detected.")

        lengths = [len(coord) for coord in coords]

        coords = [coords[np.argmax(lengths)]]

    coords = np.squeeze(coords[0], axis=(1,))

    coords = np.append(coords, coords[:1], axis=0)

    return coords

def save_image(figname, image, mask_pred, alpha=0.3):

    cmap = plt.cm.gray

    plt.figure(figsize=(8, 3.75))

    plt.subplot(1, 2, 1)

    plt.axis("off")

    plt.imshow(image, cmap=cmap)

    plt.subplot(1, 2, 2)

    plt.axis("off")

    plt.imshow(image, cmap=cmap)

    plt.imshow(mask_pred, cmap=cmap, alpha=alpha)

    plt.savefig(figname, bbox_inches='tight')

    plt.close()

def main():

    # Sort of a hack:

    # args.checkpoint = turns on saving of images

    args = opts.parse_arguments()

    args.checkpoint = False     # override for now

    glob_search = os.path.join(args.datadir, "patient*")

    patient_dirs = sorted(glob.glob(glob_search))

    if len(patient_dirs) == 0:

        raise Exception("No patient directors found in {}".format(data_dir))

    # get image dimensions from first patient

    images, _, _, _ = load_patient_images(patient_dirs[0], args.normalize)

    _, height, width, channels = images.shape

    classes = 2                 # hard coded for now

    contour_type = {'inner': 'i', 'outer': 'o'}[args.classes]

    print("Building model...")

    string_to_model = {

        "unet": models.unet,

        "dilated-unet": models.dilated_unet,

        "dilated-densenet": models.dilated_densenet,

        "dilated-densenet2": models.dilated_densenet2,

        "dilated-densenet3": models.dilated_densenet3,

    }

    model = string_to_model[args.model]

    m = model(height=height, width=width, channels=channels, classes=classes,

              features=args.features, depth=args.depth, padding=args.padding,

              temperature=args.temperature, batchnorm=args.batchnorm,

              dropout=args.dropout)

    m.load_weights(args.load_weights)

    for path in patient_dirs:

        ret = load_patient_images(path, args.normalize)

        images, patient_number, frame_indices, rotated = ret

        predictions = []

        for image in images:

            mask_pred = m.predict(image[None,:,:,:]) # feed one at a time

            predictions.append((image[:,:,0], mask_pred[0,:,:,1]))

        for (image, mask), frame_index in zip(predictions, frame_indices):

            filename = "P{:02d}-{:04d}-{}contour-auto.txt".format(

                patient_number, frame_index, contour_type)

            outpath = os.path.join(args.outdir, filename)

            print(filename)

            contour = get_contours(mask)

            if rotated:

                height, width = image.shape

                x, y = contour.T

                x, y = height - y, x

                contour = np.vstack((x, y)).T

            np.savetxt(outpath, contour, fmt='%i', delimiter=' ')

            if args.checkpoint:

                filename = "P{:02d}-{:04d}-{}contour-auto.png".format(

                    patient_number, frame_index, contour_type)

                outpath = os.path.join(args.outdir, filename)

                save_image(outpath, image, np.round(mask))

if __name__ == '__main__':

    main()