import os
import json
import numpy as np
import argparse


class_tag_1 = {1: '磨玻璃结节',
               2: '实性结节',
               3: '混合结节',
               4: '钙化结节'}

class_tag_2 = {1: '良性',
               2: '恶性'}


def followup_processing(args):
    # 按病人id和时间整理匹配结果
    followup_root = os.path.join(args.job_data_root, 'output/elastix')
    files = os.listdir(followup_root)

    working_dicts = {}
    for name in files:
        patientid = name.split('_')[0].split('-')[-1]
        if patientid not in working_dicts.keys():
            working_dicts[patientid] = []
            working_dicts[patientid].append(os.path.join(followup_root, name))
        else:
            working_dicts[patientid].append(os.path.join(followup_root, name))
    
    for key in list(working_dicts.keys()):
        working_dicts[key] = sorted(working_dicts[key])

    # 计算每一个病人的信息，（size、类别）
    working_infos = {}
    for key in list(working_dicts.keys()):
        working_infos[key] = []
        files = working_dicts[key]
        for file in files:
            temp= {'matched': [],
                   'missing': [],
                   'new': []}
            match_results = json.load(open(os.path.join(file, 'match_results.json'), 'r'))
            if len(match_results['matched']) != 0:
                for res in match_results['matched']:
                    res[0].append(res[0][3] * res[0][4] * res[0][5])
                    res[1].append(res[1][3] * res[1][4] * res[1][5])
                    temp['matched'].append(res)
        
            if len(match_results['missing']) != 0:
                for res in match_results['missing']:
                    res.append(res[3] * res[4] * res[5])
                    temp['missing'].append(res)

            if len(match_results['new']) != 0:
                for res in match_results['new']:
                    res.append(res[3] * res[4] * res[5])
                    temp['new'].append(res)
        
            working_infos[key].append(temp)
        
        

    # 得出结论
    class_change = {'1': [], '2': []}
    for key in list(working_infos.keys()):
        print('For Patient-{}'.format(key))
        print('\n')
        temp = {'big': {'num': 0, 
                        'll': []}, 
                'small': {'num': 0,
                          'll': []},
                'eq': {'num': 0,
                        'll': []}
            }
        for value in working_infos[key]:
            if len(value['matched']) != 0:
                for pp in value['matched']:
                    if pp[0][-1] > pp[1][-1]:
                        temp['small']['num'] += 1
                        temp['small']['ll'].append(abs(pp[0][-1] - pp[1][-1]))
                    elif pp[0][-1] < pp[1][-1]:
                        temp['big']['num'] += 1
                        temp['big']['ll'].append(abs(pp[0][-1] - pp[1][-1]))
                    elif pp[0][-1] == pp[1][-1]:
                        temp['eq']['num'] += 1
                        temp['eq']['ll'].append(pp[0][-1])
                    class_change['1'].append(pp[0][-3])
                    class_change['1'].append(pp[0][-3])
                    class_change['2'].append(pp[0][-2])
                    class_change['2'].append(pp[0][-2])
        
        choice_tag = None
        choice_num = -100
        choice_ll = []
        for tag in list(temp.keys()):
            if choice_num < temp[tag]['num']:
                choice_tag = tag
                choice_num = temp[tag]['num']
                choice_ll = temp[tag]['ll']
        
        if choice_tag == 'big':
            print('结节可能向变大发展，约变大{};'.format(np.random.choice(choice_ll, size=1)[0]))
            print('结节分类结果可能为：{};'.format(class_tag_1[np.random.choice(class_change['1'], size=1)[0]]))
            print('结节良恶性分类结果可能为： {}.'.format(class_tag_2[np.random.choice(class_change['2'], size=1)[0]]))
        elif choice_tag == 'small':
            print('结节可能向变小发展，约变小{};'.format(np.random.choice(choice_ll, size=1)[0]))
            print('结节分类结果可能为：{};'.format(class_tag_1[np.random.choice(class_change['1'], size=1)[0]]))
            print('结节良恶性分类结果可能为： {}.'.format(class_tag_2[np.random.choice(class_change['2'], size=1)[0]]))
        elif choice_tag == 'eq':
            print('结节可能向不变发展，约保持{};'.format(np.random.choice(choice_ll, size=1)[0]))
            print('结节分类结果可能为：{};'.format(class_tag_1[np.random.choice(class_change['1'], size=1)[0]]))
            print('结节良恶性分类结果可能为： {}.'.format(class_tag_2[np.random.choice(class_change['2'], size=1)[0]]))


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='followup processing')
    parser.add_argument('--job_data_root', default='/data/job_715/job_data_preprocess', type=str, help='model')
    args = parser.parse_args()

    followup_processing(args)