"""
Motif discovery, incorporating position sensitivity and motif degeneracy.
The basic model is similar to the original BIRT:  for a given (degenerate) motif,
find the window of positions that maximize the significance by the binomial test.

Here, the degeneracy is modeled as one canonical k-mer, and some or all of the 3k+9
"neighboring" k-mers that differ from it by just one base.
The intuition is that each substitution carries an increasing cost;  being one mutation
from consensus isn't as bad as being 3 or 4 mutations away from the consensus sequence.

This machinery is derived from my original "Vermit" implementation of the cERMIT algorithm,
which focused on log-odds "evidence" for each gene, with the in-group/out-group version
as a special case.

However, in this code I'm using binomial distribution p-values, which come only from in-group/
out-group count data (evidence = 0 or 1).  So I kept the machinery, but the nominclature is
not what I would have done if I started with this approach!
"""
import copy
import numpy as np
from scipy import weave

from pygr import util, dnaseq, rmath

def kmer_matches(kmer, suf, seq_offsets, midgap):
    """
    locs    0-based start position of matches to `kmer` within the concatenated sequences.
    matches 0-based index of corresponding sequence for each entry in `locs`.
    midgap  Number of bases between the first and second half of `kmer`
    """
    # Caching speeds the single-processor algorithm by ~7%,
    # but eliminating it frees up a LOT of memory to allow multithreading.
    if midgap == 0:
        locs = suf.find_positions(kmer) # offsets from beginning of file
        matches = np.searchsorted(seq_offsets, locs) # sequence index for each match
    else:
        half = len(kmer) // 2
        locs_left = suf.find_positions(kmer[:half])
        locs_right = suf.find_positions(kmer[half:]) - (half+midgap)
        locs = np.intersect1d(locs_left, locs_right, assume_unique=True)
        matches = np.searchsorted(seq_offsets, locs)
        matches_right = np.searchsorted(seq_offsets, locs + (half+midgap))
        same_seq = (matches == matches_right) # only valid when two halves of the match are in the same sequence!
        locs = locs[same_seq]
        matches = matches[same_seq]
    return (locs, matches)

def kmers_matches(kmers, suf, seq_offsets, midgap):
    # Gather together all matches/locations for all kmers.
    # Sort them, first by sequence ID (implicitly), then by position within the sequence.
    matches = [] # [np.array], will be joined all at once
    locs = [] # [np.array], will be joined all at once
    for kmer in kmers:
        ls, ms = kmer_matches(kmer, suf, seq_offsets, midgap)
        matches.append(ms)
        locs.append(ls)
    matches = np.concatenate(matches)
    locs = np.concatenate(locs)
    # There cannot be any duplicates in locs (assuming all kmers are unique, which they should be)
    # because there is exactly one distinct kmer at each location in the original sequences.
    o = np.argsort(locs)
    return locs[o], matches[o]

def filelocs_to_seqlocs(locs, matches, seq_lens, seq_offsets, align_5p):
    """
    Converts file-based kmer locations to sequence-based kmer locations.
    `locs` is altered *IN PLACE*, so there is no return value.
    """
    max_seq_len = seq_lens.max() - 1 # b/c they're really seq. len. + 1 for the newline
    # seq_offsets[matches] - locs == distance from right edge, <= 0
    locs -= seq_offsets[matches]
    if align_5p:
        # sequences are left-justified and 0 is the leftmost base:
        locs += seq_lens[matches]
        locs -= 1 # b/c they're really seq. len. + 1 for the newline
    else:
        # sequences are right-justified and 0 is the leftmost base:
        #locs = max_seq_len - (seq_offsets[matches] - locs)
        locs += max_seq_len
    #if not (locs.max() < max_seq_len):
    #    import pdb
    #    pdb.set_trace()
    assert locs.max() < max_seq_len
    assert locs.min() >= 0
    return max_seq_len

class SeqBins(object):
    def __init__(self, fasta, evidence, ev_wts, n_bins=1):
        self.n_bins = n_bins
        if self.n_bins <= 1:
            self.bin_ids = np.array([0]*len(evidence))
        else:
            self.bin_ids = self._cluster(fasta, n_bins)
        # The genes are divided into n bins. Let Bi and Ti be the BG and target set genes, respectively,
        # in the ith bin, and denote by bi the subset of genes from Bi whose sequence contains a hit. The
        # goal of this score is to account for cases where the fraction of targets is uneven across bins.
        # Suppose that targets within each bin are selected uniformly. Then, in bin i the probability that
        # a selected gene will contain a hit (i.e., belong to bi) is |bi|/|Bi|. Since the fraction of
        # targets in bin i is |Ti|/|T|, it follows that the probability that a selected gene will contain
        # a hit is p_m = sum_over_i[ (Ti/T) * (bi/Bi) ]
        Bi = np.array([ev_wts[self.bin_ids == cid].sum() for cid in xrange(self.n_bins)]) # total genes in each bin
        Ti = np.array([evidence[self.bin_ids == cid].sum() for cid in xrange(self.n_bins)]) # in-group genes in each bin
        T = evidence.sum() # total in-group genes across all bins
        Bi[Bi == 0] == 1 # to avoid divide-by-zero errors when Ti == Bi == 0 (no genes in bin)
        self.multiplier = (Ti / T / Bi)[None,None,:] # to be multiplied by bi, which is called wt_table in improve_window
    def _get_features(self, seq):
        features = []
        features.append(len(seq))
        bases = util.make_bag(seq)
        L = float(max(1, len(seq)))
        features.extend([bases[b]/L for b in 'ACGT'])
        return features
    def _cluster(self, fasta, n_bins):
        print "Clustering sequences by composition and length..."
        # Compute numeric features for each sequence.
        features = [self._get_features(s) for i,n,s in fasta]
        # Normalize distances by dividing by std. deviation.
        features = np.array(features, np.float)
        stddev = np.apply_along_axis(np.std, 0, features)
        print "S.d. of features:", stddev
        features /= stddev
        # For now, we just run the clustering once, even though it may be sub-optimal
        clusters, cost = util.kmeans(features, n_bins, maxiter=10)
        return np.asarray(clusters)

def improve_window(kmers, bins, evidence, ev_wts, G, mu, suf, seq_lens, seq_offsets, options, mask=None):
    """
    Evaluate the score for a set of k-mers, restricted to various possible sequence windows.
    Choose the window giving the highest Z score for these k-mers.

    kmers       Iterable of 1+ k-mers to find the optimal cERMIT-score window for
    bins        a SeqBins object defining how sequences are grouped into bins/clusters
    evidence    NumPy array of log-odds or log-enrichment scores for each gene in the pool (floats)
    ev_wts      NumPy array of weights on (0,1] for each piece of evidence, 1/(# gene models)
    G           total "number" of genes, ev_wts.sum() [to avoid recomputing every time]
    mu          average evidence, evidence.sum()/G [to avoid recomputing every time]
    suf         the suffix array searcher object
    seq_lens    length of each sequence in characters, including the newline separator
    seq_offsets position in the concatenated sequence file at which each sequence ends
    options     command line options object
    mask        Numpy array with length >= max location index; counts the number of times each bp is "used" in a motif.
                This is optional, for masking out kmers used in one motif so another motif can't claim them.
                If present, it will be updated to mask out the evidence used for these kmers.

    Raw evidence numbers should have been pre-multiplied by the provided weights (for efficiency).
    Scoring follows M.A.Newton et al, Annals of Applied Statistics 2007, 1:85-106
    rather than Georgiev et al. because the latter appears to have some typos in the formulae.
    Higher scores are better.
    """
    # Hard to find a good word for the SeqBins, where sequences have been grouped/clustered/binned
    # by sequence composition (percent A/C/G/T).  In this function "bin" is already used for
    # the sequence position windows (typically of ~25bp length).
    # Elsewhere in POWRS, there is some post-process "clustering" that already uses that name.
    # But within this function, I'll call the SeqBins "clusters" to distinguish them from
    # the positional bins.
    cluster_ids = bins.bin_ids # for access by Weave
    n_clust = bins.n_bins
    # Gather together all matches/locations for all kmers.
    # Sort them, first by sequence ID (implicitly), then by position within the sequence.
    # The algorithm below expects that locations occur in sorted order.
    #T = time.time()
    locs, matches = kmers_matches(kmers, suf, seq_offsets, options.midgap)
    # "Masking" to reduce the score of alternative motifs that closely match a previously evaluated motif.
    if mask is not None:
        motif_len = max(len(kmer) for kmer in kmers) # should all be the same length, but just in case...
        assert len(mask) >= locs[-1] + motif_len # locs are sorted, so -1 is max element
        loc_mask = np.zeros(locs.shape, dtype=np.int) # each pos. varies from 0 (totally free) to motif_len (totally masked)
        mask_limit = options.mask_limit
        weave.inline(r"""
            int n = Nlocs[0]; // see Numpy book; == locs.shape in Python
            int klen = motif_len;
            int maskmax = mask_limit;
            int ii, jj, kk, loc;
            for(ii = 0; ii < n; ii++) { // for each location...
                loc = locs[ii];
                for(jj = 0; jj < klen; jj++) {
                    if(mask[loc+jj] >= maskmax) loc_mask[ii]++; // if it's already maxed out, we can't use it (maybe)
                    else mask[loc+jj]++; // otherwise, note that we've used one count worth
                }
            }
        """, ['locs','mask','motif_len','loc_mask','mask_limit'])
        loc_mask = (loc_mask <= motif_len//2) # mask out locations that are more than half used by someone else
        locs = locs[loc_mask]
        matches = matches[loc_mask]
    # locs/matches could be empty, possibly due to masking
    if len(locs) == 0: return 0, 0, -util.Inf, 0
    # This sort of complicated indexing creates copies (unlike simple slices):
    # We need a copy (of locs, at least) because we're going to modify it in place, and it's cached. [not true anymore]
    assert locs.base == matches.base == None
    # convert file-based offsets into sequence-based offsets
    max_seq_len = filelocs_to_seqlocs(locs, matches, seq_lens, seq_offsets, options.align_5p) # locs altered in place!
    #print "Finding %i matches: %f" % (len(locs), time.time()-T); T = time.time()
    # Tally kmer locations in fixed-width bins, then build a table
    # of start,end windows showing the total (weighted) evidence
    # of genes with 1+ copies of the kmer located within that window.
    W = options.window_width
    max_W = max(1, max_seq_len // W) # in case W > max_seq_len, so we don't get a 0x0 table
    ev_table = np.zeros((max_W,max_W), dtype=np.float) # sum evidence for selected genes
    wt_table = np.zeros((max_W,max_W,n_clust), dtype=np.float) # "number" of selected genes, a.k.a. "m"
    locs //= W # convert base-pair locations into binned locations
    # now we'll subtract out evidence for genes that don't have a kmer in that window:
        ## This loop is by far the slowest part of this algorithm, so...
        #last_loc, last_seqid = None, None
        #for locbin, seqid in zip(locs, matches):
        #    if seqid != last_seqid: # first kmer of a new sequence
        #        ev_table += evidence[seqid]
        #        wt_table += ev_wts[seqid]
        #        if last_seqid is not None: # last kmer of the old sequence
        #            ev_table[last_loc:max_W,last_loc:max_W] -= evidence[last_seqid]
        #            wt_table[last_loc:max_W,last_loc:max_W] -= ev_wts[last_seqid]
        #        ev_table[0:locbin,0:locbin] -= evidence[seqid]
        #        wt_table[0:locbin,0:locbin] -= ev_wts[seqid]
        #        last_seqid = seqid
        #    else: # non-first kmer
        #        ev_table[last_loc:locbin,last_loc:locbin] -= evidence[seqid]
        #        wt_table[last_loc:locbin,last_loc:locbin] -= ev_wts[seqid]
        #    # We add one so that only bins *between* kmers will be penalized,
        #    # not bins *containing* kmers. In NumPy if start > end, it's just an empty array.
        #    last_loc = locbin + 1
        #if last_seqid is not None: # last kmer of the last sequence
        #    ev_table[last_loc:max_W,last_loc:max_W] -= evidence[last_seqid]
        #    wt_table[last_loc:max_W,last_loc:max_W] -= ev_wts[last_seqid]
    ## After re-writing in C++, the time to actually find the k-mers dominates!
    weave.inline(r"""
        int maxw = max_W; // need this to avoid int/float ambiguity error
        int nclust = n_clust;
        int last_loc = -1, last_seqid = -1, last_clust = -1;
        int n = Nlocs[0]; // see Numpy book; == locs.shape in Python
        int locbin, seqid, clust;
        float evid, evwt, last_evid, last_evwt;
        int ii, jj, kk;
        for(ii = 0; ii < n; ii++) {
            locbin = locs[ii];
            seqid = matches[ii];
            clust = cluster_ids[seqid];
            evid = evidence[seqid];
            evwt = ev_wts[seqid];
            if(seqid != last_seqid) {
                // Only need to do the diagonal and above...
                // Finish up last sequence:
                if(last_seqid != -1) {
                    for(jj = last_loc; jj < maxw; jj++) { for(kk = jj; kk < maxw; kk++) {
                        ev_table[jj*maxw + kk] -= last_evid;
                        wt_table[(jj*maxw + kk)*nclust + last_clust] -= last_evwt;
                    }}
                }
                // Now start this new sequence:
                // Add in evidence everywhere (happens ONCE per SEQUENCE)
                for(jj = 0; jj < maxw; jj++) { for(kk = jj; kk < maxw; kk++) {
                    ev_table[jj*maxw + kk] += evid;
                    wt_table[(jj*maxw + kk)*nclust + clust] += evwt;
                }}
                // Remove evidence between left edge and this occurrence
                for(jj = 0; jj < locbin; jj++) { for(kk = jj; kk < locbin; kk++) {
                    ev_table[jj*maxw + kk] -= evid;
                    wt_table[(jj*maxw + kk)*nclust + clust] -= evwt;
                }}
                last_seqid = seqid;
            } else {
                // Continuation of previous sequence;  evidence has already been added.
                // Just remove evidence between last occurrence and this occurrence.
                for(jj = last_loc; jj < locbin; jj++) { for(kk = jj; kk < locbin; kk++) {
                    ev_table[jj*maxw + kk] -= evid;
                    wt_table[(jj*maxw + kk)*nclust + clust] -= evwt;
                }}
            }
            last_loc = locbin + 1;
            last_clust = clust;
            last_evid = evid;
            last_evwt = evwt;
        }
        if(last_seqid != -1) {
            clust = cluster_ids[last_seqid];
            evid = evidence[last_seqid];
            evwt = ev_wts[last_seqid];
            for(jj = last_loc; jj < maxw; jj++) { for(kk = jj; kk < maxw; kk++) {
                ev_table[jj*maxw + kk] -= evid;
                wt_table[(jj*maxw + kk)*nclust + clust] -= evwt;
            }}
        }
    """, ['locs', 'matches', 'cluster_ids', 'evidence', 'ev_wts', 'max_W', 'n_clust', 'ev_table', 'wt_table'])
    #print "For loop over %i matches: %f" % (len(locs), time.time()-T); T = time.time()
    # Instead of the cERMIT Z-score, we can compute the binomial p-value (using in-group/out-group evidence only):
    # Trying to use non-integer valued floats in pbinom() produces NaNs
    score_table = np.around(ev_table)
    prob_table = (wt_table * bins.multiplier).sum(axis=2) # probability weighted by bin membership
    wt_table = wt_table.sum(axis=2) # condense clusters, for normal use by later code
    wt_table = np.around(wt_table)
    # This code used to compute pbinom(Ak, Nk, A/N)
    # In Amadeus and historically (tally_all), we use pbinom(Ak, A, Nk/N) instead.
    # It would be easy to use either one here, really, without changing much (see commented-out lines below).
    # P-values are quite similar, although pbinom(Ak, A, Nk/N) seems to be a little stronger in each case.
    # I'm not sure which one is better justified statistically, but top hits are nearly identical empirically.
    # We use the Amadeus version now so that we can use their binned enrichment correction.
    weave.inline(r"""
        int maxw = max_W; // need this to avoid int/float ambiguity error
        int ii, jj, kk;
        double prob = mu;
        double N = G;
        double A = (int)(prob * N);
        for(jj = 0; jj < maxw; jj++) { for(kk = jj; kk < maxw; kk++) {
            ii = jj*maxw + kk; // linear index into score_table / wt_table
            // pbinom(Ak, Nk, A/N):
            //score_table[ii] = -0.4342945 * pbinom(score_table[ii]-1, wt_table[ii], prob, 0, 1); // -log_10(P): bigger is better
            // pbinom(Ak, A, Nk/N):
            score_table[ii] = -0.4342945 * pbinom(score_table[ii]-1, A, prob_table[ii], 0, 1); // -log_10(P): bigger is better
        }}
    """, ['mu', 'max_W', 'score_table', 'prob_table', 'G'],
    **rmath.weave_inline_kwargs)
    # Disqualify bins with too few / too many hits:
    min_genes = options.min_genes
    max_genes = options.max_genes_frac * G
    score_table[wt_table < min_genes] = -util.Inf
    score_table[wt_table > max_genes] = -util.Inf
    # Almost there. Find the highest scoring, longest, rightmost window in which start <= end.
    # (Start and end are both inclusive when we read them out from the evidence table).
    olderr = np.seterr(all='ignore') # I think triu() is implemented with a multiplication...
    score_table = np.triu(score_table, k=0) # zero out everything below the diagonal (ensures start <= end) (except NaNs!)
    np.seterr(**olderr)
    max_Z = np.nanmax(score_table)
    # This was caused originally by non-integer weightings on genes -- leaving it here just in case.
    if np.isnan(max_Z):
        print "Failed with %i locs for %i kmers with -- %s" % (len(locs), len(kmers), kmers)
        return 0, 0, -util.Inf, 0
    peaks = [(end-start, end, start) for start, end in zip(*np.nonzero(score_table == max_Z))]
    peaks.sort(reverse=True)
    first_bin, last_bin = peaks[0][2], peaks[0][1]
    #score_table[~np.isfinite(score_table)] = 0 # to make plotting nicer
    #print "Post processing: %f" % (time.time()-T); T = time.time()
    # Return [inclusive, exclusive) base pair indices starting from 0 == left edge
    return first_bin*W, (last_bin+1)*W, max_Z, ev_table[first_bin,last_bin]

class Motif(object):
    def __init__(self, kmer, both_strands, data, parent=None):
        """
        `data` is a dictionary containing the following keys:
            evidence    NumPy array of 1's and 0's marking in-group and out-group, respectively (floats)
            ev_wts      NumPy array of weights on (0,1] for each piece of evidence, 1/(# gene models)
            G           total "number" of genes, ev_wts.sum() [to avoid recomputing every time]
            mu          average evidence, evidence.sum()/G [to avoid recomputing every time]
            suf         the suffix array searcher object
            seq_lens    length of each sequence in characters, including the newline separator
            seq_offsets position in the concatenated sequence file at which each sequence ends
            options     command line options object
        """
        if parent:
            self.center = parent.center # the central kmer
            self.friends = parent.friends | set([kmer]) # other included neighboring kmers
            self.other_nbrs = parent.other_nbrs - set([kmer]) # neighboring kmers not yet included
        else:
            self.center = kmer
            self.friends = set()
            self.other_nbrs = generate_neighbors(kmer)
        self.both_strands = both_strands
        # Don't want to save whole data object or pickled Motifs will be huge...
        self.options_length = data['options'].length
        self._pattern = None
        # Scoring
        self.all_kmers = set([self.center]) | self.friends
        if self.both_strands: self.all_kmers |= set(dnaseq.reverse_complement(m) for m in self.all_kmers)
        self.start, self.end, self.score, self.evidence = improve_window(kmers=self.all_kmers, **data)
        # Using center_score means that the preferred strand version will sort to the top at the end!
        if parent: self.center_score = parent.center_score
        else: self.center_score = self.score
        # When sorting, prefer both-strands version when scores are equal (e.g. palindromic k-mers)
        self.sortkey = (self.score, self.both_strands, self.center_score, self.end, self.start, self.pattern)
    @property
    def start_user(self): return self.start - self.options_length
    @property
    def end_user(self): return self.end - self.options_length
    def __hash__(self): # Motifs are uniquely identified by their string pattern, but may get different scores due to re-scoring
        return hash(self.sortkey)
    def __cmp__(self, other):
        return cmp(self.sortkey, other.sortkey)
    def __str__(self):
        return self.as_str()
    def as_str(self, shift=0, max_shift=0):
        s1 = " "*shift
        s2 = " "*(max_shift - shift)
        if self.both_strands: rc_seed = dnaseq.reverse_complement(self.center)
        else: rc_seed = "-" * len(self.center)
        return "%8.2f    [%6.1f]    %s%s%s / %s%s%s    %6i    %6i    %50s" % (
            self.score, self.evidence,
            s1, self.center, s2, s2, rc_seed, s1, self.start_user, self.end_user, self.pattern)
    @property
    def revcomp(self):
        """Returns the reverse complement of this motif (or None if not on both strands)."""
        if not self.both_strands: return None
        elif self.center == dnaseq.reverse_complement(self.center): return self # palindrome
        elif not hasattr(self, "_revcomp"):
            rc = self._revcomp = copy.copy(self)
            rc.center = dnaseq.reverse_complement(self.center)
            rc.friends = set(dnaseq.reverse_complement(f) for f in self.friends)
            rc.other_nbrs = set(dnaseq.reverse_complement(o) for o in self.other_nbrs)
            rc._pattern = None
            rc.sortkey = (rc.score, rc.both_strands, rc.center_score, rc.end, rc.start, rc.pattern)
        return self._revcomp
    @property
    def pattern(self):
        if self._pattern is None:
            self._pattern = ""
            for pos, base in enumerate(self.center):
                alts = set(f[pos] for f in self.friends)
                alts.discard(base)
                if not alts:
                    self._pattern += base
                else:
                    self._pattern += "[" + base + "".join(sorted(a.lower() for a in alts)) + "]"
        return self._pattern

def generate_neighbors(kmer):
    nbrs = set()
    for pos, base in enumerate(kmer):
        for alt in "ACGT":
            if alt != base:
                assert base != "N"
                nbrs.add(kmer[:pos] + alt + kmer[pos+1:])
    return nbrs

class Cluster(object):
    """A group of similar Motif objects."""
    def __init__(self, motif, data):
        self.motifs = [motif]
        self.lost_motifs = [] # used only in reverse complement operations...
        self.offsets = [0]
        self.all_kmers = motif.all_kmers
        self.start = motif.start
        self.end = motif.end
        self.score = motif.score
        self.evidence = motif.evidence
        # Need to keep `data` for merge scoring and revcomp scoring,
        # but it could make pickles holding Clusters quite large...
        self.data = data
    @property
    def start_user(self): return self.start - self.data['options'].length
    @property
    def end_user(self): return self.end - self.data['options'].length
    def try_merge(self, other):
        """Returns either a new Cluster or None."""
        if not overlaps(self, other): return None
        kmers_linked = False
        # First try exact alignment
        for m1, o1 in zip(self.motifs, self.offsets):
            if kmers_linked: break
            for m2, o2 in zip(other.motifs, other.offsets):
                # `or` is too permissive; `and` is better
                if m1.center in m2.friends and m2.center in m1.friends:
                    kmers_linked = True
                    offset = (o1 - o2)
                    break
        # Second try shifted left and shifted right
        # Most conservative version:  seeds overlap by N-1 bases
        # Overlapping seeds with variants turns out to be too permissive.
        for m1, o1 in zip(self.motifs, self.offsets):
            if kmers_linked: break
            for m2, o2 in zip(other.motifs, other.offsets):
                if m1.center[1:] == m2.center[:-1]:
                    kmers_linked = True
                    offset = (o1 - o2) + 1
                    break
                elif m2.center[1:] == m1.center[:-1]:
                    kmers_linked = True
                    offset = (o1 - o2) - 1
                    break
        if not kmers_linked: return None
        new = copy.copy(self)
        new.motifs = self.motifs + other.motifs
        new.offsets = self.offsets + [o+offset for o in other.offsets]
        offset = min(new.offsets)
        new.offsets = [o-offset for o in new.offsets] # set min offset to 0
        new.all_kmers = self.all_kmers | other.all_kmers
        new.start, new.end, new.score, new.evidence = improve_window(kmers=new.all_kmers, **new.data)
        if new.score > max(self.score, other.score) and overlaps(self, new) and overlaps(other, new):
            return new
        else: return None
    @property
    def revcomp(self):
        if not hasattr(self, "_revcomp"):
            # Motifs that are single stranded will return None.
            # If all our motifs are single stranded, we'll return None also.
            rc = self._revcomp = copy.copy(self)
            rc.motifs = []
            rc.lost_motifs = [] # single stranded motifs that can't be used in the revcomp
            rc.offsets = []
            rc.all_kmers = set()
            for motif, offset in zip(self.motifs, self.offsets):
                if motif.revcomp:
                    rc.motifs.append(motif.revcomp)
                    rc.offsets.append(offset) # will be reversed shortly
                    rc.all_kmers.update(motif.revcomp.all_kmers)
                else:
                    rc.lost_motifs.append(motif)
            if rc.motifs:
                max_off = max(rc.offsets)
                rc.offsets = [max_off - o for o in rc.offsets] # have to reverse direction of offsets
                rc.start, rc.end, rc.score, rc.evidence = improve_window(kmers=rc.all_kmers, **rc.data)
                rc._revcomp = self
            else:
                self._revcomp = None
        return self._revcomp
    def __cmp__(self, other):
        return cmp(self.score, other.score)
    def __len__(self):
        return len(self.motifs)
    def __str__(self):
        return "<%.2f, %i to %i, %i motifs: %s ...>" % (self.score, self.start_user, self.end_user, len(self.motifs), self.motifs[0].center)

def overlaps(mc1, mc2):
    """Compare two motifs and/or clusters to see if their location ranges overlap."""
    return (mc1.start < mc2.end) and (mc1.end > mc2.start)

def extract_positive_seqs(cluster, evidence, suf, seq_lens, seq_offsets, align_5p, midgap):
    """
    For the "positive" (in-group) sequences, extract all the unique regions
    that contain 1+ of the kmers in cluster.
    We do this to avoid the multiple-counting problem of tallying kmers individually.
    """
    regions = set() # set([(start,end,revcomp)]) within suf.corpus (the concatenated sequences)
    max_offset = max(cluster.offsets)
    kmer_len = len(cluster.motifs[0].center) + midgap # k-mer length, including "gap" residues in the middle
    region_len = max_offset + kmer_len
    for motif, offset in zip(cluster.motifs, cluster.offsets):
        # First, the positive strand:
        fwd_kmers = set([motif.center]) | motif.friends
        locs, matches = kmers_matches(fwd_kmers, suf, seq_offsets, midgap)
        seq_locs = locs.copy()
        max_seq_len = filelocs_to_seqlocs(seq_locs, matches, seq_lens, seq_offsets, align_5p)
        # matches within the positive set of sequences and the allowed region
        keep = (evidence[matches] > 0) & (cluster.start <= seq_locs) & (seq_locs < cluster.end)
        starts = locs[keep] - offset
        regions.update((start, start+region_len, False) for start in starts)
        # Then, the reverse strand
        rev_kmers = motif.all_kmers - fwd_kmers # avoid double-counting palindromes
        if rev_kmers:
            locs, matches = kmers_matches(rev_kmers, suf, seq_offsets, midgap)
            seq_locs = locs.copy()
            max_seq_len = filelocs_to_seqlocs(seq_locs, matches, seq_lens, seq_offsets, align_5p)
            # matches within the positive set of sequences and the allowed region
            keep = (evidence[matches] > 0) & (cluster.start <= seq_locs) & (seq_locs < cluster.end)
            starts = locs[keep] - (max_offset - offset)
            regions.update((start, start+region_len, True) for start in starts)
    seqs = []
    for start, end, revcomp in regions:
        subseq = suf.corpus[start:end]
        # Because of the padding at the start and/or end due to max_offset,
        # subseq may inadvertently span two different genomic sequences!
        nl = subseq.find("\n")
        if nl >= 0:
            if nl < kmer_len: subseq = subseq[nl+1:].rjust(region_len, "N") # left side shorter than k
            elif nl >= max_offset: subseq = subseq[:nl].ljust(region_len, "N") # right side shorter than k
            else: continue # could have been either side, too much trouble to figure out, skip it
        if revcomp: seqs.append(dnaseq.reverse_complement(subseq))
        else: seqs.append(subseq)
    return seqs


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
@author: shahab Sotudian
"""

import numpy as np
import pandas as pd
import math
import six
from six.moves import range
from scipy.stats import rankdata
import sys
import time
from sklearn import preprocessing
from scipy.stats import spearmanr
import random
start_time = time.time()

# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
# @@@@                         Functions
# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@

Which_Type='ANTIBODY' # 'ANTIBODY'  'ENZYME'   'OTHERS'
Num_Features= 30   # Total number of features

if Which_Type == 'ANTIBODY':
    Test_Proteins= ['3rvw','4dn4','4g6j','3hi6','3l5w','2vxt','2w9e','4g6m','3g6d','3eo1','3v6z','3hmx','3eoa','3mxw']

    Train_Proteins=['1mlc','1iqd','2jel','1nca','1ahw','1e6j','1kxq','1wej','1dqj','2fd6','2i25','1jps','2hmi','1k4c',
                    '1i9r','1bj1','1bgx','1qfw','2vis','1nsn','1bvk','1fsk','1vfb']

    Features_Ordered= ['mincomp_fa_rep','mincomp_score','mincomp_total_score','piper_fa_rep','piper_score','piper_total_score',
                       'movedcomp_total_score','movedcomp_score','var_mem_Elec','var_mem_Born','kurt_mem_Avdw','SPPS','skew_mem_Rvdw',
                       'kurt_mem_Born','avg_mem_Elec','mincomp_fa_dun','PREMIN_COMP_Torsions','movedcomp_fa_dun','POSTMIN_COMP_Torsions',
                       'PREMIN_SING_Torsions','POSTMIN_SING_Torsions','mincomp_rama_prepro','cen_Elec','movedcomp_rama_prepro','size',
                       'piper_fa_dun','mincomp_fa_atr','cen_Born','avg_mem_Born','piper_rama_prepro','mincomp_fa_sol',
                       'piper_fa_intra_sol_xover4','mincomp_omega','POSTMIN_SING_Impropers','piper_yhh_planarity','movedcomp_omega',
                       'movedcomp_fa_intra_rep','mincomp_pro_close','var_mem_Avdw','POSTMIN_SING_Bonded','piper_omega',
                       'movedcomp_pro_close','PREMIN_COMP_Impropers','mincomp_fa_intra_rep','PREMIN_SING_Impropers','POSTMIN_COMP_Impropers',
                       'PREMIN_COMP_Bonded','PREMIN_SING_Bonded','POSTMIN_COMP_Bonded','PREMIN_COMP_Angles','piper_pro_close',
                       'POSTMIN_SING_Angles','movedcomp_fa_intra_sol_xover4','PREMIN_SING_Angles','POSTMIN_COMP_Angles',
                       'mincomp_fa_intra_sol_xover4','skew_mem_Born','avg_mem_dist','movedcomp_yhh_planarity','cen_Avdw',
                       'mincomp_yhh_planarity','avg_mem_Avdw','var_mem_dist','movedcomp_ref','mincomp_ref','piper_fa_atr',
                       'movedcomp_fa_atr','var_mem_DARS','movedcomp_fa_sol','avg_mem_DARS','piper_fa_sol','var_mem_Rvdw','cen_DARS',
                       'piper_dslf_fa13','piper_ref','SPAR','var_mem_Teng','movedcomp_dslf_fa13','cen_Rvdw','movedcomp_hbond_lr_bb',
                       'mincomp_hbond_lr_bb','movedcomp_hbond_bb_sc','piper_hbond_bb_sc','movedcomp_hbond_sr_bb','mincomp_hbond_sr_bb',
                       'avg_mem_Rvdw','mincomp_hbond_bb_sc','piper_hbond_sr_bb','movedcomp_hbond_sc','mincomp_fa_elec','piper_hbond_sc',
                       'movedcomp_fa_elec','SPAS','mincomp_hbond_sc','cen_Teng','piper_fa_elec','avg_mem_Teng','piper_hbond_lr_bb','piper_time']

    DATA_ALL = pd.read_csv('/home/Final_DATA_Antibody.csv')

elif Which_Type == 'ENZYME':
    Test_Proteins= ['2gaf','4hx3','3a4s','2yvj','4fza','3fn1','3pc8','3lvk','3k75','4iz7','3vlb','4h03','3h11','1jtd','4lw4','2a1a']

    Train_Proteins=['3sgq','1fq1','2o3b','1d6r','2pcc','2mta','1kkl','1oc0','1mah','1m10','1jzd','1cgi','1ijk','1f51','2ayo','7cei',
                    '1jiw','2uuy','2oul','1us7','1jk9','1tmq','2sni','1avx','1gxd','2oob','1hia','1dfj','2nz8','2j0t','2z0e','1nw9',
                    '1ay7','1jtg','2oor','2ot3','1jwh','1bvn','1oph','1ezu','1e6e','1acb','1f6m','2abz','1f34','1jmo','1gl1','1udi',
                    '1r6q','1zm4','1eaw','1ewy','2b42','1buh','1pxv','2sic','2a9k','4cpa','1z5y','1gla','2o8v','1zli','1r0r','1fle',
                    '1oyv','1ppe','1wdw','1yvb','1clv','2ido']

    Features_Ordered= ['movedcomp_fa_rep','movedcomp_score','movedcomp_total_score','piper_lk_ball_wtd','POSTMIN_COMP_VWD03',
                       'POSTMIN_SING_VWD03','mincomp_lk_ball_wtd','var_mem_Born','mincomp_hbond_lr_bb','piper_pro_close',
                       'movedcomp_hbond_lr_bb','movedcomp_lk_ball_wtd','movedcomp_pro_close','piper_hbond_lr_bb','var_mem_Elec',
                       'POSTMIN_COMP_Angles','POSTMIN_COMP_Impropers','piper_dslf_fa13','movedcomp_dslf_fa13','mincomp_pro_close',
                       'POSTMIN_SING_Bonded','movedcomp_omega','piper_fa_rep','size','PREMIN_SING_Bonded','PREMIN_COMP_Bonded',
                       'POSTMIN_SING_Impropers','POSTMIN_COMP_Bonded','mincomp_fa_rep','mincomp_hbond_sr_bb','piper_hbond_sr_bb',
                       'piper_total_score','piper_score','piper_omega','POSTMIN_SING_Angles','avg_mem_DARS','cen_DARS',
                       'movedcomp_hbond_sr_bb','PREMIN_SING_Impropers','PREMIN_COMP_Impropers','PREMIN_SING_Angles','mincomp_omega',
                       'PREMIN_COMP_Angles','movedcomp_fa_dun','SPPS','mincomp_fa_dun','mincomp_total_score','mincomp_score',
                       'mincomp_fa_intra_rep','movedcomp_yhh_planarity','piper_rama_prepro','mincomp_rama_prepro',
                       'movedcomp_rama_prepro','piper_yhh_planarity','kurt_mem_dist','movedcomp_hbond_sc','cen_Elec',
                       'mincomp_yhh_planarity','mincomp_hbond_sc','mincomp_dslf_fa13','var_mem_dist','movedcomp_fa_intra_rep',
                       'movedcomp_hbond_bb_sc','mincomp_hbond_bb_sc','kurt_mem_Born','avg_mem_dist','avg_mem_Elec','PREMIN_SING_VWD03',
                       'PREMIN_COMP_VWD03','SPAR','PREMIN_COMP_Torsions','PREMIN_SING_Torsions','piper_fa_dun','mincomp_p_aa_pp',
                       'movedcomp_p_aa_pp','POSTMIN_COMP_Torsions','mincomp_fa_elec','piper_hbond_bb_sc','POSTMIN_SING_Torsions',
                       'var_mem_DARS','movedcomp_fa_sol','movedcomp_fa_elec','piper_fa_atr','piper_fa_elec','mincomp_fa_atr',
                       'piper_fa_sol','mincomp_fa_sol','movedcomp_fa_atr','SPAS','piper_hbond_sc','piper_fa_intra_sol_xover4',
                       'movedcomp_time','movedcomp_fa_intra_sol_xover4','mincomp_fa_intra_sol_xover4','piper_fa_intra_rep',
                       'piper_ref','movedcomp_ref','mincomp_ref','skew_mem_Born','kurt_mem_Teng','skew_mem_Teng','skew_mem_DARS',
                       'skew_mem_dist','cen_Teng','mincomp_time','cen_Rvdw','var_mem_Teng','var_mem_Rvdw','avg_mem_Avdw','avg_mem_Rvdw']

    DATA_ALL = pd.read_csv('/home//Final_DATA_Enzyme.csv')

elif Which_Type == 'OTHERS':
    Test_Proteins=['3szk','1m27','2x9a','baad','3aad','3s9d','3daw','cp57','3h2v','3bx7',
                   '3biw','3p57','4m76','3aaa','3f1p','bp57','2gtp','1rke','3r9a','3l89',
                   '1exb','4jcv']

    Train_Proteins=['1rv6','1eer','1e4k','2fju','1z0k','1ofu','1a2k','3cph','1ibr','1ib1',
                    '1k5d','1k74','1syx','3d5s','1i4d','1rlb','1he8','1r8s','1pvh','1fcc',
                    '3bp8','1azs','2hqs','1gpw','2c0l','2i9b','2hle','1de4','1qa9','1fqj',
                    '1h1v','1e96','1j2j','1n2c','1t6b','2hrk','1kac','1y64','2b4j','1ghq',
                    '2g77','1ml0','1gcq','1efn','1ffw','1bkd','1zhi','1klu','1xd3','2j7p',
                    '2cfh','1hcf','2vdb','1atn','1grn','1lfd','1b6c','1ak4','1wq1','1ira',
                    '2btf','1sbb','1xqs','1mq8','1xu1','1i2m','1he1','1s1q','2oza','1zhh',
                    '1ktz','1fak','1fc2','1kxp','1gp2','2ajf','1akj','1h9d','2a5t']

    Features_Ordered=['mincomp_score','mincomp_total_score','mincomp_fa_rep','piper_total_score',
                      'piper_score','piper_fa_rep','cen_Avdw','avg_mem_Avdw','var_mem_Teng',
                      'var_mem_Avdw','avg_mem_Teng','mincomp_lk_ball_wtd','POSTMIN_COMP_Impropers',
                      'size','POSTMIN_SING_Impropers','piper_p_aa_pp','mincomp_p_aa_pp',
                      'movedcomp_lk_ball_wtd','mincomp_pro_close','movedcomp_p_aa_pp','piper_rama_prepro',
                      'mincomp_rama_prepro','movedcomp_rama_prepro','PREMIN_COMP_Angles','PREMIN_SING_Angles',
                      'piper_yhh_planarity','cen_DARS','movedcomp_hbond_bb_sc','avg_mem_DARS','POSTMIN_COMP_VWD03',
                      'POSTMIN_COMP_Angles','POSTMIN_SING_Angles','SPPS','mincomp_hbond_bb_sc','piper_hbond_bb_sc',
                      'movedcomp_hbond_sc','mincomp_yhh_planarity','mincomp_fa_dun','mincomp_hbond_sc','movedcomp_pro_close',
                      'piper_pro_close','movedcomp_ref','mincomp_ref','piper_ref','movedcomp_yhh_planarity','SPAR',
                      'movedcomp_fa_dun','piper_fa_elec','movedcomp_fa_elec','piper_hbond_sc','movedcomp_fa_atr',
                      'piper_fa_intra_sol_xover4','movedcomp_fa_sol','mincomp_fa_elec','movedcomp_omega','piper_fa_atr',
                      'piper_fa_sol','mincomp_fa_intra_sol_xover4','SPAS','movedcomp_fa_intra_sol_xover4','mincomp_fa_atr',
                      'mincomp_fa_intra_rep','movedcomp_dslf_fa13','piper_fa_intra_rep','POSTMIN_SING_VWD03','PREMIN_COMP_VWD03',
                      'mincomp_fa_sol','PREMIN_SING_VWD03','piper_dslf_fa13','POSTMIN_SING_Torsions','PREMIN_SING_Torsions',
                      'PREMIN_COMP_Torsions','POSTMIN_COMP_Torsions','piper_fa_dun','movedcomp_fa_intra_rep','piper_hbond_sr_bb',
                      'movedcomp_hbond_sr_bb','mincomp_hbond_sr_bb','mincomp_dslf_fa13','mincomp_omega','PREMIN_COMP_Impropers',
                      'piper_omega','movedcomp_hbond_lr_bb','PREMIN_SING_Impropers','mincomp_hbond_lr_bb','piper_hbond_lr_bb',
                      'skew_mem_DARS','cen_Elec','avg_mem_Elec','avg_mem_Born','cen_Born','piper_time','kurt_mem_Teng','skew_mem_Teng',
                      'mincomp_time','movedcomp_time']

    DATA_ALL = pd.read_csv('/home/Others_Protein_DATA.csv')

X=DATA_ALL.iloc[:][Features_Ordered[0:Num_Features]]
Y=pd.DataFrame(DATA_ALL.iloc[:]['dockq'])   # ['dockq','fnat','lrms','irms']

# Normalazation
Normalization_scaler = preprocessing.MinMaxScaler()#MinMaxScaler  StandardScaler
Normalized_X=pd.DataFrame(   Normalization_scaler.fit_transform(X.values)    )
Normalized_X.columns=X.columns

# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
# @@@@                         Functions
# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@

def dcg_at_k(r, k):
    """
    Args:
        r: Relevance scores (list or numpy) in rank order
            (first element is the first item)
        k: Number of results to consider

    Returns:
        Discounted cumulative gain
    """
    r = np.asfarray(r)[:k]
    return np.sum(r / np.log2(np.arange(2, r.size + 2)))

def ndcg_at_k(r, k):

    dcg_max = dcg_at_k(sorted(r, reverse=True), k)
    if not dcg_max:
        return 0.
    return dcg_at_k(r, k) / dcg_max

def Pi_Hat(Big_X_features,x,Theta, Alpha):
    PI=1
    PI_Hat=1
    for i in range(len(Big_X_features)):
        S_x_y= np.inner( Theta, x ) - np.inner( Theta, Big_X_features[i,:] )
        # print(S_x_y)
        if S_x_y != 0:
            PI_Hat= PI_Hat + np.exp(-Alpha*S_x_y)/(1+np.exp(-Alpha*S_x_y))
        # if S_x_y<0:
        #     PI=PI+1
    return PI_Hat

def Gradient_Update(Big_X_Relevance,Big_X_features,Theta, Alpha):
        # Part1
    N_n= 1/dcg_at_k(sorted(Big_X_Relevance, reverse=True), len(Big_X_Relevance))
    if N_n == float('inf'):          # Revise N_n=0
        N_n=1
    Eq_41_Sumation= np.zeros(len(Big_X_features[1,:]))

    for Ind_x in range(len(Big_X_features)):
        global S_x_y
        x= Big_X_features[Ind_x,:]
        rx=Big_X_Relevance[Ind_x]

        # Part2 Eq_42
        Eq_42= np.zeros(len(Big_X_features[1,:]))
        for i in range(len(Big_X_features)):
            S_x_y= np.inner( Theta, x ) - np.inner( Theta, Big_X_features[i,:] )
            if S_x_y != np.array([0.]):

                Grad_Diff= x - Big_X_features[i,:]  # This is for linear score functions
                P1=np.exp(Alpha*S_x_y)/((1+np.exp(Alpha*S_x_y))**2)
                Eq_42=Eq_42+ (P1*Grad_Diff)

        Eq_42=-Alpha*Eq_42

        # Part3 Eq_43
        PP1= -(((2**rx)-1)/(math.log2(1+Pi_Hat(Big_X_features,x,Theta, Alpha)))**2)
        PP2= 1/((1+Pi_Hat(Big_X_features,x,Theta, Alpha))*np.log(2))
        Eq_43=PP1*PP2

        # Part 4 Eq_42_43

        Eq_41_Sumation= Eq_41_Sumation + (Eq_43 * Eq_42)

    # Final gradient
    return  N_n*Eq_41_Sumation

def iter_lines(lines, has_targets=True, one_indexed=True, missing=0.0):
    """Transforms an iterator of lines to an iterator of LETOR rows.
    Each row is represented by a (x, y, qid, comment) tuple.
    Parameters
    ----------
    lines : iterable of lines
        Lines to parse.
    has_targets : bool, optional
        Whether the file contains targets. If True, will expect the first token
        of every line to be a real representing the sample's target (i.e.
        score). If False, will use -1 as a placeholder for all targets.
    one_indexed : bool, optional
        Whether feature ids are one-indexed. If True, will subtract 1 from each
        feature id.
    missing : float, optional
        Placeholder to use if a feature value is not provided for a sample.
    Yields
    ------
    x : array of floats
        Feature vector of the sample.
    y : float
        Target value (score) of the sample, or -1 if no target was parsed.
    qid : object
        Query id of the sample. This is currently guaranteed to be a string.
    comment : str
        Comment accompanying the sample.
    """
    for line in lines:

        data, _, comment = line.rstrip().partition('#')
        toks = data.split()

        num_features = 0
        x = np.repeat(missing, 8)
        y = -1.0
        if has_targets:
            # print(toks)
            # print("#########################")
            y = float(toks[0])
            toks = toks[1:]

        qid = _parse_qid_tok(toks[0])

        for tok in toks[1:]:
            fid, _, val = tok.partition(':')
            fid = int(fid)
            val = float(val)
            if one_indexed:
                fid -= 1
            assert fid >= 0
            while len(x) <= fid:
                orig = len(x)
                x.resize(len(x) * 2)
                x[orig:orig * 2] = missing

            x[fid] = val
            num_features = max(fid + 1, num_features)

        assert num_features > 0
        x.resize(num_features)

        yield (x, y, qid, comment)

def read_dataset(source, has_targets=True, one_indexed=True, missing=0.0):
    """Parses a LETOR dataset from `source`.
    Parameters
    ----------
    source : string or iterable of lines
        String, file, or other file-like object to parse.
    has_targets : bool, optional
        See `iter_lines`.
    one_indexed : bool, optional
        See `iter_lines`.
    missing : float, optional
        See `iter_lines`.
    Returns
    -------
    X : array of arrays of floats
        Feature matrix (see `iter_lines`).
    y : array of floats
        Target vector (see `iter_lines`).
    qids : array of objects
        Query id vector (see `iter_lines`).
    comments : array of strs
        Comment vector (see `iter_lines`).
    """
    if isinstance(source, six.string_types):
        source = source.splitlines()

    max_width = 0
    xs, ys, qids, comments = [], [], [], []
    it = iter_lines(source, has_targets=has_targets,
                    one_indexed=one_indexed, missing=missing)
    for x, y, qid, comment in it:
        xs.append(x)
        ys.append(y)
        qids.append(qid)
        comments.append(comment)
        max_width = max(max_width, len(x))

    assert max_width > 0
    X = np.ndarray((len(xs), max_width), dtype=np.float64)
    X.fill(missing)
    for i, x in enumerate(xs):
        X[i, :len(x)] = x
    ys = np.array(ys) if has_targets else None
    qids = np.array(qids)
    comments = np.array(comments)

    return (X, ys, qids, comments)

def _parse_qid_tok(tok):
    assert tok.startswith('qid:')
    return tok[4:]

def Performance_Metrics(S_in):
    RANKS=S_in[S_in['RANK_True']==1]
    ACC_All=S_in[S_in['dockq']>0.23]
    Med_All=S_in[S_in['dockq']>0.49]
    # Spearman
    Spearman=spearmanr(S_in['RANK_Pred'], S_in['RANK_True'])[0]
    Spearman_C=spearmanr(S_in['RANK_ClusPro'], S_in['RANK_True'])[0]

    # Highest Quality
    T_1_star= int(RANKS['RANK_Pred']<=1)
    T_5_star= int(RANKS['RANK_Pred']<=5)
    T_10_star= int(RANKS['RANK_Pred']<=10)

    CT_1_star= int(RANKS['RANK_ClusPro']<=1)
    CT_5_star= int(RANKS['RANK_ClusPro']<=5)
    CT_10_star= int(RANKS['RANK_ClusPro']<=10)

    # ACC
    if ACC_All.empty:
        ACCT_1=ACCT_5=ACCT_10=ACCCT_1=ACCCT_5=ACCCT_10=0
        Have_Acc=0
    else:
        Have_Acc=1
        ACCT_1=int(np.where( sum(x<=1 for x in ACC_All['RANK_Pred'])  > 0.5 , 1, 0))
        ACCT_5=int(np.where( sum(x<=5 for x in ACC_All['RANK_Pred'])  > 0.5 , 1, 0))
        ACCT_10= int(np.where( sum(x<=10 for x in ACC_All['RANK_Pred'])  > 0.5 , 1, 0))

        ACCCT_1=int(np.where( sum(x<=1 for x in ACC_All['RANK_ClusPro'])  > 0.5 , 1, 0))
        ACCCT_5=int(np.where( sum(x<=5 for x in ACC_All['RANK_ClusPro'])  > 0.5 , 1, 0))
        ACCCT_10=int(np.where( sum(x<=10 for x in ACC_All['RANK_ClusPro'])  > 0.5 , 1, 0))

    # Med
    if Med_All.empty:
        MedT_1=MedT_5=MedT_10=MedCT_1=MedCT_5=MedCT_10=0
        Have_Med=0
    else:
        Have_Med=1
        MedT_1=int(np.where( sum(x<=1 for x in Med_All['RANK_Pred'])  > 0.5 , 1, 0))
        MedT_5=int(np.where( sum(x<=5 for x in Med_All['RANK_Pred'])  > 0.5 , 1, 0))
        MedT_10=int(np.where( sum(x<=10 for x in Med_All['RANK_Pred'])  > 0.5 , 1, 0))

        MedCT_1=int(np.where( sum(x<=1 for x in Med_All['RANK_ClusPro'])  > 0.5 , 1, 0))
        MedCT_5=int(np.where( sum(x<=5 for x in Med_All['RANK_ClusPro'])  > 0.5 , 1, 0))
        MedCT_10=int(np.where( sum(x<=10 for x in Med_All['RANK_ClusPro'])  > 0.5 , 1, 0))

    return [Spearman,Spearman_C,T_1_star,CT_1_star,T_5_star,CT_5_star,T_10_star,CT_10_star,ACCT_1,ACCCT_1,ACCT_5,ACCCT_5,ACCT_10,ACCCT_10,Have_Acc,MedT_1,MedCT_1,MedT_5,MedCT_5,MedT_10,MedCT_10,Have_Med]

# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
# @@@@                         Main Algorithm ApproxNDCG
# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@

# =========-------------------------------------------------------------------
#  Data Preprocessing  -   Train
# =========-------------------------------------------------------------------

# Load data

# X_Train=Normalized_X[DATA_ALL['pdbid'].isin(Train_Proteins)]
# Y_Train=Y[DATA_ALL['pdbid'].isin(Train_Proteins)]

# X_Train=X_Train.reset_index(drop=True)
# Y_Train=Y_Train.reset_index(drop=True)

Num_Queries=len(Train_Proteins)
Num_Features= Normalized_X.shape[1]

Weight_Y=1
Y=Y*Weight_Y
Num_itr=20
Learning_Rate = 0.1
Alpha=10

# Initialize Theta
Theta_t_1 = np.zeros([(Num_itr+1),Num_Features])
#Theta_t_1[0,:]=np.random.random([1,Num_Features])
Theta_t_1[0,:]=np.ones([1,Num_Features])

for t in range(Num_itr):
    print('**** Iteration ', t+1, '  *******************************')
    THETA= Theta_t_1[t,:]
    for i in range(Num_Queries):
        Big_X_features= Normalized_X[DATA_ALL['pdbid'].isin([Train_Proteins[i]])]
        Big_X_Relevance= Y[DATA_ALL['pdbid'].isin([Train_Proteins[i]])]

        # Compute gradient for i-th query
        Delta_Theta=  Gradient_Update(np.array(Big_X_Relevance),np.array(Big_X_features),THETA, Alpha)

        THETA = THETA + (Learning_Rate * Delta_Theta)

    Theta_t_1[(t+1),:]= THETA
    # Shuffle queries
    random.seed(t)
    random.shuffle(Train_Proteins)

# =========-------------------------------------------------------------------
#  Data Preprocessing  -   Validation
# =========-------------------------------------------------------------------

# with open("/home/vali.txt") as validationfile:
#     VX, Vy, Vqids, Vc = read_dataset(validationfile, has_targets=True , one_indexed=True)

# =========-------------------------------------------------------------------
#  Data Preprocessing  -   Test
# =========-------------------------------------------------------------------

# X_Test=Normalized_X[DATA_ALL['pdbid'].isin(Test_Proteins)]
# Y_Test=Y[DATA_ALL['pdbid'].isin(Test_Proteins)]
# X_Test=X_Test.reset_index(drop=True)
# Y_Test=Y_Test.reset_index(drop=True)

Num_Test_Queries=len(Test_Proteins)

Predicted_NDCGatk = np.array([0,1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
for pp in range(Num_itr):
    NDCG_all=[0,0,0,0,0,0,0,0,0,0,0]
    for tt in range(Num_Test_Queries):
        Test_X_features= Normalized_X[DATA_ALL['pdbid'].isin([Test_Proteins[tt]])]
        Pred_Q= np.matmul(np.array(Test_X_features)  , Theta_t_1[pp]  )     # Prediction in iteration pp-th
        Ground_Q = np.array( Y[DATA_ALL['pdbid'].isin([Test_Proteins[tt]])] )
        Pred_Ranks=(rankdata(-Pred_Q, method='ordinal'))    # Decreasing order
        Pred_Ranks = np.expand_dims(Pred_Ranks, axis=1)
        Concat= np.concatenate((Ground_Q,Pred_Ranks), axis=1)
        sorted_array = Concat[np.argsort(Concat[:, 1])]
        RGT= sorted_array[:,0]
        # NDCG @ k
        NDCG_all = np.add(NDCG_all, [pp, ndcg_at_k(RGT, 1),ndcg_at_k(RGT, 2),ndcg_at_k(RGT, 3),ndcg_at_k(RGT, 4),ndcg_at_k(RGT, 5),ndcg_at_k(RGT, 6),ndcg_at_k(RGT, 7),ndcg_at_k(RGT, 8),ndcg_at_k(RGT, 9),ndcg_at_k(RGT, 10)])

    Predicted_NDCG=NDCG_all/Num_Test_Queries
    print(Predicted_NDCG)
    Predicted_NDCGatk=np.vstack([Predicted_NDCGatk,Predicted_NDCG])

print("--- %s seconds ---" % (time.time() - start_time))

# save numpy array as csv file
from numpy import asarray
from numpy import savetxt

# save to csv file

SAVE_Name_Theta= Which_Type + '_NFeature_' + str(Num_Features) +'_Learned_Thetas_itr'+str(Num_itr)+'_Alpha'+str(Alpha)+'.csv'

SAVE_Name_Results= Which_Type + '_NFeature_' + str(Num_Features) + '_Results_itr'+str(Num_itr)+'_Alpha'+str(Alpha)+'.csv'

# save to csv file
savetxt(SAVE_Name_Theta, Theta_t_1, delimiter=',')

savetxt(SAVE_Name_Results, Predicted_NDCGatk, delimiter=',')


#!/usr/bin/env python
# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

import numpy
from pyscf import lib
from pyscf.ao2mo import _ao2mo
from pyscf.ao2mo.incore import iden_coeffs, _conc_mos
from pyscf.pbc.df.fft_ao2mo import _format_kpts
from pyscf.pbc.df import df_ao2mo
from pyscf.pbc.df import aft_ao2mo
from pyscf.pbc.lib import kpts_helper
from pyscf.pbc.lib.kpts_helper import is_zero, gamma_point, member, unique
from pyscf import __config__

def get_eri(mydf, kpts=None,
            compact=getattr(__config__, 'pbc_df_ao2mo_get_eri_compact', True)):
    if mydf._cderi is None:
        mydf.build()

    kptijkl = _format_kpts(kpts)
    eri = aft_ao2mo.get_eri(mydf, kptijkl, compact=compact)
    eri += df_ao2mo.get_eri(mydf, kptijkl, compact=compact)
    return eri

def general(mydf, mo_coeffs, kpts=None,
            compact=getattr(__config__, 'pbc_df_ao2mo_general_compact', True)):
    if mydf._cderi is None:
        mydf.build()

    kptijkl = _format_kpts(kpts)
    if isinstance(mo_coeffs, numpy.ndarray) and mo_coeffs.ndim == 2:
        mo_coeffs = (mo_coeffs,) * 4
    eri_mo = aft_ao2mo.general(mydf, mo_coeffs, kptijkl, compact=compact)
    eri_mo += df_ao2mo.general(mydf, mo_coeffs, kptijkl, compact=compact)
    return eri_mo

def ao2mo_7d(mydf, mo_coeff_kpts, kpts=None, factor=1, out=None):
    cell = mydf.cell
    if kpts is None:
        kpts = mydf.kpts
    nkpts = len(kpts)

    if isinstance(mo_coeff_kpts, numpy.ndarray) and mo_coeff_kpts.ndim == 3:
        mo_coeff_kpts = [mo_coeff_kpts] * 4
    else:
        mo_coeff_kpts = list(mo_coeff_kpts)

    # Shape of the orbitals can be different on different k-points. The
    # orbital coefficients must be formatted (padded by zeros) so that the
    # shape of the orbital coefficients are the same on all k-points. This can
    # be achieved by calling pbc.mp.kmp2.padded_mo_coeff function
    nmoi, nmoj, nmok, nmol = [x.shape[2] for x in mo_coeff_kpts]
    eri_shape = (nkpts, nkpts, nkpts, nmoi, nmoj, nmok, nmol)
    if gamma_point(kpts):
        dtype = numpy.result_type(*mo_coeff_kpts)
    else:
        dtype = numpy.complex128

    if out is None:
        out = numpy.empty(eri_shape, dtype=dtype)
    else:
        assert(out.shape == eri_shape)

    kptij_lst = numpy.array([(ki, kj) for ki in kpts for kj in kpts])
    kptis_lst = kptij_lst[:,0]
    kptjs_lst = kptij_lst[:,1]
    kpt_ji = kptjs_lst - kptis_lst
    uniq_kpts, uniq_index, uniq_inverse = unique(kpt_ji)
    ngrids = numpy.prod(mydf.mesh)
    nao = cell.nao_nr()
    max_memory = max(2000, mydf.max_memory-lib.current_memory()[0]-nao**4*16/1e6) * .5

    fswap = lib.H5TmpFile()
    tao = []
    ao_loc = None
    kconserv = kpts_helper.get_kconserv(cell, kpts)
    for uniq_id, kpt in enumerate(uniq_kpts):
        q = uniq_kpts[uniq_id]
        adapted_ji_idx = numpy.where(uniq_inverse == uniq_id)[0]

        kptjs = kptjs_lst[adapted_ji_idx]
        coulG = mydf.weighted_coulG(q, False, mydf.mesh)
        coulG *= factor

        moij_list = []
        ijslice_list = []
        for ji, ji_idx in enumerate(adapted_ji_idx):
            ki = ji_idx // nkpts
            kj = ji_idx % nkpts
            moij, ijslice = _conc_mos(mo_coeff_kpts[0][ki], mo_coeff_kpts[1][kj])[2:]
            moij_list.append(moij)
            ijslice_list.append(ijslice)
            fswap.create_dataset('zij/'+str(ji), (ngrids,nmoi*nmoj), 'D')

        for aoaoks, p0, p1 in mydf.ft_loop(mydf.mesh, q, kptjs,
                                           max_memory=max_memory):
            for ji, aoao in enumerate(aoaoks):
                ki = adapted_ji_idx[ji] // nkpts
                kj = adapted_ji_idx[ji] %  nkpts
                buf = aoao.transpose(1,2,0).reshape(nao**2,p1-p0)
                zij = _ao2mo.r_e2(lib.transpose(buf), moij_list[ji],
                                  ijslice_list[ji], tao, ao_loc)
                zij *= coulG[p0:p1,None]
                fswap['zij/'+str(ji)][p0:p1] = zij

        mokl_list = []
        klslice_list = []
        for kk in range(nkpts):
            kl = kconserv[ki, kj, kk]
            mokl, klslice = _conc_mos(mo_coeff_kpts[2][kk], mo_coeff_kpts[3][kl])[2:]
            mokl_list.append(mokl)
            klslice_list.append(klslice)
            fswap.create_dataset('zkl/'+str(kk), (ngrids,nmok*nmol), 'D')

        ki = adapted_ji_idx[0] // nkpts
        kj = adapted_ji_idx[0] % nkpts
        kptls = kpts[kconserv[ki, kj, :]]
        for aoaoks, p0, p1 in mydf.ft_loop(mydf.mesh, q, -kptls,
                                           max_memory=max_memory):
            for kk, aoao in enumerate(aoaoks):
                buf = aoao.conj().transpose(1,2,0).reshape(nao**2,p1-p0)
                zkl = _ao2mo.r_e2(lib.transpose(buf), mokl_list[kk],
                                  klslice_list[kk], tao, ao_loc)
                fswap['zkl/'+str(kk)][p0:p1] = zkl

        for ji, ji_idx in enumerate(adapted_ji_idx):
            ki = ji_idx // nkpts
            kj = ji_idx % nkpts

            moij, ijslice = _conc_mos(mo_coeff_kpts[0][ki], mo_coeff_kpts[1][kj])[2:]
            zij = []
            for LpqR, LpqI, sign in mydf.sr_loop(kpts[[ki,kj]], max_memory, False, mydf.blockdim):
                zij.append(_ao2mo.r_e2(LpqR+LpqI*1j, moij, ijslice, tao, ao_loc))

            for kk in range(nkpts):
                kl = kconserv[ki, kj, kk]
                eri_mo = lib.dot(numpy.asarray(fswap['zij/'+str(ji)]).T,
                                 numpy.asarray(fswap['zkl/'+str(kk)]))

                for i, (LrsR, LrsI, sign) in \
                        enumerate(mydf.sr_loop(kpts[[kk,kl]], max_memory, False, mydf.blockdim)):
                    zkl = _ao2mo.r_e2(LrsR+LrsI*1j, mokl_list[kk],
                                      klslice_list[kk], tao, ao_loc)
                    lib.dot(zij[i].T, zkl, sign*factor, eri_mo, 1)

                if dtype == numpy.double:
                    eri_mo = eri_mo.real
                out[ki,kj,kk] = eri_mo.reshape(eri_shape[3:])
        del(fswap['zij'])
        del(fswap['zkl'])

    return out


# Copyright 2021 AlQuraishi Laboratory
# Copyright 2021 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Constants used in AlphaFold."""

import collections
import functools
from typing import Mapping, List, Tuple

import numpy as np
import tree

# Internal import (35fd).

# Distance from one CA to next CA [trans configuration: omega = 180].
ca_ca = 3.80209737096

# Format: The list for each AA type contains chi1, chi2, chi3, chi4 in
# this order (or a relevant subset from chi1 onwards). ALA and GLY don't have
# chi angles so their chi angle lists are empty.
chi_angles_atoms = {
    "ALA": [],
    # Chi5 in arginine is always 0 +- 5 degrees, so ignore it.
    "ARG": [
        ["N", "CA", "CB", "CG"],
        ["CA", "CB", "CG", "CD"],
        ["CB", "CG", "CD", "NE"],
        ["CG", "CD", "NE", "CZ"],
    ],
    "ASN": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "OD1"]],
    "ASP": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "OD1"]],
    "CYS": [["N", "CA", "CB", "SG"]],
    "GLN": [
        ["N", "CA", "CB", "CG"],
        ["CA", "CB", "CG", "CD"],
        ["CB", "CG", "CD", "OE1"],
    ],
    "GLU": [
        ["N", "CA", "CB", "CG"],
        ["CA", "CB", "CG", "CD"],
        ["CB", "CG", "CD", "OE1"],
    ],
    "GLY": [],
    "HIS": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "ND1"]],
    "ILE": [["N", "CA", "CB", "CG1"], ["CA", "CB", "CG1", "CD1"]],
    "LEU": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD1"]],
    "LYS": [
        ["N", "CA", "CB", "CG"],
        ["CA", "CB", "CG", "CD"],
        ["CB", "CG", "CD", "CE"],
        ["CG", "CD", "CE", "NZ"],
    ],
    "MET": [
        ["N", "CA", "CB", "CG"],
        ["CA", "CB", "CG", "SD"],
        ["CB", "CG", "SD", "CE"],
    ],
    "PHE": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD1"]],
    "PRO": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD"]],
    "SER": [["N", "CA", "CB", "OG"]],
    "THR": [["N", "CA", "CB", "OG1"]],
    "TRP": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD1"]],
    "TYR": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD1"]],
    "VAL": [["N", "CA", "CB", "CG1"]],
}

# If chi angles given in fixed-length array, this matrix determines how to mask
# them for each AA type. The order is as per restype_order (see below).
chi_angles_mask = [
    [0.0, 0.0, 0.0, 0.0],  # ALA
    [1.0, 1.0, 1.0, 1.0],  # ARG
    [1.0, 1.0, 0.0, 0.0],  # ASN
    [1.0, 1.0, 0.0, 0.0],  # ASP
    [1.0, 0.0, 0.0, 0.0],  # CYS
    [1.0, 1.0, 1.0, 0.0],  # GLN
    [1.0, 1.0, 1.0, 0.0],  # GLU
    [0.0, 0.0, 0.0, 0.0],  # GLY
    [1.0, 1.0, 0.0, 0.0],  # HIS
    [1.0, 1.0, 0.0, 0.0],  # ILE
    [1.0, 1.0, 0.0, 0.0],  # LEU
    [1.0, 1.0, 1.0, 1.0],  # LYS
    [1.0, 1.0, 1.0, 0.0],  # MET
    [1.0, 1.0, 0.0, 0.0],  # PHE
    [1.0, 1.0, 0.0, 0.0],  # PRO
    [1.0, 0.0, 0.0, 0.0],  # SER
    [1.0, 0.0, 0.0, 0.0],  # THR
    [1.0, 1.0, 0.0, 0.0],  # TRP
    [1.0, 1.0, 0.0, 0.0],  # TYR
    [1.0, 0.0, 0.0, 0.0],  # VAL
]

# The following chi angles are pi periodic: they can be rotated by a multiple
# of pi without affecting the structure.
chi_pi_periodic = [
    [0.0, 0.0, 0.0, 0.0],  # ALA
    [0.0, 0.0, 0.0, 0.0],  # ARG
    [0.0, 0.0, 0.0, 0.0],  # ASN
    [0.0, 1.0, 0.0, 0.0],  # ASP
    [0.0, 0.0, 0.0, 0.0],  # CYS
    [0.0, 0.0, 0.0, 0.0],  # GLN
    [0.0, 0.0, 1.0, 0.0],  # GLU
    [0.0, 0.0, 0.0, 0.0],  # GLY
    [0.0, 0.0, 0.0, 0.0],  # HIS
    [0.0, 0.0, 0.0, 0.0],  # ILE
    [0.0, 0.0, 0.0, 0.0],  # LEU
    [0.0, 0.0, 0.0, 0.0],  # LYS
    [0.0, 0.0, 0.0, 0.0],  # MET
    [0.0, 1.0, 0.0, 0.0],  # PHE
    [0.0, 0.0, 0.0, 0.0],  # PRO
    [0.0, 0.0, 0.0, 0.0],  # SER
    [0.0, 0.0, 0.0, 0.0],  # THR
    [0.0, 0.0, 0.0, 0.0],  # TRP
    [0.0, 1.0, 0.0, 0.0],  # TYR
    [0.0, 0.0, 0.0, 0.0],  # VAL
    [0.0, 0.0, 0.0, 0.0],  # UNK
]

# Atoms positions relative to the 8 rigid groups, defined by the pre-omega, phi,
# psi and chi angles:
# 0: 'backbone group',
# 1: 'pre-omega-group', (empty)
# 2: 'phi-group', (currently empty, because it defines only hydrogens)
# 3: 'psi-group',
# 4,5,6,7: 'chi1,2,3,4-group'
# The atom positions are relative to the axis-end-atom of the corresponding
# rotation axis. The x-axis is in direction of the rotation axis, and the y-axis
# is defined such that the dihedral-angle-definiting atom (the last entry in
# chi_angles_atoms above) is in the xy-plane (with a positive y-coordinate).
# format: [atomname, group_idx, rel_position]
rigid_group_atom_positions = {
    "ALA": [
        ["N", 0, (-0.525, 1.363, 0.000)],
        ["CA", 0, (0.000, 0.000, 0.000)],
        ["C", 0, (1.526, -0.000, -0.000)],
        ["CB", 0, (-0.529, -0.774, -1.205)],
        ["O", 3, (0.627, 1.062, 0.000)],
    ],
    "ARG": [
        ["N", 0, (-0.524, 1.362, -0.000)],
        ["CA", 0, (0.000, 0.000, 0.000)],
        ["C", 0, (1.525, -0.000, -0.000)],
        ["CB", 0, (-0.524, -0.778, -1.209)],
        ["O", 3, (0.626, 1.062, 0.000)],
        ["CG", 4, (0.616, 1.390, -0.000)],
        ["CD", 5, (0.564, 1.414, 0.000)],
        ["NE", 6, (0.539, 1.357, -0.000)],
        ["NH1", 7, (0.206, 2.301, 0.000)],
        ["NH2", 7, (2.078, 0.978, -0.000)],
        ["CZ", 7, (0.758, 1.093, -0.000)],
    ],
    "ASN": [
        ["N", 0, (-0.536, 1.357, 0.000)],
        ["CA", 0, (0.000, 0.000, 0.000)],
        ["C", 0, (1.526, -0.000, -0.000)],
        ["CB", 0, (-0.531, -0.787, -1.200)],
        ["O", 3, (0.625, 1.062, 0.000)],
        ["CG", 4, (0.584, 1.399, 0.000)],
        ["ND2", 5, (0.593, -1.188, 0.001)],
        ["OD1", 5, (0.633, 1.059, 0.000)],
    ],
    "ASP": [
        ["N", 0, (-0.525, 1.362, -0.000)],
        ["CA", 0, (0.000, 0.000, 0.000)],
        ["C", 0, (1.527, 0.000, -0.000)],
        ["CB", 0, (-0.526, -0.778, -1.208)],
        ["O", 3, (0.626, 1.062, -0.000)],
        ["CG", 4, (0.593, 1.398, -0.000)],
        ["OD1", 5, (0.610, 1.091, 0.000)],
        ["OD2", 5, (0.592, -1.101, -0.003)],
    ],
    "CYS": [
        ["N", 0, (-0.522, 1.362, -0.000)],
        ["CA", 0, (0.000, 0.000, 0.000)],
        ["C", 0, (1.524, 0.000, 0.000)],
        ["CB", 0, (-0.519, -0.773, -1.212)],
        ["O", 3, (0.625, 1.062, -0.000)],
        ["SG", 4, (0.728, 1.653, 0.000)],
    ],
    "GLN": [
        ["N", 0, (-0.526, 1.361, -0.000)],
        ["CA", 0, (0.000, 0.000, 0.000)],
        ["C", 0, (1.526, 0.000, 0.000)],
        ["CB", 0, (-0.525, -0.779, -1.207)],
        ["O", 3, (0.626, 1.062, -0.000)],
        ["CG", 4, (0.615, 1.393, 0.000)],
        ["CD", 5, (0.587, 1.399, -0.000)],
        ["NE2", 6, (0.593, -1.189, -0.001)],
        ["OE1", 6, (0.634, 1.060, 0.000)],
    ],
    "GLU": [
        ["N", 0, (-0.528, 1.361, 0.000)],
        ["CA", 0, (0.000, 0.000, 0.000)],
        ["C", 0, (1.526, -0.000, -0.000)],
        ["CB", 0, (-0.526, -0.781, -1.207)],
        ["O", 3, (0.626, 1.062, 0.000)],
        ["CG", 4, (0.615, 1.392, 0.000)],
        ["CD", 5, (0.600, 1.397, 0.000)],
        ["OE1", 6, (0.607, 1.095, -0.000)],
        ["OE2", 6, (0.589, -1.104, -0.001)],
    ],
    "GLY": [
        ["N", 0, (-0.572, 1.337, 0.000)],
        ["CA", 0, (0.000, 0.000, 0.000)],
        ["C", 0, (1.517, -0.000, -0.000)],
        ["O", 3, (0.626, 1.062, -0.000)],
    ],
    "HIS": [
        ["N", 0, (-0.527, 1.360, 0.000)],
        ["CA", 0, (0.000, 0.000, 0.000)],
        ["C", 0, (1.525, 0.000, 0.000)],
        ["CB", 0, (-0.525, -0.778, -1.208)],
        ["O", 3, (0.625, 1.063, 0.000)],
        ["CG", 4, (0.600, 1.370, -0.000)],
        ["CD2", 5, (0.889, -1.021, 0.003)],
        ["ND1", 5, (0.744, 1.160, -0.000)],
        ["CE1", 5, (2.030, 0.851, 0.002)],
        ["NE2", 5, (2.145, -0.466, 0.004)],
    ],
    "ILE": [
        ["N", 0, (-0.493, 1.373, -0.000)],
        ["CA", 0, (0.000, 0.000, 0.000)],
        ["C", 0, (1.527, -0.000, -0.000)],
        ["CB", 0, (-0.536, -0.793, -1.213)],
        ["O", 3, (0.627, 1.062, -0.000)],
        ["CG1", 4, (0.534, 1.437, -0.000)],
        ["CG2", 4, (0.540, -0.785, -1.199)],
        ["CD1", 5, (0.619, 1.391, 0.000)],
    ],
    "LEU": [
        ["N", 0, (-0.520, 1.363, 0.000)],
        ["CA", 0, (0.000, 0.000, 0.000)],
        ["C", 0, (1.525, -0.000, -0.000)],
        ["CB", 0, (-0.522, -0.773, -1.214)],
        ["O", 3, (0.625, 1.063, -0.000)],
        ["CG", 4, (0.678, 1.371, 0.000)],
        ["CD1", 5, (0.530, 1.430, -0.000)],
        ["CD2", 5, (0.535, -0.774, 1.200)],
    ],
    "LYS": [
        ["N", 0, (-0.526, 1.362, -0.000)],
        ["CA", 0, (0.000, 0.000, 0.000)],
        ["C", 0, (1.526, 0.000, 0.000)],
        ["CB", 0, (-0.524, -0.778, -1.208)],
        ["O", 3, (0.626, 1.062, -0.000)],
        ["CG", 4, (0.619, 1.390, 0.000)],
        ["CD", 5, (0.559, 1.417, 0.000)],
        ["CE", 6, (0.560, 1.416, 0.000)],
        ["NZ", 7, (0.554, 1.387, 0.000)],
    ],
    "MET": [
        ["N", 0, (-0.521, 1.364, -0.000)],
        ["CA", 0, (0.000, 0.000, 0.000)],
        ["C", 0, (1.525, 0.000, 0.000)],
        ["CB", 0, (-0.523, -0.776, -1.210)],
        ["O", 3, (0.625, 1.062, -0.000)],
        ["CG", 4, (0.613, 1.391, -0.000)],
        ["SD", 5, (0.703, 1.695, 0.000)],
        ["CE", 6, (0.320, 1.786, -0.000)],
    ],
    "PHE": [
        ["N", 0, (-0.518, 1.363, 0.000)],
        ["CA", 0, (0.000, 0.000, 0.000)],
        ["C", 0, (1.524, 0.000, -0.000)],
        ["CB", 0, (-0.525, -0.776, -1.212)],
        ["O", 3, (0.626, 1.062, -0.000)],
        ["CG", 4, (0.607, 1.377, 0.000)],
        ["CD1", 5, (0.709, 1.195, -0.000)],
        ["CD2", 5, (0.706, -1.196, 0.000)],
        ["CE1", 5, (2.102, 1.198, -0.000)],
        ["CE2", 5, (2.098, -1.201, -0.000)],
        ["CZ", 5, (2.794, -0.003, -0.001)],
    ],
    "PRO": [
        ["N", 0, (-0.566, 1.351, -0.000)],
        ["CA", 0, (0.000, 0.000, 0.000)],
        ["C", 0, (1.527, -0.000, 0.000)],
        ["CB", 0, (-0.546, -0.611, -1.293)],
        ["O", 3, (0.621, 1.066, 0.000)],
        ["CG", 4, (0.382, 1.445, 0.0)],
        # ['CD', 5, (0.427, 1.440, 0.0)],
        ["CD", 5, (0.477, 1.424, 0.0)],  # manually made angle 2 degrees larger
    ],
    "SER": [
        ["N", 0, (-0.529, 1.360, -0.000)],
        ["CA", 0, (0.000, 0.000, 0.000)],
        ["C", 0, (1.525, -0.000, -0.000)],
        ["CB", 0, (-0.518, -0.777, -1.211)],
        ["O", 3, (0.626, 1.062, -0.000)],
        ["OG", 4, (0.503, 1.325, 0.000)],
    ],
    "THR": [
        ["N", 0, (-0.517, 1.364, 0.000)],
        ["CA", 0, (0.000, 0.000, 0.000)],
        ["C", 0, (1.526, 0.000, -0.000)],
        ["CB", 0, (-0.516, -0.793, -1.215)],
        ["O", 3, (0.626, 1.062, 0.000)],
        ["CG2", 4, (0.550, -0.718, -1.228)],
        ["OG1", 4, (0.472, 1.353, 0.000)],
    ],
    "TRP": [
        ["N", 0, (-0.521, 1.363, 0.000)],
        ["CA", 0, (0.000, 0.000, 0.000)],
        ["C", 0, (1.525, -0.000, 0.000)],
        ["CB", 0, (-0.523, -0.776, -1.212)],
        ["O", 3, (0.627, 1.062, 0.000)],
        ["CG", 4, (0.609, 1.370, -0.000)],
        ["CD1", 5, (0.824, 1.091, 0.000)],
        ["CD2", 5, (0.854, -1.148, -0.005)],
        ["CE2", 5, (2.186, -0.678, -0.007)],
        ["CE3", 5, (0.622, -2.530, -0.007)],
        ["NE1", 5, (2.140, 0.690, -0.004)],
        ["CH2", 5, (3.028, -2.890, -0.013)],
        ["CZ2", 5, (3.283, -1.543, -0.011)],
        ["CZ3", 5, (1.715, -3.389, -0.011)],
    ],
    "TYR": [
        ["N", 0, (-0.522, 1.362, 0.000)],
        ["CA", 0, (0.000, 0.000, 0.000)],
        ["C", 0, (1.524, -0.000, -0.000)],
        ["CB", 0, (-0.522, -0.776, -1.213)],
        ["O", 3, (0.627, 1.062, -0.000)],
        ["CG", 4, (0.607, 1.382, -0.000)],
        ["CD1", 5, (0.716, 1.195, -0.000)],
        ["CD2", 5, (0.713, -1.194, -0.001)],
        ["CE1", 5, (2.107, 1.200, -0.002)],
        ["CE2", 5, (2.104, -1.201, -0.003)],
        ["OH", 5, (4.168, -0.002, -0.005)],
        ["CZ", 5, (2.791, -0.001, -0.003)],
    ],
    "VAL": [
        ["N", 0, (-0.494, 1.373, -0.000)],
        ["CA", 0, (0.000, 0.000, 0.000)],
        ["C", 0, (1.527, -0.000, -0.000)],
        ["CB", 0, (-0.533, -0.795, -1.213)],
        ["O", 3, (0.627, 1.062, -0.000)],
        ["CG1", 4, (0.540, 1.429, -0.000)],
        ["CG2", 4, (0.533, -0.776, 1.203)],
    ],
}

# A list of atoms (excluding hydrogen) for each AA type. PDB naming convention.
residue_atoms = {
    "ALA": ["C", "CA", "CB", "N", "O"],
    "ARG": ["C", "CA", "CB", "CG", "CD", "CZ", "N", "NE", "O", "NH1", "NH2"],
    "ASP": ["C", "CA", "CB", "CG", "N", "O", "OD1", "OD2"],
    "ASN": ["C", "CA", "CB", "CG", "N", "ND2", "O", "OD1"],
    "CYS": ["C", "CA", "CB", "N", "O", "SG"],
    "GLU": ["C", "CA", "CB", "CG", "CD", "N", "O", "OE1", "OE2"],
    "GLN": ["C", "CA", "CB", "CG", "CD", "N", "NE2", "O", "OE1"],
    "GLY": ["C", "CA", "N", "O"],
    "HIS": ["C", "CA", "CB", "CG", "CD2", "CE1", "N", "ND1", "NE2", "O"],
    "ILE": ["C", "CA", "CB", "CG1", "CG2", "CD1", "N", "O"],
    "LEU": ["C", "CA", "CB", "CG", "CD1", "CD2", "N", "O"],
    "LYS": ["C", "CA", "CB", "CG", "CD", "CE", "N", "NZ", "O"],
    "MET": ["C", "CA", "CB", "CG", "CE", "N", "O", "SD"],
    "PHE": ["C", "CA", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ", "N", "O"],
    "PRO": ["C", "CA", "CB", "CG", "CD", "N", "O"],
    "SER": ["C", "CA", "CB", "N", "O", "OG"],
    "THR": ["C", "CA", "CB", "CG2", "N", "O", "OG1"],
    "TRP": [
        "C",
        "CA",
        "CB",
        "CG",
        "CD1",
        "CD2",
        "CE2",
        "CE3",
        "CZ2",
        "CZ3",
        "CH2",
        "N",
        "NE1",
        "O",
    ],
    "TYR": [
        "C",
        "CA",
        "CB",
        "CG",
        "CD1",
        "CD2",
        "CE1",
        "CE2",
        "CZ",
        "N",
        "O",
        "OH",
    ],
    "VAL": ["C", "CA", "CB", "CG1", "CG2", "N", "O"],
}

# Naming swaps for ambiguous atom names.
# Due to symmetries in the amino acids the naming of atoms is ambiguous in
# 4 of the 20 amino acids.
# (The LDDT paper lists 7 amino acids as ambiguous, but the naming ambiguities
# in LEU, VAL and ARG can be resolved by using the 3d constellations of
# the 'ambiguous' atoms and their neighbours)
# TODO: ^ interpret this
residue_atom_renaming_swaps = {
    "ASP": {"OD1": "OD2"},
    "GLU": {"OE1": "OE2"},
    "PHE": {"CD1": "CD2", "CE1": "CE2"},
    "TYR": {"CD1": "CD2", "CE1": "CE2"},
}

# Van der Waals radii [Angstroem] of the atoms (from Wikipedia)
van_der_waals_radius = {
    "C": 1.7,
    "N": 1.55,
    "O": 1.52,
    "S": 1.8,
}

Bond = collections.namedtuple(
    "Bond", ["atom1_name", "atom2_name", "length", "stddev"]
)
BondAngle = collections.namedtuple(
    "BondAngle",
    ["atom1_name", "atom2_name", "atom3name", "angle_rad", "stddev"],
)

@functools.lru_cache(maxsize=None)
def load_stereo_chemical_props() -> Tuple[
    Mapping[str, List[Bond]],
    Mapping[str, List[Bond]],
    Mapping[str, List[BondAngle]],
]:
    """Load stereo_chemical_props.txt into a nice structure.

    Load literature values for bond lengths and bond angles and translate
    bond angles into the length of the opposite edge of the triangle
    ("residue_virtual_bonds").

    Returns:
      residue_bonds:  dict that maps resname --> list of Bond tuples
      residue_virtual_bonds: dict that maps resname --> list of Bond tuples
      residue_bond_angles: dict that maps resname --> list of BondAngle tuples
    """
    # TODO: this file should be downloaded in a setup script
    stereo_chemical_props_path = "openfold/resources/stereo_chemical_props.txt"
    with open(stereo_chemical_props_path, "rt") as f:
        stereo_chemical_props = f.read()
    lines_iter = iter(stereo_chemical_props.splitlines())
    # Load bond lengths.
    residue_bonds = {}
    next(lines_iter)  # Skip header line.
    for line in lines_iter:
        if line.strip() == "-":
            break
        bond, resname, length, stddev = line.split()
        atom1, atom2 = bond.split("-")
        if resname not in residue_bonds:
            residue_bonds[resname] = []
        residue_bonds[resname].append(
            Bond(atom1, atom2, float(length), float(stddev))
        )
    residue_bonds["UNK"] = []

    # Load bond angles.
    residue_bond_angles = {}
    next(lines_iter)  # Skip empty line.
    next(lines_iter)  # Skip header line.
    for line in lines_iter:
        if line.strip() == "-":
            break
        bond, resname, angle_degree, stddev_degree = line.split()
        atom1, atom2, atom3 = bond.split("-")
        if resname not in residue_bond_angles:
            residue_bond_angles[resname] = []
        residue_bond_angles[resname].append(
            BondAngle(
                atom1,
                atom2,
                atom3,
                float(angle_degree) / 180.0 * np.pi,
                float(stddev_degree) / 180.0 * np.pi,
            )
        )
    residue_bond_angles["UNK"] = []

    def make_bond_key(atom1_name, atom2_name):
        """Unique key to lookup bonds."""
        return "-".join(sorted([atom1_name, atom2_name]))

    # Translate bond angles into distances ("virtual bonds").
    residue_virtual_bonds = {}
    for resname, bond_angles in residue_bond_angles.items():
        # Create a fast lookup dict for bond lengths.
        bond_cache = {}
        for b in residue_bonds[resname]:
            bond_cache[make_bond_key(b.atom1_name, b.atom2_name)] = b
        residue_virtual_bonds[resname] = []
        for ba in bond_angles:
            bond1 = bond_cache[make_bond_key(ba.atom1_name, ba.atom2_name)]
            bond2 = bond_cache[make_bond_key(ba.atom2_name, ba.atom3name)]

            # Compute distance between atom1 and atom3 using the law of cosines
            # c^2 = a^2 + b^2 - 2ab*cos(gamma).
            gamma = ba.angle_rad
            length = np.sqrt(
                bond1.length ** 2
                + bond2.length ** 2
                - 2 * bond1.length * bond2.length * np.cos(gamma)
            )

            # Propagation of uncertainty assuming uncorrelated errors.
            dl_outer = 0.5 / length
            dl_dgamma = (
                2 * bond1.length * bond2.length * np.sin(gamma)
            ) * dl_outer
            dl_db1 = (
                2 * bond1.length - 2 * bond2.length * np.cos(gamma)
            ) * dl_outer
            dl_db2 = (
                2 * bond2.length - 2 * bond1.length * np.cos(gamma)
            ) * dl_outer
            stddev = np.sqrt(
                (dl_dgamma * ba.stddev) ** 2
                + (dl_db1 * bond1.stddev) ** 2
                + (dl_db2 * bond2.stddev) ** 2
            )
            residue_virtual_bonds[resname].append(
                Bond(ba.atom1_name, ba.atom3name, length, stddev)
            )

    return (residue_bonds, residue_virtual_bonds, residue_bond_angles)

# Between-residue bond lengths for general bonds (first element) and for Proline
# (second element).
between_res_bond_length_c_n = [1.329, 1.341]
between_res_bond_length_stddev_c_n = [0.014, 0.016]

# Between-residue cos_angles.
between_res_cos_angles_c_n_ca = [-0.5203, 0.0353]  # degrees: 121.352 +- 2.315
between_res_cos_angles_ca_c_n = [-0.4473, 0.0311]  # degrees: 116.568 +- 1.995

# This mapping is used when we need to store atom data in a format that requires
# fixed atom data size for every residue (e.g. a numpy array).
atom_types = [
    "N",
    "CA",
    "C",
    "CB",
    "O",
    "CG",
    "CG1",
    "CG2",
    "OG",
    "OG1",
    "SG",
    "CD",
    "CD1",
    "CD2",
    "ND1",
    "ND2",
    "OD1",
    "OD2",
    "SD",
    "CE",
    "CE1",
    "CE2",
    "CE3",
    "NE",
    "NE1",
    "NE2",
    "OE1",
    "OE2",
    "CH2",
    "NH1",
    "NH2",
    "OH",
    "CZ",
    "CZ2",
    "CZ3",
    "NZ",
    "OXT",
]
atom_order = {atom_type: i for i, atom_type in enumerate(atom_types)}
atom_type_num = len(atom_types)  # := 37.

# A compact atom encoding with 14 columns
# pylint: disable=line-too-long
# pylint: disable=bad-whitespace
restype_name_to_atom14_names = {
    "ALA": ["N", "CA", "C", "O", "CB", "", "", "", "", "", "", "", "", ""],
    "ARG": [
        "N",
        "CA",
        "C",
        "O",
        "CB",
        "CG",
        "CD",
        "NE",
        "CZ",
        "NH1",
        "NH2",
        "",
        "",
        "",
    ],
    "ASN": [
        "N",
        "CA",
        "C",
        "O",
        "CB",
        "CG",
        "OD1",
        "ND2",
        "",
        "",
        "",
        "",
        "",
        "",
    ],
    "ASP": [
        "N",
        "CA",
        "C",
        "O",
        "CB",
        "CG",
        "OD1",
        "OD2",
        "",
        "",
        "",
        "",
        "",
        "",
    ],
    "CYS": ["N", "CA", "C", "O", "CB", "SG", "", "", "", "", "", "", "", ""],
    "GLN": [
        "N",
        "CA",
        "C",
        "O",
        "CB",
        "CG",
        "CD",
        "OE1",
        "NE2",
        "",
        "",
        "",
        "",
        "",
    ],
    "GLU": [
        "N",
        "CA",
        "C",
        "O",
        "CB",
        "CG",
        "CD",
        "OE1",
        "OE2",
        "",
        "",
        "",
        "",
        "",
    ],
    "GLY": ["N", "CA", "C", "O", "", "", "", "", "", "", "", "", "", ""],
    "HIS": [
        "N",
        "CA",
        "C",
        "O",
        "CB",
        "CG",
        "ND1",
        "CD2",
        "CE1",
        "NE2",
        "",
        "",
        "",
        "",
    ],
    "ILE": [
        "N",
        "CA",
        "C",
        "O",
        "CB",
        "CG1",
        "CG2",
        "CD1",
        "",
        "",
        "",
        "",
        "",
        "",
    ],
    "LEU": [
        "N",
        "CA",
        "C",
        "O",
        "CB",
        "CG",
        "CD1",
        "CD2",
        "",
        "",
        "",
        "",
        "",
        "",
    ],
    "LYS": [
        "N",
        "CA",
        "C",
        "O",
        "CB",
        "CG",
        "CD",
        "CE",
        "NZ",
        "",
        "",
        "",
        "",
        "",
    ],
    "MET": [
        "N",
        "CA",
        "C",
        "O",
        "CB",
        "CG",
        "SD",
        "CE",
        "",
        "",
        "",
        "",
        "",
        "",
    ],
    "PHE": [
        "N",
        "CA",
        "C",
        "O",
        "CB",
        "CG",
        "CD1",
        "CD2",
        "CE1",
        "CE2",
        "CZ",
        "",
        "",
        "",
    ],
    "PRO": ["N", "CA", "C", "O", "CB", "CG", "CD", "", "", "", "", "", "", ""],
    "SER": ["N", "CA", "C", "O", "CB", "OG", "", "", "", "", "", "", "", ""],
    "THR": [
        "N",
        "CA",
        "C",
        "O",
        "CB",
        "OG1",
        "CG2",
        "",
        "",
        "",
        "",
        "",
        "",
        "",
    ],
    "TRP": [
        "N",
        "CA",
        "C",
        "O",
        "CB",
        "CG",
        "CD1",
        "CD2",
        "NE1",
        "CE2",
        "CE3",
        "CZ2",
        "CZ3",
        "CH2",
    ],
    "TYR": [
        "N",
        "CA",
        "C",
        "O",
        "CB",
        "CG",
        "CD1",
        "CD2",
        "CE1",
        "CE2",
        "CZ",
        "OH",
        "",
        "",
    ],
    "VAL": [
        "N",
        "CA",
        "C",
        "O",
        "CB",
        "CG1",
        "CG2",
        "",
        "",
        "",
        "",
        "",
        "",
        "",
    ],
    "UNK": ["", "", "", "", "", "", "", "", "", "", "", "", "", ""],
}
# pylint: enable=line-too-long
# pylint: enable=bad-whitespace

# This is the standard residue order when coding AA type as a number.
# Reproduce it by taking 3-letter AA codes and sorting them alphabetically.
restypes = [
    "A",
    "R",
    "N",
    "D",
    "C",
    "Q",
    "E",
    "G",
    "H",
    "I",
    "L",
    "K",
    "M",
    "F",
    "P",
    "S",
    "T",
    "W",
    "Y",
    "V",
]
restype_order = {restype: i for i, restype in enumerate(restypes)}
restype_num = len(restypes)  # := 20.
unk_restype_index = restype_num  # Catch-all index for unknown restypes.

restypes_with_x = restypes + ["X"]
restype_order_with_x = {restype: i for i, restype in enumerate(restypes_with_x)}

def sequence_to_onehot(
    sequence: str, mapping: Mapping[str, int], map_unknown_to_x: bool = False
) -> np.ndarray:
    """Maps the given sequence into a one-hot encoded matrix.

    Args:
      sequence: An amino acid sequence.
      mapping: A dictionary mapping amino acids to integers.
      map_unknown_to_x: If True, any amino acid that is not in the mapping will be
        mapped to the unknown amino acid 'X'. If the mapping doesn't contain
        amino acid 'X', an error will be thrown. If False, any amino acid not in
        the mapping will throw an error.

    Returns:
      A numpy array of shape (seq_len, num_unique_aas) with one-hot encoding of
      the sequence.

    Raises:
      ValueError: If the mapping doesn't contain values from 0 to
        num_unique_aas - 1 without any gaps.
    """
    num_entries = max(mapping.values()) + 1

    if sorted(set(mapping.values())) != list(range(num_entries)):
        raise ValueError(
            "The mapping must have values from 0 to num_unique_aas-1 "
            "without any gaps. Got: %s" % sorted(mapping.values())
        )

    one_hot_arr = np.zeros((len(sequence), num_entries), dtype=np.int32)

    for aa_index, aa_type in enumerate(sequence):
        if map_unknown_to_x:
            if aa_type.isalpha() and aa_type.isupper():
                aa_id = mapping.get(aa_type, mapping["X"])
            else:
                raise ValueError(
                    f"Invalid character in the sequence: {aa_type}"
                )
        else:
            aa_id = mapping[aa_type]
        one_hot_arr[aa_index, aa_id] = 1

    return one_hot_arr

restype_1to3 = {
    "A": "ALA",
    "R": "ARG",
    "N": "ASN",
    "D": "ASP",
    "C": "CYS",
    "Q": "GLN",
    "E": "GLU",
    "G": "GLY",
    "H": "HIS",
    "I": "ILE",
    "L": "LEU",
    "K": "LYS",
    "M": "MET",
    "F": "PHE",
    "P": "PRO",
    "S": "SER",
    "T": "THR",
    "W": "TRP",
    "Y": "TYR",
    "V": "VAL",
}

# NB: restype_3to1 differs from Bio.PDB.protein_letters_3to1 by being a simple
# 1-to-1 mapping of 3 letter names to one letter names. The latter contains
# many more, and less common, three letter names as keys and maps many of these
# to the same one letter name (including 'X' and 'U' which we don't use here).
restype_3to1 = {v: k for k, v in restype_1to3.items()}

# Define a restype name for all unknown residues.
unk_restype = "UNK"

resnames = [restype_1to3[r] for r in restypes] + [unk_restype]
resname_to_idx = {resname: i for i, resname in enumerate(resnames)}

# The mapping here uses hhblits convention, so that B is mapped to D, J and O
# are mapped to X, U is mapped to C, and Z is mapped to E. Other than that the
# remaining 20 amino acids are kept in alphabetical order.
# There are 2 non-amino acid codes, X (representing any amino acid) and
# "-" representing a missing amino acid in an alignment.  The id for these
# codes is put at the end (20 and 21) so that they can easily be ignored if
# desired.
HHBLITS_AA_TO_ID = {
    "A": 0,
    "B": 2,
    "C": 1,
    "D": 2,
    "E": 3,
    "F": 4,
    "G": 5,
    "H": 6,
    "I": 7,
    "J": 20,
    "K": 8,
    "L": 9,
    "M": 10,
    "N": 11,
    "O": 20,
    "P": 12,
    "Q": 13,
    "R": 14,
    "S": 15,
    "T": 16,
    "U": 1,
    "V": 17,
    "W": 18,
    "X": 20,
    "Y": 19,
    "Z": 3,
    "-": 21,
}

# Partial inversion of HHBLITS_AA_TO_ID.
ID_TO_HHBLITS_AA = {
    0: "A",
    1: "C",  # Also U.
    2: "D",  # Also B.
    3: "E",  # Also Z.
    4: "F",
    5: "G",
    6: "H",
    7: "I",
    8: "K",
    9: "L",
    10: "M",
    11: "N",
    12: "P",
    13: "Q",
    14: "R",
    15: "S",
    16: "T",
    17: "V",
    18: "W",
    19: "Y",
    20: "X",  # Includes J and O.
    21: "-",
}

restypes_with_x_and_gap = restypes + ["X", "-"]
MAP_HHBLITS_AATYPE_TO_OUR_AATYPE = tuple(
    restypes_with_x_and_gap.index(ID_TO_HHBLITS_AA[i])
    for i in range(len(restypes_with_x_and_gap))
)

def _make_standard_atom_mask() -> np.ndarray:
    """Returns [num_res_types, num_atom_types] mask array."""
    # +1 to account for unknown (all 0s).
    mask = np.zeros([restype_num + 1, atom_type_num], dtype=np.int32)
    for restype, restype_letter in enumerate(restypes):
        restype_name = restype_1to3[restype_letter]
        atom_names = residue_atoms[restype_name]
        for atom_name in atom_names:
            atom_type = atom_order[atom_name]
            mask[restype, atom_type] = 1
    return mask

STANDARD_ATOM_MASK = _make_standard_atom_mask()

# A one hot representation for the first and second atoms defining the axis
# of rotation for each chi-angle in each residue.
def chi_angle_atom(atom_index: int) -> np.ndarray:
    """Define chi-angle rigid groups via one-hot representations."""
    chi_angles_index = {}
    one_hots = []

    for k, v in chi_angles_atoms.items():
        indices = [atom_types.index(s[atom_index]) for s in v]
        indices.extend([-1] * (4 - len(indices)))
        chi_angles_index[k] = indices

    for r in restypes:
        res3 = restype_1to3[r]
        one_hot = np.eye(atom_type_num)[chi_angles_index[res3]]
        one_hots.append(one_hot)

    one_hots.append(np.zeros([4, atom_type_num]))  # Add zeros for residue `X`.
    one_hot = np.stack(one_hots, axis=0)
    one_hot = np.transpose(one_hot, [0, 2, 1])

    return one_hot

chi_atom_1_one_hot = chi_angle_atom(1)
chi_atom_2_one_hot = chi_angle_atom(2)

# An array like chi_angles_atoms but using indices rather than names.
chi_angles_atom_indices = [chi_angles_atoms[restype_1to3[r]] for r in restypes]
chi_angles_atom_indices = tree.map_structure(
    lambda atom_name: atom_order[atom_name], chi_angles_atom_indices
)
chi_angles_atom_indices = np.array(
    [
        chi_atoms + ([[0, 0, 0, 0]] * (4 - len(chi_atoms)))
        for chi_atoms in chi_angles_atom_indices
    ]
)

# Mapping from (res_name, atom_name) pairs to the atom's chi group index
# and atom index within that group.
chi_groups_for_atom = collections.defaultdict(list)
for res_name, chi_angle_atoms_for_res in chi_angles_atoms.items():
    for chi_group_i, chi_group in enumerate(chi_angle_atoms_for_res):
        for atom_i, atom in enumerate(chi_group):
            chi_groups_for_atom[(res_name, atom)].append((chi_group_i, atom_i))
chi_groups_for_atom = dict(chi_groups_for_atom)

def _make_rigid_transformation_4x4(ex, ey, translation):
    """Create a rigid 4x4 transformation matrix from two axes and transl."""
    # Normalize ex.
    ex_normalized = ex / np.linalg.norm(ex)

    # make ey perpendicular to ex
    ey_normalized = ey - np.dot(ey, ex_normalized) * ex_normalized
    ey_normalized /= np.linalg.norm(ey_normalized)

    # compute ez as cross product
    eznorm = np.cross(ex_normalized, ey_normalized)
    m = np.stack(
        [ex_normalized, ey_normalized, eznorm, translation]
    ).transpose()
    m = np.concatenate([m, [[0.0, 0.0, 0.0, 1.0]]], axis=0)
    return m

# create an array with (restype, atomtype) --> rigid_group_idx
# and an array with (restype, atomtype, coord) for the atom positions
# and compute affine transformation matrices (4,4) from one rigid group to the
# previous group
restype_atom37_to_rigid_group = np.zeros([21, 37], dtype=np.int)
restype_atom37_mask = np.zeros([21, 37], dtype=np.float32)
restype_atom37_rigid_group_positions = np.zeros([21, 37, 3], dtype=np.float32)
restype_atom14_to_rigid_group = np.zeros([21, 14], dtype=np.int)
restype_atom14_mask = np.zeros([21, 14], dtype=np.float32)
restype_atom14_rigid_group_positions = np.zeros([21, 14, 3], dtype=np.float32)
restype_rigid_group_default_frame = np.zeros([21, 8, 4, 4], dtype=np.float32)

def _make_rigid_group_constants():
    """Fill the arrays above."""
    for restype, restype_letter in enumerate(restypes):
        resname = restype_1to3[restype_letter]
        for atomname, group_idx, atom_position in rigid_group_atom_positions[
            resname
        ]:
            atomtype = atom_order[atomname]
            restype_atom37_to_rigid_group[restype, atomtype] = group_idx
            restype_atom37_mask[restype, atomtype] = 1
            restype_atom37_rigid_group_positions[
                restype, atomtype, :
            ] = atom_position

            atom14idx = restype_name_to_atom14_names[resname].index(atomname)
            restype_atom14_to_rigid_group[restype, atom14idx] = group_idx
            restype_atom14_mask[restype, atom14idx] = 1
            restype_atom14_rigid_group_positions[
                restype, atom14idx, :
            ] = atom_position

    for restype, restype_letter in enumerate(restypes):
        resname = restype_1to3[restype_letter]
        atom_positions = {
            name: np.array(pos)
            for name, _, pos in rigid_group_atom_positions[resname]
        }

        # backbone to backbone is the identity transform
        restype_rigid_group_default_frame[restype, 0, :, :] = np.eye(4)

        # pre-omega-frame to backbone (currently dummy identity matrix)
        restype_rigid_group_default_frame[restype, 1, :, :] = np.eye(4)

        # phi-frame to backbone
        mat = _make_rigid_transformation_4x4(
            ex=atom_positions["N"] - atom_positions["CA"],
            ey=np.array([1.0, 0.0, 0.0]),
            translation=atom_positions["N"],
        )
        restype_rigid_group_default_frame[restype, 2, :, :] = mat

        # psi-frame to backbone
        mat = _make_rigid_transformation_4x4(
            ex=atom_positions["C"] - atom_positions["CA"],
            ey=atom_positions["CA"] - atom_positions["N"],
            translation=atom_positions["C"],
        )
        restype_rigid_group_default_frame[restype, 3, :, :] = mat

        # chi1-frame to backbone
        if chi_angles_mask[restype][0]:
            base_atom_names = chi_angles_atoms[resname][0]
            base_atom_positions = [
                atom_positions[name] for name in base_atom_names
            ]
            mat = _make_rigid_transformation_4x4(
                ex=base_atom_positions[2] - base_atom_positions[1],
                ey=base_atom_positions[0] - base_atom_positions[1],
                translation=base_atom_positions[2],
            )
            restype_rigid_group_default_frame[restype, 4, :, :] = mat

        # chi2-frame to chi1-frame
        # chi3-frame to chi2-frame
        # chi4-frame to chi3-frame
        # luckily all rotation axes for the next frame start at (0,0,0) of the
        # previous frame
        for chi_idx in range(1, 4):
            if chi_angles_mask[restype][chi_idx]:
                axis_end_atom_name = chi_angles_atoms[resname][chi_idx][2]
                axis_end_atom_position = atom_positions[axis_end_atom_name]
                mat = _make_rigid_transformation_4x4(
                    ex=axis_end_atom_position,
                    ey=np.array([-1.0, 0.0, 0.0]),
                    translation=axis_end_atom_position,
                )
                restype_rigid_group_default_frame[
                    restype, 4 + chi_idx, :, :
                ] = mat

_make_rigid_group_constants()

def make_atom14_dists_bounds(
    overlap_tolerance=1.5, bond_length_tolerance_factor=15
):
    """compute upper and lower bounds for bonds to assess violations."""
    restype_atom14_bond_lower_bound = np.zeros([21, 14, 14], np.float32)
    restype_atom14_bond_upper_bound = np.zeros([21, 14, 14], np.float32)
    restype_atom14_bond_stddev = np.zeros([21, 14, 14], np.float32)
    residue_bonds, residue_virtual_bonds, _ = load_stereo_chemical_props()
    for restype, restype_letter in enumerate(restypes):
        resname = restype_1to3[restype_letter]
        atom_list = restype_name_to_atom14_names[resname]

        # create lower and upper bounds for clashes
        for atom1_idx, atom1_name in enumerate(atom_list):
            if not atom1_name:
                continue
            atom1_radius = van_der_waals_radius[atom1_name[0]]
            for atom2_idx, atom2_name in enumerate(atom_list):
                if (not atom2_name) or atom1_idx == atom2_idx:
                    continue
                atom2_radius = van_der_waals_radius[atom2_name[0]]
                lower = atom1_radius + atom2_radius - overlap_tolerance
                upper = 1e10
                restype_atom14_bond_lower_bound[
                    restype, atom1_idx, atom2_idx
                ] = lower
                restype_atom14_bond_lower_bound[
                    restype, atom2_idx, atom1_idx
                ] = lower
                restype_atom14_bond_upper_bound[
                    restype, atom1_idx, atom2_idx
                ] = upper
                restype_atom14_bond_upper_bound[
                    restype, atom2_idx, atom1_idx
                ] = upper

        # overwrite lower and upper bounds for bonds and angles
        for b in residue_bonds[resname] + residue_virtual_bonds[resname]:
            atom1_idx = atom_list.index(b.atom1_name)
            atom2_idx = atom_list.index(b.atom2_name)
            lower = b.length - bond_length_tolerance_factor * b.stddev
            upper = b.length + bond_length_tolerance_factor * b.stddev
            restype_atom14_bond_lower_bound[
                restype, atom1_idx, atom2_idx
            ] = lower
            restype_atom14_bond_lower_bound[
                restype, atom2_idx, atom1_idx
            ] = lower
            restype_atom14_bond_upper_bound[
                restype, atom1_idx, atom2_idx
            ] = upper
            restype_atom14_bond_upper_bound[
                restype, atom2_idx, atom1_idx
            ] = upper
            restype_atom14_bond_stddev[restype, atom1_idx, atom2_idx] = b.stddev
            restype_atom14_bond_stddev[restype, atom2_idx, atom1_idx] = b.stddev
    return {
        "lower_bound": restype_atom14_bond_lower_bound,  # shape (21,14,14)
        "upper_bound": restype_atom14_bond_upper_bound,  # shape (21,14,14)
        "stddev": restype_atom14_bond_stddev,  # shape (21,14,14)
    }

restype_atom14_ambiguous_atoms = np.zeros((21, 14), dtype=np.float32)
restype_atom14_ambiguous_atoms_swap_idx = np.tile(
    np.arange(14, dtype=np.int), (21, 1)
)

def _make_atom14_ambiguity_feats():
    for res, pairs in residue_atom_renaming_swaps.items():
        res_idx = restype_order[restype_3to1[res]]
        for atom1, atom2 in pairs.items():
            atom1_idx = restype_name_to_atom14_names[res].index(atom1)
            atom2_idx = restype_name_to_atom14_names[res].index(atom2)
            restype_atom14_ambiguous_atoms[res_idx, atom1_idx] = 1
            restype_atom14_ambiguous_atoms[res_idx, atom2_idx] = 1
            restype_atom14_ambiguous_atoms_swap_idx[
                res_idx, atom1_idx
            ] = atom2_idx
            restype_atom14_ambiguous_atoms_swap_idx[
                res_idx, atom2_idx
            ] = atom1_idx

_make_atom14_ambiguity_feats()