import os
import pickle
import random
import argparse
import torch
import numpy as np


def coretest_split(index_path, test_path, val_ratio=0.1, val_num=None):
    with open(index_path, 'rb') as f:
        index = pickle.load(f)
    test_ids = [f for f in os.listdir(test_path) if len(f) == 4]
    all_ids = [os.path.basename(i[0])[:4] for i in index]
    print('Original test set size: ', len(test_ids))
    test_index = [all_ids.index(test_id) for test_id in test_ids if test_id in all_ids]
    train_val_index = list(set(range(len(all_ids))) - set(test_index))
    assert len(train_val_index) == len(all_ids) - len(test_index)
    random.shuffle(train_val_index)
    if val_num is not None:
        n_val = val_num
    else:
        n_val = int(len(train_val_index) * val_ratio)
    val_index = train_val_index[:n_val]
    train_index = train_val_index[n_val:]
    return train_index, val_index, test_index


def time_split(index_path):
    valid_ids = np.loadtxt("./data/pdbbind_v2020/timesplit_no_lig_overlap_val", dtype=str)
    test_ids = np.loadtxt("./data/pdbbind_v2020/timesplit_test", dtype=str)

    with open(index_path, 'rb') as f:
        index = pickle.load(f)
    all_ids = [os.path.basename(i[0])[:4] for i in index]
    val_index = [all_ids.index(val_id) for val_id in valid_ids if val_id in all_ids]
    test_index = [all_ids.index(test_id) for test_id in test_ids if test_id in all_ids]
    train_index = list(set(range(len(all_ids))) - set(test_index) - set(val_index))
    return train_index, val_index, test_index


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--index_path', type=str, default='./data/pdbbind_v2016/pocket_10/index.pkl')
    parser.add_argument('--split_mode', type=str, choices=['coreset', 'time'], default='coreset')
    parser.add_argument('--test_path', type=str, default='./data/pdbbind_v2016/coreset')
    parser.add_argument('--val_ratio', type=float, default=0.1)
    parser.add_argument('--val_num', type=int, default=None)
    parser.add_argument('--save_path', type=str, default='./data/pdbbind_v2016/pocket_10/split.pt')
    parser.add_argument('--seed', type=int, default=2021)
    args = parser.parse_args()
    random.seed(args.seed)

    if args.split_mode == 'coreset':
        train_index, val_index, test_index = coretest_split(args.index_path, args.test_path,
                                                            args.val_ratio, args.val_num)
    elif args.split_mode == 'time':
        train_index, val_index, test_index = time_split(args.index_path)
    else:
        raise ValueError(args.split_mode)
    torch.save({
        'train': train_index,
        'val': val_index,
        'test': test_index
    }, args.save_path)
    print('Train %d, Validation %d, Test %d.' % (len(train_index), len(val_index), len(test_index)))
    print('Done.')