"""
@Desc  : Deep Deterministic Policy Gradient algorithm
"""
from __future__ import division, print_function

import os

import numpy as np

from xt.algorithm import Algorithm
from xt.algorithm.ddpg.default_config import BATCH_SIZE, BUFFER_SIZE, TARGET_UPDATE_FREQ
from xt.algorithm.replay_buffer import ReplayBuffer
from xt.framework.register import Registers
from xt.model import model_builder
from xt.util.common import import_config
from xt.algorithm.ddpg.ddpg import DDPG

os.environ["KERAS_BACKEND"] = "tensorflow"

@Registers.algorithm.register
class DDPGAD(DDPG):
    """Deep Deterministic Policy Gradient algorithm"""

    def predict(self, state):
        """predict for ddpg"""
        inputs = state
        outputs = self.actor.predict(inputs)
        # action = self.actor.predict(state.reshape(1, state.shape[0]))
        return outputs


#!/usr/bin/env python3

import os
import glob
import numpy as np
import scipy
from ffindex import *
from parsers import parse_hhr, parse_pdb_lines
from pyrosetta import *

init("-mute all")

MAXSEQ=20000

#============================================
# Global variables such as data directory path
#============================================
LABEL_dir = "/net/scratch/minkbaek/SStor/labels"
TEMPL_dir = "/net/scratch/minkbaek/trRefine/added/templ/sub_0/templates"
#TEMPL_dir = "/net/scratch/minkbaek/SStor/templs"
DECOY_DIR = "/projects/casp/dldata"

blosum_mtx = np.array([[  4, -1, -2, -2,  0, -1, -1,  0, -2, -1, -1, -1, -1, -2, -1,  1,  0, -3, -2,  0],
                       [ -1,  5,  0, -2, -3,  1,  0, -2,  0, -3, -2,  2, -1, -3, -2, -1, -1, -3, -2, -3],
                       [ -2,  0,  6,  1, -3,  0,  0,  0,  1, -3, -3,  0, -2, -3, -2,  1,  0, -4, -2, -3],
                       [ -2, -2,  1,  6, -3,  0,  2, -1, -1, -3, -4, -1, -3, -3, -1,  0, -1, -4, -3, -3],
                       [  0, -3, -3, -3,  9, -3, -4, -3, -3, -1, -1, -3, -1, -2, -3, -1, -1, -2, -2, -1],
                       [ -1,  1,  0,  0, -3,  5,  2, -2,  0, -3, -2,  1,  0, -3, -1,  0, -1, -2, -1, -2],
                       [ -1,  0,  0,  2, -4,  2,  5, -2,  0, -3, -3,  1, -2, -3, -1,  0, -1, -3, -2, -2],
                       [  0, -2,  0, -1, -3, -2, -2,  6, -2, -4, -4, -2, -3, -3, -2,  0, -2, -2, -3, -3],
                       [ -2,  0,  1, -1, -3,  0,  0, -2,  8, -3, -3, -1, -2, -1, -2, -1, -2, -2,  2, -3],
                       [ -1, -3, -3, -3, -1, -3, -3, -4, -3,  4,  2, -3,  1,  0, -3, -2, -1, -3, -1,  3],
                       [ -1, -2, -3, -4, -1, -2, -3, -4, -3,  2,  4, -2,  2,  0, -3, -2, -1, -2, -1,  1],
                       [ -1,  2,  0, -1, -3,  1,  1, -2, -1, -3, -2,  5, -1, -3, -1,  0, -1, -3, -2, -2],
                       [ -1, -1, -2, -3, -1,  0, -2, -3, -2,  1,  2, -1,  5,  0, -2, -1, -1, -1, -1,  1],
                       [ -2, -3, -3, -3, -2, -3, -3, -3, -1,  0,  0, -3,  0,  6, -4, -2, -2,  1,  3, -1],
                       [ -1, -2, -2, -1, -3, -1, -1, -2, -2, -3, -3, -1, -2, -4,  7, -1, -1, -4, -3, -2],
                       [  1, -1,  1,  0, -1,  0,  0,  0, -1, -2, -2,  0, -1, -2, -1,  4,  1, -3, -2, -2],
                       [  0, -1,  0, -1, -1, -1, -1, -2, -2, -1, -1, -1, -1, -2, -1,  1,  5, -2, -2,  0],
                       [ -3, -3, -4, -4, -2, -2, -3, -2, -2, -3, -2, -3, -1,  1, -4, -3, -2, 11,  2, -3],
                       [ -2, -2, -2, -3, -2, -1, -2, -3,  2, -1, -1, -2, -1,  3, -3, -2, -2,  2,  7, -1],
                       [  0, -3, -3, -3, -1, -2, -2, -3, -3,  3,  1, -2,  1, -1, -2, -2,  0, -3, -1,  4]]).astype(np.float32)
blosum_mtx /= 10.0 # same as Nao's

tip_atoms = {"ALA": "CB", "CYS": "SG", "ASP": "CG", "GLU": "CD",
             "PHE": "CZ", "GLY": "CA", "HIS":"NE2", "ILE":"CD1",
             "LYS": "NZ", "LEU": "CG", "MET": "SD", "ASN": "CG",
             "PRO": "CG", "GLN": "CD", "ARG": "CZ", "SER": "OG",
             "THR":"OG1", "VAL": "CB", "TRP":"CH2", "TYR": "OH"}

def f(X, cutoff=3, scaling=3.0):
    X_prime = np.maximum(X, np.zeros_like(X) + cutoff) - cutoff
    return np.arcsinh(X_prime)/scaling

def blosum(seq):
    return blosum_mtx[seq]

def t2q(tmplt,nres):

    dmin = 2.0
    dmax = 20.0

    qmap = tmplt['qmap']
    L = qmap.shape[0]

    # bin distance matrix
    nbins = 18
    bins = np.linspace(dmin, dmax, nbins+1)
    dist = tmplt['dist6d']
    dist[dist<0.001] = 999.9
    dbin = np.digitize(dist, bins).astype(np.float32)
    dbin[dbin > nbins] = 0

    f1d = np.zeros((nres,4), dtype=np.float32)
    f1d[qmap] = np.concatenate([
        tmplt['phi'][:,None],
        tmplt['psi'][:,None],
        tmplt['scores'][:,None],
        tmplt['conf'][:,None]
    ], axis=-1)

    f2d = np.zeros((nres,nres,6), dtype=np.float32)
    f2d[np.ix_(qmap,qmap)] = np.concatenate([
        dbin[:,:,None],
        tmplt['omega6d'][:,:,None],
        tmplt['theta6d'][:,:,None],
        tmplt['phi6d'][:,:,None],
        np.tile(tmplt['stats'][None,None,:],[L,L,1])
    ], axis=-1)

    return [f1d,f2d]

def subsample_msa(msa):
    n_sub = int(msa.shape[0]/2)
    if n_sub < 5:
        return msa
    seq = msa[0]
    tmp = msa[1:]
    np.random.shuffle(tmp)
    return np.concatenate([seq[np.newaxis,:], tmp[:n_sub-1,:]], axis=0)

def get_seqsep(n_res):
    pos = np.arange(n_res)
    pos = pos.astype(np.float32)
    tiled_pos = np.tile(pos, (n_res, 1))
    seqsep = np.abs(tiled_pos - tiled_pos.T)/100.0 - 1.0
    return seqsep

def get_labels(pdb, noise_level=0.0):
    label_fn = os.path.join(LABEL_dir, pdb + ".npz")
    data = np.load(label_fn)
    #
    msa = subsample_msa(data['msa'])
    seq = blosum(msa[0])
    #
    dssp = data['dssp'].astype(np.float32)
    phi = data['phi'].astype(np.float32)
    psi = data['psi'].astype(np.float32)
    omg = data['omg'].astype(np.float32)
    omg = np.stack((1-omg, omg), axis=-1)

    return dssp, phi, psi, omg, msa, seq

def get_input_features(filename, is_train=True, mask_diag=8):
    fn_s = glob.glob("%s/%s/*.npz"%(DECOY_DIR, filename))
    np.random.shuffle(fn_s)
    fn = fn_s[0]
    #
    data = np.load(fn)
    #
    seqsep = get_seqsep(data['maps'].shape[0])
    #
    # structure-based features
    d_CB = f(data['maps'][:,:,0].astype(np.float32))
    o_map = np.stack([data["omega6d"], data["theta6d"], data["phi6d"]], axis=-1).astype(np.float32)
    o_map = np.concatenate([np.sin(o_map), np.cos(o_map)], axis=-1)
    feat_2d = np.concatenate((d_CB[:,:,None], o_map, seqsep[:,:,None]), axis=-1)
    if is_train: # add gaussian noise in training stage to avoid overfitting
        scale = np.random.uniform(low=0.05, high=0.15)
        noise = np.random.normal(scale=scale, size=feat_2d.shape).astype(np.float32)
        feat_2d += noise
    feat_2d = np.nan_to_num(feat_2d)
    #
    if mask_diag > -1:
        for i in range(feat_2d.shape[-1]):
            for k in range(mask_diag):
                np.fill_diagonal(feat_2d[k:,:,i], 0.0)
                np.fill_diagonal(feat_2d[:,k:,i], 0.0)
    #
    # Template information
    nres = data['maps'].shape[0]
    tlist = glob.glob("%s/%s/*.npz"%(TEMPL_dir, filename))
    np.random.shuffle(tlist)
    n_templ = np.random.randint(10, 26)
    tlist = tlist[:n_templ]
    pthreads = [t2q(np.load(f), nres) for f in tlist]
    pth_1d = np.stack([p[0] for p in pthreads], axis=0)
    pth_2d = np.stack([p[1] for p in pthreads], axis=0)
    #
    return feat_2d, pth_1d, pth_2d

def load_train_data(pdb, noise_level=0.0, tag=None, is_train=True, mask_diag=8):
    #
    SS_label, phi_label, psi_label, omg_label, msa, seq = get_labels(pdb, noise_level=noise_level)
    feat_2d, pth_1d, pth_2d = get_input_features(pdb, is_train=is_train, mask_diag=mask_diag)
    #

    return seq, msa, feat_2d, pth_1d, pth_2d, SS_label, phi_label, psi_label, omg_label

# TODO
# input feature generation for prediction
def read_msa(a3m_fn):
    import string

    seqs = []
    table = str.maketrans(dict.fromkeys(string.ascii_lowercase))

    # read file
    with open(a3m_fn) as fp:
        for line in fp:
            if line[0] == '>': continue
            seqs.append(line.rstrip().translate(table)) # remove lowercase letters
    # convert letters into numbers
    alphabet = np.array(list("ARNDCQEGHILKMFPSTWYV-"), dtype='|S1').view(np.uint8)
    msa = np.array([list(s) for s in seqs], dtype='|S1').view(np.uint8)
    for i in range(alphabet.shape[0]):
        msa[msa == alphabet[i]] = i

    msa[msa > 20] = 20
    return msa

def get_dihedrals(a, b, c, d):
    b0 = -1.0*(b - a)
    b1 = c - b
    b2 = d - c

    b1 /= np.linalg.norm(b1, axis=-1)[:,None]

    v = b0 - np.sum(b0*b1, axis=-1)[:,None]*b1
    w = b2 - np.sum(b2*b1, axis=-1)[:,None]*b1

    x = np.sum(v*w, axis=-1)
    y = np.sum(np.cross(b1, v)*w, axis=-1)

    return np.arctan2(y, x)

def get_angles(a, b, c):
    v = a - b
    v /= np.linalg.norm(v, axis=-1)[:,None]

    w = c - b
    w /= np.linalg.norm(w, axis=-1)[:,None]

    x = np.sum(v*w, axis=1)

    return np.arccos(np.clip(x, -1.0, 1.0))

def get_dist_ori_maps(model, resNo_s, n_res, dmax=20.0, return_sep=False):
    atm_names = list()
    for i in range(1, model.size()+1):
        if model.residue(i).name()[:3] == "GLY":
            atm_names.append('CA')
        else:
            atm_names.append('CB')
    #
    CA  = np.zeros((n_res, 3), dtype=np.float32)
    N   = np.zeros((n_res, 3), dtype=np.float32)
    O   = np.zeros((n_res, 3), dtype=np.float32)
    C   = np.zeros((n_res, 3), dtype=np.float32)
    CB  = np.zeros((n_res, 3), dtype=np.float32)
    #
    for idx, resNo in enumerate(resNo_s):
        i = idx + 1 # resNo in pose
        CA[resNo, :] = np.array(model.residue(i).atom("CA").xyz())
        N[resNo, :] = np.array(model.residue(i).atom("N").xyz())
        O[resNo, :] = np.array(model.residue(i).atom("O").xyz())
        C[resNo, :] = np.array(model.residue(i).atom("C").xyz())
        CB[resNo, :] = np.array(model.residue(i).atom(atm_names[idx]).xyz())
    #
    mask = np.zeros((n_res, n_res), dtype=np.float32)
    for i in resNo_s:
        for j in resNo_s:
            mask[i,j] = 1.0
    #
    dist_CB_CB = scipy.spatial.distance.cdist(CB, CB, metric='euclidean')
    #
    dist_maps = f(dist_CB_CB[:,:,None])
    #
    # recreate Cb from given N, Ca, C
    b = CA - N
    c = C - CA
    a = np.cross(b, c)
    Cb = -0.58273431*a + 0.56802827*b - 0.54067466*c + CA
    #
    # neighbor search
    kdCb = scipy.spatial.cKDTree(Cb)
    indices = kdCb.query_ball_tree(kdCb, dmax)
    #
    # indices of contacting residues
    idx = np.array([[i,j] for i in range(len(indices)) for j in indices[i] if (i != j) and (i in resNo_s) and (j in resNo_s)]).T
    idx0 = idx[0]
    idx1 = idx[1]
    #
    # matrix of Ca-Cb-Cb-Ca dihedrals
    omega6d = np.zeros((n_res, n_res))
    ang = get_dihedrals(CA[idx0], CB[idx0], CB[idx1], CA[idx1])
    omega6d[idx0, idx1] = ang
    #
    # matrix of polar coord theta
    theta6d = np.zeros((n_res, n_res))
    ang = get_dihedrals(N[idx0], CA[idx0], Cb[idx0], Cb[idx1])
    theta6d[idx0, idx1] = ang
    #
    # matrix of polar coord phi
    phi6d = np.zeros((n_res, n_res))
    ang = get_angles(CA[idx0], Cb[idx0], Cb[idx1])
    phi6d[idx0, idx1] = ang
    #
    if return_sep:
        dist6d = np.zeros((nres, nres), dtype=np.float32)
        dist6d[idx0, idx1] = np.linalg.norm(Cb[idx1]-Cb[idx0], axis=-1)
        return dist6d, omega6d, theta6d, phi6d
    #
    ori_map = np.stack([omega6d, theta6d, phi6d], axis=-1).astype(np.float32)
    ori_map = np.concatenate([np.sin(ori_map), np.cos(ori_map)], axis=-1)
    #
    ori_map = ori_map * mask[:,:,None]
    dist_maps = dist_maps * mask[:,:,None]
    #
    return dist_maps, ori_map

def get_phipsi(xyz, tmap):
    nres = xyz.shape[1]

    N = xyz[0]
    Ca = xyz[1]
    C = xyz[2]

    C_prev = np.insert(C, 0, np.nan, axis=0)
    N_next = np.insert(N, nres, np.nan, axis=0)

    phi = get_dihedrals(C_prev[:-1,:], N, Ca, C)
    psi = get_dihedrals(N, Ca, C, N_next[1:,:])
    phi[0] = 0.0
    psi[-1] = 0.0
    tor = np.stack((phi, psi), axis=-1)[tmap]
    return tor

# get 6d coordinates from x,y,z coords of N,Ca,C atoms
def get_coords6d(xyz, dmax):

    nres = xyz.shape[1]

    # three anchor atoms
    N  = xyz[0]
    Ca = xyz[1]
    C  = xyz[2]

    # recreate Cb given N,Ca,C
    b = Ca - N
    c = C - Ca
    a = np.cross(b, c)
    Cb = -0.58273431*a + 0.56802827*b - 0.54067466*c + Ca

    # fast neighbors search to collect all
    # Cb-Cb pairs within dmax
    kdCb = scipy.spatial.cKDTree(Cb)
    indices = kdCb.query_ball_tree(kdCb, dmax)

    # indices of contacting residues
    idx = np.array([[i,j] for i in range(len(indices)) for j in indices[i] if i != j]).T
    idx0 = idx[0]
    idx1 = idx[1]

    # Cb-Cb distance matrix
    dist6d = np.zeros((nres, nres))
    dist6d[idx0,idx1] = np.linalg.norm(Cb[idx1]-Cb[idx0], axis=-1)

    # matrix of Ca-Cb-Cb-Ca dihedrals
    omega6d = np.zeros((nres, nres))
    omega6d[idx0,idx1] = get_dihedrals(Ca[idx0], Cb[idx0], Cb[idx1], Ca[idx1])

    # matrix of polar coord theta
    theta6d = np.zeros((nres, nres))
    theta6d[idx0,idx1] = get_dihedrals(N[idx0], Ca[idx0], Cb[idx0], Cb[idx1])

    # matrix of polar coord phi
    phi6d = np.zeros((nres, nres))
    phi6d[idx0,idx1] = get_angles(Ca[idx0], Cb[idx0], Cb[idx1])

    return dist6d, omega6d, theta6d, phi6d

def get_2D_str_feat(model, n_res):
    resNo_s = list()
    for i in range(1, model.size()+1):
        resNo = model.pdb_info().number(i) # resNo in PDB file
        resNo_s.append(resNo-1) # -1 to make it as python array index
    resNo_s = np.array(resNo_s)
    d_map, o_map = get_dist_ori_maps(model, resNo_s, n_res)
    seqsep = get_seqsep(n_res)
    return np.concatenate((d_map, o_map, seqsep[:,:,None]), axis=-1)

def make_input_features(a3m_fn, pdb_fn=None, hhr_fn=None, ffdb=None, mask_diag=8, seqID=999.0):
    if isinstance(a3m_fn, str):
        msa = read_msa(a3m_fn)
    else:
        msa = a3m_fn
    seq = blosum(msa[0])
    if len(msa) > MAXSEQ:
        msa = msa[:MAXSEQ]
    #
    if pdb_fn == None: return seq, msa
    model = pose_from_file(pdb_fn)
    feat_2d = get_2D_str_feat(model, len(msa[0]))
    for i in range(feat_2d.shape[-1]):
        for k in range(mask_diag):
            np.fill_diagonal(feat_2d[k:,:,i], 0.0)
            np.fill_diagonal(feat_2d[:,k:,i], 0.0)

    pth_1d = list()
    pth_2d = list()
    if hhr_fn != None:
        L0 = len(msa[0])
        dmin = 2.0
        dmax = 20.0
        nbins = 18
        dbins = np.linspace(dmin, dmax, nbins+1)
        hhr = parse_hhr(hhr_fn, ffdb.index, seqID)
        for i, hit in enumerate(hhr[:10]):
            # extract template from FFindexDB
            entry = get_entry_by_name(hit[0], ffdb.index)
            data = read_entry_lines(entry, ffdb.data)
            xyz_tmp, idx_tmp = parse_pdb_lines(data)
            idx_tmp = idx_tmp-1 # change residue indices to start from 0

            # 6d coordinates
            d, o, t, p = get_coords6d(xyz_tmp, 20.0)
            db = np.digitize(d,dbins)

            # residue indices in the query and the template
            sel = np.intersect1d(hit[1][:,1], idx_tmp, return_indices=True)[1]
            qmap = hit[1][sel,0]
            tmap = hit[1][sel,1]

            # format 1d features
            f1d_a = np.zeros((L0,2), dtype=np.float32)
            f1d_a[qmap] = get_phipsi(xyz_tmp, tmap)
            f1d_k = np.zeros((L0,2), dtype=np.float32)
            f1d_k[qmap] = hit[1][sel,2:]
            pth_1d.append(np.concatenate((f1d_a, f1d_k), axis=-1))

            # format 2d features
            f2d_k = np.zeros((L0,L0,6), dtype=np.float32)
            ij = np.ix_(qmap,qmap.T)
            f2d_k[:,:,0][ij] = db[tmap][:,tmap]
            f2d_k[:,:,1][ij] = o[tmap][:,tmap]
            f2d_k[:,:,2][ij] = t[tmap][:,tmap]
            f2d_k[:,:,3][ij] = p[tmap][:,tmap]
            f2d_k[:,:,4:][ij] = np.array(hit[2:])#s0d[i][None,None,:]
            pth_2d.append(f2d_k)

        pth_1d = np.array(pth_1d)
        pth_2d = np.array(pth_2d)
    #
    return seq, msa, feat_2d, pth_1d, pth_2d


from ops import conv, relu, conv_, max_pooling, fully_connected, batchnorm, prelu, leaky_relu, B_residual_blocks, pixelshuffler
import tensorflow as tf
import numpy as np

class generator:
    def __init__(self, name, B):
        self.name = name
        self.B = B

    def __call__(self, inputs, train_phase):
        with tf.variable_scope(self.name):
            inputs = conv("conv1", inputs, 64, 9)
            inputs = prelu("alpha1", inputs)
            skip_connection = tf.identity(inputs)
            #The paper has 16 residual blocks
            for b in range(1, self.B + 1):
                inputs = B_residual_blocks("B"+str(b), inputs, train_phase)
            # inputs = B_residual_blocks("B2", inputs, train_phase)
            # inputs = B_residual_blocks("B3", inputs, train_phase)
            # inputs = B_residual_blocks("B4", inputs, train_phase)
            # inputs = B_residual_blocks("B5", inputs, train_phase)
            inputs = conv("conv2", inputs, 64, 3)
            inputs = batchnorm(inputs, train_phase, "BN")
            inputs = inputs + skip_connection
            inputs = conv("conv3", inputs, 256, 3)
            inputs = pixelshuffler(inputs, 2)
            inputs = prelu("alpha2", inputs)
            inputs = conv("conv4", inputs, 256, 3)
            inputs = pixelshuffler(inputs, 2)
            inputs = prelu("alpha3", inputs)
            inputs = conv("conv5", inputs, 3, 9)
        return tf.nn.tanh(inputs)

    def var_list(self):
        return tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope=self.name)

class discriminator:
    def __init__(self, name):
        self.name = name

    def __call__(self, inputs, train_phase):
        with tf.variable_scope(self.name, reuse=tf.AUTO_REUSE):
            # inputs = tf.random_crop(inputs, [-1, 70, 70, 3])
            inputs = conv("conv1_1", inputs, 64, 3, 2)
            inputs = leaky_relu(inputs, 0.2)
            # inputs = conv("conv1_2", inputs, 64, 3, is_SN=True)
            # inputs = leaky_relu(inputs, 0.2)
            inputs = conv("conv2_1", inputs, 128, 3, 2)
            inputs = batchnorm(inputs, train_phase, "BN1")
            inputs = leaky_relu(inputs, 0.2)
            # inputs = conv("conv2_2", inputs, 128, 3, is_SN=True)
            # inputs = leaky_relu(inputs, 0.2)
            inputs = conv("conv3_1", inputs, 256, 3, 2)
            inputs = batchnorm(inputs, train_phase, "BN2")
            inputs = leaky_relu(inputs, 0.2)
            # inputs = conv("conv3_2", inputs, 256, 3, is_SN=True)
            # inputs = leaky_relu(inputs, 0.2)
            inputs = conv("conv4_1", inputs, 512, 3, 2)
            inputs = batchnorm(inputs, train_phase, "BN3")
            inputs = leaky_relu(inputs, 0.2)
            # inputs = fully_connected("fc", inputs, 512, is_SN=True)
            output = fully_connected("output", inputs, 1)
        return output

    def var_list(self):
        return tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope=self.name)

def vggnet(inputs, vgg_path):
    inputs = (inputs + 1) * 127.5
    inputs = tf.reverse(inputs, [-1]) - np.array([103.939, 116.779, 123.68])
    para = np.load(vgg_path+"vgg19.npy", encoding="latin1").item()
    inputs = relu(conv_(inputs, para["conv1_1"][0], para["conv1_1"][1]))
    inputs = relu(conv_(inputs, para["conv1_2"][0], para["conv1_2"][1]))
    inputs = max_pooling(inputs)
    inputs = relu(conv_(inputs, para["conv2_1"][0], para["conv2_1"][1]))
    inputs = relu(conv_(inputs, para["conv2_2"][0], para["conv2_2"][1]))
    F = inputs
    inputs = max_pooling(inputs)
    inputs = relu(conv_(inputs, para["conv3_1"][0], para["conv3_1"][1]))
    inputs = relu(conv_(inputs, para["conv3_2"][0], para["conv3_2"][1]))
    inputs = relu(conv_(inputs, para["conv3_3"][0], para["conv3_3"][1]))
    inputs = max_pooling(inputs)
    inputs = relu(conv_(inputs, para["conv4_1"][0], para["conv4_1"][1]))
    inputs = relu(conv_(inputs, para["conv4_2"][0], para["conv4_2"][1]))
    inputs = relu(conv_(inputs, para["conv4_3"][0], para["conv4_3"][1]))
    return F


#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# Copyright (C) 2018 Brno University of Technology FIT
# Author: Jan Profant <jan.profant@phonexia.com>
# All Rights Reserved

import h5py
import numpy as np
from scipy.io.idl import AttrDict
from scipy.sparse import coo_matrix

from vbdiar.utils.utils import Utils

class PLDA(object):

    V = None
    U = None
    D = None
    mc = None
    mu = None
    _R = None
    _T = None
    _y = None
    _x = None
    Syy = None
    vdim = None
    srank = None
    crank = None
    noise_type = None

    def init(self, vdim, srank, crank, diag="isotropic"):
        """

        Args:
            vdim:
            srank:
            crank:
            diag:
        """
        self.vdim = vdim
        self.srank = srank
        self.crank = crank
        self.noise_type = diag

        np.random.seed(7)
        self.V = np.array(np.random.randn(vdim, srank), np.float) * 1e0
        self.U = np.array(np.random.randn(vdim, crank), np.float) * 1e0
        if diag == "zero":
            self.D = np.zeros((vdim, 1), np.float)
        else:
            self.D = np.ones((vdim, 1), np.float)

        self.mu = np.zeros((vdim, 1), np.float)

        # expectations
        self._R = np.eye(self.srank + self.crank, self.srank + self.crank)
        self._T = np.zeros((self.srank + self.crank, self.vdim))
        self.Syy = np.zeros((self.srank, self.srank))
        # point estimates
        self._y = None
        self._x = None

    def __str__(self):
        return "= PLDA model d=%d,s=%d,c=%d,type=%s" % (self.vdim, self.srank, self.crank, self.noise_type)

    def __init__(self, model_file):
        """

        Args:
            model_file:

        Returns:

        """
        model = h5py.File(model_file, 'r')
        v, u, d, noise_type, mu = model['V'][:], model['U'][:], model['D'][:], model['noise_type'], model['mu'][:]
        vdim, srank = v.shape
        crank = u.shape[1]
        self.init(vdim, srank, crank, noise_type)
        self.V = v
        self.U = u
        self.D = d
        self.mu = mu
        self.cache_statistics()

    @staticmethod
    def symmetrize(a):
        a = a + a.T
        a *= 0.5
        return a

    @staticmethod
    def logdet_chol(chol_m):
        """Return the log determinant of a matrix given the cholesky decomposition"""
        return 2 * np.sum(np.log(np.diagonal(chol_m)))

    @staticmethod
    def invhandle(m, func_only=True):
        """Return a function handle on multiplying a matrix
           by the inverse of the marix provided in this function
           if func_only=False then return: [logdet,chol,handle]
           (Niko's style)
           This is used when scoring the PLDA system
           """
        cm = np.linalg.cholesky(m)

        def h(a):
            return np.linalg.solve(cm.T, np.linalg.solve(cm, a.T))

        if func_only:
            return h
        else:
            logdet = PLDA.logdet_chol(cm)
            return [logdet, cm, h]

    def cache_statistics(self):
        """
            various values independant of the data in the cache object
        """
        # values independant of the data
        mc = AttrDict()
        mc.UD = self.U*self.D
        mc.VD = self.V*self.D
        mc.J = np.dot(self.U.T, mc.VD)

        mc.K = np.sqrt(self.D) * self.U
        mc.K = PLDA.symmetrize(np.dot(mc.K.T, mc.K)+np.eye(self.crank))

        mc.iK = np.linalg.inv(mc.K)
        mc.P0 = mc.VD.T-np.dot(np.dot(mc.J.T, mc.iK), mc.UD.T)
        self.mc = mc
        return mc

    def score_with_constant_n(self, nt, ft, nt2, ft2):
        """  Compute verification score
             This changes for every pair nT nt and can be done only once for 1 session training
             fT = P*y_head = (V'D-J'iK*U'D)*f, where f is the first order stats
        """
        # Covariance and first order stats preparation
        Py = PLDA.symmetrize(np.dot(self.mc.P0, self.V))  # (DV-iKJDU)V=VDV-JiKJ

        # P in eq. 31, EM for PLDA by Niko
        Py_Train = nt * Py + np.eye(self.srank)
        Py_Test = nt2 * Py + np.eye(self.srank)
        Py_SameSpk = (nt2 + nt) * Py + np.eye(self.srank)

        [logdetQ1, cholQ1, hQ1] = PLDA.invhandle(Py_Train, func_only=False)
        [logdetQ2, cholQ2, hQ2] = PLDA.invhandle(Py_Test, func_only=False)
        [logdetQ12, cholQ12, hQ12] = PLDA.invhandle(Py_SameSpk, func_only=False)
        # fT*hQ1(fT.T) is terms y_head*P*y_head in eq 30, EM for PLDA by Niko
        Q1 = 0.5 * (logdetQ1 + sum(ft * hQ12(ft.T)) - sum(ft * hQ1(ft.T)))
        Q2 = 0.5 * (logdetQ2 + sum(ft2 * hQ12(ft2.T)) - sum(ft2 * hQ2(ft2.T)))
        # Q2=Q2[np.newaxis,:] (can be dangerous)
        Q1 = Q1[:, np.newaxis]
        scores = np.dot(ft.T, hQ12(ft2.T))
        Q2 -= 0.5 * logdetQ12
        scores += Q1
        scores += Q2
        return scores

    def prepare_stats(self, data):
        """ Convert data (e.g. i-vectors) to statistics useful for scoring
            Input:
                data:      ivect-dim x nb_ex matrix
                seg2model: defines of multisession enrollment (or test). It is
                           2 column array of integer indices mapping rows of data
                           (1st column) to rows of output statistics (2nd column).
                           By default (seg2model=None) each vector in data has its
                           own output vector of statistics (i.e. single session
                           enrollment or test). seg2model can be also represented
                           by coo_matrix or by 1D array of labels if each input
                           vector belongs to exactly one enrollment.
            Output: a AttrDict object with
                    N: vector of counts for each speaker
                    F: array of sum of ivector for each speaker transformed by P0
                       to srank dimensionality
        """
        seg2model = np.arange(len(data))
        if seg2model.ndim == 1:
            seg2model = np.c_[np.arange(len(data)), seg2model]
        seg2model = coo_matrix((np.ones(len(seg2model)), (seg2model[:, 0], seg2model[:, 1])))

        stats = AttrDict()
        stats.N = np.array(seg2model.T.sum(1))
        stats.F = np.dot(self.mc.P0, seg2model.T.dot(data - self.mu).T).T
        return stats

    def score(self, test, enroll):
        """
            Score ivectors based on the PLDA model
            Input:
                PLDA object. PLDA.V and PLDA.U gives the subspaces
                stats objects: enroll and test
            Output:
                2D array of scores of all possibilities
        """
        test = test - self.mu
        enroll = enroll - self.mu
        test = Utils.l2_norm(test)
        enroll = Utils.l2_norm(enroll)
        Tstats = self.prepare_stats(test)
        tstats = self.prepare_stats(enroll)
        # Create scores
        scores = np.zeros((len(Tstats.N), len(tstats.N)), 'f')
        (a, b) = scores.shape

        # Score for each uniq combination of N enroll and M test sessions (only enroll for now)
        for n_enroll_sessions in np.unique(Tstats.N):
            idxs_T = np.where(Tstats.N == n_enroll_sessions)[0]
            for n_test_sessions in np.unique(tstats.N):
                idxs_t = np.where(tstats.N == n_test_sessions)[0]
                scores[np.ix_(idxs_T, idxs_t)] = self.score_with_constant_n(n_enroll_sessions, Tstats.F[idxs_T, :].T,
                                                                            n_test_sessions, tstats.F[idxs_t, :].T)
        return scores.T


#!/usr/bin/env python
# -*- coding: utf-8 -*-
#

import sys
import os
import time
import itertools

from heapq import heappush, heappop

# Compatibility with Python < 2.6
#
try:
  from heapq import heappushpop
except ImportError:
  def heappushpop(heap, item):
    heappush(heap, item)
    return heappop(heap)

import numpy

from prosci.util.protein import Pdb, ResidueList, Protein
from prosci.util.pdb3d import dist
from prosci.util.geohash import GeometricHash

from prosci.loops import ccd
from prosci.loops.loopmodel import ANCHOR_LENGTH, get_max_loop_length, describe_anchors, iterate_database, get_loop_structure, add_oxygens, is_loop_length_in_database, get_min_loop_length_in_database, relabel_loop, is_clash, calculate_rmsd, is_consecutive, make_contact_gh, find_contacts, get_contact_class, find_contacts_simple, get_native_contacts
from prosci.loops.esst import load_esst

# Define consecutive residues very loosely, for purposes of checking the framework
def is_framework_consecutive(a, b):
  # Set max distance for consecutive C and N atoms to 2.8 Angstroms
  # ... which is very loose
  #
  return is_consecutive(a, b, maxlen_sq=7.84)

def _heapadd(heap, item, maxitems):
  if len(heap) < maxitems:
    heappush(heap, item)
  else:
    heappushpop(heap, item)

class FreadError(RuntimeError):
  pass

class BadInputError(FreadError):
  pass

class UnsupportedLoopLengthError(BadInputError):
  pass

class LoopStretchError(BadInputError):
  pass

class NonConsecutiveAnchorError(BadInputError):
  pass

class Fread(object):
  def __init__(self, db=None, subst_table_path=None, score_cutoff=25, open_rmsd_cutoff=1.0, closed_rmsd_cutoff=0.3, vdw_factor=0.7, close=True, verbose=False, errstream=sys.stderr, meld=True, max_melding_rmsd=1.0, nostruc=False, mutate=False, ccd_target_rmsd=0.15, ccd_iterations=5000, max_decoys=100, first_decoys=0, extension_size=sys.maxint, extension_minimum=0, calculate_contacts=False, contact_distance=6.0, contact_identity=0.8):
    self.score_cutoff = score_cutoff
    self.open_rmsd_cutoff = open_rmsd_cutoff
    self.closed_rmsd_cutoff = closed_rmsd_cutoff
    self.vdw_factor = vdw_factor
    self.close = close
    self.verbose = verbose
    self.errstream = errstream
    self.meld = meld
    self.max_melding_rmsd = max_melding_rmsd
    self.min_database_loop_length = 1
    self.no_structure_output = nostruc
    self.mutate = mutate
    self.ccd_target_rmsd = ccd_target_rmsd
    self.ccd_iterations = ccd_iterations
    self.max_decoys = max_decoys
    self.first_decoys = first_decoys
    self.extension_size = extension_size
    self.extension_minimum = extension_minimum
    self.min_b_factor = 0.0

    self.calculate_contacts = calculate_contacts
    self.contact_distance = contact_distance
    self.contact_identity = contact_identity

    self.warnings = []

    if db:
      self.set_db(db)

    self.set_subst_tables(subst_table_path)

    self.results = []
    self.counter = itertools.count(1)
    self.idecoy = 0

  def set_options(self, opts):
    for o in opts:
      if o in self.__dict__ and not o.startswith('_'):
        self.__dict__[o] = opts[o]

  def set_silent(self):
    #self.verbose = False
    self.errstream = open(os.devnull, 'w')

  def note(self, message, force=False):
    if force or self.verbose:
      self.errstream.write(str(message)+"\n")

  def warn(self, exception, force=False):
    if isinstance(exception, str):
      exception = FreadError(exception)
    self.note("WARNING: "+str(exception), force)
    self.warnings.append(exception)
    return exception

  def die(self, exception, force=True):
    raise self.warn(exception, force)

  def set_subst_tables(self, subst_table_path):
    self.subst_tables = load_esst(subst_table_path)

  def set_db(self, db):
    self.min_database_loop_length = get_min_loop_length_in_database(db)
    self.db = db

  def set_structure(self, structure):
    self.clear_results()
    p = Pdb(structure)
    self.ligand_atoms = p.ligands
    self.residues = ResidueList(p)
    add_oxygens(self.residues)

  def get_loop_index(self, residue_number, inscode="", chain=None):
    if inscode:
      inscode = inscode.strip()
    if chain:
      chain = chain.strip()

    if isinstance(residue_number, str) and residue_number[-1].isalpha():
        ires = int(residue_number[:-1])
        inscode = residue_number[-1]
    else:
        ires = int(residue_number)

    start_of_loop=None
    for i,r in enumerate(self.residues):
      if r.ires == ires and r.inscode == inscode and (not chain or chain == r.chain):
          start_of_loop = i+1
          break

    return start_of_loop

  def get_structure_sequence(self, startindex=0, endindex=None):
    return self.residues[startindex:endindex].get_seq()

  def set_loop(self, start_of_loop, loop_sequence, resnum=False, chain=None, structure_known=None):
    self.clear_results()
    if chain is not None:
      chain = chain.strip()

    # We've been given a residue number instead of a start index
    if resnum:
      if isinstance(start_of_loop, str):
        if start_of_loop[-1].isalpha():
          ires = int(start_of_loop[:-1])
          inscode = start_of_loop[-1]
        else:
          ires = int(start_of_loop)
          inscode = ""
      else:
        ires = int(start_of_loop)
        inscode = ""

      found=False
      for i,r in enumerate(self.residues):
        if r.ires == ires and r.inscode == inscode and (not chain or chain == r.chain):
            start_of_loop = i+1
            found = True
            self.note("Located residue left of loop to be modelled: %d%s. Loop starting index: %d. Sequence surrounding start of loop: %s|%s" % (r.ires, r.inscode, start_of_loop, self.residues[max(0,start_of_loop-8):start_of_loop].get_seq(), self.residues[start_of_loop:start_of_loop+8].get_seq()))
            break

      if not found:
        msg = "Residue before start of loop (residue number %d%s" % (ires, inscode)
        if chain:
          msg += ", chain '%s'" % chain
        msg += ") not found in query structure"
        self.die(BadInputError(msg))

    if start_of_loop < ANCHOR_LENGTH or start_of_loop > len(self.residues) - ANCHOR_LENGTH:
      self.die(BadInputError("Cannot model loop closer than %d residues to the terminus (need a complete anchor on both sides)." % (ANCHOR_LENGTH)))

    # Assuming no coordinates are present in file, start and end of the loop are the same
    end_of_loop = start_of_loop

    if structure_known is not None:
      # We were told if loop coordinates are present in input file.
      # Skip auto-detection.
      #
      self.loop_structure_is_known = structure_known
      if structure_known:
        self.note("User tells me loop structure is present.")
      else:
        self.note("User tells me loop structure is not present.")
    # Auto-detect if loop coordinates are present in input file.
    else:
      # Are the coordinates are in the file?
      self.loop_structure_is_known = is_framework_consecutive(self.residues[start_of_loop-1],
                                                    self.residues[start_of_loop])
      self.note("Is loop structure present: "+str(self.loop_structure_is_known))

    # If we have the native loop structure, adjust end_of_loop.
    # Also do some additional sanity checks
    #
    if self.loop_structure_is_known:
      end_of_loop += len(loop_sequence)
      strucseq = ""
      for i,r in enumerate(self.residues[start_of_loop:end_of_loop]):
          strucseq += r.get_seq()
      if (not self.mutate) and (loop_sequence != strucseq):
        self.die(BadInputError("Residues differ between sequence and structure input: %s, %s" % (loop_sequence, strucseq)))
      self.note("Loop sequence of given ATOM co-ordinates: %s" % (strucseq))
      del strucseq

    if end_of_loop != start_of_loop:
      if len(loop_sequence) != end_of_loop - start_of_loop:
        self.die(BadInputError("Loop co-ordinates present in input, but number of residues (%d) does not match length of input loop sequence (%d)." % (end_of_loop - start_of_loop, len(loop_sequence))))
      if end_of_loop > len(self.residues) - ANCHOR_LENGTH:
        self.die(BadInputError("Cannot model loop closer than %d residues to the terminus (need a complete anchor on both sides)." % (ANCHOR_LENGTH)))

    # Set ourselves up for loop extension
    #
    self.seq = self.loop_sequence = loop_sequence
    self.s = self.start_of_loop = start_of_loop
    self.e = self.end_of_loop = end_of_loop

    self.verify_anchors()

    while len(self.seq) < self.min_database_loop_length:
      try:
        self.verify_stretch()
      except LoopStretchError:
        pass
      self.extend_loop()

    while True:
      try:
        self.verify_stretch()
        break
      except LoopStretchError as e:
        try:
          self.extend_loop()
        except BadInputError:
          self.die(LoopStretchError(str(e)+" Cannot extend gap further."))

  def extend_loop(self):
      """Extend loop by one residue on each side. Raises NonConsecutiveAnchorError() or UnsupportedLoopLengthError() upon failure.
      """
      start_of_loop = self.start_of_loop
      end_of_loop = self.end_of_loop
      s = self.s - 1
      e = self.e + 1

      length = (start_of_loop - s) + (e - end_of_loop) + len(self.loop_sequence)

      if s<ANCHOR_LENGTH or not is_framework_consecutive(self.residues[s-ANCHOR_LENGTH], self.residues[s-ANCHOR_LENGTH+1]) or e>len(self.residues)-ANCHOR_LENGTH or not is_framework_consecutive(self.residues[e+ANCHOR_LENGTH-2], self.residues[e+ANCHOR_LENGTH-1]):
        self.die(NonConsecutiveAnchorError("Cannot extend loop to length %d, due to gaps in the query structure or proximity to the termini"%(length)))

      if not is_loop_length_in_database(self.db, length):
        self.die(UnsupportedLoopLengthError("Cannot extend loop to length %d, due to database limitations"%(length)))

      seq = ""
      for i in xrange(s, start_of_loop):
        seq += self.residues[i].get_seq()
      seq += self.loop_sequence
      for i in xrange(end_of_loop, e):
        seq += self.residues[i].get_seq()

      assert len(seq) == length, str([s, start_of_loop, end_of_loop, e, length, len(seq), seq, dbdir, pdb_file, start_of_loop, loop_sequence])

      self.s = s
      self.e = e
      self.seq = seq

      self.note("Extending loop to length %d" % (length))

  def verify_anchors(self):
    # Ensure anchors are consecutive stretches of amino acids
    for x in xrange(self.s-ANCHOR_LENGTH+1, self.s):
      if not is_framework_consecutive(self.residues[x-1], self.residues[x]):
        self.die(NonConsecutiveAnchorError("Anchor residues not consecutive in framework structure: residue index %s, %s"%(self.s, self.residues.code)))
    for x in xrange(self.e, self.e+ANCHOR_LENGTH-1):
      if not is_framework_consecutive(self.residues[x], self.residues[x+1]):
        self.die(NonConsecutiveAnchorError("Anchor residues not consecutive in framework structure: residue index %s, %s"%(self.s, self.residues.code)))

  def verify_stretch(self):
    # Ensure anchors are not too far apart for loop to stretch gap
    anchor_distance = dist(self.residues[self.s-1].C, self.residues[self.e].N)
    if anchor_distance > get_max_loop_length(len(self.seq)) * 1.05:
      self.die(LoopStretchError("Loop span (%.2f Angstrom) too large to be closed using %d residues. Trying to extend gap." % (anchor_distance, len(self.seq))))

  def verify_length(self):
    # Ensure loop is not too long for the database
    if not is_loop_length_in_database(self.db, len(self.seq)):
      self.die(UnsupportedLoopLengthError("Cannot model loop of length %d, due to database limitations"%(len(self.seq))))

  def verify(self):
    self.verify_anchors()
    self.verify_stretch()
    self.verify_length()

  def model(self, top=None, stop_after=None, f_rank_decoy=None, f_stop_search=None, f_filter=None):
    """Model loop using the FREAD algorithm. This is the method handling most of the work. Raises UnsupportedLoopLengthError if loop length is not supported.
    """
    if top is None:
      top = self.max_decoys

    if stop_after is None:
      stop_after = self.first_decoys

    if top <= 0:
      top = sys.maxint

    if stop_after > 0 and self.idecoy >= stop_after:
      return self.results

    if not f_rank_decoy:
      f_rank_decoy = FREAD_RANKING
    if not f_stop_search:
      f_stop_search = lambda x: False
    if not f_filter:
      f_filter = lambda x: True

    while len(self.seq) < self.min_database_loop_length:
      self.extend_loop()

    self.verify()

    meld_anchors = self.meld  # and (self.open_rmsd_cutoff <= self.max_melding_rmsd)
    close_loop = self.close

    verbose = self.verbose
    residues = self.residues
    start_of_loop = self.s
    end_of_loop = self.e
    loop_sequence = self.seq

    start = start_of_loop - ANCHOR_LENGTH
    end = end_of_loop + ANCHOR_LENGTH
    loop_length = len(loop_sequence)
    total_length = loop_length + 2*ANCHOR_LENGTH

    self.note("Loop region: N(%4d, %4d)  C(%4d,%4d)  Loop(%3d,'%s')" % (start, start_of_loop, end_of_loop, end, loop_length, loop_sequence))

    ############################################################################

                # Get query anchors and prepare for clash checking #

    # Get anchor coordinates
    #
    anchor_N = residues[start:start_of_loop]
    anchor_C = residues[end_of_loop:end]

    # Describe anchors in query structure
    #
    anchor_description, query_transform = describe_anchors(anchor_N, anchor_C, loop_length)

    # Build a GeometricHash of the query structure (without the loop region), for
    # clash checking
    #
    coords = []
    gh_atoms = []
    for r in residues[:start]+residues[end:]:
      for a in r:
        #if a.atom in ("N", "CA", "C", "O", "CB"):
        coords.append(a.xyz)
        gh_atoms.append(a)
    gh = GeometricHash(numpy.array(coords))
    del coords

    # Inter-residue contacts
    if self.calculate_contacts:
      p = Protein((residues[:start]+residues[end:]).split_chains())
      p.ligands = Protein(ResidueList(self.ligand_atoms).split_chains())
      contact_gh, contact_gh_atoms = make_contact_gh(p)

    ############################################################################

                                # Search the database #

    results = [] # Heap Queue of the top-ranking decoys
    for decoy in self.results:

      _heapadd(results, (f_rank_decoy(decoy), decoy.idecoy, decoy), top)

    for decoy in iterate_database(self.db, loop_length, self.subst_tables, anchor_description, loop_sequence, self.open_rmsd_cutoff, self.score_cutoff):

        if len(results) >= top:
          if decoy.internal_rmsd > results[0][-1].anchor_rmsd_open:
            continue

        # Retrieve loop structure from database
        decoy_residues = get_loop_structure(self.db, decoy.struc, decoy.start, total_length)

        assert len(decoy_residues) == total_length

        # Superimpose anchors and check anchor RMSD before starting
        anchor_rmsd_open = ccd.superimpose(decoy_residues[:ANCHOR_LENGTH]+decoy_residues[-ANCHOR_LENGTH:], anchor_N+anchor_C, decoy_residues)

        if anchor_rmsd_open > self.open_rmsd_cutoff:
          self.note("%s_%d_%d : Anchor RMSD too large: %.3f"%(decoy.struc, decoy.start, loop_length, anchor_rmsd_open))
          continue

        if len(results) >= top:
          if anchor_rmsd_open > results[0][-1].anchor_rmsd_open:
            continue

        # Save start residue number of loop, in database structure
        decoy.startres = decoy_residues[ANCHOR_LENGTH].ires
        decoy.startinscode = decoy_residues[ANCHOR_LENGTH].inscode

        # Relabel residues and discard non-matching atoms
        relabel_loop(decoy_residues[ANCHOR_LENGTH:-ANCHOR_LENGTH], loop_sequence, prevatom=anchor_N[-1].CA, nextatom=anchor_C[0].CA)

        if self.loop_structure_is_known:
          loop_rmsd_open = calculate_rmsd(decoy_residues[ANCHOR_LENGTH:-ANCHOR_LENGTH], residues[start_of_loop:end_of_loop])

        # Are we allowed to meld this decoy?
        #
        meld_this_decoy = meld_anchors and (anchor_rmsd_open <= self.max_melding_rmsd) #and not self.no_structure_output
        close_this_decoy = close_loop and anchor_rmsd_open > self.ccd_target_rmsd and (not meld_anchors or anchor_rmsd_open > self.max_melding_rmsd)

        #if not self.no_structure_output or close_this_decoy:
        decoy.nanchor = anchor_N.deep_copy()
        decoy.canchor = anchor_C.deep_copy()

        if meld_this_decoy:
          meld(decoy.nanchor, decoy_residues[:ANCHOR_LENGTH])
          meld(decoy.canchor, decoy_residues[-ANCHOR_LENGTH:], invertWeights=True)

        if not close_this_decoy:
          anchor_rmsd_closed = anchor_rmsd_open
          iterations = 0
        else:
          # Loop closure
          anchor_rmsd_closed, iterations = ccd.close_loop(decoy.nanchor, decoy_residues, decoy.canchor, target_rmsd=self.ccd_target_rmsd, iterations=self.ccd_iterations)
          if anchor_rmsd_closed > self.closed_rmsd_cutoff:
            self.note("Failed to close loop: %s_%d_%d"%(decoy.struc, decoy.start, loop_length))
            continue

        # Cut off the decoy loop's anchors
        decoy_residues = decoy_residues[ANCHOR_LENGTH:-ANCHOR_LENGTH]

        #if not self.no_structure_output:
        # Restore main chain oxygens, which got lost during the melding/closing procedure
        decoy_residues[-1].O = None # This oxygen is wrong anyway, so delete it
        add_oxygens(decoy.nanchor+decoy_residues[:1], force=True)
        add_oxygens(decoy_residues[-1:]+decoy.canchor+residues[end:end+1], force=True)

        # Clash check
        is_clashing = is_clash(gh, gh_atoms, decoy_residues, self.vdw_factor)
        if is_clashing:
          self.note("Clash detected in decoy: %s_%d_%d"%(decoy.struc, decoy.start, loop_length))
          continue

        if self.loop_structure_is_known:
          decoy.loop_rmsd_open = loop_rmsd_open
          if not iterations:
            decoy.loop_rmsd_closed = loop_rmsd_open
          else:
            decoy.loop_rmsd_closed = calculate_rmsd(decoy_residues, residues[start_of_loop:end_of_loop])

        decoy.length = loop_length
        decoy.anchor_rmsd_open = anchor_rmsd_open
        decoy.iterations = iterations

        #if not self.no_structure_output:
        decoy.loop = decoy_residues

        if self.calculate_contacts:
          contacts = ""
          for i, r in enumerate(decoy_residues):
            contact_atoms = find_contacts(contact_gh, r, maxdist=self.contact_distance)
            contact_atoms.extend(find_contacts_simple(decoy.nanchor+decoy_residues+decoy.canchor, i, maxdist=self.contact_distance))
            contacts += get_contact_class(r.chain, contact_atoms)

          decoy.native_contacts = get_native_contacts(self.db, decoy.struc, decoy.start, total_length)

          id = 0
          for x,y in zip(decoy.native_contacts, contacts):
            if x == y:
              id += 1
          decoy.contact_identity = float(id) / len(contacts)
          decoy.contacts = contacts

          if (self.contact_identity > 0) and (decoy.contact_identity < self.contact_identity):
            continue

        # Get per-residue substitution scores
        #
        tables = self.subst_tables.tables
        ascii2index = self.subst_tables.ascii2index
        seqmap = tuple([ascii2index[ord(s)] for s in loop_sequence])
        dihed = decoy.dihedrals
        score_list = []
        for i,x in enumerate(decoy.seq):
          score_list.append(tables[int(dihed[i])][seqmap[i]][ascii2index[ord(x)]])
        assert sum(score_list) == decoy.score
        decoy.residue_scores = score_list

        # Save per-residue scores in the occupancy column and
        # the total score in the B factor column
        #
        #
        for i,r in enumerate(decoy.loop):
          for a in r:
            a.b = self.min_b_factor + max(0, 30 - decoy.score)
            # Average score over a 3 residue window...
            # because that's how MEDELLER does it
            sc = score_list[max(0,i-1):i+2]
            a.occup = float(sum(sc))/len(sc)

        if not f_filter(decoy):
          self.note("User-defined filter is excluding decoy: %s_%d_%d"%(decoy.struc, decoy.start, loop_length))
          continue

        self.idecoy = decoy.idecoy = next(self.counter)
        # Save result
        _heapadd(results, (f_rank_decoy(decoy), decoy.idecoy, decoy), top)

        if stop_after > 0 and self.idecoy >= stop_after:
          break

        if f_stop_search(decoy):
          break

    self.note("%d decoys found"%(len(results)))

    self.results = []
    while results:
      r = heappop(results)
      self.results.append(r[-1])
    self.results.reverse()
    return self.results

  def clear_results(self):
    self.results = []
    self.warnings = []
    self.counter = itertools.count(1)
    self.idecoy = 0

  def write_summary(self, outstream=sys.stdout, write_decoy_sequence=False):
    """Write summary information to specified output stream.
    """
    for decoy in self.results:
      outstream.write("%s_%d%s_%d\t%d\t%.3f\t%.3f\t%d\t%s" % (decoy.struc, decoy.startres, decoy.startinscode, decoy.length, decoy.score, decoy.internal_rmsd, decoy.anchor_rmsd_open, decoy.iterations,decoy.seq))
      if write_decoy_sequence:
        outstream.write("\t%s" % (decoy.seq))
      if self.calculate_contacts:
        outstream.write("\t%s" % (decoy.contacts))
        outstream.write("\t%s" % (decoy.native_contacts))
        outstream.write("\t%.3f" % (decoy.contact_identity))
      if self.loop_structure_is_known:
        outstream.write("\t%.3f\t%.3f" % (decoy.loop_rmsd_open, decoy.loop_rmsd_closed))
      outstream.write("\n")

  def write_decoy_structures(self, out_dir, top=0, suffix=".loop.atm", idfilter=[]):
    """Write decoy structure files (PDB format) to specified directory.
    """
    if self.no_structure_output:
      self.die("Cannot write decoy structures as structure output is disabled")

    if out_dir:
      if not os.path.exists(out_dir):
        os.mkdir(out_dir)

    for i, decoy in enumerate(self.results):
      if top > 0 and i >= top:
        break

      decoyname = "%s_%d%s_%d"%(decoy.struc, decoy.startres, decoy.startinscode, decoy.length)

      if idfilter and (decoyname not in idfilter):
        continue

      decoyfile = os.path.join(out_dir, decoyname+suffix)
      f = open(decoyfile, "w")
      try:
        for r in decoy.nanchor:
          f.write(str(r))
        for r in decoy.loop:
          f.write(str(r))
        for r in decoy.canchor:
          f.write(str(r))
      finally:
        f.close()

  def assemble_model(self, decoy):
    # Start and end of loop region in model
    start = self.s - ANCHOR_LENGTH
    end = self.e + ANCHOR_LENGTH

    if isinstance(decoy, int):
      decoy = self.results[decoy]
      if self.no_structure_output:
        self.die("Cannot write decoy structures as structure output is disabled")

    model = ResidueList([])
    for i in xrange(start):
      model.append(self.residues[i].copy())
    model.extend(decoy.nanchor)
    model.extend(decoy.loop)
    model.extend(decoy.canchor)
    for i in xrange(end, len(self.residues)):
      model.append(self.residues[i].copy())

    # Replace main chain oxygen before the loop, to ensure the peptide bond there is planar
    model[start-1].O = None
    add_oxygens(model, start=start-1, end=start, force=True)

    return model

  def write_model_structures(self, out_dir, top=0, suffix=".model.atm", idfilter=[]):
    """Write model structure files (PDB format), including decoy residues, to specified directory.
    """

    if self.no_structure_output:
      self.die("Cannot write model structures as structure output is disabled")

    if out_dir:
      if not os.path.exists(out_dir):
        os.mkdir(out_dir)
    else:
      out_dir = ""

    for i, decoy in enumerate(self.results):
      if top > 0 and i >= top:
        break

      decoyname = "%s_%d%s_%d"%(decoy.struc, decoy.startres, decoy.startinscode, decoy.length)

      if idfilter and (decoyname not in idfilter):
        continue

      decoyfile = os.path.join(out_dir, decoyname+suffix)

      model = self.assemble_model(decoy)

      f = open(decoyfile, "w")
      try:
        f.write(str(model))
      finally:
        f.close()

  def open_errstream(self, filename):
      if filename == "-":
        self.errstream = sys.stderr
      else:
        self.errstream = open(filename, "w")

  def close_errstream(self):
      if self.errstream != sys.stderr:
        self.errstream.close()
        self.set_silent()

  def automodel_loop(self, start_of_loop, loop_sequence,
                      loopsuffix = "",
                      dbdir = '.',
                      strucdir = 'decoys',
                      summary_file_name = 'summary',
                      write_decoys = True,
                      write_models = True,
                      write_decoy_sequence = False,
                      resnum = False,
                      chain = None,
                      structure_known = None,
                      f_rank_decoy = None,
                      f_stop_search = None,
                      f_filter = None,
                      idfilter = [],
                      **kwargs):
      """Search the given database and return a list of decoys for the given loop.
      """

      # Set options defined in the Fread constructor
      self.set_options(kwargs)

      if self.extension_size < 0:
        max_extension = sys.maxint
      else:
        max_extension = self.extension_size

      if strucdir is not None:
        strucdir = strucdir.strip()
      if (not strucdir) or (not write_decoys and not write_models):
          strucdir = None
          self.no_structure_output = True

      for db_group in dbdir.split(":"): # try these databases until we find something
        databases = db_group.split("|") # combine results from all these databases
        self.set_db(databases[0])

        try:

          self.set_loop(start_of_loop, loop_sequence.upper(), resnum=resnum, chain=chain, structure_known=structure_known)
        except FreadError:
          break

        # Extend loop if extension_minimum is set
        #
        for i in xrange(self.extension_minimum):
          try:
            self.extend_loop()
          except FreadError:
            break

        # Try to model, and extend loop if nothing was found
        anyresults = False
        for db in databases:

          self.set_db(db)
          try:

            anyresults |= bool(self.model(f_rank_decoy=f_rank_decoy, f_stop_search=f_stop_search, f_filter=f_filter))
          except FreadError:
            break

        for i in xrange(max_extension):
          if anyresults:
            break
          try:
            self.extend_loop()
          except FreadError:
            break
          for db in databases:
            self.set_db(db)

            anyresults |= bool(self.model(f_rank_decoy=f_rank_decoy, f_stop_search=f_stop_search, f_filter=f_filter))

        if anyresults:
          break

      # Write results to STDOUT or a file, if one was specified
      if summary_file_name:
        if summary_file_name == "-":
          self.write_summary(sys.stdout, write_decoy_sequence=write_decoy_sequence)
        else:
          root, ext = os.path.splitext(summary_file_name)
          outstream = open(root + loopsuffix + ext, "w")
          self.write_summary(outstream, write_decoy_sequence=write_decoy_sequence)
          outstream.close()

      if not self.no_structure_output:
        if write_decoys:
          self.write_decoy_structures(strucdir + loopsuffix, suffix=".loop.pdb", idfilter = idfilter)
        if write_models:
          self.write_model_structures(strucdir + loopsuffix, suffix=".model.pdb", idfilter = idfilter)

      return self.results

def meld(fragScaffold, fragPrediction, invertWeights=False):
    """Averages the coordinates of the two Pdb arguments. Move the first object's co-ordinates onto the averaged position.

    By default, the first object is assumed to be part of the N-terminal fragment, the second is part of the C-terminal fragment. This can be reversed by setting invertWeights=True."""

    resS = fragScaffold
    resP = fragPrediction
    L = len(resS)

    assert len(resS) == len(resP)

    def averageCoord(P, S):
      # P = coordinate of prediction
      # S = coordinate of scaffold
      # D = distance (in residues) from (loop+anchor)-fragment end
      # L = anchor length
      return 1.0/(L+1) * (D*P + (L+1-D)*S)

    for i, (rS, rP) in enumerate(zip(resS, resP)):
      if invertWeights:
        D = len(resS)-i
      else:
        D = i+1

      newN = averageCoord(rP.N.xyz, rS.N.xyz)
      newCA = averageCoord(rP.CA.xyz, rS.CA.xyz)
      newC = averageCoord(rP.C.xyz, rS.C.xyz)

      T_from, T_to, rotmat = ccd.get_rotmat([rS.N.xyz, rS.CA.xyz, rS.C.xyz], [newN, newCA, newC])

      rS.O = None # Remove main chain oxygen - can regenerate this later

      for a in rS:
        a.xyz = numpy.dot(a.xyz - T_from, rotmat) + T_to

def FREAD_RANKING(decoy):
  """Returns a comparable object (a tuple of scores) used for ranking a given decoy.

  Bigger values will be ranked higher."""
  return (-decoy.anchor_rmsd_open, decoy.score, -decoy.internal_rmsd, -decoy.iterations)

def FREAD_RANKING_BY_ESSS(decoy):
  """Returns a comparable object (a tuple of scores) used for ranking a given decoy.

  Bigger values will be ranked higher."""
  return (decoy.score, -decoy.anchor_rmsd_open, -decoy.internal_rmsd, -decoy.iterations)


import sys
from datetime import datetime
import math
from collections import Counter
from transformers import BertForSequenceClassification, Trainer, TrainingArguments, AdamW, RobertaForTokenClassification, RobertaTokenizer, AutoTokenizer
from transformers.modeling_outputs import SequenceClassifierOutput
from sklearn.metrics import accuracy_score, precision_recall_fscore_support, classification_report
import unicodedata
from sklearn.model_selection import train_test_split
import torch
from sklearn.model_selection import KFold
import numpy as np
import collections
import argparse
import jsonlines
import json
from tqdm import tqdm
import itertools
import os
from cleantext import clean
from urllib.parse import unquote
import torch

from feverous.utils.annotation_processor import AnnotationProcessor, EvidenceType
from feverous.utils.wiki_page import WikiPage, get_wikipage_by_id
from feverous.database.feverous_db import FeverousDB
from feverous.utils.log_helper import LogHelper

LogHelper.setup()
logger = LogHelper.get_logger(__name__)

class FEVEROUSDataset(torch.utils.data.Dataset):
    def __init__(self, encodings, labels, use_labels = True):
        self.encodings = encodings
        self.labels = labels
        self.use_labels = use_labels

    def __getitem__(self, idx):
        item = {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}
        if self.use_labels:
            item['labels'] = torch.tensor(self.labels[idx])
        return item

    def __len__(self):
        return len(self.labels)

def process_data(claim_verdict_list):

    text = [x[0] for x in claim_verdict_list]#["I love Pixar.", "I don't care for Pixar."]
    labels = [x[1] for x in claim_verdict_list] #get value from enum

    return text, labels

def compute_metrics(pred):
    labels = list(itertools.chain(*pred.label_ids))
    preds = list(itertools.chain(*pred.predictions.argmax(-1)))
    preds = [pred for i, pred in enumerate(preds) if labels[i] != -100]
    labels = [lab for lab in labels if lab != -100]

    precision, recall, f1, _ = precision_recall_fscore_support(labels, preds, average='micro')
    acc = accuracy_score(labels, preds)
    # map_verdict_to_index = {0:'Not enough Information', 1: 'Supported', 2:'Refuted'}
    class_rep = classification_report(labels, preds, output_dict=True)
    return {
        'accuracy': acc,
        'f1': f1,
        'precision': precision,
        'recall': recall,
        'class_rep': class_rep
    }

def model_trainer(args, test_dataset):
    # model = BertForSequenceClassification.from_pretrained('bert-base-uncased', num_labels =4)
    model = RobertaForTokenClassification.from_pretrained(args.model_path, num_labels = 3, return_dict=True)

    #/anfs/bigdisc/rmya2/faiss_data/results_table_to_cell2/checkpoint-1400/'
    training_args = TrainingArguments(
    per_device_train_batch_size=16,  # batch size per device during training
    per_device_eval_batch_size=16,   # batch size for evaluation
    # warmup_steps=0,                # number of warmup steps for learning rate scheduler
    logging_dir='./logs',
    output_dir='./model_output'
    )

    trainer = Trainer(
    model=model,                         # the instantiated 🤗 Transformers model to be trained
    args=training_args,                  # training arguments, defined above
    eval_dataset=test_dataset,          # evaluation dataset
    compute_metrics = compute_metrics,
    )
    return trainer, model

def report_average(reports):
    mean_dict = dict()
    for label in reports[0].keys():
        dictionary = dict()

        if label in 'accuracy':
            mean_dict[label] = sum(d[label] for d in reports) / len(reports)
            continue

        for key in reports[0][label].keys():
            dictionary[key] = sum(d[label][key] for d in reports) / len(reports)
        mean_dict[label] = dictionary

    return mean_dict

def convert_evidence_into_tables_gold(annotation, db):
    evidence = annotation.get_evidence()[0] #only bother with first set for now
    all_cells = []
    all_cells_id = []
    for piece in evidence:
        if '_cell_' in piece:
            all_cells.append(piece)
            all_cells_id.append('_'.join(piece.split('_')[1:]))
    cell_ids_by_table = {}
    table_by_id = {}
    for i,cell in enumerate(all_cells):
        pa, page = get_wikipage_by_id(cell, db)
        if pa == None:
            continue
        table = pa.get_table_from_cell(all_cells_id[i])
        tab_id = page + "_"  + table.get_id().replace('t_', 'table_')
        if tab_id in cell_ids_by_table:
            cell_ids_by_table[tab_id].append(all_cells_id[i])
        else:
            cell_ids_by_table[tab_id] = [all_cells_id[i]]
            table_by_id[tab_id] = table
    output_tables = []
    output_tables_ids = []
    output_labels = []
    for id, cells in cell_ids_by_table.items():
        curr_tab = table_by_id[id]
        table_flat = [annotation.get_claim(), '</s>']
        table_flat_ids = [0, 0]
        labels_flat = [0, 0]
        for row in curr_tab.rows:
            for j, cell in enumerate(row.row):
                table_flat.append(str(cell))
                # table_flat_ids.append(id.split('_')[0] + cell.get_id().replace('hc_', '_header_cell_').replace('c_','_cell_'))
                table_flat_ids.append(id.split('_')[0] + '_' + cell.get_id().replace('hc_', 'header_cell_').replace('c_','cell_'))
                # cell_header_ele = [ele.split('_')[0] + '_' +  el.get_id().replace('hc_', 'header_cell_') for el in wiki_page.get_context('_'.join(ele.split('_')[1:])) if "header_cell_" in el.get_id()]
                if cell.get_id() in cells:
                    labels_flat.append(1)
                else:
                    labels_flat.append(0)
                if j < len(row.row) -1:
                    table_flat.append('|')
                    table_flat_ids.append(0)
                    labels_flat.append(0)
            table_flat.append('</s>')
            table_flat_ids.append(0)
            labels_flat.append(0)
        output_tables.append(table_flat)
        output_tables_ids.append(table_flat_ids)
        output_labels.append(labels_flat)
    return (output_tables, output_tables_ids, output_labels)

def convert_evidence_into_tables(annotation, db, args):
    predicted_sentences= [ele[0] + '_sentence_' + ele[1].split('_')[1] for ele in annotation.predicted_evidence  if ele[1].split('_')[0] == 'sentence']
    predicted_tables =  [ele[0] + '_table_' + ele[1].split('_')[1] for ele in annotation.predicted_evidence if ele[1].split('_')[0] == 'table']
    evidence = set(list(predicted_sentences[:3] + predicted_tables[:2]))
    tables = [ele for ele in evidence if '_table_' in ele]
    output_tables = []
    output_tables_ids = []
    for tab_id in tables:
        pa, page = get_wikipage_by_id(tab_id, db)
        if pa == None:
            continue
        curr_tab = pa.page_items['_'.join(tab_id.split('_')[1:])]
        table_flat = [annotation.get_claim(), '</s>']
        table_flat_ids = [0, 0]
        for row in curr_tab.rows:
            for j, cell in enumerate(row.row):
                table_flat.append(str(cell))
                table_flat_ids.append(page + '_' + cell.get_id().replace('hc_', 'header_cell_').replace('c_','cell_'))
                if j < len(row.row) -1:
                    table_flat.append('|')
                    table_flat_ids.append(0)
            table_flat.append('</s>')
            table_flat_ids.append(0)
        output_tables.append(table_flat)
        output_tables_ids.append(table_flat_ids)
    return (output_tables, output_tables_ids)

def tokenize_and_align_labels(text_in, labels_in, all_input_ids, tokenizer):
    label_all_tokens = False
    tokenized_inputs = tokenizer(
        text_in,
        padding=True,
        truncation=True,
        # We use this argument because the texts in our dataset are lists of words (with a label for each word).
        is_split_into_words=True,
    )
    labels = []
    all_input_ids_new = []
    for i, label in enumerate(labels_in):
        word_ids = tokenized_inputs.word_ids(batch_index=i)
        previous_word_idx = None
        label_ids = []
        input_ids = []
        for word_idx in word_ids:
            # Special tokens have a word id that is None. We set the label to -100 so they are automatically
            # ignored in the loss function.
            if word_idx is None:
                label_ids.append(-100)
                input_ids.append(-100)
            # We set the label for the first token of each word.
            elif word_idx != previous_word_idx:
                label_ids.append(label[word_idx])
                input_ids.append(all_input_ids[i][word_idx])
            # For the other tokens in a word, we set the label to either the current label or -100, depending on
            # the label_all_tokens flag.
            else:
                label_ids.append(label[word_idx] if label_all_tokens else -100)
                input_ids.append(all_input_ids[i][word_idx] if label_all_tokens else -100)
            previous_word_idx = word_idx
        labels.append(label_ids)
        all_input_ids_new.append(input_ids)
    tokenized_inputs["labels"] = labels
    return tokenized_inputs, all_input_ids_new, labels

def extract_cells_from_tables(annotations, args):
    db = FeverousDB(args.wiki_path)
    all_input = []
    all_input_ids = []
    all_labels = []
    anno_ids = []
    for i, anno in enumerate(tqdm(annotations)):
        tabs,tabs_ids = convert_evidence_into_tables(anno, db, args)
        all_input += tabs
        all_input_ids += tabs_ids
        all_labels += [[0] * len(el) for el in tabs] #dummy labels
        anno_ids += len(tabs) * [anno.get_id()]

    logger.info('Sample entry: {}'.format(all_input[0]))
    # logger.info('Sample label: {}'.format(all_labels[0]))
    logger.info('Sample id: {}'.format(anno_ids[0]))

    if not args.trivial_baseline:
        tokenizer = AutoTokenizer.from_pretrained('roberta-base', do_lower_case=True, add_prefix_space=True)

        text_test, all_input_ids, labels_test = tokenize_and_align_labels(all_input, all_labels, all_input_ids, tokenizer) #dummy labels

        # lab_after = [[pred for i, pred in enumerate(label) if labels_test[j][i] != -100] for j,label in enumerate(labels_test)]

        test_dataset = FEVEROUSDataset(text_test, labels_test)

        trainer, model = model_trainer(args, test_dataset)
        model_output = trainer.predict(test_dataset)

        # predictions = model_output.predictions.argmax(-1)
        predictions =  (model_output.predictions > 0.25).astype(int)
        # print(predictions)
        predictions = [[1 if ele[1]==1 else 0 for ele in e] for e in predictions]
        # print(predictions)
        # print([list(pred).count(1) for pred in predictions])
        labels = model_output.label_ids

        # preds = pred.predictions.argmax(-1)

        predictions = [[pred for i, pred in enumerate(pred_instance) if labels[j][i] != -100] for j,pred_instance in enumerate(predictions)]

        all_input_ids = [[pred for i, pred in enumerate(label) if labels[j][i] != -100] for j,label in enumerate(all_input_ids)]

        # print(predictions)
        predictions_map = {anno_ids[i]:[all_input_ids[i][j] for j,el in enumerate(predictions[i]) if predictions[i][j] == 1 and all_input_ids[i][j] != 0] for i,annota in enumerate(anno_ids)}
    else:
        predictions_map = {value: [e for e in all_input_ids[i] if e !=0][:25] for i, value in enumerate(anno_ids)}
    # print(predictions_map.keys())

    with jsonlines.open(os.path.join(args.input_path.split('.jsonl')[0] + '.cells.jsonl') , 'w') as writer:
        with jsonlines.open(os.path.join(args.input_path)) as f:
             for i,line in enumerate(f.iter()):
                 if i == 0:
                     writer.write({'header':''}) # skip header line
                     continue
                 # if len(line['evidence'][0]['content']) == 0: continue
                 predicted_sentences= [ele[0] + '_sentence_' + ele[1].split('_')[1] for ele in line['predicted_evidence'] if ele[1].split('_')[0] == 'sentence']
                 predicted_tables =  [ele[0] + '_table_' + ele[1].split('_')[1] for ele in line['predicted_evidence']  if ele[1].split('_')[0] == 'table' ]
                 line['predicted_evidence'] = set(list(predicted_sentences + predicted_tables))
                 if not args.trivial_baseline:
                     line['predicted_evidence'] = [ele for ele in line['predicted_evidence'] if ('_table_' not in ele and '_table_caption_' not in ele)]
                 else:
                     line['predicted_evidence'] = [ele for ele in line['predicted_evidence'] if ('_table_' not in ele and '_table_caption_' not in ele and '_sentence_' not in ele)]
                 # print(line['id'])
                 for key, value in predictions_map.items():
                     if key == line['id']:
                         line['predicted_evidence'] += predictions_map[key]
                 writer.write(line)

    logger.info('Finished extracting cells...')

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--input_path', type=str, help='/path/to/data')
    parser.add_argument('--model_path', type=str, help='/path/to/data')
    parser.add_argument('--trivial_baseline', action='store_true', default=False)
    parser.add_argument('--max_sent', type=int, default=5, help='/path/to/data')
    parser.add_argument('--wiki_path', type=str)

    args = parser.parse_args()

    anno_processor =AnnotationProcessor(args.input_path)
    annotations = [annotation for annotation in anno_processor]
    # annotations.sort(key=lambda x: x.source, reverse=True)
    logger.info('Start extracting cells from Tables...')
    extract_cells_from_tables(annotations, args)

if __name__ == "__main__":
    main()


import numpy as np
import cv2
from picamera.array import PiRGBArray as pi_rgb
from picamera import PiCamera as picam
from vibration import Haptic
import time

height = 320
width = 640

font = cv2.FONT_HERSHEY_PLAIN
color = [255, 255, 255]
left_boundary = int(width * 3 / 8)
right_boundary = int(width * 5 / 8)
upper_boundary = int(height * 2 / 3)

class PiCam:
    def __init__(self):
        self.cam = picam()
        self.cam.resolution = (width, height)
        self.cam.framerate = 10
        self.raw_cap = pi_rgb(picam, size=self.cam.resolution)
        time.sleep(0.1)

def execute_warning(haptic_obj, position_vector):
    haptic_obj.vibrate(100 * position_vector[0])
    haptic_obj.vibrate(100 * position_vector[1])
    haptic_obj.vibrate(100 * position_vector[2])

def position_warning(haptic_obj, x_pos, y_pos):
    # pos_mat : [upper, left, right]
    position_mat = [0, 0, 0]
    if y_pos > upper_boundary:
        position_mat[0] = 1
    if x_pos < left_boundary:
        position_mat[1] = 1
    elif x_pos > right_boundary:
        position_mat[2] = 1
    execute_warning(haptic_obj, position_mat)
    return position_mat

if __name__ == '__main__':
    camera = PiCam()
    vibration_motor = Haptic()
    time.sleep(0.1)
    net = cv2.dnn.readNet("mppi_yolo/yolov3.weights",
                          "mppi_yolo/yolov3.cfg")

    with open("mppi_yolo/coco.names", "r") as f:
        classes = f.read().splitlines()

    for frame in camera.cam.capture_continuous(camera.raw_cap, format="bgr", use_video_port=True):
        img = frame.array
        blob = cv2.dnn.blobFromImage(img, np.divide(1, 255), (width, height), (0, 0, 0), swapRB=True, crop=False)
        net.setInput(blob)
        output_layers_names = net.getUnconnectedOutLayersNames()
        layer_outputs = net.forward(output_layers_names)

        boxes = []
        confidences = []
        class_ids = []

        for output in layer_outputs:
            for detection in output:
                scores = detection[5:]
                class_id = np.argmax(scores)
                confidence = scores[class_id]
                if confidence > 0.5:
                    center_x = int(np.multiply(detection[0], width))
                    center_y = int(np.multiply(detection[1], height))

                    w = int(np.multiply(detection[2], width))
                    h = int(np.multiply(detection[3], height))

                    x = int(np.subtract(center_x, np.divide(w, 2)))
                    y = int(np.subtract(center_y, np.divide(h, 2)))

                    boxes.append([x, y, w, h])
                    confidences.append(float(confidence))
                    class_ids.append(class_id)

        indexes = cv2.dnn.NMSBoxes(boxes, confidences, 0.6, 0.4)

        if len(indexes) > 0:
            for i in indexes.flatten():
                x, y, w, h = boxes[i]
                center_x = x + w // 2
                center_y = y + h // 2
                label = str(classes[class_ids[i]])
                confidence = str(np.round_(confidences[i], 2))
                pos_mat = position_warning(vibration_motor, center_x, center_y)
                cv2.circle(img, (center_x, center_y), 3, color, 2)
                cv2.putText(img, f"{label} {confidence}", (x, y + 20), font, 2, color)

        cv2.imshow("image", img)
        key = cv2.waitKey(1)
        if key == 27:
            break

        camera.raw_cap.truncate(0)

    cv2.destroyAllWindows()


import time
import torchvision
import cv2
import numpy as np
from robot import *
from image_folder2 import *
from config import *
from alset_ddqn import *
import torch.nn.functional as F
import time
import torch
import torch.optim as optim
import torch.nn.functional as F
import torchvision.datasets as datasets
import torchvision.models as models
import torchvision.transforms as transforms

class ALSETNN():
    def __init__(self,alset_robot,outputs, nn_name, app_type):
        self.robot = alset_robot
        self.app_name = None
        self.nn_name = nn_name
        self.app_name = nn_name
        self.app_type = app_type
        self.alset_dqn = None
        self.device = None
        self.model = None
        self.optimizer = None
        self.outputs = outputs   # outputs may be much smaller than full_action_set for a function
        self.num_outputs = len(outputs)
        # self.joystick_actions = ["REWARD","PENALTY"]
        self.dsu = DatasetUtils(self.app_name, self.app_type, self.nn_name)
        self.cfg = Config()
        self.automatic_mode = False
        self.auto_func = None
        self.func_flow_model = None

    def nn_init(self, gather_mode=False):
        self.model = torchvision.models.alexnet(pretrained=True)
        self.model.classifier[6] = torch.nn.Linear(self.model.classifier[6].in_features, self.num_outputs)
        model_path = self.dsu.best_model(mode=self.app_type, nn_name=self.nn_name)
        try:
          self.model.load_state_dict(torch.load(model_path))
          print("Loaded model state.")
        except:
          print("Starting from new Alexnet model.")
          torch.save(self.model.state_dict(), model_path)
        self.device = torch.device('cuda')
        self.model = self.model.to(self.device)
        robot_dirs = []
        self.nn_dir = self.dsu.dataset_path(mode=self.app_type, nn_name=self.nn_name)
        robot_dirs.append(self.nn_dir)
        func_classifier = self.cfg.get_func_value(self.nn_name, "CLASSIFIER")
        if func_classifier is not None:
            outputs = func_classifier[0]
            self.cfg.full_action_set = outputs
            self.func_flow_model = FunctionalApp(self.robot, self.nn_name, "FUNC")
            self.func_flow_model.nn_init()
            action_name = self.func_flow_model.eval_func_flow_model(reward_penalty=None,init=True)
        else:
            outputs = self.cfg.full_action_set
        for dir_name in outputs:
          robot_dirs.append(self.nn_dir + dir_name)
        self.dsu.mkdirs(robot_dirs)
        print("nn_init: " , robot_dirs)
        # Class can be used as a single-NN app that can do any action
        if self.is_automated_function():
          self.auto_func = AutomatedFuncs(self.robot)
          self.auto_func.set_automatic_function(self.nn_name)
        return self.is_automated_function(), outputs

    def preprocess(self, camera_value):
         global device, normalize
         mean = 255.0 * np.array([0.485, 0.456, 0.406])
         stdev = 255.0 * np.array([0.229, 0.224, 0.225])
         normalize = torchvision.transforms.Normalize(mean, stdev)

         x = camera_value
         x = cv2.cvtColor(x, cv2.COLOR_BGR2RGB)
         x = x.transpose((2, 0, 1))
         x = torch.from_numpy(x).float()
         x = normalize(x)
         x = x.to(self.device)
         x = x[None, ...]
         return x

    def nn_process_image(self, NN_name=None, image=None, reward_penalty=None):
        if self.cfg.get_func_value(NN_name, "CLASSIFIER") is not None:
          return self.nn_process_image_cnn(NN_name=NN_name, image=image, reward_penalty=reward_penalty)
        elif image is None:
          print("NN process image None", NN_name, reward_penalty)
          return "NOOP"
        else:
          return self.nn_process_image_dqn(NN_name, image, reward_penalty)

    def nn_process_image_dqn(self, NN_name=None, image=None, reward_penalty=None):
        print("NN process image")
        if self.alset_dqn is None:
          self.alset_dqn = ALSET_DDQN(alset_robot=self.robot, initialize_model=False, do_train_model=False, app_name=NN_name, app_type="FUNC")
          self.alset_dqn.nn_init(gather_mode=False)
        return self.alset_dqn.nn_process_image(NN_name, image, reward_penalty)
        # NN_name=None, image=None, reward_penalty=None)

    def nn_process_image_cnn(self, NN_name=None, image=None, reward_penalty=None):
        # if image == None:
        if image is None:
          print("NN process image None")
          return
        print("NN process image")
        x = image
        x = self.preprocess(x)
        y = self.model(x)

        # we apply the `softmax` function to normalize the output vector
        # so it sums to 1 (which makes it a probability distribution)
        y = F.softmax(y, dim=1)

        print(y.flatten)
        max_prob = 0
        best_action = -1
        # for i in range(self.num_outputs):
        # for i in range(len(self.cfg.full_action_set)):
        for i in range(len(self.outputs)):
            prob = float(y.flatten()[i])
            # print("PROB", i, self.cfg.full_action_set[i], prob)
            print("PROB", i, self.outputs[i], prob)
            if max_prob < prob:
                # for j, name in enumerate(self.cfg.full_action_set):
                for j, name in enumerate(self.outputs):
                    # if name == self.cfg.full_action_set[i]:
                    if name == self.outputs[i]:
                        if (reward_penalty is not None or
                            name not in ["REWARD1", "PENALTY1", "REWARD2", "PENALTY2"]):
                          max_prob = prob
                          best_action = j
                          break
                if best_action == -1:
                    # print("invalid action " + self.cfg.full_action_set[i] + "not in " + self.cfg.full_action_set)
                    print("invalid action " + self.outputs[i] + " not in " + self.outputs)
                    exit()
        # action_name = self.cfg.full_action_set[best_action]
        if self.func_flow_model is None:
          action_name = self.outputs[best_action]
        else:
          print("eval func flow input:", self.outputs[best_action])
          curr_nn, action_name = self.func_flow_model.eval_func_flow_model(self.outputs[best_action])
          print("eval func flow:", curr_nn, self.outputs[best_action], action_name)
        return action_name
#        if action_name == "FORWARD":
#            print("NN FORWARD")
#            self.robot.forward()
#        elif action_name == "REVERSE":
#            print("NN REVERSE")
#            self.robot.reverse()
#        elif action_name == "LEFT":
#            print("NN LEFT")
#            self.robot.left()
#        elif action_name == "RIGHT":
#            print("NN RIGHT")
#            self.robot.right()
#            # self.robot.left() # for simplified tabletop NN
#        elif action_name == "LOWER_ARM_DOWN":
#            print("NN LOWER ARM DOWN")
#            self.robot.lower_arm("DOWN")
#        elif action_name == "LOWER_ARM_UP":
#            print("NN LOWER ARM UP")
#            self.robot.lower_arm("UP")
#        elif action_name == "UPPER_ARM_DOWN":
#            print("NN UPPER ARM DOWN")
#            self.robot.upper_arm("DOWN")
#        elif action_name == "UPPER_ARM_UP":
#            print("NN UPPER ARM UP")
#            self.robot.upper_arm("UP")
#        elif action_name == "GRIPPER_CLOSE":
#            print("NN GRIPPER_CLOSE")
#            self.robot.gripper("CLOSE")
#        elif action_name == "GRIPPER_OPEN":
#            print("NN GRIPPER_OPEN")
#            self.robot.gripper("OPEN")
#        else:
#            print("Action:", best_action, max_prob, action_name)
#            time.sleep(0.08)
#            self.robot.stop()

    def is_automated_function(self):
        try:
          idx1 = self.cfg.func_registry.index(self.nn_name)
          for [func,parent_func] in self.cfg.func_automated:
            if func == self.nn_name:
              print("is_automated_func True")
              return True
          print("is_automated_func False")
          return False
        except:
          print("is_automated_func False")
          return False

    def nn_set_automatic_mode(self, TF):
        if not self.is_automated_function:
          print("not an automated function: ", self.nn_name)
          exit()
        self.automatic_mode = TF

    def nn_automatic_mode(self):
        return self.automatic_mode

    def nn_automatic_action(self, NN_name, frame, feedback):
        if not self.is_automated_function():
          print("not an automated function: ", NN_name)
          exit()
        return self.auto_func.automatic_function(frame, feedback)

    def nn_upon_penalty(self, penalty):
        # goto teleop mode
        self.nn_set_automatic_mode(False)

    def nn_upon_reward(self, reward):
        exit()

    # train the NN to do imitation learning starting from pre-trained imagenet
    #
    # Train a single NN based upon all 8 of the functional NN from TT_FUNC.
    # Handle new data gathered for TT_FUNC.
    # Use a split of the data so that we know the test accuracy.
    # This test accuracy might help with determining whether there is enough
    # data for TT_DQN.
    def train(self, dataset_root_list=None, best_model_path=None, full_action_set=None, noop_remap=None, only_new_images=True):
        if self.cfg.get_func_value(self.nn_name, "CLASSIFIER") is not None:
          self.train_cnn(dataset_root_list, best_model_path, full_action_set, noop_remap, only_new_images)
        else:
          self.train_dqn(dataset_root_list, best_model_path, full_action_set, noop_remap, only_new_images)

    def train_dqn(self, dataset_root_list=None, best_model_path=None, full_action_set=None, noop_remap=None, only_new_images=True):
        if dataset_root_list is None:
          # num_NN = self.cfg.func_registry.index(self.app_name)
          dataset_root_list = []
          dsp = self.dsu.dataset_path("FUNC", self.nn_name)
          dataset_root_list.append(dsp)
        for ds_root in dataset_root_list:
           func_name = self.dsu.get_func_name_from_full_path(ds_root)
           alset_dqn = ALSET_DDQN(initialize_model=False, do_train_model=False, app_name=func_name, app_type="FUNC")
           alset_dqn.nn_init(gather_mode=False)
           alset_dqn.parse_func_dataset(func_name, False, "FUNC")
           # alset_dqn.parse_func_dataset(init=False, app_name=ds_root, app_mode="FUNC")
           #   def parse_func_dataset(self, init=False, app_name=None, app_mode="FUNC"):
           # TypeError: parse_func_dataset() got an unexpected keyword argument 'app_name'

    def train_cnn(self, dataset_root_list=None, best_model_path=None, full_action_set=None, noop_remap=None, only_new_images=True):
        # nn.py is a primitive NN that can be used by higher layers like tabletop_func_app.
        # Set default parameters if unset by caller
        if dataset_root_list is None:
          # num_NN = self.cfg.func_registry.index(self.app_name)
          dataset_root_list = []
          dsp = self.dsu.dataset_path("FUNC", self.nn_name)
          dataset_root_list.append(dsp)
        best_model = self.dsu.best_model("FUNC", self.nn_name, classifier=True)
        if full_action_set is None:
           full_action_set = self.cfg.full_action_set
        if noop_remap is None:
           pass

        #######################
        # idx = self.dsu.dataset_indices(mode="FUNC", nn_name=self.nn_name, position="NEXT")
        # idx_processed = self.dsu.dataset_idx_processed(mode="FUNC", nn_name=self.nn_name)
        #######################

        # if self.robot.initialize:
        # elif self.robot.train_new_data:

        # Dataset transforms
        #
        # Problems to deal with:
        # - The robot actions must be consistent for all apps.
        # - automatic mode: Determine which images are NOOP operations and
        #   train accordingly.  Automatic mode has only REWARD/PENALTY/NOOP operations.
        #   These need to be mapped from the underlying operations: LEFT, UPPER_ARM_{UP/DOWN}
        # - DQN has 2 additional joystick modes: CUBE_OFF_TABLE_REWARD, ROBOT_OFF_TABLE_PENALTY
        #
        # nn_init needs to store info about additional mappings:
        #
        # self.robot_actions = ( "FORWARD", "REVERSE", "LEFT", "RIGHT", "LOWER_ARM_DOWN", "LOWER_ARM_UP", "UPPER_ARM_DOWN", "UPPER_ARM_UP", "GRIPPER_OPEN", "GRIPPER_CLOSE")
        # self.joystick_actions = ("REWARD","PENALTY","CUBE_OFF_TABLE_REWARD","ROBOT_OFF_TABLE_PENALTY")
        # self.automatic_mode() returns True/False. If True:
        #   self.automatic_actions = ( "NOOP", "REWARD", "PENALTY")
        #   self.automatic_mode_noop_mapping = ("LEFT", "LOWER_ARM_UP", "LOWER_ARM_DOWN")
        #
        # model = models.alexnet(pretrained=True)
        # model.classifier[6] = torch.nn.Linear(model.classifier[6].in_features, len(full_action_set))
        # self.model is already loaded from nn_init() = which is a prerequisite
        if self.model is None:
            print("nn_init() not called before training.")
            exit
        NUM_EPOCHS = 30
        # NUM_EPOCHS = 3   # for debugging
        best_accuracy = 0.0
        optimizer = optim.SGD(self.model.parameters(), lr=0.00005, momentum=0.9)
        for root in dataset_root_list:
            while True:
              print("nn train: ", root, best_model_path)
              func_classifier = self.cfg.get_func_value(self.nn_name, "CLASSIFIER")
              if func_classifier is not None:
                full_action_set2 = func_classifier[0]
              else:
                full_action_set2 = full_action_set
              print("fas: ", full_action_set2)
              dataset = ImageFolder2(
                  root,
                  self.nn_name,
                  self.app_type,
                  transforms.Compose([
                      transforms.ColorJitter(0.1, 0.1, 0.1, 0.1),
                      transforms.Resize((224, 224)),
                      transforms.ToTensor(),
                      transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
                  ]),
                  full_action_set = full_action_set2,
                  remap_to_noop = noop_remap,
                  only_new_images = only_new_images
              )
              if len(dataset.imgs) == 0:
                if dataset.all_images_processed("FUNC", self.nn_name):
                  print("no more images")
                  break
                print("next index")
                dataset.save_images_processed("FUNC", self.nn_name)
                continue
              print("num dataset imgs:", len(dataset.imgs))
              print("len dataset     :", len(dataset))
              if len(dataset) > len(dataset.imgs):
                print("dataset[0]:")
                print(dataset[0])
                print("datasetimg:")
                print(dataset.imgs[0])
              # Attributes:
              # classes (list): List of the class names sorted alphabetically.
              # class_to_idx (dict): Dict with items (class_name, class_index).
              # imgs (list): List of (image path, class_index) tuples

              # apply transforms
              # full_action_set => revises classes, class_to_idx
              # noop_remap      => revises imgs' class_index
              #
              # classes = [d.name for d in os.scandir(dir) if d.is_dir()]
              # classes.sort()
              # class_to_idx = {cls_name: i for i, cls_name in enumerate(classes)}
              # imgs= [image_path, class_index:

              # problem: we're currently training incrementally. So, dataset may not be big
              # enough to split into train/test.
              # However, we're starting on a pretrained alexnet model and trying to
              # specialize the model from there.  The original code that this was based
              # upon was training from scratch.
              if len(dataset) > 500:
                train_dataset, test_dataset = torch.utils.data.random_split(dataset, [len(dataset) - 50, 50])
              elif (len(dataset) > 50):
                  # test on a subset of the dataset
                  train_dataset, test_dataset = torch.utils.data.random_split(dataset, [len(dataset) - 50, 50])
                  train_dataset = dataset    # train on full incremental dataset
              else:
                  train_dataset = dataset    # train on full incremental dataset
                  test_dataset = dataset     # test on full incremental dataset

              print("len train_dataset:", len(train_dataset))
              print("len test_dataset :", len(test_dataset))

              train_loader = torch.utils.data.DataLoader(
                train_dataset,
                batch_size=8,
                shuffle=True,
                num_workers=0
              )
              test_loader = torch.utils.data.DataLoader(
                test_dataset,
                batch_size=8,
                shuffle=True,
                num_workers=0
              )

              for epoch in range(NUM_EPOCHS):

                  for images, labels in iter(train_loader):
                      images = images.to(self.device)
                      labels = labels.to(self.device)
                      optimizer.zero_grad()
                      outputs = self.model(images)
                      loss = F.cross_entropy(outputs, labels)
                      loss.backward()
                      optimizer.step()

                  # we may not want to do test counts. Take all the data and
                  # apply it here.
                  test_error_count = 0.0
                  for images, labels in iter(test_loader):
                      images = images.to(self.device)
                      labels = labels.to(self.device)
                      outputs = self.model(images)
                      test_error_count += float(torch.sum(torch.abs(labels - outputs.argmax(1))))

                  test_accuracy = 1.0 - float(test_error_count) / float(len(test_dataset))
                  print('%d: %f' % (epoch, test_accuracy))
                  if test_accuracy > best_accuracy:
                      print("bm, bmp", best_model, best_model_path)
                      best_accuracy = test_accuracy
              torch.save(self.model.state_dict(), best_model)
              dataset.save_images_processed("FUNC", self.nn_name)

    # wipe_memory(self)
    def cleanup(self):
        del self.model
        if self.optimizer is not None:
          self._optimizer_to(torch.device('cpu'))
          del self.optimizer
          gc.collect()
        torch.cuda.empty_cache()
        # you can check that memory was released using nvidia-smi

    def _optimizer_to(self, device):
      for param in self.optimizer.state.values():
        # Not sure there are any global tensors in the state dict
        if isinstance(param, torch.Tensor):
            param.data = param.data.to(self.device)
            if param._grad is not None:
                param._grad.data = param._grad.data.to(self.device)
        elif isinstance(param, dict):
            for subparam in param.values():
                if isinstance(subparam, torch.Tensor):
                    subparam.data = subparam.data.to(self.device)
                    if subparam._grad is not None:
                        subparam._grad.data = subparam._grad.data.to(device)