from os.path import dirname, realpath

import sys

sys.path.append(dirname(dirname(realpath(__file__))))

import argparse

from argparse import Namespace

import pickle

from sandstone.learn.metrics.factory import get_metric

import numpy as np

import os

LOGGER_KEYS = ['censors', 'golds'] 

def make_logging_dict(results):

    logging_dict = {}

    logging_dict = {k: np.array( results[0].get('test_{}'.format(k), 0) ) for k in LOGGER_KEYS}

    logging_dict['probs'] = []

    prob_key = [k.split('_probs')[0] for k in results[0].keys() if 'probs' in k][0]

    logging_dict['probs'] = np.mean([np.array(r['{}_probs'.format(prob_key)])  for r in results ], axis = 0)

    return logging_dict

parser = argparse.ArgumentParser()

parser.add_argument('--parent_dir', type = str, default = '/Mounts/rbg-storage1/logs/lung_ct/')

parser.add_argument('--result_file_names', type = str, nargs = '+', default =  ["7a07ee56c93e2abd100a47542e394bed","1e94034923b44462203cfcf29ae29061","9ff7a8ba3b1f9eb7216f35222b3b8524","18490328d1790f7b6b8c86d97f25103c"])

parser.add_argument('--test_suffix', type = str, default = 'test')

parser.add_argument('--metric_name', nargs= '*', type = str, default = 'survival')

if __name__ == '__main__':

    args = parser.parse_args()

    result_args = [Namespace(**pickle.load(open(os.path.join(args.parent_dir, '{}.results'.format(f)), 'rb'))) for f in args.result_file_names]

    test_full_paths = [os.path.join(args.parent_dir, '{}.results.{}_.predictions'.format(f, args.test_suffix)) for f in args.result_file_names]

    test_results = [pickle.load(open(f, 'rb')) for f in  test_full_paths]

    logging_dict = make_logging_dict(test_results)

    performance_dict = {}

    metrics = [get_metric(m) for m in args.metric_name]

    for m in metrics:

        performance_dict.update( m(logging_dict, result_args[0]) )

    print(performance_dict)