import os

import sys

import numpy as np

import nibabel as nib

import time

import torch

import torch.nn as nn

import torch.nn.functional as F

import _init_paths

from libs.network import U_Net, U_NetDF

from utils.image_list import to_image_list

import libs.datasets.augment as standard_augment

import libs.datasets.joint_augment as joint_augment

def progress_bar(curr_idx, max_idx, time_step, repeat_elem = "_"):

    max_equals = 55

    step_ms = int(time_step*1000)

    num_equals = int(curr_idx*max_equals/float(max_idx))

    len_reverse =len('Step:%d ms| %d/%d ['%(step_ms, curr_idx, max_idx)) + num_equals

    sys.stdout.write("Step:%d ms|%d/%d [%s]" %(step_ms, curr_idx, max_idx, " " * max_equals,))

    sys.stdout.flush()

    sys.stdout.write("/b" * (max_equals+1))

    sys.stdout.write(repeat_elem * num_equals)

    sys.stdout.write("/b"*len_reverse)

    sys.stdout.flush()

    if curr_idx == max_idx:

        print('/n')

def load_nii(img_path):

    """

    Function to load a 'nii' or 'nii.gz' file, The function returns

    everyting needed to save another 'nii' or 'nii.gz'

    in the same dimensional space, i.e. the affine matrix and the header

    Parameters

    ----------

    img_path: string

    String with the path of the 'nii' or 'nii.gz' image file name.

    Returns

    -------

    Three element, the first is a numpy array of the image values,

    the second is the affine transformation of the image, and the

    last one is the header of the image.

    """

    nimg = nib.load(img_path)

    return nimg.get_fdata(), nimg.affine, nimg.header

def save_nii(vol, affine, hdr, path, prefix, suffix):

    vol = nib.Nifti1Image(vol, affine, hdr)

    vol.set_data_dtype(np.uint8)

    nib.save(vol, os.path.join(path, prefix+'_'+suffix + ".nii.gz"))

def get_person_names(root_path=None):

    persons_name = os.listdir(root_path)

    persons_name = [pn for pn in persons_name if "patient" in pn]

    persons_name.sort()

    return persons_name

def get_patient_data(patient, root_path):

    patient_data = {}

    infocfg_p = os.path.join(root_path, patient, "Info.cfg")

    with open(infocfg_p) as f_in:

        for line in f_in:

            l = line.rstrip().split(": ")

            patient_data[l[0]] = l[1]

    ed_path = os.path.join(root_path, patient, "%s_frame%02d.nii.gz" % (patient, int(patient_data['ED'])))

    es_path = os.path.join(root_path, patient, "%s_frame%02d.nii.gz" % (patient, int(patient_data['ES'])))

    img_4d_path = os.path.join(root_path, patient, "{}_4d.nii.gz".format(patient))

    # ed_gt_path = os.path.join(root_path, patient, "%s_frame%02d_gt.nii.gz" % (patient, int(patient_data['ED'])))

    # es_gt_path = os.path.join(root_path, patient, "%s_frame%02d_gt.nii.gz" % (patient, int(patient_data['ES'])))

    ed, affine, hdr = load_nii(ed_path)

    patient_data['ED_VOL'] = np.swapaxes(ed, 0, -1)

    patient_data['3D_affine'] = affine

    patient_data['3D_hdr'] = hdr

    es, _, _ = load_nii(es_path)  # (w, h, slices)

    patient_data['ES_VOL'] = np.swapaxes(es, 0, -1)

    img_4d, affine_4d, hdr_4d = load_nii(img_4d_path)  # (w, h, slices, times)

    patient_data['4D'] = np.swapaxes(img_4d, 0, 1)

    patient_data['4D_affine'] = affine_4d

    patient_data['4D_hdr'] = hdr_4d

    patient_data['size'] = img_4d.shape[:2][::-1]

    patient_data['pid'] = patient

    # ed_gt = load_nii(ed_gt_path)

    # patient_data['ED_GT'] = np.swapaxes(ed_gt, 0, 1)

    # es_gt = load_nii(es_gt_path)

    # patient_data['ES_GT'] = np.swapaxes(es_gt, 0, 1)

    return patient_data

def test_it(model, data, device="cuda", used_df=True):

    model.eval()

    imgs = data

    imgs = imgs.to(device)

    # gts = gts.to(device)

    net_out = model(imgs)

    if used_df:

        preds_out = net_out[0]

        preds_df = net_out[1]

    else:

        preds_out = net_out[0]

        preds_df = None

    preds_out = nn.functional.softmax(preds_out, dim=1)

    _, preds = torch.max(preds_out, 1)

    preds = preds.unsqueeze(1)  # (N, 1, *)

    return preds, preds_df

def transform(imgs):

    mean = 63.19523533061758

    std = 70.74166957523165

    trans = standard_augment.Compose([standard_augment.To_PIL_Image(),

                                    # joint_augment.RandomAffine(0,translate=(0.125, 0.125)),

                                    # joint_augment.RandomRotate((-180,180)),

                                    # joint_augment.FixResize(224),

                                    standard_augment.to_Tensor(),

                                    standard_augment.normalize([mean], [std]),

                                      ])

    return trans(imgs)

def test_voxel(model, imgs, used_df, multi_batches=False, resize=None):

    """ imgs: (slices, H, W)

        preds: (slices, 1, H, W)

    """

    imgs = imgs[..., None].astype(np.float32)

    B, _, _, C = imgs.shape

    if multi_batches:

        data, origin_shape = to_image_list(imgs, size_divisible=32, return_size=True)

        preds, _ = test_it(model, data)

        # for j in range(imgs.shape[0]):

        #     preds[j, ...] = pred.cpu().numpy()[j, :, :origin_shape[j][0], :origin_shape[j][1]]

    else:

        preds = torch.zeros(B, C, resize[0], resize[1])

        for j, pt in enumerate(imgs):

            data = [transform(pt)]

            data, origin_shape = to_image_list(data, size_divisible=32, return_size=True)

            pred, _ = test_it(model, data, used_df=used_df)

            preds[j, ...] = pred[0, 0, :origin_shape[0][0], :origin_shape[0][1]]

    if resize is not None:

        # preds = F.interpolate(preds, size=resize, mode='nearest')

        preds = preds[..., :resize[0], :resize[1]]

    return preds.cpu().numpy()[:, 0, ...]

def create_model(model_name, selfeat):

    if model_name == 'U_NetDF':

        model = U_NetDF(selfeat=selfeat, auxseg=True)

    elif model_name == 'U_Net':

        model = U_Net()

    # elif model_name == 'Resnet18_DfUnet':

    #     model = Resnet18_DfUnet()

    # elif model_name == 'DenseNet':

    #     model = DenseNet()

    # elif model_name == 'DenseNet_DF':

    #     model = DenseNet_DF(selfeat=selfeat)

    return model

def test(mgpus, model_name, model_path, used_df, selfeat, log_path):

    model = create_model(model_name, selfeat=selfeat)

    if mgpus is not None and len(mgpus) > 2:

        model = nn.DataParallel(model)

    model.cuda()

    checkpoint = torch.load(model_path, map_location='cpu')

    model.load_state_dict(checkpoint['model_state'])

    root_path = "MICCAIACDC2017/ACDC_DataSet/testing/testing/"

    root_path = "/root/ACDC_DataSet/testing/testing/"

    persons_name = get_person_names(root_path)

    for j, pn in enumerate(persons_name):

        s_time = time.time()

        patient_data = get_patient_data(pn, root_path)

        # (slices, h, w)

        es_pred = test_voxel(model, patient_data['ES_VOL'], used_df=used_df, resize=patient_data['size'])

        ed_pred = test_voxel(model, patient_data['ED_VOL'], used_df=used_df, resize=patient_data['size'])

        es_pred = np.transpose(es_pred, (2, 1, 0))

        ed_pred = np.transpose(ed_pred, (2, 1, 0))

        img_4D = patient_data['4D']

        h, w, s, t = img_4D.shape

        pred_4D = np.zeros((w, h, s, t))

        for i in range(img_4D.shape[-1]):

            pred = test_voxel(model, np.transpose(img_4D[...,i], (2, 0, 1)), used_df=used_df, resize=patient_data['size'])

            pred_4D[..., i] = np.transpose(pred, (2, 1, 0))

        save_path = os.path.join(log_path, "all_predictions")

        os.makedirs(save_path, exist_ok=True)

        CheckSizeAndSaveVolume(pred_4D, patient_data, save_path)

        progress_bar(j%(len(persons_name)+1), len(persons_name),time.time()-s_time)

def CheckSizeAndSaveVolume(seg_4D, patient_data, save_path):

    """

    TODO:

    """ 

    prefix = patient_data['pid']

    suffix = '4D'

    # save_nii(seg_4D, patient_data['4D_affine'], patient_data['4D_hdr'], save_path, prefix, suffix)

    suffix = 'ED'

    ED_phase_n = int(patient_data['ED'])

    ED_pred = seg_4D[:,:,:,ED_phase_n]

    save_nii(ED_pred.astype(np.uint8), patient_data['3D_affine'], patient_data['3D_hdr'], save_path, prefix, suffix)

    suffix = 'ES'

    ES_phase_n = int(patient_data['ES'])

    ES_pred = seg_4D[:,:,:,ES_phase_n]

    save_nii(ES_pred.astype(np.uint8), patient_data['3D_affine'], patient_data['3D_hdr'], save_path, prefix, suffix)

    # ED_GT = patient_data.get('ED_GT', None)

    results = []

    return results

if __name__ == "__main__":

    # get_person_names()

    import argparse

    parser = argparse.ArgumentParser(description="arg parser")

    parser.add_argument('--mgpus', type=str, default='0', required=False, help='whether to use multiple gpu')

    parser.add_argument('--used_df', type=str, default='True', help='whether to use df')

    parser.add_argument('--model', type=str, default='', help='whether to use df')

    parser.add_argument('--selfeat', type=bool, default=True, help='whether to use feature select')

    parser.add_argument('--model_path', type=str, default=None, help='whether to train with evaluation')

    parser.add_argument('--output_dir', type=str, default="logs/acdc_logs/logs_supcat_auxseg/predtions", required=False, help='specify an output directory if needed')

    parser.add_argument('--log_file', type=str, default="../log_predtion.txt", help="the file to write logging")

    args = parser.parse_args()

    if args.mgpus is not None:

        os.environ["CUDA_VISIBLE_DEVICES"] = args.mgpus

    model_path = "logs/acdc_logs/logs_supcat_auxseg/ckpt/checkpoint_epoch_70.pth"

    test(args.mgpus, "U_NetDF", model_path, args.used_df, args.selfeat, args.output_dir)