import csv
caffe_root = '/home/toanhoi/caffe/build/tools/caffe/'
import sys
import os
sys.path.insert(0, caffe_root + 'python')
os.environ["GLOG_minloglevel"] = "1"
import caffe
import argparse
import numpy as np
from medpy.io import load, save

#0:     Background (everything outside the brain)
#10:   Cerebrospinal fluid (CSF)
#150: Gray matter (GM)
#250: White matter (WM)
def convert_label_submit(label_img):
    label_processed=np.zeros(label_img.shape[0:]).astype(np.uint8)
    for i in range(label_img.shape[2]):
        label_slice=label_img[:, :, i]
        label_slice[label_slice == 1] = 10
        label_slice[label_slice == 2] = 150
        label_slice[label_slice == 3] = 250
        label_processed[:, :, i]=label_slice
    return label_processed

def convert_label(label_img):
    label_processed=np.zeros(label_img.shape[0:]).astype(np.uint8)
    for i in range(label_img.shape[2]):
        label_slice=label_img[:, :, i]
        label_slice[label_slice == 10] = 1
        label_slice[label_slice == 150] = 2
        label_slice[label_slice == 250] = 3
        label_processed[:, :, i]=label_slice
    return label_processed
#Reference https://github.com/ginobilinie/infantSeg
def dice(im1, im2,tid):
    im1=im1==tid
    im2=im2==tid
    im1=np.asarray(im1).astype(np.bool)
    im2=np.asarray(im2).astype(np.bool)
    if im1.shape != im2.shape:
        raise ValueError("Shape mismatch: im1 and im2 must have the same shape.")
    # Compute Dice coefficient
    intersection = np.logical_and(im1, im2)
    dsc=2. * intersection.sum() / (im1.sum() + im2.sum())
    return dsc

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='makes a plot from Caffe output')
    parser.add_argument("-start")
    parser.add_argument("-end")

    if (os.environ.get('CAFFE_CPU_MODE')):
        caffe.set_mode_cpu()
    else:
        caffe.set_mode_gpu()

    data_path = '/media/toanhoi/Study/databaseSeg/ISeg/iSeg-2017-Training'

    subject_id=9
    subject_name = 'subject-%d-' % subject_id

    f_T1 = os.path.join(data_path, subject_name + 'T1.hdr')
    inputs_T1, header_T1 = load(f_T1)
    inputs_T1 = inputs_T1.astype(np.float32)

    f_T2 = os.path.join(data_path, subject_name + 'T2.hdr')
    inputs_T2, header_T2 = load(f_T2)
    inputs_T2 = inputs_T2.astype(np.float32)

    f_l = os.path.join(data_path, subject_name + 'label.hdr')
    labels, header_label = load(f_l)
    labels = labels.astype(np.uint8)
    labels=convert_label(labels)

    mask = inputs_T1 > 0
    mask = mask.astype(np.bool)
    # ======================normalize to 0 mean and 1 variance====
    # Normalization
    inputs_T1_norm =(inputs_T1 - inputs_T1[mask].mean()) / inputs_T1[mask].std()
    inputs_T2_norm = (inputs_T2 - inputs_T2[mask].mean()) / inputs_T2[mask].std()

    inputs_T1_norm = inputs_T1_norm[:, :, :, None]
    inputs_T2_norm = inputs_T2_norm[:, :, :, None]

    inputs = np.concatenate((inputs_T1_norm, inputs_T2_norm), axis=3)
    inputs = inputs[None, :, :, :, :]
    inputs = inputs.transpose(0, 4, 3, 1, 2)
    num_class=4
    num_paches=0

    model_def='./deploy_3d_denseseg.prototxt'
    model_weights = "./snapshot/3d_denseseg_iseg_iter_200000.caffemodel"
    net = caffe.Net(model_def, model_weights,caffe.TEST)
    patch_input = [64, 64, 64]
    xstep = 16
    ystep = 8#16
    zstep = 16#16
    deep_slices    = np.arange(patch_input[0] // 2, inputs.shape[2] - patch_input[0] // 2 + xstep, xstep)
    height_slices  = np.arange(patch_input[1] // 2, inputs.shape[3] - patch_input[1] // 2 + ystep, ystep)
    width_slices   = np.arange(patch_input[2] // 2, inputs.shape[4] - patch_input[2] // 2 + zstep, zstep)
    output = np.zeros((num_class,) + inputs.shape[2:])
    count_used=np.zeros((inputs.shape[2],inputs.shape[3],inputs.shape[4]))+1e-5

    total_patch=len(deep_slices)*len(height_slices)*len(width_slices)
    for i in range(len(deep_slices)):
        for j in range(len(height_slices)):
            for k in range(len(width_slices)):
                num_paches=num_paches+1
                deep   = deep_slices[i]
                height = height_slices[j]
                width  = width_slices[k]
                raw_patches= inputs[:,:,deep - patch_input[0] // 2:deep + patch_input[0] // 2,
                                         height - patch_input[1] // 2:height + patch_input[1] // 2,
                                         width - patch_input[2] // 2:width + patch_input[2] // 2]
                print "Processed: ",num_paches ,"/", total_patch
                net.blobs['data'].data[...] = raw_patches
                net.forward()

                #Major voting https://github.com/ginobilinie/infantSeg
                temp_predic=net.blobs['softmax'].data[0].argmax(axis=0)
                for labelInd in range(4):  # note, start from 0
                    currLabelMat = np.where(temp_predic == labelInd, 1, 0)  # true, vote for 1, otherwise 0
                    output[labelInd, deep - patch_input[0] // 2:deep + patch_input[0] // 2,
                        height - patch_input[1] // 2:height + patch_input[1] // 2,
                        width - patch_input[2] // 2:width + patch_input[2] // 2] += currLabelMat
                #Average
                # output[slice(None),deep - patch_input[0] // 2:deep + patch_input[0] // 2,
                #         height - patch_input[1] // 2:height + patch_input[1] // 2,
                #         width - patch_input[2] // 2:width + patch_input[2] // 2]+=net.blobs['softmax'].data[0]

                count_used[deep - patch_input[0] // 2:deep + patch_input[0] // 2,
                        height - patch_input[1] // 2:height + patch_input[1] // 2,
                        width - patch_input[2] // 2:width + patch_input[2] // 2]+=1

    output=output/count_used
    y = np.argmax(output, axis=0)
    out_label=y.transpose(1,2,0)
    dsc_0 = dice(out_label , labels, 0)
    dsc_1 = dice(out_label , labels, 1)
    dsc_2 = dice(out_label , labels, 2)
    dsc_3 = dice(out_label , labels, 3)
    dsc = np.mean([dsc_1, dsc_2, dsc_3])  # ignore Background
    print dsc_1, dsc_2, dsc_3, dsc
    with open('result_3d_dense_seg.csv', 'a+') as csvfile:
        datacsv = csv.writer(csvfile, delimiter=',', quoting=csv.QUOTE_MINIMAL)
        datacsv.writerow([dsc_1, dsc_2, dsc_3, dsc])

    out_label=out_label.astype(np.uint8)
    out_label = convert_label_submit(out_label)
    save(out_label, '{}/{}'.format("./", "3d_dense_seg_result.hdr"), header_T1)