# -*- coding: utf-8 -*-

import sys

import random

import os

import numpy as np

from collections import defaultdict

sys.path.append(os.getcwd()) #add the env path

from sklearn.model_selection import train_test_split,StratifiedKFold

from main import train

from config import DRUG_EXAMPLE, RESULT_LOG, PROCESSED_DATA_DIR, LOG_DIR, MODEL_SAVED_DIR, ENTITY2ID_FILE, KG_FILE, \

    EXAMPLE_FILE,  DRUG_VOCAB_TEMPLATE, ENTITY_VOCAB_TEMPLATE, \

    RELATION_VOCAB_TEMPLATE, SEPARATOR, THRESHOLD, TRAIN_DATA_TEMPLATE, DEV_DATA_TEMPLATE, \

    TEST_DATA_TEMPLATE, ADJ_ENTITY_TEMPLATE, ADJ_RELATION_TEMPLATE, ModelConfig, NEIGHBOR_SIZE

from utils import pickle_dump, format_filename,write_log,pickle_load

def read_entity2id_file(file_path: str, drug_vocab: dict, entity_vocab: dict):

    print(f'Logging Info - Reading entity2id file: {file_path}' )

    assert len(drug_vocab) == 0 and len(entity_vocab) == 0

    with open(file_path, encoding='utf8') as reader:

        count=0

        for line in reader:

            if(count==0):

                count+=1

                continue

            drug, entity = line.strip().split('\t')

            drug_vocab[entity]=len(drug_vocab) 

            entity_vocab[entity] = len(entity_vocab)

def read_example_file(file_path:str,separator:str,drug_vocab:dict):

    print(f'Logging Info - Reading example file: {file_path}')

    assert len(drug_vocab)>0

    examples=[]

    with open(file_path,encoding='utf8') as reader:

        for idx,line in enumerate(reader):

            d1,d2,flag=line.strip().split(separator)[:3]

            if d1 not in drug_vocab or d2 not in drug_vocab:

                continue

            if d1 in drug_vocab and d2 in drug_vocab:

                examples.append([drug_vocab[d1],drug_vocab[d2],int(flag)])

    examples_matrix=np.array(examples)

    print(f'size of example: {examples_matrix.shape}')

    X=examples_matrix[:,:2]

    y=examples_matrix[:,2:3]

    train_data_X, valid_data_X,train_y,val_y = train_test_split(X,y, test_size=0.2,stratify=y)

    train_data=np.c_[train_data_X,train_y]

    valid_data_X, test_data_X,val_y,test_y = train_test_split(valid_data_X,val_y, test_size=0.5)

    valid_data=np.c_[valid_data_X,val_y]

    test_data=np.c_[test_data_X,test_y]

    return examples_matrix

def read_kg(file_path: str, entity_vocab: dict, relation_vocab: dict, neighbor_sample_size: int):

    print(f'Logging Info - Reading kg file: {file_path}')

    kg = defaultdict(list)

    with open(file_path, encoding='utf8') as reader:

        count=0

        for line in reader:

            if count==0:

                count+=1

                continue

            head, tail, relation = line.strip().split(' ') 

            if head not in entity_vocab:

                entity_vocab[head] = len(entity_vocab)

            if tail not in entity_vocab:

                entity_vocab[tail] = len(entity_vocab)

            if relation not in relation_vocab:

                relation_vocab[relation] = len(relation_vocab)

            # undirected graph

            kg[entity_vocab[head]].append((entity_vocab[tail], relation_vocab[relation]))

            kg[entity_vocab[tail]].append((entity_vocab[head], relation_vocab[relation]))

    print(f'Logging Info - num of entities: {len(entity_vocab)}, '

          f'num of relations: {len(relation_vocab)}')

    print('Logging Info - Constructing adjacency matrix...')

    n_entity = len(entity_vocab)

    adj_entity = np.zeros(shape=(n_entity, neighbor_sample_size), dtype=np.int64)

    adj_relation = np.zeros(shape=(n_entity, neighbor_sample_size), dtype=np.int64)

    for entity_id in range(n_entity):

        all_neighbors = kg[entity_id]

        n_neighbor = len(all_neighbors)

        sample_indices = np.random.choice(

            n_neighbor,

            neighbor_sample_size,

            replace=False if n_neighbor >= neighbor_sample_size else True

        )

        adj_entity[entity_id] = np.array([all_neighbors[i][0] for i in sample_indices])

        adj_relation[entity_id] = np.array([all_neighbors[i][1] for i in sample_indices])

    return adj_entity, adj_relation

def process_data(dataset: str, neighbor_sample_size: int,K:int):

    drug_vocab = {}

    entity_vocab = {}

    relation_vocab = {}

    read_entity2id_file(ENTITY2ID_FILE[dataset], drug_vocab, entity_vocab)

    pickle_dump(format_filename(PROCESSED_DATA_DIR, DRUG_VOCAB_TEMPLATE, dataset=dataset),drug_vocab)

    pickle_dump(format_filename(PROCESSED_DATA_DIR, ENTITY_VOCAB_TEMPLATE, dataset=dataset),entity_vocab)

    examples_file=format_filename(PROCESSED_DATA_DIR, DRUG_EXAMPLE, dataset=dataset)

    examples = read_example_file(EXAMPLE_FILE[dataset], SEPARATOR[dataset],drug_vocab)

    np.save(examples_file,examples)

    adj_entity_file = format_filename(PROCESSED_DATA_DIR, ADJ_ENTITY_TEMPLATE, dataset=dataset)

    adj_relation_file = format_filename(PROCESSED_DATA_DIR, ADJ_RELATION_TEMPLATE, dataset=dataset)

    adj_entity, adj_relation = read_kg(KG_FILE[dataset], entity_vocab, relation_vocab,

                                       neighbor_sample_size)

    pickle_dump(format_filename(PROCESSED_DATA_DIR, DRUG_VOCAB_TEMPLATE, dataset=dataset),

                drug_vocab)

    pickle_dump(format_filename(PROCESSED_DATA_DIR, ENTITY_VOCAB_TEMPLATE, dataset=dataset),

                entity_vocab)

    pickle_dump(format_filename(PROCESSED_DATA_DIR, RELATION_VOCAB_TEMPLATE, dataset=dataset),

                relation_vocab)

    adj_entity_file = format_filename(PROCESSED_DATA_DIR, ADJ_ENTITY_TEMPLATE, dataset=dataset)

    np.save(adj_entity_file, adj_entity)

    print('Logging Info - Saved:', adj_entity_file)

    adj_relation_file = format_filename(PROCESSED_DATA_DIR, ADJ_RELATION_TEMPLATE, dataset=dataset)

    np.save(adj_relation_file, adj_relation)

    print('Logging Info - Saved:', adj_entity_file)

    cross_validation(K,examples,dataset,neighbor_sample_size)

def cross_validation(K_fold,examples,dataset,neighbor_sample_size):

    subsets=dict()

    n_subsets=int(len(examples)/K_fold)

    remain=set(range(0,len(examples)-1))

    for i in reversed(range(0,K_fold-1)):

        subsets[i]=random.sample(remain,n_subsets)

        remain=remain.difference(subsets[i])

    subsets[K_fold-1]=remain

    aggregator_types=['sum','concat','neigh']

    for t in aggregator_types:

        count=1

        temp={'dataset':dataset,'aggregator_type':t,'avg_auc':0.0,'avg_acc':0.0,'avg_f1':0.0,'avg_aupr':0.0}

        for i in reversed(range(0,K_fold)):

            test_d=examples[list(subsets[i])]

            val_d,test_data=train_test_split(test_d,test_size=0.5)

            train_d=[]

            for j in range(0,K_fold):

                if i!=j:

                    train_d.extend(examples[list(subsets[j])])

            train_data=np.array(train_d)               

            train_log=train(

            kfold=count,

            dataset=dataset,

            train_d=train_data,

            dev_d=val_d,

            test_d=test_data,

            neighbor_sample_size=neighbor_sample_size,

            embed_dim=32,

            n_depth=2,

            l2_weight=1e-7,

            lr=2e-2,

            #lr=5e-3,

            optimizer_type='adam',

            batch_size=2048,

            aggregator_type=t,

            n_epoch=50,

            callbacks_to_add=['modelcheckpoint', 'earlystopping']

            )     

            count+=1

            temp['avg_auc']=temp['avg_auc']+train_log['test_auc']

            temp['avg_acc']=temp['avg_acc']+train_log['test_acc']

            temp['avg_f1']=temp['avg_f1']+train_log['test_f1']

            temp['avg_aupr']=temp['avg_aupr']+train_log['test_aupr']

        for key in temp:

            if key=='aggregator_type' or key=='dataset':

                continue

            temp[key]=temp[key]/K_fold

        write_log(format_filename(LOG_DIR, RESULT_LOG[dataset]),temp,'a')

        print(f'Logging Info - {K_fold} fold result: avg_auc: {temp["avg_auc"]}, avg_acc: {temp["avg_acc"]}, avg_f1: {temp["avg_f1"]}, avg_aupr: {temp["avg_aupr"]}')

if __name__ == '__main__':

    if not os.path.exists(PROCESSED_DATA_DIR):

        os.makedirs(PROCESSED_DATA_DIR)

    if not os.path.exists(LOG_DIR):

        os.makedirs(LOG_DIR)

    if not os.path.exists(MODEL_SAVED_DIR):

        os.makedirs(MODEL_SAVED_DIR)

    model_config = ModelConfig()

    process_data('kegg',NEIGHBOR_SIZE['kegg'],5)