import csv, pickle, json, os, math, sys

from collections import OrderedDict

from bioservices import UniProt

import numpy as np

import pickle

PROT_FILE = "proteins.fasta"

CHEM_FILE = "ligands.tab"

AFF_FILE = "Y.tab"

def prepare_new_data(fpath, test=True):

    prots = read_proteins(fpath)

    chems = read_chemicals(fpath)

    Y = np.zeros((len(chems), len(prots)))

    if  os.path.exists(fpath + AFF_FILE):

        Y = np.loadtxt(fpath + AFF_FILE)

    pickle.dump(Y, open(fpath + "Y","wb"), protocol=pickle.HIGHEST_PROTOCOL)

    label_row_inds, label_col_inds = np.where(np.isnan(Y)==False)

    #json.dump(linepos, open(FLAGS.test_path  + "csv_pos_match.txt","w"))

    indic = set(range(len(label_row_inds)))

    indic = sorted(indic, key=os.urandom)

    if not os.path.exists(fpath + "folds/"):

        os.makedirs(fpath + "folds/")

    if test:

        json.dump(indic, open(fpath + "folds/test_fold.txt","w"))

    else:

        json.dump(indic, open(fpath + "folds/train_fold.txt","w"))

def read_chemicals(datafolder):

    counter =0

    filepath = datafolder + CHEM_FILE

    chemicals = {}

    with open(filepath) as file:

         next(file)

         for row in file:

            chem_id = row.split('\t')[0]

            smiles = (row.split('\t')[1]).strip()

            chemicals[chem_id] = smiles

            counter +=1

    print("%d number(s) of chemical(s)" % counter)

    json.dump(chemicals, open(datafolder + 'ligands.txt', 'w'))

    return chemicals

def read_proteins(datafolder):

    proteins = {}

    counter =0

    fa=""

    filename = datafolder + PROT_FILE

    with open(filename) as f:

        fa = f.readlines()

    idindex=[]

    for i, line in enumerate(fa):

        if ">" in line:

            idindex.append(i)

    idindex.append(i)

    for i, idx in enumerate(idindex):

        if i < len(idindex)-1:

            idx1 = idindex[i+1]

            info = fa[idx].split()

            pid = info[0][4:10]

            seq = "".join(fa[idx+1:idx1])

            seq = seq.replace("\n","")

            proteins[pid] = seq

            counter +=1

    print("%d number(s) of protein(s)" % counter)

    json.dump(proteins, open(datafolder + 'proteins.txt', 'w'))

    return proteins