import os
import torch
import sys
sys.path.append("..")
from torch.utils.data.dataloader import DataLoader
from esm.data import ESMStructuralSplitDataset

# for split_level in ['family', 'superfamily', 'fold']:
#     for cv_partition in ['0', '1', '2', '3', '4']:
#         esm_structural_train = ESMStructuralSplitDataset(
#             split_level=split_level, 
#             cv_partition=cv_partition, 
#             split='train', 
#             root_path = os.path.expanduser('.'),
#             download=False
#         )
#         esm_structural_valid = ESMStructuralSplitDataset(
#             split_level=split_level, 
#             cv_partition=cv_partition, 
#             split='valid', 
#             root_path = os.path.expanduser('.'),
#             download=False
#         )

esm_structural_train = ESMStructuralSplitDataset(
    split_level='fold', 
    cv_partition='0', 
    split='train', 
    root_path = os.path.expanduser('.'),
    download=False
)
esm_structural_valid = ESMStructuralSplitDataset(
    split_level='fold', 
    cv_partition='0', 
    split='valid', 
    root_path = os.path.expanduser('.'),
    download=False
)
#name=esm_structural_train['names'][0]

train_loader = DataLoader(dataset=esm_structural_train,
                        batch_size=1,
                        shuffle=False,
                        #collate_fn=esm_structural_train.__collate_fn__
                        )
valid_loader = DataLoader(dataset=esm_structural_valid,
                        batch_size=1,
                        shuffle=False,
                        #collate_fn=esm_structural_train.__collate_fn__
                        )
#aa=next(iter(train_loader))
#batch_labels, batch_strs, batch_tokens = batch_converter(data)
f=open(os.path.join(os.path.expanduser('.'), 'valid_seq.txt'), 'w')
# for idx, batch in enumerate(train_loader):
#     print(batch[0][0],end=' ',file=f)
#     print(batch[1]['seq'][0],file=f)

for idx, batch in enumerate(valid_loader):
    print('>'+batch[0][0],end='\n',file=f)
    print(batch[1]['seq'][0],file=f)


