import settings
import gdv_utils

def get_node_labels(network_name,ground_truth_type):
    if network_name == 'aging':
        return get_aging_node_labels(ground_truth_type)
    elif 'enron' in network_name:
        return get_enron_node_labels(ground_truth_type)
    return None


def get_node_label_title(network_name,ground_truth_type,label):
    if network_name == 'aging':
        return 'aging' if label==1 else 'non-aging'
    elif 'enron' in network_name and ground_truth_type == 'manager':
        return 'manager' if label==1 else 'other'
    return label

def get_aging_node_labels(ground_truth_type = 'sequenceage'):
    """
    Gets ground truth node labels for 'aging' dataset (1 - age-related, 0 - otherwise)
    :return:
    """
    all_nodes = gdv_utils.getNodeNames('aging')
    res = dict(zip(all_nodes, [0]*len(all_nodes)))

    file_dir = settings.BASE_INPUT_DIR + 'aging/ground_truth/'
    if ground_truth_type == 'sequenceage':
        files = ['SequenceAge']
    elif ground_truth_type == 'unionallbutdynetage':
        files = ['SequenceAge','ADExpressionAge','BrainExpression2004Age','BrainExpression2008Age']
    elif ground_truth_type == 'unionall':
        files = ['SequenceAge','ADExpressionAge','BrainExpression2004Age','BrainExpression2008Age','DyNetAge']
    elif ground_truth_type == 'unionall':
        files = ['SequenceAge','ADExpressionAge','BrainExpression2004Age','BrainExpression2008Age','DyNetAge']
    elif ground_truth_type == 'brain2004age':
        files = ['BrainExpression2004Age']
    elif ground_truth_type == 'brain2008age':
        files = ['BrainExpression2008Age']
    elif ground_truth_type == 'dynetage':
        files = ['DyNetAge']
    else:
        raise Exception("Unknown ground truth type: %s" % ground_truth_type)
    missing_nodes_count = 0
    present_nodes_count = 0
    for f in files:
        for line in open(file_dir + f,'r').readlines():
            node_id = line.strip()
            if res.has_key(node_id):
                if res[node_id]==0:
                    present_nodes_count += 1
                res[node_id]=1
            else:
                missing_nodes_count += 1
    if missing_nodes_count>0:
        print '%s nodes from \'%s\' ground truth files not in the network (but %s nodes are)' % (missing_nodes_count, ground_truth_type, present_nodes_count)

    return res

def get_enron_node_labels(ground_truth_type = 'manager'):

    if not ground_truth_type in ['manager','profession']:
        raise Exception("Unknown ground truth type: %s" % ground_truth_type)

    res = {}
    for line in open(settings.BASE_INPUT_DIR + 'enron/aggregated/enron_aggregated_1.txt','r').readlines():
        values = line.strip().split(' ')
        res[values[0]]=0
        res[values[1].replace('\n','')]=0

    i=0
    for line in open(settings.BASE_INPUT_DIR + 'enron/ground_truth/employees.txt','r').readlines():
        values = line.strip().split('\t')

        if ground_truth_type == 'manager':
            if len(values)>2 and ('manag' in values[2].lower() or 'ceo' in values[2].lower() or 'president' in values[2].lower() or 'director' in values[2].lower()):
                res[str(i)] = 1
        elif ground_truth_type == 'profession':
            if len(values)<3 or values[2].strip() == '':
                res[str(i)] = 'employee'
            else:
                v = values[2].strip().lower()
                if 'legal' in v:
                    v = v.replace('legal','').replace('department','').strip()
                elif 'lawyer' in v or 'trader' in v or 'n/a' in v:
                    v = 'employee'
                res[str(i)] = v

        i+=1

    return res