import numpy as np
import jax.numpy as jnp
import pickle
import sys
import random
from alphafold.model import model
from alphafold.model import config
model_name = 'model_1'
model_config = config.model_config(model_name)
model_runner = model.RunModel(model_config)

def features2parallel(features_processed):
    data_res0 = dict()
    for key in features_processed[0].keys():
        features0 = []
        for i in range(len(features_processed)):
            features0.append(jnp.array(features_processed[i][key]))
        features0 = jnp.array(features0)
        data_res0[key] = features0
    return data_res0

def process_path(path):
    feature=[]
    labels=[]
    for j, index in enumerate(path):
        feat_paths = np.char.decode(index, 'utf-8')
        fs = open(str(feat_paths), "rb")
        feature_dict = pickle.load(fs)
        labels.append(feature_dict['flu'])
        feature.append(model_runner.process_features(
            feature_dict, random_seed=random.randrange(sys.maxsize)))
    processed_feature_dict = features2parallel(feature)
    aa=np.array(labels).astype(np.float32) 
    processed_feature_dict['flu']=jnp.array(aa)
    #print(processed_feature_dict['druglabels'])
    return processed_feature_dict


def process_path_save(path):
    feature=[]
    for j, index in enumerate(path):
        feat_paths = np.char.decode(index, 'utf-8')
        fs = open(str(feat_paths), "rb")
        feature_dict = pickle.load(fs)
        feature=model_runner.process_features(
            feature_dict, random_seed=random.randrange(sys.maxsize))
        feature['flu']=jnp.array(float(feature_dict['flu']))
    return feature

def process_path_load(path):
    feature=[]
    labels=[]
    for j, index in enumerate(path):
        feat_paths = np.char.decode(index, 'utf-8')
        fs = open(str(feat_paths), "rb")
        feature_dict = pickle.load(fs)
        labels.append(feature_dict['flu'])
        feature.append(feature_dict)
    processed_feature_dict = features2parallel(feature)
    aa=np.array(labels).astype(np.float32) 
    processed_feature_dict['flu']=jnp.array(aa)
    #print(processed_feature_dict['druglabels'])
    return processed_feature_dict