#!/usr/bin/python

"""
Calculate dN/dS using Nei Gojobori method with Juke-Cantor's multiple-substitution correction (optional), and whole sequence or sliding window. Tested against MATLAB's dnds().
USAGE:

    Run this from command-line, from path (../../):
    hpcleap_dnds/py/scripts/dnds.py <query sequence> <reference sequence>
        For more information on query and reference sequence command-line arguments see ** in DETAILS section below.

REQUIREMENTS:
    scripts:
        - hpcleap_dnds/py/scripts/changes.py (hpcleap_dnds/py/scripts/)
    data:
        - observed_changes.p, potential_changes.p: returned by changes.py
        seq1, seq2: @todo: make cmd or file input arg
DETAILS:
** Query/Reference Input sequences:  both sequences are DNA CDS-transcripts
of interest (exons-only). To give an example, say the "query" sequence
corresponds to the Anopheles gambiae TEP1 protein's transcript id
(VectorBase: AGAP010815), and the other "reference" sequence corresponds
to the orthologous Aedes aegypti TEP1 protein's trancript id (VectorBase:
AAEL001802). The orthologous reference sequence is required, along with
the query sequence of interest, in order to estimate dN/dS values (sliding
 or whole), which requires simulations/approximations of ancestral DNA
substitution mutations (synonymous, non-synonymous changes).
In the main webpage: vg-genes.html, the query sequence is what the user
will request aggregate informations from (after clicking "GO!"), whereas
the orthologous reference sequence is only used to calculate dnds.py,
without the aggregate information requested.
"""
############
# IMPORTS: #
############

# default modules
import sys
import pickle
import math
import warnings

#import matplotlib.pyplot as plt
import numpy as np
# andy-developed modules: imported from py/scripts/ (i.e. "." relative to where this file is)
import changes as codon_pair_data  # /hpcleap_dnds/py/scripts/changes.py
import align as align_then_trim    # /hpcleap_dnds/py/scripts/align.py

# just for testing, @todo: remove testing imports
import pdb
import time
start_time = time.time()  # @time

# @todo: make a check to ensure the input seqs are divisible by 3 (i.e. n_aa_residues = n_dna_residues/3)

# @todo:REMOVE: \/ and \/\/
# with open('../data/observed_changes_dict.p','rb') as f_observed:
#     changes_observed = pickle.load(f_observed)

# @todo:REMOVE: \/ and \/\/: the final script cannot use absolut paths, note: the path from which this script is executed is where current working directory is, in this case it should be a .html that is in {root} calling this script residing in {root}/scripts directory
 # changes_observed = pickle.load(open('/home/qiime/Desktop/hpcleap_wp6_compbio/hpcleap_bioinf/data/observed_changes_dict.p','rb'))
# changes_potential= pickle.load(open('/home/qiime/Desktop/hpcleap_wp6_compbio/hpcleap_bioinf/data/potential_changes_dict.p','rb'))

#############
# FUNCTIONS #
#############

def dnds( seq1, seq2, changes_potential, changes_observed, msCorrect='approximate', sliding=False, windowLength=3, stepLength=1):
    """ Perform dN/dS analysis, using the 'NG' algoritm, includes both whole sequence or sliding window, and either an approximate or exact multiple-substiution correction method. (@todo: make sure it actually is exact... it could be
             something else)

    ARGS:
        seq1,  a DNA sequence as string of letters, AGTC. Seq1 must be equal in length
            to, and aligned with, seq2, with gaps trimmed away. @todo: how on earth can
            we reliably make this work on the web service?

            e.g. seq1 = 'ATGCGCAAATACTCCCCCTTCCGAAATGGATACATGGAACCCACCCTTGGGCAGCACCTCCCAACCCTGTCTTTTCCAGACCCCGGACTCCGGCCCCAAAACCTGTACACCCTCTGGGGAGGCTCCGTTGTCTGCATGTACCTCTACCAGCTTTCCCCCCCCATCACCTGGCCCCTCCTGCCCCATGTGATTTTTTGCCACCCCGGCCAGCTCGGGGCCTTCCTCACCAATGTTCCCTACAAACGAATAGAAAAACTCCTCTATAAAATTTCCCTTACCACAGGGGCCCTAATAATTCTACCCGAGGACTGTTTGCCCACCACCCTTTTCCAGCCTGCTAGGGCACCCGTCACGCTGACAGCCTGGCAAAACGGCCTCCTTCCGTTCCACTCAACCCTCACCACTCCAGGCCTTATTTGGACATTTACCGATGGCACGCCTATGATTTCCGGGCCCTGCCCTAAAGATGGCCAGCCATCTTTAGTACTACAGTCCTCCTCCTTTATATTTCACAAATTTCAAACCAAGGCCTACCACCCCTCATTTCTACTCTCACACGGCCTCATACAGTACTCTTCCTTTCATAATTTGCATCTCCTATTTGAAGAATACACCAACATCCCCATTTCTCTACTTTTTAACGAAAAAGAGGCAGATGACAATGACCATGAGCCCCAAATATCCCCCGGGGGCTTAGAGCCTCTCAGTGAAAAACATTTCCGTGAAACAGAAGTC'

        seq2,  a DNA sequence similar to seq1 but with differences (substitutions),
            representing a CDS orthologue of seq1 from a different species. Read
            description of seq1 for other required similarities to avoid errors.

            e.g. seq2 = 'ATGCGCAAGTACTCCCCCTTCCGAAACGGATACATGGAACCCACCCTTGGGCAACACCTCCCAACCCTGTCTTTTCCAGACCCCGGCCTCCGGCCCCAAAACCTGTACACCCTCTGGGGAGACTCTGTTGTCTGCCTGTACCTCTACCAGCTCTCCCCCCCCATCACCTGGCCCCTCCCGCCCCATGTGATTTTTTGCCACCCCGGCCAGCTCGGGGCCTTCCTCACCAATGTTCCCTACAAGCGTATGGAAGAACTCCTCTATAAAATTTCCCTTACCACAGGGGCCCTAATAATTCTACCCGAGGACTGTTTACCAACCACCCTTTTCCAGCCTGCTAGGGCCCCCGTCACGTTGACCGCCTGGCAGAACGGCCTCCTTCCGTTCCACTCAACCCTCACCACTCCAGGCCTTATTTGGACATTTACCGATGGCACGCCTATGGTTTCCGGACCCTGCCCCAAAGATGGCCAGCCATCTTTAGTACTACAGTCCTCCTCATTTATATTTCACAAATTTCAAACCAAGGCCTACCACCCTTCATTTCTACTCTCACACGGCCTCATACAGTACTCCTCCTTTCACAATTTACATCTCCTTTTTGAAGAATACACCAACATCCCCGTTTCTCTACTTTTTAACGAAAAAGAGGCAAATGACACTGACCATGAGCCCCAAATATCCCCCGGGGGCTTAGAGCCTCCCGCTGAAAAACATTTCCGCGAAACAGAAGTC'
        changes_potential, a dict, with key=pair of codons tuple, e.g. ('ATG','ATG'), and value=['S':<S>,'N':<N>]. Where <S> is the number of potential synonmyous sites for each codon (averaged between the two codons), and <N> is the same but for non-synonymous sites.
            e.g. changes.potential_changes_dict(...)  (see: ./changes.py)
        changes_observed, @todo
        msCorrect, a string to toggle between multiple-substitution correction methods:
            "approximate", "exact" (@todo: make sure it actually is exact... it could be
             something else)

            e.g. msCorrect = 'approximate'
        sliding, a boolean to toggle between sliding window analysis (vector of dN/dS values at successive chunks of sequence) or whole sequence analysis (a single
            dN/dS value for the given pair of input sequences), either: True, False

            e.g. sliding = False
        windowLength, an integer specifying the width of the sliding window, measured in no. of codons in the window to measure dN/dS over, from 1-to-length(seq1)

            e.g. windowLength = 50
        stepLength, an integer specifying no. of codons to shift the sliding window with each iteration. If stepLength < windowLength then windows will overlap, overlapping is dealt with prior to plotting (acts to smooth values, averages along the overlaps are taken as a dN/dS value for any codon).
            e.g. stepLength = 1
    NOTES:
        Sources of formulae:
            http://www.megasoftware.net/mega4/WebHelp/part_iv___evolutionary_analysis/computing_evolutionary_distances/distance_models/synonymouse_and_nonsynonymous_substitution_models/hc_nei_gojobori_method.html
    """

    def chunks(l, n):
        """ Yield successive n-sized chunks from l. """
        for i in xrange(0, len(l), n):
            yield l[i:i+n]

    # todo: stop codons to deal with, reject
    # todo: ambiguous bases to deal with:
        # gaps,  @done
        # Ns,    @todo
        # Xs     @todo

    warning_count = 0

    # STATS per CODON-PAIR:
    codons_seq1   = [codon for codon in chunks(seq1,3)]  #splits
    codons_seq2   = [codon for codon in chunks(seq2,3)]
    codons_paired = [pair for pair in zip(codons_seq1,codons_seq2) if (len(pair[0])+len(pair[1]))==6] # aligned codons are paired into tuples, excess codons are truncated, @todo: in main example, we lose 5 bps of data

    # @done: the next for loop is extremely innefficient, I should set the structure of the changes_potential and changes_observed dicts to how I want it to look, a priori, i.e. when it's instantiated in changes.py.
    # OR just remove this chunk and access the data as they come in the args

    changes_all = {'observed':{'S':[],'N':[]},'potential':{'S':[],'N':[]}}

    for pair in codons_paired:
        changes_all['potential']['S'].append(changes_potential[pair]['S'])
        changes_all['potential']['N'].append(changes_potential[pair]['N'])
        changes_all['observed']['S'].append(changes_observed[pair]['S'])
        changes_all['observed']['N'].append(changes_observed[pair]['N'])

    list_S  = changes_all['potential']['S']
    list_Sd = changes_all['observed']['S']

    list_N  = changes_all['potential']['N']
    list_Nd = changes_all['observed']['N']

    if sliding:
        # STATS for each WINDOW seq
        intervals    = range(0,len(codons_paired)-windowLength+1,stepLength)
        windows      = zip(intervals,[i + windowLength - 1 for i in intervals])

        window_stats = {}

        #window_stats_list = []

        # @done: test against matlab's sliding window, also @todo: find out what stepLength does, @todo: try to plot the sliding window version

        for window_i,window in enumerate(windows):

            start = window[0]
            end   = window[1]+1

            window_stats[window] = {    'S':sum(list_S[start:end]),
                                        'Sd':sum(list_Sd[start:end]),
                                        'N': sum(list_N[start:end]),
                                        'Nd':sum(list_Nd[start:end])    }

            pS = window_stats[window]['Sd']/window_stats[window]['S']
            pN = window_stats[window]['Nd']/window_stats[window]['N']

            try:
                if msCorrect=='approximate':
                    dN = -(3./4.)*math.log(1.-(4./3.)*pN)
                    dS = -(3./4.)*math.log(1.-(4./3.)*pS)

                # @todo: what is this commented code? I don't remember...
                # elif msCorrect=='exact':
                #     d=ln(1-p*4/3)/ln(1-3/(4*N))

                else: # msCorrect=='????'  # @todo: is this the exact one? Or something else?
                    dN = pN
                    dS = pS
                window_stats[window]['dNdS'] = dN/dS
            # @todo: I'm not sure I'm treating the following exceptions in the right way...
            # technically it woud be best to exclude these from downstream analyses?
            # e.g. missing value/datapoint on a plot of dN/dS (y-axis) vs. window interval (x-axis)
            except ZeroDivisionError:
                warning_count += 1
                #warn_msg = "Query and Reference sequences are too divergent. Approximate multiple-substitutions correction cannot be achieved, for window: "+str(window_i)+", dS is zero, leading to a division error when trying dN/dS... try alternative value for argument: msCorrect (e.g. 'exact') OR alternative value for argument: windowLength (e.g. "+str(windowLength+20)+") ...\n"   # # @TODO: uncomment for verbose warning message prints // @ANDY-2017-01-30
                #warnings.warn(warn_msg)  # @TODO: uncomment for verbose warning message prints // @ANDY-2017-01-30
                window_stats[window]['dNdS'] = float('Inf')
            except ValueError:
                warning_count += 1
                #warn_msg="Query and Reference sequences are too divergent. Approximate multiple-substitutions correction cannot be achieved, for window: "+str(window_i)+",  SYNONYMOUS changes per synonymous site, pS>=3/4, log() operation will yeild return undefined... try alternative value for argument: msCorrect (e.g. 'exact') OR alternative value for argument: windowLength (e.g. "+str(windowLength+20)+") ...\n"  # # @TODO: uncomment for verbose warning message prints // @ANDY-2017-01-30
                #warnings.warn(warn_msg)  # @TODO: uncomment for verbose warning message prints // @ANDY-2017-01-30
                window_stats[window]['dNdS'] = float('nan')

        return window_stats,warning_count  # list of dnds per window interval // dict of dnds, key=(<from #base pair>,<to #base pair>), value=<dN/dS of the window specified in the key>
    else:
        # STATS for WHOLE SEQ
        S   = sum(list_S)
        Sd  = sum(list_Sd)
        pS  = Sd/S
        N   = sum(list_N)
        Nd  = sum(list_Nd)
        pN  = Nd/N

        try:
            if msCorrect=='approximate':

                if (pS>=3./4.):
                    raise ValueError("Query and reference sequences are too divergent. Approximate multiple-substitutions correction cannot be achieved, SYNONYMOUS changes per synonymous site, pS>=3/4, log() operation will yeild return undefined. Try alternative value for argument: msCorrect (e.g. 'exact')...")

                if (pN>=3./4.):
                    raise ValueError("Query and reference sequences are too divergent. Approximate multiple-substitutions correction cannot be achieved, NON-SYNONYMOUS changes per synonymous site, pN>=3/4, log() operation will yeild return undefined. Try alternative value for argument: msCorrect (e.g. 'exact')...")

                dS  = -(3./4.)*math.log(1.-((4./3.)*pS))
                dN  = -(3./4.)*math.log(1.-((4./3.)*pN))
                dN_dS = dN/dS

            else: # @todo: is this the exact one? Or something else?

                # @DONE: one day the following three lines of code will error, giving a ZeroDivisionError, this needs to be handled with try
                dS = pS  # i.e. dS = Sd/S
                dN = pN
                dN_dS = dN/dS
        except ValueError:
            warning_count += 1
            warnings.warn("Query and reference sequencea are too divergent. ValueError: Approximate multiple-substitutions correction cannot be achieved: UNKNOWN reason, probably due to illegal numbers in a log() function...\n")
            dN_dS = float("nan")
        except ZeroDivisionError:
            warning_count += 1
            warnings.warn("Query and reference sequences are too divergent. ZeroDiviSionError: Approximate multiple-substitutions correction cannot be achieved: UNKNOWN reason, probably due to illegal numbers in a log() function...\n")
            dN_dS = float('Inf')

        return dN_dS, warning_count  # i.e. omega = dN/dS = (Nd/N)/(Sd/S)

def plot_dnds_sliding(dnds_slide_dict):

    """ Plots sliding dN/dS values (y-axis) along the input sequence's aligned codons (x-axis). If the sliding windows overlap (see below), then plot_dnds_sliding will also average the overlapping dN/dS values for each codon.
        ----     window 1
         ----    window 2
          ----   ...
        ^^^^^^^  take average along each column (codon)
    ARGS:
        dnds_slide_dict,    the output of dnds() if the 'sliding' optional argument is set to True
            e.g. dnds_slide_dict = dnds( s1, s2, potential_changes, observed_changes, msCorrect='approximate', sliding=True, windowLength=50, stepLength=1 )
    RETURNS:
        None,   a plot is generated and written to: py/data/dnds_sliding_test.png
    """
    #import pickle

    # #
    # # dnds_slide_dict has overlapping windows of dnds calculated, windows overlap by stepLength bps, and are uniformly windowLength wide
    # #
    # with open("py/data/dnds_slide_dict.p","r") as fi:
    #     dnds_slide_dict=pickle.load(fi)

    #window_intervals = sorted(dnds_slide_dict.keys()) # @todo: I dont think sorting the windows will make a difference to final result, but it will make it slower, @todo: test this just in case
    window_intervals = dnds_slide_dict.keys() # @todo: I dont think sorting the windows will make a difference to final result, but it will make it slower, @todo: test this just in case

    #
    # vectorize and find max of the window_intervals
    #
    max_window_position = np.amax(window_intervals) # e.g. 243 @done: a whole order of magnitude faster than: max_window_interval = max(window_intervals, key=lambda x: x[1])
    # @todo: ^ doing equivalent operations on np.array() (instead of list) is much faster, so maybe this means we should  np.arrays() in the rest of the code in dnds.py and changes.py
    # @DONE: Q: would finding the max() be faster than sorting then indexing? A: yes

    #
    # Initialise matrix with NaN, each row is the values for a specific window, these will cascade
    #       and overlap to various degrees depending on stepLength and windowLength
    #

    # ----     window 1
    #  ----    window 2
    #   ----   ...
    # ^^^^^^^  take average along each column
    #
    overlap_matrix      = np.empty((len(window_intervals),max_window_position+1))  # @todo: are you sure it's +1? initialize empty matrix, note: entries are not actually NaN, just near-zero
    overlap_matrix[:]   = np.NAN # initiate empty np array with NaN, so later we can mask

    for window_i,window in enumerate(window_intervals):

        start = window[0] # e.g. 0
        end   = window[1] # e.g. 49

        # in the i-th row, fill all elements from the window[0]-th to window[1]-th with the dN/dS value for this window
        overlap_matrix[window_i,start:end+1] = dnds_slide_dict[window]['dNdS'] # @todo: are you sure it's +1? test, keep in mind for these indices it does -1 for the "to" part

    #
    # Mask all non-finite values, to allow proper plotting
    #
    nan_masker              = ~np.isfinite(overlap_matrix) # boolean matrix, True if element is finite, False if element is Inf or NaN
    overlap_matrix_masked   = np.ma.masked_array(overlap_matrix,mask=nan_masker)

    #
    # Columwise averages (i.e. avg all values in each column, leading to a 1D vector of averages)
    #
    overlap_matrix_avg      = overlap_matrix_masked.mean(axis=0)
    #overlap_matrix_avg.mean() # sanity: 0.15584966052233765 (whole seq average: 0.152307100775) not bad!

    #
    # Plot dN/dS along the input sequences
    #
    # plt.plot(overlap_matrix_avg)
    # plt.show()
    # plt.savefig("py/data/dnds_sliding_test.png")
    # avg_matrix = overlap_matrix.mean(axis=1)
    # print overlap_matrix_avg
    return list(overlap_matrix_avg), overlap_matrix_avg.mean()

# @ANDY:code redundancy, can uncomment if we need to load this whole pipeline as a function

def dnds_pipeline(qry_seq_in, ref_seq_in):

    """ Runs the whole dnds pipeline for sliding windows, returns the smoothed vector
    of dN/dS values along the codons of the query sequence.
    Essentially the same code as: if __name__ == "__main__": ...
    """

    #
    # TRY CACHED DATA:
    #   Open dictionaries that have cached computationally intensively
    #   produced results, these are all possible codon pairs with cached statistics for all possible pairs regardless of user input seqs

    print("\t============================================================================") # @todo:REMOVE
    print("\tdN/dS sliding analysis: pre-cached statistics for all possible codon pairs...")
    print("\t============================================================================") # @todo:REMOVE

    #
    try:
    #if (os.path.exists("./py/data/observed_changes_dict.p") and os.path.exists("./py/data/potential_changes_dict.p")):
    # LOAD CACHED (fast)

        #
        # Unpickle codonPair-to-statistics dictionaries
        #
        f1 = open('observed_changes_dict.p','rb')
        observed_changes  = pickle.load(f1)
        f1.close()

        f2 = open('potential_changes_dict.p','rb')
        potential_changes = pickle.load(f2)
        f2.close()

        print("\t\tLOADED!") # @todo:REMOVE

    except IOError:
    #else:
    # CREATE NEW (slow)

        #
        # Create dictionary of codon (DNA-triplet, e.g. "ATG") -to- amino
        #   acid residue (e.g. "M")
        #
        nt_to_aa_dict     = codon_pair_data.geneticCode("standard")

        # alternative 1 {{

        # #@todo:Urgent:debug:2016-11-28: why does the following two lines (commented) lead to silent error? For some reason I HAVE to pickle the data to get it to work... when I just return the dicts directly I get the wrong dN/dS values. The following two lines illustrate this, when uncommented in place of the "# @2:Create" and "# @2:Unpickle" blocks of code...
        # observed_changes  = codon_pair_data.potential_changes_dict(nt_to_aa_dict)
        # potential_changes = codon_pair_data.observed_changes_dict(nt_to_aa_dict)

        #}} 1 alternative 2 {{

        #
        # @2:Create the cached codonPair-to-statistics dictionaries, then pickle
        #
        codon_pair_data.potential_changes_dict(nt_to_aa_dict)
        codon_pair_data.observed_changes_dict(nt_to_aa_dict)

        #
        # @2:Unpickle codonPair-to-statistics dictionaries
        #
        #f1 = open('./py/data/observed_changes_dict.p','rb') # @todo:coderedundancey = bad, wrap these lines into a function and call the function
        f1 = open('observed_changes_dict.p','r') # @todo:coderedundancey = bad, wrap these lines into a function and call the function
        observed_changes  = pickle.load(f1)
        f1.close()

        #f2 = open('./py/data/potential_changes_dict.p','rb')
        f2 = open('potential_changes_dict.p','r')
        potential_changes = pickle.load(f2)
        f2.close()

        # }} alternative 2

        print("\t\tCREATED!") # @todo:REMOVE

    #
    #  Calculate dN/dS
    #

    # @todo:
    #   Q:These are either taken as cmd input args or from a tmp file?
    #   A: No, we take them from the output of align.align_query_vs_reference() and align.trim_gaps_from_aligned_seqs() (align. imported as align_then_trim.)

    # alternative 1 {{

    # Aligned and Trimmed HTLV-1 vs. STLV-1 proteins (MATLAB example to benhmark dnds calculations)
    # s1 = 'ATGCGCAAATACTCCCCCTTCCGAAATGGATACATGGAACCCACCCTTGGGCAGCACCTCCCAACCCTGTCTTTTCCAGACCCCGGACTCCGGCCCCAAAACCTGTACACCCTCTGGGGAGGCTCCGTTGTCTGCATGTACCTCTACCAGCTTTCCCCCCCCATCACCTGGCCCCTCCTGCCCCATGTGATTTTTTGCCACCCCGGCCAGCTCGGGGCCTTCCTCACCAATGTTCCCTACAAACGAATAGAAAAACTCCTCTATAAAATTTCCCTTACCACAGGGGCCCTAATAATTCTACCCGAGGACTGTTTGCCCACCACCCTTTTCCAGCCTGCTAGGGCACCCGTCACGCTGACAGCCTGGCAAAACGGCCTCCTTCCGTTCCACTCAACCCTCACCACTCCAGGCCTTATTTGGACATTTACCGATGGCACGCCTATGATTTCCGGGCCCTGCCCTAAAGATGGCCAGCCATCTTTAGTACTACAGTCCTCCTCCTTTATATTTCACAAATTTCAAACCAAGGCCTACCACCCCTCATTTCTACTCTCACACGGCCTCATACAGTACTCTTCCTTTCATAATTTGCATCTCCTATTTGAAGAATACACCAACATCCCCATTTCTCTACTTTTTAACGAAAAAGAGGCAGATGACAATGACCATGAGCCCCAAATATCCCCCGGGGGCTTAGAGCCTCTCAGTGAAAAACATTTCCGTGAAACAGAA'
    # s2 = 'ATGCGCAAGTACTCCCCCTTCCGAAACGGATACATGGAACCCACCCTTGGGCAACACCTCCCAACCCTGTCTTTTCCAGACCCCGGCCTCCGGCCCCAAAACCTGTACACCCTCTGGGGAGACTCTGTTGTCTGCCTGTACCTCTACCAGCTCTCCCCCCCCATCACCTGGCCCCTCCCGCCCCATGTGATTTTTTGCCACCCCGGCCAGCTCGGGGCCTTCCTCACCAATGTTCCCTACAAGCGTATGGAAGAACTCCTCTATAAAATTTCCCTTACCACAGGGGCCCTAATAATTCTACCCGAGGACTGTTTACCAACCACCCTTTTCCAGCCTGCTAGGGCCCCCGTCACGTTGACCGCCTGGCAGAACGGCCTCCTTCCGTTCCACTCAACCCTCACCACTCCAGGCCTTATTTGGACATTTACCGATGGCACGCCTATGGTTTCCGGACCCTGCCCCAAAGATGGCCAGCCATCTTTAGTACTACAGTCCTCCTCATTTATATTTCACAAATTTCAAACCAAGGCCTACCACCCTTCATTTCTACTCTCACACGGCCTCATACAGTACTCCTCCTTTCACAATTTACATCTCCTTTTTGAAGAATACACCAACATCCCCGTTTCTCTACTTTTTAACGAAAAAGAGGCAAATGACACTGACCATGAGCCCCAAATATCCCCCGGGGGCTTAGAGCCTCCCGCTGAAAAACATTTCCGCGAAACAGAA'

    # }} 1 alternative 2  {{  # @todo: debug:ever since the .p loading bug,
    #                               which is still un-resolved (see: "@todo:Urgent:debug:2016-11-28")
    #                               I think we should test a load( .p ) version just in case and check

    #
    # AGAP010815 (gambiae)
    #     https://www.vectorbase.org/Anopheles_gambiae/Gene/Sequence?db=core;g=AGAP010815;r=3L:11202091-11206882;t=AGAP010815-RA
    qry_seq_raw = qry_seq_in
    # s1_AGAP010815_RA = 'ATGTGGCAGTTCATAAGGTCACGAATATTAACGGTGATAATCTTCATAGGTGCTGCTCATGGGCTACTGGTTGTGGGTCCGAAATTTATACGGGCCAACCAGGAATACACTCTGGTGATCAGCAACTTTAACTCACAGCTAAGCAAAGTGGACCTGCTGTTAAAACTGGAAGGCGAAACTGATAATGGTTTAAGCGTTCTGAACGTTACCAAGATGGTTGACGTGCGACGTAATATGAACCGAATGATCAACTTCAATATGCCTGAGGATCTGACGGCTGGAAACTACAAAATAACTATCGATGGACAGCGTGGCTTCAGCTTTCACAAGGAGGCAGAGCTGGTGTATCTCAGCAAATCGATATCGGGGCTAATACAGGTCGATAAGCCCGTATTTAAACCTGGGGATACGGTGAACTTCCGTGTGATCGTGCTGGACACGGAGCTGAAACCGCCGGCGAGGGTCAAGTCGGTTTATGTAACTATACGAGATCCTCAGCGCAATGTGATTCGCAAATGGTCCACGGCAAAACTGTATGCCGGTGTGTTCGAGAGCGATCTACAGATAGCGCCTACTCCAATGCTCGGGGTCTGGAATATCTCGGTGGAGGTGGAAGGAGAAGAGCTTGTGTCAAAGACGTTTGAGGTGAAGGAGTACGTGTTGTCAACGTTCGACGTGCAGGTCATGCCATCGGTGATTCCACTGGAAGAGCATCAAGCTGTGAATCTTACAATCGAAGCGAACTATCACTTTGGTAAGCCAGTGCAAGGAGTGGCCAAGGTGGAGCTGTACCTAGACGACGATAAGCTAAAACTGAAAAAAGAGCTGACTGTGTACGGAAAGGGCCAGGTAGAGTTGCGCTTTGACAATTTTGCAATGGATGCGGATCAGCAGGATGTACCAGTGAAGGTGTCGTTCGTCGAGCAGTACACAAATCGTACGGTGGTCAAACAGTCACAAATCACGGTATATAGGTATGCGTACCGAGTAGAGTTGATAAAAGAGAGTCCACAGTTTCGTCCGGGACTCCCGTTCAAATGTGCGCTTCAGTTTACACACCATGATGGAACACCGGCTAAAGGCATTAGCGGTAAGGTAGAGGTATCCGATGTACGATTCGAAACGACAACAACGAGTGATAACGATGGATTGATTAAGCTCGAGCTGCAACCAAGTGAGGGTACTGAACAACTCAGTATTCACTTCAATGCTGTTGATGGATTCTTTTTTTATGAAGATGTGAATAAGGTAGAAACGGTTACGGATGCGTATATTAAACTGGAGCTGAAATCACCGCATCAAACGGAACAAATTGATGCGTTTCATGGTGACGTGCACGGAGCGCATGACATTCTTCGTGTACTATGTCATGTCAAAGGGCAATATCATCGATGCAGGATTCATGCGACCCAACAAGCAACCGAAGTACCTGTTGCAGCTGAACGCAACAGAAAAGATGATTCCGAGGGCGAAAATTCTCATCGCTACCGTAGCGGGCCGCACGGTGGTGTACGACTTCGCAGACCTCGATTTCCAAGAGCTTCGCAATAATTTTGATTTAAGCATTGACGAGCAAGAGATCAAGCCGGGACGACAAATCGAGCTGAGCATGTCTGGACGCCCAGGAGCGTACGTTGGGCTGGCCGCGTATGACAAAGCCTTGCTGCTTTTCAACAAGAACCACGACCTGTTCTGGGAGGACATTGGGCAGGTGTTTGATGGGTTCCATGCAATCAATGAGAACGAGTTTGACATATTCCACAGCTTGGGTCTGTTCGCCAGGACATTGGACGATATCTTGTTCGACAGTGCAAATGAAAAGACGGGGCGTAATGCACTGCAGTCAGGCAAGCCGATCGGCAAGCTGGTGTCGTATCGGACGAACTTCCAGGAATCGTGGTTGTGGAAAAATGTTTCCATCGGACGATCGGGAAGTCGCAAGTTGATCGAGGTAGTACCGGACACGACCACCTCCTGGTATCTGACGGGCTTCTCGATCGATCCCGTGTACGGGTTGGGTATCATCAAGAAGCCAATCCAGTTCACAACAGTCCAGCCGTTCTACATCGTAGAGAACTTACCATATTCAATCAAACGAGGCGAAGCGGTTGTGTTGCAGTTTACGCTGTTCAACAACCTTGGAGCGGAGTATATAGCCGATGTGACGCTGTACAATGTGGCCAACCAGACCGAGTTCGTCGGACGTCCAAATACGGATCTCAGCTACACCAAATCCGTGAGCGTTCCTCCAAAAGTTGGTGTGCCAATCTCGTTCCTCATCAAGGCCCGCAAGCTCGGCGAGATGGCGGTTCGTGTAAAGGCTTCGATAATGCTGGGACACGAAACGGACGCCCTGGAAAAGGTAATACGGGTGATGCCTGAAAGTTTGGTGCAGCCGAGAATGGATACACGCTTTTTCTGCTTCGACGATCACAAAAATCAAACGTTTCCGATCAACTTGGACATCAACAAGAAGGCCGACAGTGGATCGACAAAGATTGAGTTTCGACTAAATCCCAATTTGTTGACCACGGTCATCAAGAACCTGGACCATCTTCTCGGCGTTCCGACGGGATGTGGTGAGCAGAATATGGTCAAATTTGTTCCCAACATTTTGGTACTGGATTATTTGCATGCCATCGGGTCGAAAGAACAGCATCTAATCGACAAAGCTACGAATTTGTTGCGTCAAGGATATCAAAACCAGATGCGCTACCGTCAGACGGATGGTTCATTTGGTTTGTGGGAGACTACTAATGGTAGCGTGTTTCTCACCGCGTTCGTTGGCACATCGATGCAAACTGCAGTAAAATACATAAGCGATATTGATGCAGCAATGGTGGAGAAGGCATTGGATTGGTTAGCCTCGAAGCAGCATTTCTCGGGACGGTTTGACAAGGCCGGTGCAGAGTATCACAAAGAAATGCAAGGAGGGTTGCGCAATGGTGTGGCCCTCACATCATATGTGTTGATGGCATTGCTGGAGAATGACATTGCCAAAGCAAAGCACGCAGAGGTGATTCAAAAAGGAATGACCTATCTGAGCAATCAGTTTGGATCCATCAACAATGCATACGACCTATCGATAGCAACCTACGCGATGATGTTGAACGGACACACCATGAAGGAGGAGGCACTCAATAAGCTGATTGATATGTCTTTCATTGATGCTGATAAAAACGAACGGTTCTGGAACACAACGAATCCAATAGAAACCACCGCATATGCTCTGCTGTCGTTTGTGATGGCCGAGAAGTACACAGACGGTATACCGGTCATGAATTGGTTGGTGAATCAACGTTACGTTACCGGTAGCTTTCCGAGCACGCAAGACACGTTTGTGGGGCTGAAAGCGCTGACCAAAATGGCGGAAAAGATATCTCCGTCCCGAAACGACTACACCGTTCAACTGAAGTACAAGAAGAGTGCAAAATACTTCAAAATAAACTCGGAGCAAATTGATGTGGAAAACTTCGTGGATATACCGGAGGACACAAAAAAGCTCGAGATCAATGTGGGGGGCATTGGATTTGGGTTGTTAGAGGTGGTTTATCAATTTAATTTGAATCTCGTCAACTTTGAGAATAGATTCCAACTAGACCTGGAGAAACAGAACACAGGCTCTGACTACGAGCTGAGGCTGAAGGTCTGTGCCAGCTACATACCCCAGCTGACCGACAGACGATCGAACATGGCACTGATTGAGGTAACCTTACCGAGCGGTTACGTGGTTGATCGCAATCCGATCAGCGAGCAGACGAAGGTGAATCCGATTCAGAAAACTGAAATCCGTTACGGTGGCACTTCAGTCGTTTTATACTACGACAATATGGGCAGCGAGCGTAACTGTTTCACCCTGACCGCGTACAGACGCTTTAAGGTCGCATTGAAGCGTCCAGCGTATGTGGTTGTGTATGATTATTATAATACAAATCTGAACGCCATCAAAGTGTACGAAGTGGACAAGCAGAATTTGTGCGAAATCTGTGACGAAGAAGACTGTCCTGCAGAGTGCAAAAAATAG'
    #
    # AAEL001802 (aegypti)
    #     https://www.vectorbase.org/Aedes_aegypti/Gene/Sequence?db=core;g=AAEL001802;r=supercont1.43:685886-717122;t=AAEL001802-RA
    ref_seq_raw = ref_seq_in
    #s2_AAEL001802_RA = 'ATGTCGGTATTCATACAAACGGACAAACCGGTGTATACCCCGGGAGATCTGATACGTTTTCGGGTAATCGTGGTGGATGCTGACACTAGACCTGTGACTAGTATTAAAACGGTAAATATAGCGATCGACGATTCTGCAAAAAATTCCATTCGAAAGTGGCCTTATGCCAAGTTGTTAAACGGCATCTTTGAGTCACAAGTGCAATTAGCTTCTTCGCCTGTTCTTGGCACCTGGATTATCAACGTAACAGCTTCCGACGACATCATTGTCACCAAACAGATAGAAGTTAAGGAATATGTGTTGCCAAAATTTTTCGTGAAAGTTTACCCTTCGGAGGTTCTATTGGGGAAAAATAAGAAGGTTTCTCTTACCTTAGATGCCTATTACACGTTCAAAGAACCCGTCGACGGCAATTACAAAGTTGAGTTATTTTTGGACCATACCAAGAGAAAGCCTGACTTCATAAAAAGTGATCGAATCACCGGTAAAACATCACTTGAGTTTCAATTGAAAAATGAAGTAGACATTGATGGCGACGAGCAGTACACTGATGTCACGGTTGAAGTTGAAGTTGTCGAGGCATTTTCTAATCGCACAGTTAGTATAACTGAGAATATTCCGATTTATCGTCAGCCTTATACCGTGACCCTTCTTCCATCTGCACCATCATTTCGACCAGGAGTTCCATTCAATGTACAAATAGTTGTGAAAGATCAGCTTGGACACCCTCCTGCCGAAGAAAAAGCGGCATCAATTGACCTTACTGTAGAGTTCCATTTGCCCATTGACAGTGACACCAAATCTATCACTGTAGATCTGGACGAGAAAGGAACAGGTCAGCTCACATTAGAGCCCCGCCCAGACGCCCAAGAACTGAAAGTGAACGCTACATATGACTCTCAACAATACGATGTAATTCACGATCCGATACATGGTTTCAGTTCGCAAAGTAAGCAGTACATCACAGTAACTCTGAATCCAAAATACTATAACAACATTAAAGTCGATAAGGACATCGTACTGGACATCTCCTGCACTGAAACAATGACGCACTTCTCGTACATCGTTGTCACCAGAGGAAACATAGTGGAAGCATCGAACGTTCCTGTCAGGATAAAAAAGAAACATTCTCTGAGATTGAAAATGACTTCAAAAATGTCTCCGGAGTCGAGGCTTCTAGTGTACTATACAAACAGGGAGTATCTCATCTTTGATGATATTGAGCTGAAGTTCGATTCGTTCAACAACGACTTCAAATTCGATTTGAACGATGATGAGTATTTTCCAGGGCAATCAGTTTATATCGATGTATACGCTTCAAAGGATTCATACGTTGCGTTCAGTGGAATCGATGAAAGTGTACTCCTGGTAGGCAAAGAGCGCCATGACTTCAACAAAGGAGATGTGCTCAAGGAACTCGCTCTTTACGGAGCAACAAATGATGCCGAGTTTGACTTGTTCCACGTAAGTTTCATGTCAAATGGTTTGATTATTCCAGTTAATGTATCTGTAACTCGCTCACAGAATGCACGATTTGGTACTCTACTAGGAAGGACTAGGCAGCAAGCGATTGAAATTCGAACTCAATTCCTAGAATCCTGGTTATGGAAATCCTTTTCCATGGATGGTCGAAACAACTTCAAAGCAATAGAAGACTCGGTTCCGGATACTATTACAACGTATCACGTGTCAGGATTTGCTTTAAGTCCAACACTAGGTCTTGGAGTAATCCAACAACCAGTGAGTTTCACCGTTCGTAAAAAATTCTACTTGGTTGCAAATTTGCCTTACTCGATCAAACGGGGTGAAGTGGCGTTGATTCAGGTTACCGTCTTCAACTTCCTAGGAAGCAGCATAACAACCGATGTGACGCTGTTCAATAAACGCGATGAAATTGAGTTTGTCGAGAATGCATCCACTAATAATACACATCGAACAAAGGCGGTAATTGTCCCGAATAACAATGGAAAATCTGTATCATTTATGGTGAAAGCAAAGAAATTAGGACAGATTGCGATCAAATTCCAGGCGGTAAACCTGCTGGAAACGGATGCATTGGAGCACATGTTACGAGTAACCCCAGAGAGCCATCGCTATGAGAAAAATGTAGCTCGATTCGTTGAGCTACCAAAGTTTGAGACGCAAACTTTCGATGTGAAGCTGGACATTCCCAAAAATATCGACGAGGGTTCTGCTCAAATCAAATTCACGTTAGACCCGGACATTTTGGGAACAGCCATCAGCAACCTAGACGGGTTGATCCGGAAACCCTTTGGATGTGGCGAACAAAATATGCTCCATTTTGTGCCAAATATAGTCGTTTTGGATTATCTTAACGAAACCAACACAGCGGCAGAAGATGTGAGGACCAAAGCGATAAATTTTCTTAGCAGCGGATATCAAAACCAGCTACGCTACAAACGTTCGGATGGGGCCTTCAGTGTCTGGGGACAATCGTATGCTGGCAGTACATTTTTGACGGCCTTTGTGGCGAAATCATTCAAAATAGCAGCCAAATACATTCAGGTGGATAAGTCTATAGTAGACGCGGCATTCGACTGGTTAGTGAAACAACAACAATCAGATGGGCGGTTCCCAGAAGTGGGGCAAGTATTCCAAGCAGATATGCAGGGTGGGCTTCGTAATAACGGTTTTGCGCTTACCGCGTATGTTCTGATCGCTTTTGCTGAAAATAAGGAAGTATACAGAAAATACCAATCACAACTGAACAAAACTACTAACTTCATAGCAGATAGACTTGCTAATATGGAGAATCCATACGACCTCTCGCTGTCCACTTATGCGTTGATGCTAACAAATCATGGCAAGCGCACCGAGTTTCTTCACAAATTAGTCGAAAAGTCGATATTTGACCGCAATCAAACTGAGAGATATTGGGACAGCAAACCAGTTGATATTGAAGTTGCTGGATATGCTCTATTGTCATACGTAGCTGCCGGTAAATTATTGGATGCAACGCCTATCATGCGGTGGCTCAACAAGCAGCGTTATGGTCTCGGAGGCTATCCTGGAACTCAGGAAACATTCGTTGGATTGAAAGCATTGGCAACGTTCGCTGCAAATGTAACTAGTAGGAGAAACGAATATACTGTAAGGATATTCTACGAACCAAATGGTCGACGAACATTCGACGTACACATGCACAATTCGTTTAATATTCAAGAGCTTGACATTCCTAATAACATCAGAAAAATGAAGGTGGAAGTTGAAGGCATCGGCAGAGGCTTCTTCCAAGTGGCATATCAGTACTATCAAAATATGCAGGTGGCTAAGCCCAGTTTCAGCATTACAATTAATCAGCTTAACACCACGACGGAACACATGCAGCAATTGGACGTGTGTGTGAAATACATACCAAAAGAGGCTTATCAAAAATCGAATATGGCTTTGGTGGAAATATTCTTGCCTAGTGGGCTTGTAGCAGACTCAGATGCCATTACGGACAAGACTGGAGGAATTCGAAGAATTGAAAGACGTTTTTCGGACACCTCAGTAGTTATATATTATGATAATTTGGACCCCGAAGACAAGTGCTTCCGAGTGACTGCTTATCGTCGGTATAAAATTGCATTGCATTTGCCATCATATATTATAGTTTATGATTATTATAATTTTGAGCGCTTTGCCATTCAAAAGTACGAAGGAAAGGTGCTGCAGCTCTGCGATATTTGTGAAGACGAGGACTGCGAAACTTTATCATGTCAAAATAGCTCGAAATTGGCAATAATGTAA'

    ###########################
    # Benchmarking vs. MATLAB #
    ###########################

    print("\t=========================================")
    print("\tProcessing heuristic codon alignments....")
    print("\t=========================================")

    # @done:@testing:benchmark vs. matlab: @todo: place in README these details
    #   using s1_AGAP010815_RA and s2_AAEL001802_RA (after aligning/trimming)
    # MATLAB whole dnds vs. andy-wenping whole dnds (msCorrect='approximate'): dN/dS:
    #   1.0958 vs. 1.04939841193 (Elapsed time: 1.952491045, Total warnings: 0), respectively
    # MATLAB mean(sliding dnds) vs. andy-wenping mean(sliding dnds): Mean dN/dS over all 50-codon long windows:
    #   1.7254581782566112 (smoothing?) vs. 1.49715496621 (incl. smoothing, stepLength=1) (Elapsed time: 2.29854488373, Total warnings: 173) (when taking the mean of all 50 window intervals in MATLAB's sliding dnds()
    # MATLAB sliding dnds plot vs. andy-wenping sliding dnds plot: similar sliding dnds plots:
    #   see:  D:\Dropbox\00_HPC-LEAP\WP5 Thematic Cyprus Computational Biology\giannis_database

    # @benchmark:matlab's dnds()
    # dN/dS=0.15270463083614955 with msCorrect='approximate', dN/dS=0.16529268957638238 with
    # msCorrect='exact', MATLAB gives: dN/dS=0.0221=dnds(seq1,seq2, 'geneticCode', 1, 'Method', 'NG'). This is a huge error, an order of magnitude off.
    # seq1 = 'ATGCGCAAATACTCCCCCTTCCGAAATGGATACATGGAACCCACCCTTGGGCAGCACCTCCCAACCCTGTCTTTTCCAGACCCCGGACTCCGGCCCCAAAACCTGTACACCCTCTGGGGAGGCTCCGTTGTCTGCATGTACCTCTACCAGCTTTCCCCCCCCATCACCTGGCCCCTCCTGCCCCATGTGATTTTTTGCCACCCCGGCCAGCTCGGGGCCTTCCTCACCAATGTTCCCTACAAACGAATAGAAAAACTCCTCTATAAAATTTCCCTTACCACAGGGGCCCTAATAATTCTACCCGAGGACTGTTTGCCCACCACCCTTTTCCAGCCTGCTAGGGCACCCGTCACGCTGACAGCCTGGCAAAACGGCCTCCTTCCGTTCCACTCAACCCTCACCACTCCAGGCCTTATTTGGACATTTACCGATGGCACGCCTATGATTTCCGGGCCCTGCCCTAAAGATGGCCAGCCATCTTTAGTACTACAGTCCTCCTCCTTTATATTTCACAAATTTCAAACCAAGGCCTACCACCCCTCATTTCTACTCTCACACGGCCTCATACAGTACTCTTCCTTTCATAATTTGCATCTCCTATTTGAAGAATACACCAACATCCCCATTTCTCTACTTTTTAACGAAAAAGAGGCAGATGACAATGACCATGAGCCCCAAATATCCCCCGGGGGCTTAGAGCCTCTCAGTGAAAAACATTTCCGTGAAACAGAAGTC'
    # seq2 = 'ATGCGCAAGTACTCCCCCTTCCGAAACGGATACATGGAACCCACCCTTGGGCAACACCTCCCAACCCTGTCTTTTCCAGACCCCGGCCTCCGGCCCCAAAACCTGTACACCCTCTGGGGAGACTCTGTTGTCTGCCTGTACCTCTACCAGCTCTCCCCCCCCATCACCTGGCCCCTCCCGCCCCATGTGATTTTTTGCCACCCCGGCCAGCTCGGGGCCTTCCTCACCAATGTTCCCTACAAGCGTATGGAAGAACTCCTCTATAAAATTTCCCTTACCACAGGGGCCCTAATAATTCTACCCGAGGACTGTTTACCAACCACCCTTTTCCAGCCTGCTAGGGCCCCCGTCACGTTGACCGCCTGGCAGAACGGCCTCCTTCCGTTCCACTCAACCCTCACCACTCCAGGCCTTATTTGGACATTTACCGATGGCACGCCTATGGTTTCCGGACCCTGCCCCAAAGATGGCCAGCCATCTTTAGTACTACAGTCCTCCTCATTTATATTTCACAAATTTCAAACCAAGGCCTACCACCCTTCATTTCTACTCTCACACGGCCTCATACAGTACTCCTCCTTTCACAATTTACATCTCCTTTTTGAAGAATACACCAACATCCCCGTTTCTCTACTTTTTAACGAAAAAGAGGCAAATGACACTGACCATGAGCCCCAAATATCCCCCGGGGGCTTAGAGCCTCCCGCTGAAAAACATTTCCGCGAAACAGAAGTC'
    # dnds('ATGCGCAAATACTCCCCCTTCCGAAATGGATACATGGAACCCACCCTTGGGCAGCACCTCCCAACCCTGTCTTTTCCAGACCCCGGACTCCGGCCCCAAAACCTGTACACCCTCTGGGGAGGCTCCGTTGTCTGCATGTACCTCTACCAGCTTTCCCCCCCCATCACCTGGCCCCTCCTGCCCCATGTGATTTTTTGCCACCCCGGCCAGCTCGGGGCCTTCCTCACCAATGTTCCCTACAAACGAATAGAAAAACTCCTCTATAAAATTTCCCTTACCACAGGGGCCCTAATAATTCTACCCGAGGACTGTTTGCCCACCACCCTTTTCCAGCCTGCTAGGGCACCCGTCACGCTGACAGCCTGGCAAAACGGCCTCCTTCCGTTCCACTCAACCCTCACCACTCCAGGCCTTATTTGGACATTTACCGATGGCACGCCTATGATTTCCGGGCCCTGCCCTAAAGATGGCCAGCCATCTTTAGTACTACAGTCCTCCTCCTTTATATTTCACAAATTTCAAACCAAGGCCTACCACCCCTCATTTCTACTCTCACACGGCCTCATACAGTACTCTTCCTTTCATAATTTGCATCTCCTATTTGAAGAATACACCAACATCCCCATTTCTCTACTTTTTAACGAAAAAGAGGCAGATGACAATGACCATGAGCCCCAAATATCCCCCGGGGGCTTAGAGCCTCTCAGTGAAAAACATTTCCGTGAAACAGAAGTC', 'ATGCGCAAGTACTCCCCCTTCCGAAACGGATACATGGAACCCACCCTTGGGCAACACCTCCCAACCCTGTCTTTTCCAGACCCCGGCCTCCGGCCCCAAAACCTGTACACCCTCTGGGGAGACTCTGTTGTCTGCCTGTACCTCTACCAGCTCTCCCCCCCCATCACCTGGCCCCTCCCGCCCCATGTGATTTTTTGCCACCCCGGCCAGCTCGGGGCCTTCCTCACCAATGTTCCCTACAAGCGTATGGAAGAACTCCTCTATAAAATTTCCCTTACCACAGGGGCCCTAATAATTCTACCCGAGGACTGTTTACCAACCACCCTTTTCCAGCCTGCTAGGGCCCCCGTCACGTTGACCGCCTGGCAGAACGGCCTCCTTCCGTTCCACTCAACCCTCACCACTCCAGGCCTTATTTGGACATTTACCGATGGCACGCCTATGGTTTCCGGACCCTGCCCCAAAGATGGCCAGCCATCTTTAGTACTACAGTCCTCCTCATTTATATTTCACAAATTTCAAACCAAGGCCTACCACCCTTCATTTCTACTCTCACACGGCCTCATACAGTACTCCTCCTTTCACAATTTACATCTCCTTTTTGAAGAATACACCAACATCCCCGTTTCTCTACTTTTTAACGAAAAAGAGGCAAATGACACTGACCATGAGCCCCAAATATCCCCCGGGGGCTTAGAGCCTCCCGCTGAAAAACATTTCCGCGAAACAGAAGTC')

    #
    # ALIGN REF vs. QUERY seq, then TRIM GAPS:
    #     Next we want to remove all characters (trim) in each of the sequences that is aligned to a gap position**** in the top alignment
    #

    #  Alignment:
    #     ATGTGC----TAA  qry_seq
    #     ATG--A--TTTGA  ref_seq
    #        ^^ ^^^^     gap positions****
    #  Trimming:
    #     ATGCTAA        qry_seq (after trimming)
    #     ATGATGA        ref_seq (after trimming)

    # alignment scores are calculated from these two parameters, best scoring alignments have few gaps and gaps are small
    aln_gap_open = -10
    aln_gap_extend = -0.5 # @todo: make cmd args later?

    qry_seq_aln, ref_seq_aln                          = align_then_trim.align_query_vs_reference(qry_seq_raw, ref_seq_raw, aln_gap_open, aln_gap_extend)
    qry_seq_trimmed, ref_seq_trimmed, qry_seq_indices = align_then_trim.trim_gaps_from_aligned_seqs( qry_seq_aln, ref_seq_aln )

    # @Note: benhamrking: it is qry_seq_trimmed and ref_seq_trimmed that should be used as input to MATLAB's dnds() function, for benchmarking

    print("\t\tCOMPLETE!")

    #}} alternative 2

    # ///

    # alternative 1 {{

    print("\t===========================")
    print("\tSliding window analysis....")
    print("\t===========================")

    # # @NOTE:uncomment below to achieve dnds of 0.15.. or 0.164 if using exact method, interestingly
    # dnds_whole, warning_count = dnds( qry_seq_trimmed, ref_seq_trimmed, potential_changes, observed_changes, msCorrect='approximate', sliding=False)
    # print "dN/dS: "+str(dnds_whole)

    # }} 1 alternative 2 {{

    # @DONE: work with the sliding window version instead of dnds_whole
    dnds_slide_dict, warning_count = dnds( qry_seq_trimmed, ref_seq_trimmed, potential_changes, observed_changes, msCorrect='approximate', sliding=True, windowLength=25, stepLength=1 )
    #
    # Plot the sliding window values
    #
    dnds_sliding_vec, dnds_sliding_mean = plot_dnds_sliding(dnds_slide_dict)

    print("\t\tCOMPLETE!")

    #
    # Summary statistics
    #

    print("\t===============================================")
    print("\tElapsed time: "+str(time.time() - start_time))   # @time
    print("\tTotal warnings (missing values): "+str(warning_count))
    print "\tAvg. dN/dS over all windows: "+str(dnds_sliding_mean)
    print("\t===============================================")
    # @todo: show the user also the whole dnds value somewhere in the webpage
    print("\t\tdN/dS Sliding Analysis: All Jobs Complete!")

    return dnds_sliding_vec, qry_seq_indices

########
# MAIN # (still testing, cmd-args-mode not available yet)
########
if __name__ == "__main__":

    #
    # TRY CACHED DATA:
    #   Open dictionaries that have cached computationally intensively
    #   produced results
    #
    try:
    #if (os.path.exists("./py/data/observed_changes_dict.p") and os.path.exists("./py/data/potential_changes_dict.p")):
    # LOAD CACHED (fast)

        #
        # Unpickle codonPair-to-statistics dictionaries
        #
        f1 = open('observed_changes_dict.p','rb')
        observed_changes  = pickle.load(f1)
        f1.close()

        f2 = open('potential_changes_dict.p','rb')
        potential_changes = pickle.load(f2)
        f2.close()
        print "LOADED!!" # @todo:REMOVE

    except IOError:
    #else:
    # CREATE NEW (slow)

        #
        # Create dictionary of codon (DNA-triplet, e.g. "ATG") -to- amino
        #   acid residue (e.g. "M")
        #
        nt_to_aa_dict     = codon_pair_data.geneticCode("standard")

        # alternative 1 {{

        # #@todo:Urgent:debug:2016-11-28: why does the following two lines (commented) lead to silent error? For some reason I HAVE to pickle the data to get it to work... when I just return the dicts directly I get the wrong dN/dS values. The following two lines illustrate this, when uncommented in place of the "# @2:Create" and "# @2:Unpickle" blocks of code...
        # observed_changes  = codon_pair_data.potential_changes_dict(nt_to_aa_dict)
        # potential_changes = codon_pair_data.observed_changes_dict(nt_to_aa_dict)

        #}} 1 alternative 2 {{

        #
        # @2:Create the cached codonPair-to-statistics dictionaries, then pickle
        #
        codon_pair_data.potential_changes_dict(nt_to_aa_dict)
        codon_pair_data.observed_changes_dict(nt_to_aa_dict)

        #
        # @2:Unpickle codonPair-to-statistics dictionaries
        #
        f1 = open('observed_changes_dict.p','rb') # @todo:coderedundancey = bad, wrap these lines into a function and call the function
        observed_changes  = pickle.load(f1)
        f1.close()

        f2 = open('potential_changes_dict.p','rb')
        potential_changes = pickle.load(f2)
        f2.close()

        # }} alternative 2

        print "CREATED!!" # @todo:REMOVE

    #
    #  Calculate dN/dS
    #

    # @todo:
    #   Q:These are either taken as cmd input args or from a tmp file?
    #   A: No, we take them from the output of align.align_query_vs_reference() and align.trim_gaps_from_aligned_seqs() (align. imported as align_then_trim.)

    # alternative 1 {{

    # Aligned and Trimmed HTLV-1 vs. STLV-1 proteins (MATLAB example to benhmark dnds calculations)
    # s1 = 'ATGCGCAAATACTCCCCCTTCCGAAATGGATACATGGAACCCACCCTTGGGCAGCACCTCCCAACCCTGTCTTTTCCAGACCCCGGACTCCGGCCCCAAAACCTGTACACCCTCTGGGGAGGCTCCGTTGTCTGCATGTACCTCTACCAGCTTTCCCCCCCCATCACCTGGCCCCTCCTGCCCCATGTGATTTTTTGCCACCCCGGCCAGCTCGGGGCCTTCCTCACCAATGTTCCCTACAAACGAATAGAAAAACTCCTCTATAAAATTTCCCTTACCACAGGGGCCCTAATAATTCTACCCGAGGACTGTTTGCCCACCACCCTTTTCCAGCCTGCTAGGGCACCCGTCACGCTGACAGCCTGGCAAAACGGCCTCCTTCCGTTCCACTCAACCCTCACCACTCCAGGCCTTATTTGGACATTTACCGATGGCACGCCTATGATTTCCGGGCCCTGCCCTAAAGATGGCCAGCCATCTTTAGTACTACAGTCCTCCTCCTTTATATTTCACAAATTTCAAACCAAGGCCTACCACCCCTCATTTCTACTCTCACACGGCCTCATACAGTACTCTTCCTTTCATAATTTGCATCTCCTATTTGAAGAATACACCAACATCCCCATTTCTCTACTTTTTAACGAAAAAGAGGCAGATGACAATGACCATGAGCCCCAAATATCCCCCGGGGGCTTAGAGCCTCTCAGTGAAAAACATTTCCGTGAAACAGAA'
    # s2 = 'ATGCGCAAGTACTCCCCCTTCCGAAACGGATACATGGAACCCACCCTTGGGCAACACCTCCCAACCCTGTCTTTTCCAGACCCCGGCCTCCGGCCCCAAAACCTGTACACCCTCTGGGGAGACTCTGTTGTCTGCCTGTACCTCTACCAGCTCTCCCCCCCCATCACCTGGCCCCTCCCGCCCCATGTGATTTTTTGCCACCCCGGCCAGCTCGGGGCCTTCCTCACCAATGTTCCCTACAAGCGTATGGAAGAACTCCTCTATAAAATTTCCCTTACCACAGGGGCCCTAATAATTCTACCCGAGGACTGTTTACCAACCACCCTTTTCCAGCCTGCTAGGGCCCCCGTCACGTTGACCGCCTGGCAGAACGGCCTCCTTCCGTTCCACTCAACCCTCACCACTCCAGGCCTTATTTGGACATTTACCGATGGCACGCCTATGGTTTCCGGACCCTGCCCCAAAGATGGCCAGCCATCTTTAGTACTACAGTCCTCCTCATTTATATTTCACAAATTTCAAACCAAGGCCTACCACCCTTCATTTCTACTCTCACACGGCCTCATACAGTACTCCTCCTTTCACAATTTACATCTCCTTTTTGAAGAATACACCAACATCCCCGTTTCTCTACTTTTTAACGAAAAAGAGGCAAATGACACTGACCATGAGCCCCAAATATCCCCCGGGGGCTTAGAGCCTCCCGCTGAAAAACATTTCCGCGAAACAGAA'

    # }} 1 alternative 2  {{  # @todo: debug:ever since the .p loading bug,
    #                               which is still un-resolved (see: "@todo:Urgent:debug:2016-11-28")
    #                               I think we should test a load( .p ) version just in case and check

    #
    # AGAP010815 (gambiae)
    #     https://www.vectorbase.org/Anopheles_gambiae/Gene/Sequence?db=core;g=AGAP010815;r=3L:11202091-11206882;t=AGAP010815-RA
    qry_seq_raw = sys.argv[1]
    # OBP7 gambiae: v
    #qry_seq_raw = "ATGTGTGAATATTCGAATACGCGCAACAAGATGAGCAACCTGGTCGTCGTCCTCGTCCTGCTGACGATGTACATTGTGCTTTCGGCCCCATTCGAAATACCGGACCGGTACAAAAAGCCGGCTAAAATGTTGCACGAAATTTGTATCGCCGAGTCGGGCGCCTCGGAGGAGCAGCTGCGCACCTGTCTCGATGGAACCGTACCGACAGCTCCGGCCGCCAAGTGCTACATCCACTGCCTGTTCGACAAGATCGACGTGGTGGACGAGGCGACTGGGCGCATCCTGCTCGACCGACTGCTTTACATCATCCCGGACGACGTGAAGGCAGCGGTGGACCATTTAACGCGCGAATGTAGCCACATCGTAACGCCGGATAAGTGCGAAACCGCCTACGAGACGGTCAAATGTTATTTCAATGCGCACGACGAGGTGATCAAATTCTGCCACCTACTAGTGCTGGAGTGA"
    #s1_AGAP010815_RA = 'ATGTGGCAGTTCATAAGGTCACGAATATTAACGGTGATAATCTTCATAGGTGCTGCTCATGGGCTACTGGTTGTGGGTCCGAAATTTATACGGGCCAACCAGGAATACACTCTGGTGATCAGCAACTTTAACTCACAGCTAAGCAAAGTGGACCTGCTGTTAAAACTGGAAGGCGAAACTGATAATGGTTTAAGCGTTCTGAACGTTACCAAGATGGTTGACGTGCGACGTAATATGAACCGAATGATCAACTTCAATATGCCTGAGGATCTGACGGCTGGAAACTACAAAATAACTATCGATGGACAGCGTGGCTTCAGCTTTCACAAGGAGGCAGAGCTGGTGTATCTCAGCAAATCGATATCGGGGCTAATACAGGTCGATAAGCCCGTATTTAAACCTGGGGATACGGTGAACTTCCGTGTGATCGTGCTGGACACGGAGCTGAAACCGCCGGCGAGGGTCAAGTCGGTTTATGTAACTATACGAGATCCTCAGCGCAATGTGATTCGCAAATGGTCCACGGCAAAACTGTATGCCGGTGTGTTCGAGAGCGATCTACAGATAGCGCCTACTCCAATGCTCGGGGTCTGGAATATCTCGGTGGAGGTGGAAGGAGAAGAGCTTGTGTCAAAGACGTTTGAGGTGAAGGAGTACGTGTTGTCAACGTTCGACGTGCAGGTCATGCCATCGGTGATTCCACTGGAAGAGCATCAAGCTGTGAATCTTACAATCGAAGCGAACTATCACTTTGGTAAGCCAGTGCAAGGAGTGGCCAAGGTGGAGCTGTACCTAGACGACGATAAGCTAAAACTGAAAAAAGAGCTGACTGTGTACGGAAAGGGCCAGGTAGAGTTGCGCTTTGACAATTTTGCAATGGATGCGGATCAGCAGGATGTACCAGTGAAGGTGTCGTTCGTCGAGCAGTACACAAATCGTACGGTGGTCAAACAGTCACAAATCACGGTATATAGGTATGCGTACCGAGTAGAGTTGATAAAAGAGAGTCCACAGTTTCGTCCGGGACTCCCGTTCAAATGTGCGCTTCAGTTTACACACCATGATGGAACACCGGCTAAAGGCATTAGCGGTAAGGTAGAGGTATCCGATGTACGATTCGAAACGACAACAACGAGTGATAACGATGGATTGATTAAGCTCGAGCTGCAACCAAGTGAGGGTACTGAACAACTCAGTATTCACTTCAATGCTGTTGATGGATTCTTTTTTTATGAAGATGTGAATAAGGTAGAAACGGTTACGGATGCGTATATTAAACTGGAGCTGAAATCACCGCATCAAACGGAACAAATTGATGCGTTTCATGGTGACGTGCACGGAGCGCATGACATTCTTCGTGTACTATGTCATGTCAAAGGGCAATATCATCGATGCAGGATTCATGCGACCCAACAAGCAACCGAAGTACCTGTTGCAGCTGAACGCAACAGAAAAGATGATTCCGAGGGCGAAAATTCTCATCGCTACCGTAGCGGGCCGCACGGTGGTGTACGACTTCGCAGACCTCGATTTCCAAGAGCTTCGCAATAATTTTGATTTAAGCATTGACGAGCAAGAGATCAAGCCGGGACGACAAATCGAGCTGAGCATGTCTGGACGCCCAGGAGCGTACGTTGGGCTGGCCGCGTATGACAAAGCCTTGCTGCTTTTCAACAAGAACCACGACCTGTTCTGGGAGGACATTGGGCAGGTGTTTGATGGGTTCCATGCAATCAATGAGAACGAGTTTGACATATTCCACAGCTTGGGTCTGTTCGCCAGGACATTGGACGATATCTTGTTCGACAGTGCAAATGAAAAGACGGGGCGTAATGCACTGCAGTCAGGCAAGCCGATCGGCAAGCTGGTGTCGTATCGGACGAACTTCCAGGAATCGTGGTTGTGGAAAAATGTTTCCATCGGACGATCGGGAAGTCGCAAGTTGATCGAGGTAGTACCGGACACGACCACCTCCTGGTATCTGACGGGCTTCTCGATCGATCCCGTGTACGGGTTGGGTATCATCAAGAAGCCAATCCAGTTCACAACAGTCCAGCCGTTCTACATCGTAGAGAACTTACCATATTCAATCAAACGAGGCGAAGCGGTTGTGTTGCAGTTTACGCTGTTCAACAACCTTGGAGCGGAGTATATAGCCGATGTGACGCTGTACAATGTGGCCAACCAGACCGAGTTCGTCGGACGTCCAAATACGGATCTCAGCTACACCAAATCCGTGAGCGTTCCTCCAAAAGTTGGTGTGCCAATCTCGTTCCTCATCAAGGCCCGCAAGCTCGGCGAGATGGCGGTTCGTGTAAAGGCTTCGATAATGCTGGGACACGAAACGGACGCCCTGGAAAAGGTAATACGGGTGATGCCTGAAAGTTTGGTGCAGCCGAGAATGGATACACGCTTTTTCTGCTTCGACGATCACAAAAATCAAACGTTTCCGATCAACTTGGACATCAACAAGAAGGCCGACAGTGGATCGACAAAGATTGAGTTTCGACTAAATCCCAATTTGTTGACCACGGTCATCAAGAACCTGGACCATCTTCTCGGCGTTCCGACGGGATGTGGTGAGCAGAATATGGTCAAATTTGTTCCCAACATTTTGGTACTGGATTATTTGCATGCCATCGGGTCGAAAGAACAGCATCTAATCGACAAAGCTACGAATTTGTTGCGTCAAGGATATCAAAACCAGATGCGCTACCGTCAGACGGATGGTTCATTTGGTTTGTGGGAGACTACTAATGGTAGCGTGTTTCTCACCGCGTTCGTTGGCACATCGATGCAAACTGCAGTAAAATACATAAGCGATATTGATGCAGCAATGGTGGAGAAGGCATTGGATTGGTTAGCCTCGAAGCAGCATTTCTCGGGACGGTTTGACAAGGCCGGTGCAGAGTATCACAAAGAAATGCAAGGAGGGTTGCGCAATGGTGTGGCCCTCACATCATATGTGTTGATGGCATTGCTGGAGAATGACATTGCCAAAGCAAAGCACGCAGAGGTGATTCAAAAAGGAATGACCTATCTGAGCAATCAGTTTGGATCCATCAACAATGCATACGACCTATCGATAGCAACCTACGCGATGATGTTGAACGGACACACCATGAAGGAGGAGGCACTCAATAAGCTGATTGATATGTCTTTCATTGATGCTGATAAAAACGAACGGTTCTGGAACACAACGAATCCAATAGAAACCACCGCATATGCTCTGCTGTCGTTTGTGATGGCCGAGAAGTACACAGACGGTATACCGGTCATGAATTGGTTGGTGAATCAACGTTACGTTACCGGTAGCTTTCCGAGCACGCAAGACACGTTTGTGGGGCTGAAAGCGCTGACCAAAATGGCGGAAAAGATATCTCCGTCCCGAAACGACTACACCGTTCAACTGAAGTACAAGAAGAGTGCAAAATACTTCAAAATAAACTCGGAGCAAATTGATGTGGAAAACTTCGTGGATATACCGGAGGACACAAAAAAGCTCGAGATCAATGTGGGGGGCATTGGATTTGGGTTGTTAGAGGTGGTTTATCAATTTAATTTGAATCTCGTCAACTTTGAGAATAGATTCCAACTAGACCTGGAGAAACAGAACACAGGCTCTGACTACGAGCTGAGGCTGAAGGTCTGTGCCAGCTACATACCCCAGCTGACCGACAGACGATCGAACATGGCACTGATTGAGGTAACCTTACCGAGCGGTTACGTGGTTGATCGCAATCCGATCAGCGAGCAGACGAAGGTGAATCCGATTCAGAAAACTGAAATCCGTTACGGTGGCACTTCAGTCGTTTTATACTACGACAATATGGGCAGCGAGCGTAACTGTTTCACCCTGACCGCGTACAGACGCTTTAAGGTCGCATTGAAGCGTCCAGCGTATGTGGTTGTGTATGATTATTATAATACAAATCTGAACGCCATCAAAGTGTACGAAGTGGACAAGCAGAATTTGTGCGAAATCTGTGACGAAGAAGACTGTCCTGCAGAGTGCAAAAAATAG'
    #qry_seq_raw = s1_AGAP010815_RA
    #
    # AAEL001802 (aegypti)
    #     https://www.vectorbase.org/Aedes_aegypti/Gene/Sequence?db=core;g=AAEL001802;r=supercont1.43:685886-717122;t=AAEL001802-RA
    ref_seq_raw = sys.argv[2]
    # OBP7 aegypti: v
    #ref_seq_raw = "ATGATGGAACAGCTTATGCTGGCAGTTTTGCTGGCGGTTTTTCTCGGGCTCGTAGCAGATGTTACGATGGCCGCTCAAATCAAGGACAATTTGGAGCTACCCGAATATTACAAACGTCCGGCCAAAATTCTGCACAACATCTGTCTGGCAGAATCCGGTGCCATGGAGAGCAAACTAAAGCAGTGCATGGACGGAGTGCTTCATGACGACCGGGAAGTCAAGTGCTACATCCATTGTCTATTCGACAAGGTGGACGTAATCGACGAAGCAACCGGGCAGATCCTGTTGGACCGATTGGCACCACTGGCACCGGACAACGATGTGAAGGATGTGTTCAATCATTTGACCAAAGAGTGTGGTCATATCAAACTACAAGATTCCTGCGATACGGCGTACGAAGTGGCCAAATGTTACTTCGCGGCACACGATCAGGTCGTCAAATTCTGTCACCTGTTGATGGCTGATGTTACCAGCTAG"
    #s2_AAEL001802_RA = 'ATGTCGGTATTCATACAAACGGACAAACCGGTGTATACCCCGGGAGATCTGATACGTTTTCGGGTAATCGTGGTGGATGCTGACACTAGACCTGTGACTAGTATTAAAACGGTAAATATAGCGATCGACGATTCTGCAAAAAATTCCATTCGAAAGTGGCCTTATGCCAAGTTGTTAAACGGCATCTTTGAGTCACAAGTGCAATTAGCTTCTTCGCCTGTTCTTGGCACCTGGATTATCAACGTAACAGCTTCCGACGACATCATTGTCACCAAACAGATAGAAGTTAAGGAATATGTGTTGCCAAAATTTTTCGTGAAAGTTTACCCTTCGGAGGTTCTATTGGGGAAAAATAAGAAGGTTTCTCTTACCTTAGATGCCTATTACACGTTCAAAGAACCCGTCGACGGCAATTACAAAGTTGAGTTATTTTTGGACCATACCAAGAGAAAGCCTGACTTCATAAAAAGTGATCGAATCACCGGTAAAACATCACTTGAGTTTCAATTGAAAAATGAAGTAGACATTGATGGCGACGAGCAGTACACTGATGTCACGGTTGAAGTTGAAGTTGTCGAGGCATTTTCTAATCGCACAGTTAGTATAACTGAGAATATTCCGATTTATCGTCAGCCTTATACCGTGACCCTTCTTCCATCTGCACCATCATTTCGACCAGGAGTTCCATTCAATGTACAAATAGTTGTGAAAGATCAGCTTGGACACCCTCCTGCCGAAGAAAAAGCGGCATCAATTGACCTTACTGTAGAGTTCCATTTGCCCATTGACAGTGACACCAAATCTATCACTGTAGATCTGGACGAGAAAGGAACAGGTCAGCTCACATTAGAGCCCCGCCCAGACGCCCAAGAACTGAAAGTGAACGCTACATATGACTCTCAACAATACGATGTAATTCACGATCCGATACATGGTTTCAGTTCGCAAAGTAAGCAGTACATCACAGTAACTCTGAATCCAAAATACTATAACAACATTAAAGTCGATAAGGACATCGTACTGGACATCTCCTGCACTGAAACAATGACGCACTTCTCGTACATCGTTGTCACCAGAGGAAACATAGTGGAAGCATCGAACGTTCCTGTCAGGATAAAAAAGAAACATTCTCTGAGATTGAAAATGACTTCAAAAATGTCTCCGGAGTCGAGGCTTCTAGTGTACTATACAAACAGGGAGTATCTCATCTTTGATGATATTGAGCTGAAGTTCGATTCGTTCAACAACGACTTCAAATTCGATTTGAACGATGATGAGTATTTTCCAGGGCAATCAGTTTATATCGATGTATACGCTTCAAAGGATTCATACGTTGCGTTCAGTGGAATCGATGAAAGTGTACTCCTGGTAGGCAAAGAGCGCCATGACTTCAACAAAGGAGATGTGCTCAAGGAACTCGCTCTTTACGGAGCAACAAATGATGCCGAGTTTGACTTGTTCCACGTAAGTTTCATGTCAAATGGTTTGATTATTCCAGTTAATGTATCTGTAACTCGCTCACAGAATGCACGATTTGGTACTCTACTAGGAAGGACTAGGCAGCAAGCGATTGAAATTCGAACTCAATTCCTAGAATCCTGGTTATGGAAATCCTTTTCCATGGATGGTCGAAACAACTTCAAAGCAATAGAAGACTCGGTTCCGGATACTATTACAACGTATCACGTGTCAGGATTTGCTTTAAGTCCAACACTAGGTCTTGGAGTAATCCAACAACCAGTGAGTTTCACCGTTCGTAAAAAATTCTACTTGGTTGCAAATTTGCCTTACTCGATCAAACGGGGTGAAGTGGCGTTGATTCAGGTTACCGTCTTCAACTTCCTAGGAAGCAGCATAACAACCGATGTGACGCTGTTCAATAAACGCGATGAAATTGAGTTTGTCGAGAATGCATCCACTAATAATACACATCGAACAAAGGCGGTAATTGTCCCGAATAACAATGGAAAATCTGTATCATTTATGGTGAAAGCAAAGAAATTAGGACAGATTGCGATCAAATTCCAGGCGGTAAACCTGCTGGAAACGGATGCATTGGAGCACATGTTACGAGTAACCCCAGAGAGCCATCGCTATGAGAAAAATGTAGCTCGATTCGTTGAGCTACCAAAGTTTGAGACGCAAACTTTCGATGTGAAGCTGGACATTCCCAAAAATATCGACGAGGGTTCTGCTCAAATCAAATTCACGTTAGACCCGGACATTTTGGGAACAGCCATCAGCAACCTAGACGGGTTGATCCGGAAACCCTTTGGATGTGGCGAACAAAATATGCTCCATTTTGTGCCAAATATAGTCGTTTTGGATTATCTTAACGAAACCAACACAGCGGCAGAAGATGTGAGGACCAAAGCGATAAATTTTCTTAGCAGCGGATATCAAAACCAGCTACGCTACAAACGTTCGGATGGGGCCTTCAGTGTCTGGGGACAATCGTATGCTGGCAGTACATTTTTGACGGCCTTTGTGGCGAAATCATTCAAAATAGCAGCCAAATACATTCAGGTGGATAAGTCTATAGTAGACGCGGCATTCGACTGGTTAGTGAAACAACAACAATCAGATGGGCGGTTCCCAGAAGTGGGGCAAGTATTCCAAGCAGATATGCAGGGTGGGCTTCGTAATAACGGTTTTGCGCTTACCGCGTATGTTCTGATCGCTTTTGCTGAAAATAAGGAAGTATACAGAAAATACCAATCACAACTGAACAAAACTACTAACTTCATAGCAGATAGACTTGCTAATATGGAGAATCCATACGACCTCTCGCTGTCCACTTATGCGTTGATGCTAACAAATCATGGCAAGCGCACCGAGTTTCTTCACAAATTAGTCGAAAAGTCGATATTTGACCGCAATCAAACTGAGAGATATTGGGACAGCAAACCAGTTGATATTGAAGTTGCTGGATATGCTCTATTGTCATACGTAGCTGCCGGTAAATTATTGGATGCAACGCCTATCATGCGGTGGCTCAACAAGCAGCGTTATGGTCTCGGAGGCTATCCTGGAACTCAGGAAACATTCGTTGGATTGAAAGCATTGGCAACGTTCGCTGCAAATGTAACTAGTAGGAGAAACGAATATACTGTAAGGATATTCTACGAACCAAATGGTCGACGAACATTCGACGTACACATGCACAATTCGTTTAATATTCAAGAGCTTGACATTCCTAATAACATCAGAAAAATGAAGGTGGAAGTTGAAGGCATCGGCAGAGGCTTCTTCCAAGTGGCATATCAGTACTATCAAAATATGCAGGTGGCTAAGCCCAGTTTCAGCATTACAATTAATCAGCTTAACACCACGACGGAACACATGCAGCAATTGGACGTGTGTGTGAAATACATACCAAAAGAGGCTTATCAAAAATCGAATATGGCTTTGGTGGAAATATTCTTGCCTAGTGGGCTTGTAGCAGACTCAGATGCCATTACGGACAAGACTGGAGGAATTCGAAGAATTGAAAGACGTTTTTCGGACACCTCAGTAGTTATATATTATGATAATTTGGACCCCGAAGACAAGTGCTTCCGAGTGACTGCTTATCGTCGGTATAAAATTGCATTGCATTTGCCATCATATATTATAGTTTATGATTATTATAATTTTGAGCGCTTTGCCATTCAAAAGTACGAAGGAAAGGTGCTGCAGCTCTGCGATATTTGTGAAGACGAGGACTGCGAAACTTTATCATGTCAAAATAGCTCGAAATTGGCAATAATGTAA'
    #ref_seq_raw = s2_AAEL001802_RA

    ###########################
    # Benchmarking vs. MATLAB #
    ###########################

    # @done:@testing:benchmark vs. matlab: @todo: place in README these details
    #   using s1_AGAP010815_RA and s2_AAEL001802_RA (after aligning/trimming)
    # MATLAB whole dnds vs. andy-wenping whole dnds (msCorrect='approximate'): dN/dS:
    #   1.0958 vs. 1.04939841193 (Elapsed time: 1.952491045, Total warnings: 0), respectively
    # MATLAB mean(sliding dnds) vs. andy-wenping mean(sliding dnds): Mean dN/dS over all 50-codon long windows:
    #   1.7254581782566112 (smoothing?) vs. 1.49715496621 (incl. smoothing, stepLength=1) (Elapsed time: 2.29854488373, Total warnings: 173) (when taking the mean of all 50 window intervals in MATLAB's sliding dnds()
    # MATLAB sliding dnds plot vs. andy-wenping sliding dnds plot: similar sliding dnds plots:
    #   see:  D:\Dropbox\00_HPC-LEAP\WP5 Thematic Cyprus Computational Biology\giannis_database

    # @benchmark:matlab's dnds()
    # dN/dS=0.15270463083614955 with msCorrect='approximate', dN/dS=0.16529268957638238 with
    # msCorrect='exact', MATLAB gives: dN/dS=0.0221=dnds(seq1,seq2, 'geneticCode', 1, 'Method', 'NG'). This is a huge error, an order of magnitude off.
    # seq1 = 'ATGCGCAAATACTCCCCCTTCCGAAATGGATACATGGAACCCACCCTTGGGCAGCACCTCCCAACCCTGTCTTTTCCAGACCCCGGACTCCGGCCCCAAAACCTGTACACCCTCTGGGGAGGCTCCGTTGTCTGCATGTACCTCTACCAGCTTTCCCCCCCCATCACCTGGCCCCTCCTGCCCCATGTGATTTTTTGCCACCCCGGCCAGCTCGGGGCCTTCCTCACCAATGTTCCCTACAAACGAATAGAAAAACTCCTCTATAAAATTTCCCTTACCACAGGGGCCCTAATAATTCTACCCGAGGACTGTTTGCCCACCACCCTTTTCCAGCCTGCTAGGGCACCCGTCACGCTGACAGCCTGGCAAAACGGCCTCCTTCCGTTCCACTCAACCCTCACCACTCCAGGCCTTATTTGGACATTTACCGATGGCACGCCTATGATTTCCGGGCCCTGCCCTAAAGATGGCCAGCCATCTTTAGTACTACAGTCCTCCTCCTTTATATTTCACAAATTTCAAACCAAGGCCTACCACCCCTCATTTCTACTCTCACACGGCCTCATACAGTACTCTTCCTTTCATAATTTGCATCTCCTATTTGAAGAATACACCAACATCCCCATTTCTCTACTTTTTAACGAAAAAGAGGCAGATGACAATGACCATGAGCCCCAAATATCCCCCGGGGGCTTAGAGCCTCTCAGTGAAAAACATTTCCGTGAAACAGAAGTC'
    # seq2 = 'ATGCGCAAGTACTCCCCCTTCCGAAACGGATACATGGAACCCACCCTTGGGCAACACCTCCCAACCCTGTCTTTTCCAGACCCCGGCCTCCGGCCCCAAAACCTGTACACCCTCTGGGGAGACTCTGTTGTCTGCCTGTACCTCTACCAGCTCTCCCCCCCCATCACCTGGCCCCTCCCGCCCCATGTGATTTTTTGCCACCCCGGCCAGCTCGGGGCCTTCCTCACCAATGTTCCCTACAAGCGTATGGAAGAACTCCTCTATAAAATTTCCCTTACCACAGGGGCCCTAATAATTCTACCCGAGGACTGTTTACCAACCACCCTTTTCCAGCCTGCTAGGGCCCCCGTCACGTTGACCGCCTGGCAGAACGGCCTCCTTCCGTTCCACTCAACCCTCACCACTCCAGGCCTTATTTGGACATTTACCGATGGCACGCCTATGGTTTCCGGACCCTGCCCCAAAGATGGCCAGCCATCTTTAGTACTACAGTCCTCCTCATTTATATTTCACAAATTTCAAACCAAGGCCTACCACCCTTCATTTCTACTCTCACACGGCCTCATACAGTACTCCTCCTTTCACAATTTACATCTCCTTTTTGAAGAATACACCAACATCCCCGTTTCTCTACTTTTTAACGAAAAAGAGGCAAATGACACTGACCATGAGCCCCAAATATCCCCCGGGGGCTTAGAGCCTCCCGCTGAAAAACATTTCCGCGAAACAGAAGTC'
    # dnds('ATGCGCAAATACTCCCCCTTCCGAAATGGATACATGGAACCCACCCTTGGGCAGCACCTCCCAACCCTGTCTTTTCCAGACCCCGGACTCCGGCCCCAAAACCTGTACACCCTCTGGGGAGGCTCCGTTGTCTGCATGTACCTCTACCAGCTTTCCCCCCCCATCACCTGGCCCCTCCTGCCCCATGTGATTTTTTGCCACCCCGGCCAGCTCGGGGCCTTCCTCACCAATGTTCCCTACAAACGAATAGAAAAACTCCTCTATAAAATTTCCCTTACCACAGGGGCCCTAATAATTCTACCCGAGGACTGTTTGCCCACCACCCTTTTCCAGCCTGCTAGGGCACCCGTCACGCTGACAGCCTGGCAAAACGGCCTCCTTCCGTTCCACTCAACCCTCACCACTCCAGGCCTTATTTGGACATTTACCGATGGCACGCCTATGATTTCCGGGCCCTGCCCTAAAGATGGCCAGCCATCTTTAGTACTACAGTCCTCCTCCTTTATATTTCACAAATTTCAAACCAAGGCCTACCACCCCTCATTTCTACTCTCACACGGCCTCATACAGTACTCTTCCTTTCATAATTTGCATCTCCTATTTGAAGAATACACCAACATCCCCATTTCTCTACTTTTTAACGAAAAAGAGGCAGATGACAATGACCATGAGCCCCAAATATCCCCCGGGGGCTTAGAGCCTCTCAGTGAAAAACATTTCCGTGAAACAGAAGTC', 'ATGCGCAAGTACTCCCCCTTCCGAAACGGATACATGGAACCCACCCTTGGGCAACACCTCCCAACCCTGTCTTTTCCAGACCCCGGCCTCCGGCCCCAAAACCTGTACACCCTCTGGGGAGACTCTGTTGTCTGCCTGTACCTCTACCAGCTCTCCCCCCCCATCACCTGGCCCCTCCCGCCCCATGTGATTTTTTGCCACCCCGGCCAGCTCGGGGCCTTCCTCACCAATGTTCCCTACAAGCGTATGGAAGAACTCCTCTATAAAATTTCCCTTACCACAGGGGCCCTAATAATTCTACCCGAGGACTGTTTACCAACCACCCTTTTCCAGCCTGCTAGGGCCCCCGTCACGTTGACCGCCTGGCAGAACGGCCTCCTTCCGTTCCACTCAACCCTCACCACTCCAGGCCTTATTTGGACATTTACCGATGGCACGCCTATGGTTTCCGGACCCTGCCCCAAAGATGGCCAGCCATCTTTAGTACTACAGTCCTCCTCATTTATATTTCACAAATTTCAAACCAAGGCCTACCACCCTTCATTTCTACTCTCACACGGCCTCATACAGTACTCCTCCTTTCACAATTTACATCTCCTTTTTGAAGAATACACCAACATCCCCGTTTCTCTACTTTTTAACGAAAAAGAGGCAAATGACACTGACCATGAGCCCCAAATATCCCCCGGGGGCTTAGAGCCTCCCGCTGAAAAACATTTCCGCGAAACAGAAGTC')

    #
    # ALIGN REF vs. QUERY seq, then TRIM GAPS:
    #     Next we want to remove all characters (trim) in each of the sequences that is aligned to a gap position**** in the top alignment
    #

    #  Raw input:
    #     ATGTGCTAA      qry
    #     ATGATTTGA      ref
    #
    # Alignment:
    #     ATGTGC----TAA  qry_seq
    #     ATG--A--TTTGA  ref_seq
    #        ^^ ^^^^     gap positions****
    #  Trimming:
    #     ATGCTAA        qry_seq (after trimming)
    #     ATGATGA        ref_seq (after trimming)

    # alignment scores are calculated from these two parameters, best scoring alignments have few gaps and gaps are small
    aln_gap_open = -10
    aln_gap_extend = -0.5 # @todo: make cmd args later?

    qry_seq_aln, ref_seq_aln         = align_then_trim.align_query_vs_reference(qry_seq_raw, ref_seq_raw, aln_gap_open, aln_gap_extend)
    qry_seq_trimmed, ref_seq_trimmed = align_then_trim.trim_gaps_from_aligned_seqs( qry_seq_aln, ref_seq_aln )

    # @Note: benhamrking: it is qry_seq_trimmed and ref_seq_trimmed that should be used as input to MATLAB's dnds() function, for benchmarking

    #}} alternative 2

    # ///

    # alternative 1 {{

    # # @NOTE:uncomment below to achieve dnds of 0.15.. or 0.164 if using exact method, interestingly
    # dnds_whole, warning_count = dnds( qry_seq_trimmed, ref_seq_trimmed, potential_changes, observed_changes, msCorrect='approximate', sliding=False)
    # print "dN/dS: "+str(dnds_whole)

    # }} 1 alternative 2 {{

    # @DONE: work with the sliding window version instead of dnds_whole
    dnds_slide_dict, warning_count = dnds( qry_seq_trimmed, ref_seq_trimmed, potential_changes, observed_changes, msCorrect='approximate', sliding=True, windowLength=50, stepLength=1 )
    #
    # Plot the sliding window values
    #
    dnds_sliding_vec, dnds_sliding_mean = plot_dnds_sliding(dnds_slide_dict)

    print "Mean dN/dS over all windows: "+str(dnds_sliding_mean)
    print "dN/dS over sliding windows: "+str(dnds_sliding_vec)
    # @todo: show the user also the whole dnds value somewhere in the webpage

    # }} alternative 2


"""
Module with functionality to compute anomaly scores based on reconstructions
"""
import numpy as np
import scipy.constants as cst
from skimage.color import gray2rgb  # convert spectra to 3 channels

from .metrics import ReconstructionMetrics

###############################################################################
GALAXY_LINES = {
    # EMISSION
    "OII_3727": 3727.0,
    "H_delta_4102": 4102.0,
    "H_gamma_4340": 4340.0,
    "H_beta_4861": 4861.0,
    "OIII_4959": 4959.0,
    "OIII_5007": 5007.0,
    "NII_6548": 6548.0,
    "H_alpha_6563": 6563.0,
    "NII_6584": 6584.0,
    "SII_6716": 6716.0,
    "SII_6720": 6720.0,
    "SII_6731": 6731.0,
    # ABSORPTION
}
##############################################################################
class ReconstructionAnomalyScore(ReconstructionMetrics):
    """
    Class to deal with the outliers based on a generative model trained with
    tensorflow.keras
    """

    ###########################################################################
    def __init__(
        self,
        reconstruct_fucntion,
        wave: np.array,
        lines: list = None,
        velocity_filter: float = 50,
        percentage: int = 100,
        relative: bool = False,
        epsilon: float = 1e-3,
    ):
        """
        INPUTS
            reconstruct_fucntion: reconstruct method of trained
                generative model

            lines: list with lines to discard to compute anomaly_score
            velocity_filter: Doppler velocity to consider at the moment of
                line filtering. It is in units of Km/s.
                DeltaWave = (v/c) * wave
            wave: common grid to spectra

            percentage: percentage of fluxes with the highest
                reconstruction error to consider to compute
                the anomaly score
            relative: whether or not the score is weigthed by the input
            epsilon: float value to avoid division by zero
        """

        self.reconstruct = reconstruct_fucntion
        self.wave = wave

        self.lines = lines
        filter_lines = velocity_filter == 0
        self.filter_lines = filter_lines
        self.velocity_filter = velocity_filter

        super().__init__(percentage, relative, epsilon)

    ###########################################################################
    def score(
        self, observation: np.array, metric: str, p: float = 0.33
    ) -> np.array:

        # in case I pass a spectra with one dimension
        # this line converts 1D array to (1, n_wave, 3)
        # an image where each channel has the spectrun
        observation = self.spectra_to_batch_image(observation)

        assert observation.ndim == 4

        observation, reconstruction = self.reconstruct_and_filter(
            observation, self.lines, self.velocity_filter
        )

        # make compatible batch of spectra's images with metrics
        observation = observation[:, 0, :, 0]
        reconstruction = reconstruction[:, 0, :, 0]
        if metric == "lp":

            assert np.isscalar(p)

            anomaly_score = super().lp(observation, reconstruction, p)
            return anomaly_score.reshape((-1, 1))

        if metric == "mse":

            anomaly_score = super().mse(observation, reconstruction)
            return anomaly_score.reshape((-1, 1))

        if metric == "mad":

            anomaly_score = super().mad(observation, reconstruction)
            return anomaly_score.reshape((-1, 1))

        print(f"{metric} not implemented")

    ###########################################################################
    def reconstruct_and_filter(
        self, observation: np.array, lines: list, velocity_filter: float
    ) -> tuple:

        """
        PARAMETERS
            observation: array with the origin of fluxes
            lines: list with lines to discard to compute anomaly_score
            velocity_filter: Doppler velocity to consider at the moment of
                line filtering. It is in units of Km/s.
                DeltaWave = (v/c) * wave

        OUTPUTS
            observation, reconstruction:
                np.arrays with the filter if it applies
        """

        # get (n_batch, flux). This is what is compatible with
        # reconstruction method
        reconstruction = self.reconstruct(observation[:, 0, :, 0])

        if self.filter_lines is True:

            velocity_mask = self.get_velocity_filter_mask(
                lines, velocity_filter
            )

            observation = observation[:, 0, velocity_mask, 0]
            reconstruction = reconstruction[:, velocity_mask]

        observation = self.spectra_to_batch_image(observation)
        reconstruction = self.spectra_to_batch_image(reconstruction)

        assert observation.ndim == 4
        assert reconstruction.ndim == 4

        return observation, reconstruction

    ###########################################################################
    def get_velocity_filter_mask(
        self, lines: list, velocity_filter: float
    ) -> np.array:

        """
        Compute array with filters for narrow emission lines
        PARAMETERS

            lines: list with lines to discard to compute anomaly_score.
                Check VELOCITY_LINES dictionary at the begin in the document.
            velocity_filter: Doppler velocity to consider at the moment of
                line filtering. It is in units of Km/s.
                DeltaWave = (v/c) * wave

        OUTPUT

            velocity_mask: array of bools with the regions to discard
        """

        c = cst.c * 1e-3  # [km/s]
        alpha = velocity_filter / c  # filter width

        velocity_mask = self.wave.astype(bool)

        for line in lines:

            delta_wave = GALAXY_LINES[line] * alpha
            # move line to origin
            wave = self.wave - GALAXY_LINES[line]
            line_mask = (wave < -delta_wave) | (delta_wave < wave)
            # update velocity mask
            velocity_mask *= line_mask

        return velocity_mask

    ###########################################################################
    def spectra_to_batch_image(self, spectra):

        # If a 1D spec is passed
        if spectra.ndim == 1:
            # get (1, flux)
            gray_spectra = spectra[np.newaxis, ...]
            # get (1, flux, 3)
            spectra_image = gray2rgb(gray_spectra)
            # get (n_batch, 1, flux, 3)
            return spectra_image[np.newaxis]
        # array of spectra: (n_batch, flux)
        if spectra.ndim == 2:
            # get (n_bacth, flux, 3)
            gray_spectra = gray2rgb(spectra)
            # get (n_bacth, 1, flux, 3)
            return gray_spectra[:, np.newaxis, ...]
        # if already image pass to (n_batch, 1, flux, 3)
        if spectra.ndim == 3:
            return spectra[np.newaxis, ...]

        return spectra


#-*- coding:utf-8 -*-
#
# Original code is here: https://github.com/openai/guided-diffusion
#

import numpy as np
import torch as th
import torch.nn as nn
import torch.nn.functional as F

from .fp16_util import convert_module_to_f16, convert_module_to_f32
from .modules import *

NUM_CLASSES = 1

class UNetModel(nn.Module):
    """
    The full UNet model with attention and timestep embedding.
    :param in_channels: channels in the input Tensor.
    :param model_channels: base channel count for the model.
    :param out_channels: channels in the output Tensor.
    :param num_res_blocks: number of residual blocks per downsample.
    :param attention_resolutions: a collection of downsample rates at which
        attention will take place. May be a set, list, or tuple.
        For example, if this contains 4, then at 4x downsampling, attention
        will be used.
    :param dropout: the dropout probability.
    :param channel_mult: channel multiplier for each level of the UNet.
    :param conv_resample: if True, use learned convolutions for upsampling and
        downsampling.
    :param dims: determines if the signal is 1D, 2D, or 3D.
    :param num_classes: if specified (as an int), then this model will be
        class-conditional with `num_classes` classes.
    :param use_checkpoint: use gradient checkpointing to reduce memory usage.
    :param num_heads: the number of attention heads in each attention layer.
    :param num_heads_channels: if specified, ignore num_heads and instead use
                               a fixed channel width per attention head.
    :param num_heads_upsample: works with num_heads to set a different number
                               of heads for upsampling. Deprecated.
    :param use_scale_shift_norm: use a FiLM-like conditioning mechanism.
    :param resblock_updown: use residual blocks for up/downsampling.
    :param use_new_attention_order: use a different attention pattern for potentially
                                    increased efficiency.
    """

    def __init__(
        self,
        image_size,
        in_channels,
        model_channels,
        out_channels,
        num_res_blocks,
        attention_resolutions,
        dropout=0,
        channel_mult=(1, 2, 4, 8),
        conv_resample=True,
        dims=2,
        num_classes=None,
        use_checkpoint=False,
        use_fp16=False,
        num_heads=1,
        num_head_channels=-1,
        num_heads_upsample=-1,
        use_scale_shift_norm=False,
        resblock_updown=False,
        use_new_attention_order=False,
    ):
        super().__init__()

        if num_heads_upsample == -1:
            num_heads_upsample = num_heads

        self.image_size = image_size
        self.in_channels = in_channels
        self.model_channels = model_channels
        self.out_channels = out_channels
        self.num_res_blocks = num_res_blocks
        self.attention_resolutions = attention_resolutions
        self.dropout = dropout
        self.channel_mult = channel_mult
        self.conv_resample = conv_resample
        self.num_classes = num_classes
        self.use_checkpoint = use_checkpoint
        self.dtype = th.float16 if use_fp16 else th.float32
        self.num_heads = num_heads
        self.num_head_channels = num_head_channels
        self.num_heads_upsample = num_heads_upsample

        time_embed_dim = model_channels * 4
        self.time_embed = nn.Sequential(
            linear(model_channels, time_embed_dim),
            nn.SiLU(),
            linear(time_embed_dim, time_embed_dim),
        )

        if self.num_classes is not None:
            self.label_emb = nn.Embedding(num_classes, time_embed_dim)

        ch = input_ch = int(channel_mult[0] * model_channels)
        self.input_blocks = nn.ModuleList(
            [TimestepEmbedSequential(conv_nd(dims, in_channels, ch, 3, padding=1))]
        )
        self._feature_size = ch
        input_block_chans = [ch]
        ds = 1
        for level, mult in enumerate(channel_mult):
            for _ in range(num_res_blocks):
                layers = [
                    ResBlock(
                        ch,
                        time_embed_dim,
                        dropout,
                        out_channels=int(mult * model_channels),
                        dims=dims,
                        use_checkpoint=use_checkpoint,
                        use_scale_shift_norm=use_scale_shift_norm,
                    )
                ]
                ch = int(mult * model_channels)
                if ds in attention_resolutions:
                    layers.append(
                        AttentionBlock(
                            ch,
                            use_checkpoint=use_checkpoint,
                            num_heads=num_heads,
                            num_head_channels=num_head_channels,
                            use_new_attention_order=use_new_attention_order,
                        )
                    )
                self.input_blocks.append(TimestepEmbedSequential(*layers))
                self._feature_size += ch
                input_block_chans.append(ch)
            if level != len(channel_mult) - 1:
                out_ch = ch
                self.input_blocks.append(
                    TimestepEmbedSequential(
                        ResBlock(
                            ch,
                            time_embed_dim,
                            dropout,
                            out_channels=out_ch,
                            dims=dims,
                            use_checkpoint=use_checkpoint,
                            use_scale_shift_norm=use_scale_shift_norm,
                            down=True,
                        )
                        if resblock_updown
                        else Downsample(
                            ch, conv_resample, dims=dims, out_channels=out_ch
                        )
                    )
                )
                ch = out_ch
                input_block_chans.append(ch)
                ds *= 2
                self._feature_size += ch

        self.middle_block = TimestepEmbedSequential(
            ResBlock(
                ch,
                time_embed_dim,
                dropout,
                dims=dims,
                use_checkpoint=use_checkpoint,
                use_scale_shift_norm=use_scale_shift_norm,
            ),
            AttentionBlock(
                ch,
                use_checkpoint=use_checkpoint,
                num_heads=num_heads,
                num_head_channels=num_head_channels,
                use_new_attention_order=use_new_attention_order,
            ),
            ResBlock(
                ch,
                time_embed_dim,
                dropout,
                dims=dims,
                use_checkpoint=use_checkpoint,
                use_scale_shift_norm=use_scale_shift_norm,
            ),
        )
        self._feature_size += ch

        self.output_blocks = nn.ModuleList([])
        for level, mult in list(enumerate(channel_mult))[::-1]:
            for i in range(num_res_blocks + 1):
                ich = input_block_chans.pop()
                layers = [
                    ResBlock(
                        ch + ich,
                        time_embed_dim,
                        dropout,
                        out_channels=int(model_channels * mult),
                        dims=dims,
                        use_checkpoint=use_checkpoint,
                        use_scale_shift_norm=use_scale_shift_norm,
                    )
                ]
                ch = int(model_channels * mult)
                if ds in attention_resolutions:
                    layers.append(
                        AttentionBlock(
                            ch,
                            use_checkpoint=use_checkpoint,
                            num_heads=num_heads_upsample,
                            num_head_channels=num_head_channels,
                            use_new_attention_order=use_new_attention_order,
                        )
                    )
                if level and i == num_res_blocks:
                    out_ch = ch
                    layers.append(
                        ResBlock(
                            ch,
                            time_embed_dim,
                            dropout,
                            out_channels=out_ch,
                            dims=dims,
                            use_checkpoint=use_checkpoint,
                            use_scale_shift_norm=use_scale_shift_norm,
                            up=True,
                        )
                        if resblock_updown
                        else Upsample(ch, conv_resample, dims=dims, out_channels=out_ch)
                    )
                    ds //= 2
                self.output_blocks.append(TimestepEmbedSequential(*layers))
                self._feature_size += ch

        self.out = nn.Sequential(
            normalization(ch),
            nn.SiLU(),
            zero_module(conv_nd(dims, input_ch, out_channels, 3, padding=1)),
        )

    def convert_to_fp16(self):
        """
        Convert the torso of the model to float16.
        """
        self.input_blocks.apply(convert_module_to_f16)
        self.middle_block.apply(convert_module_to_f16)
        self.output_blocks.apply(convert_module_to_f16)

    def convert_to_fp32(self):
        """
        Convert the torso of the model to float32.
        """
        self.input_blocks.apply(convert_module_to_f32)
        self.middle_block.apply(convert_module_to_f32)
        self.output_blocks.apply(convert_module_to_f32)

    def forward(self, x, timesteps, y=None):
        """
        Apply the model to an input batch.
        :param x: an [N x C x ...] Tensor of inputs.
        :param timesteps: a 1-D batch of timesteps.
        :param y: an [N] Tensor of labels, if class-conditional.
        :return: an [N x C x ...] Tensor of outputs.
        """
        assert (y is not None) == (
            self.num_classes is not None
        ), "must specify y if and only if the model is class-conditional"

        hs = []
        emb = self.time_embed(timestep_embedding(timesteps, self.model_channels))

        if self.num_classes is not None:
            assert y.shape == (x.shape[0],)
            emb = emb + self.label_emb(y)

        h = x.type(self.dtype)
        for module in self.input_blocks:
            h = module(h, emb)
            hs.append(h)
        h = self.middle_block(h, emb)
        for module in self.output_blocks:
            h = th.cat([h, hs.pop()], dim=1)
            h = module(h, emb)
        h = h.type(x.dtype)
        return self.out(h)

def create_model(
    image_size,
    num_channels,
    num_res_blocks,
    channel_mult="",
    learn_sigma=False,
    class_cond=False,
    use_checkpoint=False,
    attention_resolutions="16",
    num_heads=1,
    num_head_channels=-1,
    num_heads_upsample=-1,
    use_scale_shift_norm=False,
    dropout=0,
    resblock_updown=False,
    use_fp16=False,
    use_new_attention_order=False,
):
    if channel_mult == "":
        if image_size == 512:
            channel_mult = (0.5, 1, 1, 2, 2, 4, 4)
        elif image_size == 256:
            channel_mult = (1, 1, 2, 2, 4, 4)
        elif image_size == 128:
            channel_mult = (1, 1, 2, 3, 4)
        elif image_size == 64:
            channel_mult = (1, 2, 3, 4)
        else:
            raise ValueError(f"unsupported image size: {image_size}")
    else:
        channel_mult = tuple(int(ch_mult) for ch_mult in channel_mult.split(","))

    attention_ds = []
    for res in attention_resolutions.split(","):
        attention_ds.append(image_size // int(res))

    return UNetModel(
        image_size=image_size,
        in_channels=3,
        model_channels=num_channels,
        out_channels=(3 if not learn_sigma else 6),
        num_res_blocks=num_res_blocks,
        attention_resolutions=tuple(attention_ds),
        dropout=dropout,
        channel_mult=channel_mult,
        num_classes=(NUM_CLASSES if class_cond else None),
        use_checkpoint=use_checkpoint,
        use_fp16=use_fp16,
        num_heads=num_heads,
        num_head_channels=num_head_channels,
        num_heads_upsample=num_heads_upsample,
        use_scale_shift_norm=use_scale_shift_norm,
        resblock_updown=resblock_updown,
        use_new_attention_order=use_new_attention_order,
    )


#!/usr/bin/env python

import os,sys
import numpy as np
from aqml.cheminfo import *
from aqml.cheminfo.rw.xyz import *
from representation.xb import get_nzs

T,F = True,False

pi2 = np.pi*np.pi
a2b = 1.8897261258369282  # angstrom to bohr

np.set_printoptions(precision=2)

class rkrr(object):
    """
    recursive krr
    """
    def __init__(self, ys):
        """ Note that the column of `ys represents the same property
        generated from different level of theory """
        self._ys = np.array(ys)
        self.nm = len(ys)
        self.nl = len(ys[0]) # number of levels
        # initialize test set size `n2 to 0
        self.n2 = 0

    def init_m(self, obj):
        """ initialize molecules """
        #assert type(obj) is molecules
        self._zsu = np.unique( obj.zs )
        self._zs = obj.zs
        self._nas = obj.nas
        self._nsheav = obj.nsheav
        self._coords = obj.coords
        _ias2 = np.cumsum(self._nas)
        _ias1 = np.array([0,]+list(_ias2[:-1]),np.int)
        self._ias1 = _ias1
        self._ias2 = _ias2
        # get energy of dressed atom
        self._nzs = get_nzs(_ias1,_ias2,self._zs,self._zsu)

    def get_idx(self,idx=-1,n1s=[],n1max=-1,seed1=1,namin=1,namax=7):
        """
        get training/test idx

        i) AML: idx=-n ([n] target mols)
                namax=0 -> one training set, i.e., all amons
                namax=7 -> 7 training sets, the i-th set is {N_I = i} (i=1,2,...,7)
                namax=-5 -> one training set, comprising of all amons with N_I <= -namax
        ii) random sampling: idx>0 (__not__ quite useful?!)
        iii) manually select training set: idx=[0,3,10,...],
        """

        if isinstance(idx,int):
            if idx == 0:
                # only k1 is to be calculated, then we can randomly choose training data
                # from the molecules and shuffle many times to make an averaged LC
                _idx1 = np.arange(self.nm)
                idx2 = []
            elif idx > 0:
                np.random.seed(seed1) # fix the random sequence
                tidxs = np.random.permutation(self.nm)
                _idx1 = tidxs[:-idx]; idx2 = tidxs[-idx:]
                if n1max > 0:
                    _idx1 = _idx1[:n1max] # choose a subset
            else:
                tidxs = np.arange(self.nm); _idx1 = tidxs[:idx]; idx2 = tidxs[idx:]
        elif isinstance(idx,list):
            # idx as the maximal training set
            _idx1 = np.array(idx,np.int)
            tidxs = np.arange(self.nm)
            idx2 = np.setdiff1d(tidxs,_idx1)
        else:
            print('#ERROR: unsupported type of `idx')
            raise

        n2 = len(idx2) # now test set is fixed!!
        self.n2 = n2

        # now get smaller training set sizes
        nsu_heav = np.unique( self._nsheav ) # already in ascending order
        self.nsu_heav = nsu_heav
        idx1 = _idx1
        nn1 = len(n1s)
        aml = True
        if namax == 0:
            aml = False
            if nn1 == 0:
                # use only one training set, including all amons
                idxs1 = [_idx1]
            else:
                idxs1 = [ _idx1[:n1] for n1 in n1s ]
        elif namax < 0:
            # use only one training set, {N_I <= -namax}
            idx1 = _idx1[ np.logical_and(self._nsheav[_idx1]>=namin, self._nsheav[_idx1] <= -namax) ]
            idx2 = np.setdiff1d(tidxs,idx1)
            idxs1 = [idx1]; n1s = [len(idx1)]
        else:
            # use `namax-`namin+1 training sets to generate a LC
            idxs1 = []
            n1s = []
            # Note that `nheav may not be of the same order as the mol idx in filename,
            # so we'd better sort the idx of filenames! E.g., na=5 in frag_01.xyz while na=3 in frag_05.xyz
            idx1_sorted = []
            t = self._nsheav[_idx1]
            cnt = 0
            for na in range(namin,namax+1):
                if np.any(na==nsu_heav):
                    idx_i = _idx1[t==na]
                    leni = len(idx_i)
                    idx1_sorted += list(idx_i)
                    cnt += leni
                    n1s.append(cnt)
                else:
                    if cnt == 0:
                        n1s.append(np.nan) #
                    else:
                        n1s.append(cnt)
            print(' ** initial n1s = ', n1s)
            idx1 = idx1_sorted
        self.aml = aml
        self.n1s = n1s

        tidx = np.concatenate((idx1,idx2)).astype(np.int)
        n1,n2 = len(idx1),len(idx2)
        nt = n1+n2
        self.ys = self._ys[tidx]

        null = np.array([],np.int)
        _nas = self._nas[tidx]
        _coords = []; _zs = []
        for i1 in tidx:
            ib1,ie1 = self._ias1[i1],self._ias2[i1]
            _coords += list(self._coords[ib1:ie1])
            _zs += list(self._zs[ib1:ie1])
        _coords = np.array(_coords)
        _zs = np.array(_zs,np.int)

        self.nas1 = _nas[:n1]
        self.nas2 = null if n2 == 0 else _nas[n1:]
        self.nas = _nas

        self.nzs = self._nzs[tidx]
        self.nsheav = self._nsheav[tidx]

        # atomic index
        ias_e = np.cumsum(_nas)
        ias_b = np.concatenate(([0],ias_e[:-1]))
        ias1 = np.concatenate( [ np.arange(ias_b[i],ias_e[i]) for i in range(n1) ] )
        ias2 = null if n2 == 0 else np.concatenate( [ np.arange(ias_b[i],ias_e[i]) for i in range(n1,nt) ] )
        iast = np.concatenate((ias1,ias2))
        self.coords = _coords[iast]
        self.zs1 = _zs[ias1]
        self.zs2 = _zs[ias2]
        self.zs = np.concatenate((self.zs1,self.zs2))
        self.nat1 = len(ias1)
        self.nat2 = len(ias2)

    def calc_ae_dressed(self,nzs1,ys1,nzs2,ys2):
        esb = np.linalg.lstsq(nzs1,ys1)[0]
        ys1p = np.dot(nzs1,esb)
        #print ' +++++ ys1.shape, ys1p.shape = ', ys1.shape, ys1p.shape
        dys1 = ys1 - ys1p
        ys2_base = np.dot(nzs2,esb)
        dys2 = ys2 - ys2_base
        return dys1,dys2, ys2_base

    def calc_e_base(self,nzs1,ys1,nzs2):
        esb = np.linalg.lstsq(nzs1,ys1,rcond=None)[0]
        ys1p = np.dot(nzs1,esb)
        #print ' +++++ ys1.shape, ys1p.shape = ', ys1.shape, ys1p.shape
        dys1 = ys1 - ys1p
        ys2_base = np.dot(nzs2,esb)
        return dys1, ys2_base

    def run(self, mks, usebl=T, llambda=1e-10, iprt=True):
        """do KRR training & test

        vars
        ================
        """
        if isinstance(mks, (list,tuple)):
            tmpk1,tmpk2 = mks;
        elif isinstance(mks, str):
            if os.path.exists(mks):
                dic = np.load(mks)
                tmpk1,tmpk2 = [ dic[key].copy() for key in ['k1','k2'] ]
        else:
            print('#ERROR: unknow input')
            raise
        mk1,mk2 = tmpk1.copy(), tmpk2.copy()

        n2 = self.n2
        tidxs = np.arange(self.nm)
        idxs1 = tidxs[:-n2]; idxs2 = tidxs[-n2:]

        # now train many krr models
        nl = self.nl # number of levels
        mods = []
        n1sr = self.n1s[::-1] # order reversed
        ys2p = np.zeros(n2)
        ns_nested = []
        for l in range(nl-1):
            #lvi = levels[i]
            _ims1 = list(range(n1sr[l]))
            ims1 = _ims1 #[ np.logical_not( np.isnan(self.ys[_ims1,l]) ) ] # e.g., if E_QMC is not avail, was set to NaN
            n1 = len(ims1)
            ns_nested.append(n1)
            nzs1, nzs2 = self.nzs[ims1], self.nzs[idxs2]
            k1 = mk1[ims1][:,ims1]
            k1[np.diag_indices_from(k1)] += llambda
            k2 = mk2[:,ims1]
            if l==0: # one model
                _ys1 = self.ys[ims1,l]
                ys1, ys2_base = self.calc_e_base(nzs1,_ys1,nzs2)
            else: # i <= nl-2:
                #_ys1_l = self.ys[ims1,l]; _ys1_lm1 = self.ys[ims1,l-1]
                _ys1_l = self.ys[ims1,l] - self.ys[ims1,l-1]
                ys1_l, ys2_base_l = self.calc_e_base(nzs1,_ys1_l,nzs2)
                #ys1_lm1, ys2_base_lm1 = self.calc_e_base(nzs1,_ys1_lm1,nzs2)
                ys1 = ys1_l; ys2_base = ys2_base_l
                #ys1 = ys1_l - ys1_lm1; ys2_base = 0.
            alphas = np.linalg.solve(k1,ys1)
            ys2p_l = np.dot(k2,alphas) + ys2_base
            ys2p += ys2p_l
        print(' ** ns_nested = ', ns_nested )

        # last machine
        _n1s = self.n1s
        l = nl-1
        maes = []
        n1s = []
        for i,_n1 in enumerate(_n1s):
            _ims1 = np.arange(_n1)
            #print ' ims1 = ', _ims1, self.ys[_ims1,l]
            ims1 = _ims1[ np.logical_not( np.isnan(self.ys[_ims1,l]) ) ] # e.g., if E_QMC is not avail, was set to NaN
            n1 = len(ims1)
            nzs1, nzs2 = self.nzs[ims1], self.nzs[idxs2]
            if np.linalg.matrix_rank(nzs1) < len(nzs1[0]):
                print('%6s'%('None'))
                n1s.append(np.nan)
                maes.append(np.nan) #; rmses.append(np.nan)
                continue

            k1 = mk1[ims1][:,ims1]
            k1[np.diag_indices_from(k1)] += llambda
            k2 = mk2[:,ims1]
            #_ys1_l = self.ys[ims1,l]; _ys1_lm1 = self.ys[ims1,l-1]
            _ys1_l = self.ys[ims1,l] - self.ys[ims1,l-1]
            ys1_l, ys2_base_l = self.calc_e_base(nzs1,_ys1_l,nzs2)
            #ys1_lm1, ys2_base_lm1 = self.calc_e_base(nzs1,_ys1_lm1,nzs2)
            ys1 = ys1_l
            #ys1 = ys1_l - ys1_lm1
            ys2_base = ys2_base_l
            alphas = np.linalg.solve(k1,ys1)
            #print(' --> ys2p, k2*alpha, ys2_base = ',ys2p, np.dot(k2,alphas), ys2_base)
            ys2p_final = ys2p + np.dot(k2,alphas) + ys2_base

            dys = ys2p_final - self.ys[idxs2,l]
            print('%6d %.2f'%(len(ims1), dys[0]))
            maes.append( dys[0] )
            n1s.append(n1)
        self.maes = maes
        self.n1s = n1s


#!/usr/bin/env python
import os
import sys
import numpy as np
import matplotlib
matplotlib.use('PDF')
from StringIO import StringIO
from mgplottools.mpl import get_color, set_axis, new_figure, ls, \
                            set_color_cycle

def create_figure(outfile, pulse_duration, pop_0i, pop_rr, gate_error):

    # Layout
    fig_width       = 12.5              # Total canvas (cv) width
    left_margin     = 2.25              # Left cv -> plot area
    right_margin    = 2.25              # plot area -> right cv
    top_margin      = 0.3               # top cv -> plot area
    bottom_margin   = 1.0               # bottom cv -> panel 1
    h               = 2.5               # height of each panel
    w = fig_width - (left_margin + right_margin)  # width of panel
    fig_height = bottom_margin + h + top_margin
    fig = new_figure(fig_width, fig_height)

    set_color_cycle()

    # Panel
    pos = [left_margin/fig_width, bottom_margin/fig_height,
           w/fig_width, h/fig_height]
    ax = fig.add_axes(pos)
    ax.plot(pulse_duration, pop_0i, dashes=ls['dashed'],
            label='max pop in $\Ket{0i}$')
    ax.plot(pulse_duration, pop_rr, dashes=ls['long-dashed'],
            label='max pop in $\Ket{rr}$')
    ax.plot(pulse_duration, gate_error, color='black',
            label='gate error')
    set_axis(ax, 'x', 0, 700, 100, minor=4, range=(10, 700),
             label='central pulse duration (ns)')
    ax.set_yscale('log')
    ax.set_ylim(1.0e-3, 1.0e0)
    ax.legend()

    # output
    fig.savefig(outfile, format=os.path.splitext(outfile)[1][1:])

def read_data(datfile):
    data = """
# pulse dur [ns]  max pop 0i        max pop rr      gate error
  10              0.0837            0.8827          0.54675839
  20              0.0506            0.7553          0.42540254
  30              0.0363            0.5689          0.31267201
  40              0.0283            0.3972          0.23354138
  50              0.0232            0.2699          0.17602662
  75              0.016             0.1153          0.08980501
  100             0.0122            0.0694          0.05393173
  150             0.0083            0.0343          0.0249168
  200             0.0063            0.02            0.01393326
  250             0.005             0.0128          0.00952858
  300             0.0042            0.009           0.00662097
  350             0.0036            0.0066          0.00482424
  400             0.0032            0.0051          0.00366514
  450             0.0029            0.004           0.00298152
  500             0.0025515683      0.0032745798    0.00240066
  550             0.0023239946      0.0027134376    0.00195581
  600             0.0021279147      0.0022725114    0.00175541
  650             0.0019674974      0.0019544359    0.00135133
  700             0.0018390561      0.0016939797    0.00085459
"""
    return np.genfromtxt(StringIO(data), unpack=True)

def main(argv=None):
    if argv is None:
        argv = sys.argv
    basename = os.path.splitext(__file__)[0]
    outfile = basename + '.pdf'
    data_folder = basename
    create_figure(outfile, *read_data(data_folder))

if __name__ == "__main__":
    sys.exit(main())