from jax._src.api import T
from alphafold.model import model
from alphafold.model import data as param_loader
from alphafold.model import config
from absl import logging
from absl import flags
from absl import app
import optax
import sys
import random
import pickle
import haiku as hk
import os
import jax
import jax.profiler
import numpy as np
import jax.numpy as jnp
#from alphafold.train.data_system import cast_to_precision
import time
import torch
import scipy
from typing import List, Tuple, Any
from typing import Any, Mapping
import functools
import tensorflow as tf
from process_path import process_path,process_path_save


def datasets_train():
    features_dir='/home/public/bigdata/my/datasets/fitness/sta/train/feature'
    dataset = tf.data.Dataset.list_files(
        os.path.join(features_dir, '*.pkl'))
    dataset = dataset.shuffle(buffer_size=1 * 1)
    dataset = dataset.batch(1)
    return iter(dataset.as_numpy_iterator())

def datasets_test():
    features_dir='/home/public/bigdata/my/datasets/fitness/sta/test/feature'
    dataset = tf.data.Dataset.list_files(
        os.path.join(features_dir, '*.pkl'))
    dataset = dataset.shuffle(buffer_size=1 * 1)
    dataset = dataset.batch(1)
    return iter(dataset.as_numpy_iterator())

def main(_):
    train_set = datasets_test()
    # print('start train epoch ',step)
    for idx,batch in enumerate(train_set):
        print('test ',idx)
        data=process_path_save(batch)
        pickle.dump(data, open(os.path.join('data/test',str(idx)+'.pkl'), 'wb'))
        # if(idx==4):
        #     break

if __name__ == '__main__':
    app.run(main)