from pkg_resources import add_activation_listener
import monai
import torch
import itk
import numpy as np
import os.path
import argparse
import sys
from pathlib import Path
import deep_atlas_train
import glob

ROOT_DIR = str(Path(os.getcwd()).parent.parent.absolute())
sys.path.insert(0, os.path.join(ROOT_DIR, 'deepatlas/test'))

from test import (
    seg_training_inference, load_json, reg_training_inference
)

def parse_command_line():
    parser = argparse.ArgumentParser(
        description='pipeline for deep atlas test')
    parser.add_argument('-gpu', metavar='id of gpu', type=str, default='0',
                        help='id of gpu device to use')
    parser.add_argument('-ti', metavar='task id and name', type=str,
                        help='task name and id')
    parser.add_argument('-nf', metavar='number of fold', type=int,
                        help='number of fold for testing')
    parser.add_argument('-op', metavar='output path for prediction step', type=str,
                        help="relative path of the output directory, should be same name in the registration, crop and final prediction steps")
    argv = parser.parse_args()
    return argv


def path_to_id(path):
    return os.path.basename(path).split('.')[0]


def main():
    ROOT_DIR = str(Path(os.getcwd()).parent.parent.absolute())
    args = parse_command_line()
    gpu = args.gpu
    task = args.ti
    num_fold = f'fold_{args.nf}'
    output_path = os.path.join(ROOT_DIR, 'deepatlas_results', task, 'customize_predicted_results')
    fold_path = os.path.join(output_path, num_fold)
    out_path = os.path.join(fold_path, args.op)
    json_path = os.path.join(
        ROOT_DIR, 'deepatlas_results', task, "training_results", num_fold, 'dataset.json')
    seg_model_path = os.path.join(
        ROOT_DIR, 'deepatlas_results', task, 'training_results', num_fold, 'SegNet', 'seg_net_best.pth')
    reg_model_path = os.path.join(
        ROOT_DIR, 'deepatlas_results', task, 'training_results', num_fold, 'RegNet', 'reg_net_best.pth')
    json_file = load_json(json_path)
    labels = json_file['labels']
    num_label = len(labels.keys())
    network_info = json_file['network']
    spatial_dim = network_info['spatial_dim']
    dropout = network_info['dropout']
    activation_type = network_info['activation_type']
    normalization_type = network_info['normalization_type']
    num_res = network_info['num_res']
    device = torch.device("cuda:" + gpu)
    output_seg_path = os.path.join(out_path, 'SegNet')
    output_reg_path = os.path.join(out_path, 'RegNet')
    try:
        os.mkdir(output_path)
    except:
        print(f'{output_path} is already existed !!!')

    try:
        os.mkdir(fold_path)
    except:
        print(f'{fold_path} is already existed !!!')
    
    try:
        os.mkdir(out_path)
    except:
        print(f'{out_path} is already existed !!!')

    try:
        os.mkdir(output_seg_path)
    except:
        print(f'{output_seg_path} is already existed !!!')

    try:
        os.mkdir(output_reg_path)
    except:
        print(f'{output_reg_path} is already existed !!!')

    seg_net = deep_atlas_train.get_seg_net(
        spatial_dim, num_label, dropout, activation_type, normalization_type, num_res)
    reg_net = deep_atlas_train.get_reg_net(
        spatial_dim, spatial_dim, dropout, activation_type, normalization_type, num_res)
    
    img_path = os.path.join(ROOT_DIR, 'deepatlas_preprocessed', task, 'customize_test_data', args.op, 'images')
    seg_path = os.path.join(ROOT_DIR, 'deepatlas_preprocessed', task, 'customize_test_data', args.op, 'labels')
    total_img_paths = []
    total_seg_paths = []
    for i in sorted(glob.glob(img_path + '/*.nii.gz')):
        total_img_paths.append(i)

    for j in sorted(glob.glob(seg_path + '/*.nii.gz')):
        total_seg_paths.append(j)
    
    seg_ids = list(map(path_to_id, total_seg_paths))
    img_ids = map(path_to_id, total_img_paths)
    data = []
    for img_index, img_id in enumerate(img_ids):
        data_item = {'img': total_img_paths[img_index]}
        if img_id in seg_ids:
            data_item['seg'] = total_seg_paths[seg_ids.index(img_id)]
        data.append(data_item)

    seg_training_inference(seg_net, device, seg_model_path, output_seg_path, num_label, json_path=None, data=data)
    reg_training_inference(reg_net, device, reg_model_path, output_reg_path, num_label, json_path=None, data=data)


if __name__ == '__main__':
    main()