# Matthew Smith, 12/5/19
# Shuo Gu, 02/18/20
# Shoichet Lab: Torsion Strain Rotation Project
# This Python script will contain all of the functions for using the Torsion
# Library method for estimating torsion strain energy. To use these functions,
# source this script in either another Python script or a Python commandline
# session as:
#   exec(open("TL_Functions.py").read())
# when in the same directory.

# We will use TL_2.1_VERSION_6.xml as the Torsion Library.

# Explanation of how this works from TL_Lookup_Test.py:
# Since the Torsion Library contains very specific SMARTS patterns to define
# the torsion patterns (also called torsion rules in the acutal library) and
# there are only 514 torsion patterns in TL_2.1, it is best to just search
# over all of the torsion rules in the library, instead of trying to find
# and classify every substructure in the molecule. This was Ying's suggestion.
# For each torsion rule, we will use the GetSubstructMatches function in
# the package rdkit.Chem to find the indeces of the atoms that match the
# torsion rule. If there are multiple such substructures, then this function
# should return multiple sets of indeces. We will then find the information
# we need (dihedral angle, method, estimate, lower bound, upper bound, and
# flag). Since the same torsion pattern in the molecule can have multiple
# torsion rules in the library applying to it, we need to choose the intended
# torsion rule. The authors of the original Torsion Library ordered the
# torsion rules within each subclass in decreasing specificity, so we will
# keep track of the number of the torsion rule for each matching substructure.
# We will then pick the set of information with the lowest counting index,
# assuming that this is the most specific. Since the first hierarchy class in
# the Torsion Library is the general class, we will need to loop over all of
# the specific classes first. This process is not the most efficient, since
# we are looking up and calculating information that we immediately throw away,
# but we can optimize this procedure later.

# Once we have one set of information for each set of indeces defining the
# bonds, we need to condense these by torsion pattern. We must do this because,
# for a given bond, each atom can be bonded to 2 or 3 other atoms in a torsion
# pattern. This gives between 2*2=4 and 3*3=9 possible lists of 4 atom indeces
# defining the torsion pattern, not including reversing the order of the
# indeces! Therefore we need to find all of the information sets (angle,
# estimate, etc.) with the same atom indeces defining the bond (listed
# forwards and backwards) and decide which one we will keep. Since we are
# using these energy estimates for identifying unstable conformations in
# DOCK-ing campaigns, being conservstive with our hit-picking means being
# liberal in our energy assignments. We thus keep the set of information with
# the highest energy estimate, prioritizing the estimates from specific
# hierarchy classes over ones from the general hierarchy class.

# Shuo Gu, 01/10/20, fixed a bug of input reading
# Shuo Gu, 01/12/20, Introduced interpolation between energy bins
# Shuo Gu, 02/18/20, Changed the initial energy 0 to 100 for the approximate method

import xml.etree.ElementTree as ET #For reading XML files
from rdkit import Chem #The module with the RDKit functions we will need
import os #For the function to read in the molecules from the .mol2 file
import numpy as np #For arrays and NumPy function
from math import sqrt, atan2, pi #For calculating dihedral angles
from math import ceil #Ceiling function

# Import the XML file, using this as a guide:
# https://docs.python.org/3/library/xml.etree.elementtree.html
tree = ET.parse("TL_2.1_VERSION_6.xml")
root = tree.getroot()

# We will turn the procedure for estimating torsion strain energy in
# TL_Lookup_Test.py into a function that we can call on every molecule
# in the list, which will return an object of the following class:
class TP_list(object):
    def __init__(self, indeces, angles, smarts, hc, methods, E,
    CI_l, CI_u, flags):
        self.indeces = indeces #List of lists of indeces
        self.angles = angles #List of angles
        self.smarts = smarts #List of SMARTS strings
        self.hc = hc #List of strings for hierarchy class
        self.methods = methods #List of strings of energy-estimation methods
        self.E = E #List of energy estimates
        self.CI_l = CI_l #List of lower bounds for 95% CI of energy estimates
        self.CI_u = CI_u #List of upper bounds for 95% CI of energy estimates
        self.flags = flags
        # List of Booleans flagging whether an angle whose "method" is
        # "approximate" is not observed
        self.TP_indeces = [j for j in range(len(indeces))]
        # Indeces of the torsion patterns, not the atoms!

    # Getters and setters:
    def get_indeces(self):
        return(self.indeces)
    def set_indeces(self, inds):
        self.indeces = inds
    def get_angles(self):
        return(self.angles)
    def set_angles(self, angs):
        self.angles = angs
    def get_smarts(self):
        return(self.smarts)
    def set_smarts(self, sms):
        self.smarts = sms
    def get_hc(self):
        return(self.hc)
    def set_hc(self, hcs):
        self.hc = hcs
    def get_methods(self):
        return(self.methods)
    def set_methods(self, meths):
        self.methods = meths
    def get_E(self):
        return(self.E)
    def set_E(self, Es):
        self.E = Es
    def get_CI_l(self):
        return(self.CI_l)
    def set_CI_l(self, ls):
        self.CI_l = ls
    def get_CI_u(self):
        return(self.CI_u)
    def set_CI_u(self, us):
        self.CI_u = us
    def get_flags(self):
        return(self.flags)
    def set_flags(self, fs):
        self.flags = fs
    def get_TP_indeces(self): #Don't need a setter for this
        return(self.TP_indeces)

    # A method to return the information for a subset of the torsion patterns
    def get_TPs(self, inds = None):
        # The parameter inds is a list of indeces for the torsion patterns,
        # not the atoms! We default to using all of the torsion patterns
        if inds == None:
            inds = [j for j in range(len(self.indeces))]
        # Create a list of torsion pattern info to be returned:
        tps = [] #Initialize
        for j in inds:
            tps.append([self.TP_indeces[j], self.E[j],
            self.CI_l[j], self.CI_u[j], self.indeces[j], self.angles[j],
            self.smarts[j], self.hc[j], self.methods[j], self.flags[j]])
        return(tps)

    # A method to find the sum of the energy estimates, with the 95% CI
    def sum(self, cutoff = None):
        # If at least one of the torsion patterns is flagged, then we
        # do not calculate the estimate or CI. We return -1 * (number flagged)
        flagged = sum(self.flags) #Treats True = 1 and False = 0
        if flagged > 0:
            return([-1*flagged, 0, 0]) #We do not try to calculate the CI
        else:
            if cutoff == None:
                cutoff == 0 #The default cutoff is 0, or using every angle
            ret = [0] * 3 #Initialize the returned array
            for i in range(len(self.E)):
                if self.E[i] >= cutoff:
                    # If the energy estimate is larger than the cutoff
                    ret[0] += self.E[i]
                    ret[1] += self.CI_l[i]
                    ret[2] += self.CI_u[i]
            return(ret)

# We need a function to read in .mol2 files, which RDKit does not supply to us
# We will use a function from
# https://chem-workflows.com/articles/2019/07/18/building-a-multi-molecule-mol2-reader-for-rdkit/
# Shuo has modified this function significantly to get it to work with the
# .mol2 files that this lab uses, which contain many lines of comments before
# each line "@<TRIPOS>MOLECULE"
def Mol2MolSupplier (file = None):
    names = [] #Make a list to hold the molecule names
    mols = {} #Make a dictionary
    with open(file, 'r') as f:
        fileend = os.fstat(f.fileno()).st_size
        count = 0
        line = f.readline()

        while not f.tell() == fileend:
            if line.startswith("#") or line == '\n':
                line = f.readline()
            if line.startswith("@<TRIPOS>MOLECULE"):
                count += 1
                mol = []
                mol.append(line)
                line = f.readline()
                if line != "\n" and line.split()[0].strip() not in names:
                    name = line.split()[0].strip()
                    print(name)
                else:
                    name = "mol2Number" + str(count)
                    print(name)

                while not line.startswith("@<TRIPOS>MOLECULE"):
                    mol.append(line)
                    line = f.readline()

                    if f.tell() == fileend:
                        mol.append(line)
                        break
                block = ",".join(mol).replace(',','')
                m = Chem.rdmolfiles.MolFromMol2Block(block, sanitize=False, removeHs = False)
                names.append(name)
                mols[name] = m
    return(names, mols)

# Here is an updated version to use with the output "file" string buffer object
# created by the db2_file_like function in the db2_to_mol2.py script
def db2MolSupplier(file):
    names = [] #Make a list to hold the molecule names
    mols = {} #Make a dictionary
    with file as f: #file is already opened as a string buffer
        bufferend = len(f.getvalue())
        count = 0
        line = f.readline()
        while not f.tell() == bufferend:
            if line.startswith("#") or line == '\n':
                line = f.readline()
            if line.startswith("@<TRIPOS>MOLECULE"):
                count += 1
                mol = []
                mol.append(line)
                line = f.readline()

                name = "db2Number" + str(count)

                while not line.startswith("@<TRIPOS>MOLECULE"):
                    mol.append(line)
                    line = f.readline()

                    if f.tell() == bufferend:
                        mol.append(line)
                        break
                block = ",".join(mol).replace(',','')
                m = Chem.rdmolfiles.MolFromMol2Block(block, sanitize=False, removeHs = False)
                names.append(name)
                mols[name] = m
    return(names, mols)

# The first function we will need to make TP_list objects is the dihedral angle
# calculator. I made this function (and its helper function) for my iPQB
# summer project of writing a script that creates Ramachandran plots
# (Make_Rama.py). I have changed some of the comments explaining the math
# First a helper function for creating unit vectors
def unit(a):
    # The argument should be a NumPy array with 1 axis
    return(a / sqrt(np.dot(a,a))) #Scales a by its norm

def dihedral(a_1, a_2, a_3, a_4):
    # The arguments should all be NumPy arrays with 1 axis and length 3
    # These atoms should be in order, with a_2 and a_3 defining the bond
    # of interest

    # The 3 displacement vectors:
    b_1 = a_2 - a_1
    b_2 = a_3 - a_2
    b_3 = a_4 - a_3

    n_1 = unit(np.cross(b_1, b_2))
    n_2 = unit(np.cross(b_2, b_3))

    # Imagine the first atom (a_1) is above the middle bond (from a_2 to a_3),
    # so that b_1 points downward. Then n_1 points out of the page

    m = unit(np.cross(n_1, b_2))
    # I moved the normalization to be after the cross product. Moving
    # the normalization should not change the end result because the cross
    # product commutes with scalar multiplication and ||n_1|| = 1

    # Looking down b_2, we can consider n_1 to be the x-axis and
    # m to be the y-axis. Then the dihedral angle is the angle that
    # n_2 makes with the x-axis when projected into this plane. Since dihedral
    # angles are measured going clockwise, we need to negate the angle
    # that we get back from atan
    x = np.dot(n_1, n_2) #Project n_2 onto n_1
    y = np.dot(m, n_2) #Project n_2 onto m
    return(-atan2(y,x) * 180 / pi) #Return the angle in degrees

# The next function we will need will calculate angular differences,
# in degrees. This function will calculate theta_1 - theta_2. See my notes
# from 9/19/19
def ang_diff(theta_1, theta_2):
    # (-180,180] -> [0, 360)
    if theta_1 < 0:
        theta_1 += 360
    if theta_2 < 0:
        theta_2 += 360
    del_theta = (theta_1 - theta_2) % 360 #Angular difference
    # [0, 360) -> (-180, 180]
    if del_theta > 180:
        del_theta -= 360
    return(del_theta)
    # Test this works: ang_diff(0, 179), ang_diff(0, -179),
    # and ang_diff(-179, 179)

# This function will allow us to do the matching for each torsion rule
def tp_match(tp, hc, j, mol, pos, bi):
    # tp is a torsion pattern, hc is the type of hierarchyClass ("general" or
    # "specific", and j is the current value for i
    # This function turned the global variables mol, positions, and bond_info
    # in TL_Lookup_Test.py into parameters mol, pos and bi, respectively
    smarts = tp.get("smarts")
    # Create the histograms for energy estimates and bounds of confidence
    # intervals, if available
    hist_E = [] #Initialize for energy estimates
    hist_l = [] #Initialize for lower bounds of CIs
    hist_u = [] #Initialize for upper bounds of CIs
    if tp.get("method") == "exact":
        for bin in tp.find("histogram_converted").findall("bin"):
            hist_E.append(float(bin.get("energy")))
            hist_l.append(float(bin.get("lower")))
            hist_u.append(float(bin.get("upper")))
    matches = mol.GetSubstructMatches(Chem.MolFromSmarts(smarts))
    # A list of lists
    for match in matches: #For each match
        # Some of the SMARTS for the torion patterns actually have 5 atoms.
        # We need to ingore these
        if len(match) > 4:
            continue #Go to the next match

        # First get the atom locations
        # Based on https://github.com/rdkit/rdkit/issues/1982
        r_1 = np.array(pos[match[0]])
        r_2 = np.array(pos[match[1]])
        r_3 = np.array(pos[match[2]])
        r_4 = np.array(pos[match[3]])

        theta = dihedral(r_1, r_2, r_3, r_4) #Dihedral angle

        # Changed next line from "TP.get" to "tp.get"
        if tp.get("method") == "exact": #If using the exact method
            # First figure out what bin we are in in the histogram.
            # We define the bins by the right endpoints, like we did
            # when we made the plots of the energy profiles
            bin_num = ceil(theta / 10) + 17
            # Starting with the -170 deg bin
            energy = (hist_E[bin_num]-hist_E[(bin_num+35)%36])/10.0*(theta-(bin_num-17)*10)+hist_E[bin_num]
            lower = (hist_l[bin_num]-hist_l[(bin_num+35)%36])/10.0*(theta-(bin_num-17)*10)+hist_l[bin_num]
            upper = (hist_u[bin_num]-hist_u[(bin_num+35)%36])/10.0*(theta-(bin_num-17)*10)+hist_u[bin_num]

            bi.append(
            [
            list(match), #Convert tuple to list
            theta,
            smarts,
            hc, #"general" or "specific"
            "exact",
            energy,
            lower,
            upper,
            False, #This only could apply for the approximate method
            j #We will take this out when we create the final object
            ]
            )

        else: #If using the approximate method
            not_observed = True
            # Initialize the flag for not observing that angle
            energy = 100
            # Initialize the energy, in case we cannot estimate it
            # Loop over all the possible angle peaks
            for angle in tp.find("angleList").findall("angle"):
                theta_0 = float(angle.get("theta_0")) #Peak location
                delta = ang_diff(theta, theta_0) #Angular displacement
                if abs(delta) <= float(angle.get("tolerance2")):
                    # If within the tolerance range for that peak
                    beta_1 = float(angle.get("beta_1"))
                    beta_2 = float(angle.get("beta_2"))
                    # The coefficients for the regression
                    energy = beta_1*(delta ** 2) + beta_2*(delta ** 4)
                    # Using the "not-as-small angle approximation"
                    not_observed = False
                    break
                    # Break the for loop to avoid problems if the
                    # observed angle sits at the border between two
                    # peaks
            bi.append(
            [
            list(match), #Convert tuple to list
            theta,
            smarts,
            hc, #"general" or "specific"
            "approximate",
            energy,
            energy, #Lower bound for CI, which we cannot find for approx. method
            energy, #Upper bound for CI, which we cannot find for approx. method
            not_observed,
            j #We will take this out when we create the final object
            ]
            )

# This most general function will automate what we did in TL_Lookup_Test.py
def TL_lookup(mol): #mol is read in from the .mol2 file
    positions = mol.GetConformer().GetPositions()
    # List of lists of atom coordinates. Luckily RDKit starts indexing at 0
    bond_info = []
    # Initialize an empty list that will hold the information for each bond
    i = 0 #Initialize count of torsion rules

    # Loop over all of the specific hierarchy classes
    for HC in root.findall("hierarchyClass"):
        if HC.get("name") != "GG": #Not the general class
            for TP in HC.iter("torsionRule"): #Loop over each torsion rule
                tp_match(TP, "specific", i, mol, positions, bond_info)
                i += 1 #Increase the count for the torsion rule

    # Now for the general method:
    for TP in root.find("hierarchyClass[@name='GG']").iter("torsionRule"):
        tp_match(TP, "general", i, mol, positions, bond_info)
        i += 1 #Increase the count for the torsion rule

    # Now that we have all of the torsion patterns, we need to be able to find
    # duplicates. The first such way is if the entire pattern is reversed. We
    # can fix this problem by making sure that all of the lists of indeces have
    # the second index (the first in the bond of interest) lower than the third
    # index (the second in the bond of interest)
    for  bond in bond_info: #Loop over every bond
        if bond[0][1] > bond[0][2]:
            bond[0].reverse() #Reverse this list
            bond.append(True) #Mark that we reversed this bond's indeces
        else:
            bond.append(False) #Mark that we did not reverse this bond's
            # indeces. We will remove this marking later

    # Next we condense the bond_info by the lists of 4 atoms defining the bonds.
    # We will pick the entry of bond_info that has the lowest value for i for
    # each match, since the torsion rules in the Torsion Library are arranged
    # (within each hierarchy class or hierarchy subclass) in decreasing
    # specificity, and we loop over all of the specific hierarchy classes
    # before the general one
    bond_info_red = [bond_info[0]] #Initialize a list for the reduced bond info
    # This reduced list needs at least one element for checking subelements
    for j in range(1, len(bond_info)):
        # Skip the first bond, which is already in the reduced list
        atom_0 = bond_info[j][0][0] #First atom index
        atom_1 = bond_info[j][0][1] #Second atom index
        atom_2 = bond_info[j][0][2] #Third bond index
        atom_3 = bond_info[j][0][3] #Fourth bond index

        unmatched = True #Initialize not finding a match
        for k in range(len(bond_info_red)):
            # Check against everything in the growing reduced list
            if bond_info_red[k][0][0] == atom_0 \
            and bond_info_red[k][0][1] == atom_1 \
            and bond_info_red[k][0][2] == atom_2 \
            and bond_info_red[k][0][3] == atom_3:
                # If there is a match in ALL of the atom indeces
                unmatched = False
                if bond_info[j][9] < bond_info_red[k][9]:
                    # The ninth index gives the value for i
                    # If the new bond has a lower value, then we use it
                    # to replace the current one
                    bond_info_red[k] = bond_info[j]
                    break
                    # No need to continue looking for matches, since there
                    # should be no more than 1

        if unmatched: #If no match
            bond_info_red.append(bond_info[j]) #Append the current bond

    # Now we condense the bond_info_red by the 2 atoms actually defining the
    # bond. We will pick the entry of bond_info_red for each match that has
    # the highest energy estimate, prioritizing torsion patterns from
    # "specific" hierarchy classes over ones from the "general" hierarchy class
    b_i_r = [bond_info_red[0]] #Initialize a list for the further reduced bond info
    # I used the name b_i_r over bond_info_red_2
    for j in range(1, len(bond_info_red)):
        # Skip the first bond, which is already in the reduced list
        atom_1 = bond_info_red[j][0][1] #First atom index of the bond
        atom_2 = bond_info_red[j][0][2] #Second atom index of the bond

        unmatched = True #Initialize not finding a match
        for k in range(len(b_i_r)):
            # Check against everything in the growing list
            if b_i_r[k][0][1] == atom_1 and b_i_r[k][0][2] == atom_2:
                # If there is a match in the two atom indeces defining the bond
                unmatched = False
                if bond_info_red[j][3][0] > b_i_r[k][3][0] \
                or (
                bond_info_red[j][5] > b_i_r[k][5] \
                and
                bond_info_red[j][3][0] == b_i_r[k][3][0]
                ):
                    # The third index gives "general" or "specific", and we
                    # access the first char in that string. We use:
                    #   's' > 's' gives False
                    #   'g' > 'g' gives False
                    #   'g' > 's' gives False
                    #   's' > 'g' gives True
                    # to prioritize using the specific classe over the general one.
                    # The fifth index gives the energy estimate
                    b_i_r[k] = bond_info_red[j] #Replace the current bond info
                    break
                    # No need to continue looking for matches, since there
                    # should be no more than 1

        if unmatched: #If no match
            b_i_r.append(bond_info_red[j]) #Append the current bond

    # Now that we have only the bond information that we want, we need to
    # un-do any times we reversed the atom indeces
    for  bond in b_i_r:
        if bond[10]:
            # The tenth index gives whether or not we reversed the indeces
            bond[0].reverse() #Reverse this list

    # Now that we have all the information we need, we can return the
    # object that contains it
    return(
    TP_list([bond[0] for bond in b_i_r], [bond[1] for bond in b_i_r],
    [bond[2] for bond in b_i_r], [bond[3] for bond in b_i_r],
    [bond[4] for bond in b_i_r], [bond[5] for bond in b_i_r],
    [bond[6] for bond in b_i_r], [bond[7] for bond in b_i_r],
    [bond[8] for bond in b_i_r]) #Using Python list comprehension
    )


#!/usr/bin/env python
"""
radarData.py
This file holds the RadarData class that hold the radar data and processes it.

@author: John Swoboda
"""
import ipdb
import scipy as sp
# My modules
from isrutilities.physConstants import v_C_0, v_Boltz
from SimISR import Path
from SimISR.IonoContainer import IonoContainer
from SimISR.utilFunctions import CenteredLagProduct, MakePulseDataRepLPC,dict2h5,h52dict,readconfigfile, BarkerLag, update_progress
#from SimISR import specfunctions
#from SimISR.analysisplots import plotspecsgen

class RadarDataFile(object):
    """
        This class will will take the ionosphere class and create radar data both
        at the IQ and fitted level.

        Variables
        simparams - A dictionary that holds simulation parameters the keys are the
            following
            'angles': Angles (in degrees) for the simulation.  The az and el angles
                are stored in a list of tuples [(az_0,el_0),(az_1,el_1)...]
            'IPP' : Interpulse period in seconds, this is the IPP from position to
                position.  In other words if a IPP is 10 ms and there are 10 beams it will
                take 100 ms to go through all of the beams.
            'Tint'  This is the integration time in seconds.
            'Timevec': This is the time vector for each frame.  It is referenced to
                the begining of the frame.  Also in seconds
            'Pulse': The shape of the pulse.  This is a numpy array.
            'TimeLim': The length of time that the experiment will run, which cooresponds
                to how many pulses per position are used.
        sensdict - A dictionary that holds the sensor parameters.
        rawdata: This is a NbxNpxNr numpy array that holds the raw IQ data.
        rawnoise: This is a NbxNpxNr numpy array that holds the raw noise IQ data.
    """

    def __init__(self, Ionodict, inifile, outdir, outfilelist=None):
        """
            This function will create an instance of the RadarData class.  It will
            take in the values and create the class and make raw IQ data.
            Inputs:
            sensdict - A dictionary of sensor parameters
            angles - A list of tuples which the first position is the az angle
                and the second position is the el angle.
            IPP - The interpulse period in seconds represented as a float.
            Tint - The integration time in seconds as a float.  This will be the
            integration time of all of the beams.
            time_lim - The length of time of the simulation the number of time points
                will be calculated.
            pulse - A numpy array that represents the pulse shape.
            rng_lims - A numpy array of length 2 that holds the min and max range
                that the radar will cover.
        """
        (sensdict, simparams) = readconfigfile(inifile)
        self.simparams = simparams
        N_angles = len(self.simparams['angles'])

        NNs = int(self.simparams['NNs'])
        self.sensdict = sensdict
        Npall = sp.floor(self.simparams['TimeLim']/self.simparams['IPP'])
        Npall = int(sp.floor(Npall/N_angles)*N_angles)
        Np = Npall/N_angles

        print("All spectrums created already")
        filetimes = Ionodict.keys()
        filetimes.sort()
        ftimes = sp.array(filetimes)
        simdtype = self.simparams['dtype']
        pulsetimes = sp.arange(Npall)*self.simparams['IPP'] +ftimes.min()
        pulsefile = sp.array([sp.where(itimes-ftimes >= 0)[0][-1] for itimes in pulsetimes])
        # differentiate between phased arrays and dish antennas
        if sensdict['Name'].lower() in ['risr', 'pfisr', 'risr-n']:

            beams = sp.tile(sp.arange(N_angles), Npall/N_angles)
        else:

            # for dish arrays
            brate = simparams['beamrate']
            beams2 = sp.repeat(sp.arange(N_angles), brate)
            beam3 = sp.concatenate((beams2, beams2[::-1]))
            ntile = int(sp.ceil(Npall/len(beam3)))
            leftover = int(Npall-ntile*len(beam3))
            if ntile > 0:
                beams = sp.tile(beam3, ntile)
                beams = sp.concatenate((beams, beam3[:leftover]))
            else:
                beams = beam3[:leftover]

        pulsen = sp.repeat(sp.arange(Np), N_angles)
        pt_list = []
        pb_list = []
        pn_list = []
        fname_list = []
        self.datadir = outdir
        self.maindir = outdir.parent
        self.procdir = self.maindir/'ACF'
        Nf = len(filetimes)
        progstr = 'Data from {:d} of {:d} being processed Name: {:s}.'
        if outfilelist is None:
            print('\nData Now being created.')

            Noisepwr = v_Boltz*sensdict['Tsys']*sensdict['BandWidth']
            self.outfilelist = []
            for ifn, ifilet in enumerate(filetimes):

                outdict = {}
                ifile = Ionodict[ifilet]
                ifilename = Path(ifile).name
                update_progress(float(ifn)/Nf, progstr.format(ifn, Nf, ifilename))
                curcontainer = IonoContainer.readh5(ifile)
                if ifn == 0:
                    self.timeoffset = curcontainer.Time_Vector[0, 0]
                pnts = pulsefile == ifn
                pt = pulsetimes[pnts]
                pb = beams[pnts]
                pn = pulsen[pnts].astype(int)
                rawdata = self.__makeTime__(pt, curcontainer.Time_Vector,
                                            curcontainer.Sphere_Coords,
                                            curcontainer.Param_List, pb)
                d_shape = rawdata.shape
                n_tempr = sp.random.randn(*d_shape).astype(simdtype)
                n_tempi = 1j*sp.random.randn(*d_shape).astype(simdtype)
                noise = sp.sqrt(Noisepwr/2)*(n_tempr+n_tempi)
                outdict['AddedNoise'] = noise
                outdict['RawData'] = rawdata+noise
                outdict['RawDatanonoise'] = rawdata
                outdict['NoiseData'] = sp.sqrt(Noisepwr/2)*(sp.random.randn(len(pn), NNs).astype(simdtype)+
                                                            1j*sp.random.randn(len(pn), NNs).astype(simdtype))
                outdict['Pulses'] = pn
                outdict['Beams'] = pb
                outdict['Time'] = pt
                fname = '{0:d} RawData.h5'.format(ifn)
                newfn = self.datadir/fname
                self.outfilelist.append(str(newfn))
                dict2h5(str(newfn), outdict)

                #Listing info
                pt_list.append(pt)
                pb_list.append(pb)
                pn_list.append(pn)
                fname_list.append(fname)
            infodict = {'Files':fname_list, 'Time':pt_list, 'Beams':pb_list, 'Pulses':pn_list}
            dict2h5(str(outdir.joinpath('INFO.h5')), infodict)

        else:
            infodict = h52dict(str(outdir.joinpath('INFO.h5')))
            alltime = sp.hstack(infodict['Time'])
            self.timeoffset = alltime.min()
            self.outfilelist = outfilelist

#%% Make functions
    def __makeTime__(self,pulsetimes,spectime,Sphere_Coords,allspecs,beamcodes):
        """This is will make raw radar data for a given time period and set of
        spectrums. This is an internal function called by __init__.
        Inputs-
        self - RadarDataFile object.
        pulsetimes - The time the pulses are sent out in reference to the spectrums.
        spectime - The times for the spectrums.
        Sphere_Coords - An Nlx3 array that holds the spherical coordinates of the spectrums.
        allspecs - An NlxNdtimexNspec array that holds the spectrums.
        beamcodes - A NBx4 array that holds the beam codes along with the beam location
        in az el along with a system constant in the array. """

        range_gates = self.simparams['Rangegates']
        pulse = self.simparams['Pulse']
        sensdict = self.sensdict
        pulse2spec = sp.array([sp.where(itimes-spectime >= 0)[0][-1] for itimes in pulsetimes])
        Np = len(pulse2spec)
        lp_pnts = len(pulse)
        samp_num = sp.arange(lp_pnts)
        #N_rg = len(range_gates)# take the size
        N_samps = len(range_gates)
        angles = self.simparams['angles']
        Nbeams = len(angles)
        rho = Sphere_Coords[:, 0]
        Az = Sphere_Coords[:, 1]
        El = Sphere_Coords[:, 2]
        rng_len = self.sensdict['t_s']*v_C_0*1e-3/2.
        speclen = allspecs.shape[-1]
        simdtype = self.simparams['dtype']
        out_data = sp.zeros((Np, N_samps), dtype=simdtype)
        weights = {ibn:self.sensdict['ArrayFunc'](Az, El, ib[0], ib[1], sensdict['Angleoffset'])
                   for ibn, ib in enumerate(angles)}
        for istn in range(len(spectime)):
            for ibn in range(Nbeams):
                #print('\t\t Making Beam {0:d} of {1:d}'.format(ibn, Nbeams))
                weight = weights[ibn]
                for isamp in range(N_samps):
                    range_g = range_gates[isamp]

                    if range_g <= 0:
                        continue
                    range_m = range_g*1e3
                    rnglims = [range_g-rng_len/2.0, range_g+rng_len/2.0]
                    rangelog = (rho >= rnglims[0])&(rho < rnglims[1])
                    # Get the number of points covered
                    cur_pnts = samp_num+isamp
                    keep_pnt = sp.logical_and(cur_pnts >= 0, cur_pnts < N_samps)
                    cur_pnts = cur_pnts[keep_pnt]
                    # This is a nearest neighbors interpolation for the
                    # spectrums in the range domain
                    if sp.sum(rangelog) == 0:
                        minrng = sp.argmin(sp.absolute(range_g-rho))
                        rangelog[minrng] = True

                    #create the weights and weight location based on the beams pattern.
                    weight_cur = weight[rangelog]
                    weight_cur = weight_cur/weight_cur.sum()
                    specsinrng = allspecs[rangelog]
                    if specsinrng.ndim == 3:
                        specsinrng = specsinrng[:, istn]
                    elif specsinrng.ndim == 2:
                        specsinrng = specsinrng[istn]
                    specsinrng = specsinrng*sp.tile(weight_cur[:, sp.newaxis], (1, speclen))
                    cur_spec = specsinrng.sum(0)
                    # based off new way of calculating
                    pow_num = sensdict['Pt']*sensdict['Ksys'][ibn]*sensdict['t_s']
                    pow_den = range_m**2
                    curdataloc = sp.where(sp.logical_and((pulse2spec == istn),
                                                         (beamcodes == ibn)))[0]
                    # create data
                    if len(curdataloc) == 0:
                        print('\t\t No data for {0:d} of {1:d} in this time period'.format(ibn, Nbeams))
                        continue
                    cur_pulse_data = MakePulseDataRepLPC(pulse, cur_spec, 20,
                                                         len(curdataloc), numtype=simdtype)
                    cur_pulse_data = cur_pulse_data*sp.sqrt(pow_num/pow_den)
                    for idatn, idat in enumerate(curdataloc):
                        d_pnts = cur_pnts[::-1]-isamp
                        cursum = cur_pulse_data[idatn, d_pnts]+out_data[idat, cur_pnts]
                        out_data[idat, cur_pnts] = cursum
        return out_data
        #%% Processing
    def processdataiono(self):
        """ This will perform the the data processing and create the ACF estimates
        for both the data and noise but put it in an Ionocontainer.
        Inputs:

        Outputs:
        Ionocontainer- This is an instance of the ionocontainer class that will hold the acfs.
        """
        (datalags, noiselags) = self.processdata()
        return lagdict2ionocont(datalags, noiselags, self.sensdict, self.simparams,
                                datalags['Time'])

    def processdata(self):
        """ This will perform the the data processing and create the ACF estimates
        for both the data and noise.
        Inputs:
        timevec - A numpy array of times in seconds where the integration will begin.
        inttime - The integration time in seconds.
        lagfunc - A function that will make the desired lag products.
        Outputs:
        DataLags: A dictionary with keys 'Power' 'ACF','RG','Pulses' that holds
        the numpy arrays of the data.
        NoiseLags: A dictionary with keys 'Power' 'ACF','RG','Pulses' that holds
        the numpy arrays of the data.
        """
        timevec = self.simparams['Timevec'] +self.timeoffset
        inttime = self.simparams['Tint']
        # Get array sizes

        NNs = int(self.simparams['NNs'])
        range_gates = self.simparams['Rangegates']
        N_samps = len(range_gates)# take the size
        pulse = self.simparams['Pulse']
        Pulselen = len(pulse)
        N_rg = N_samps - Pulselen+1
        simdtype = self.simparams['dtype']
        Ntime = len(timevec)

        if 'outangles' in self.simparams.keys():
            Nbeams = len(self.simparams['outangles'])
            inttime = inttime
        else:
            Nbeams = len(self.simparams['angles'])

        # Choose type of processing
        if self.simparams['Pulsetype'].lower() == 'barker':
            lagfunc = BarkerLag
            Nlag = 1
        else:
            lagfunc = CenteredLagProduct
            Nlag = Pulselen
        # initialize output arrays
        outdata = sp.zeros((Ntime, Nbeams, N_rg, Nlag), dtype=simdtype)
        outaddednoise = sp.zeros((Ntime, Nbeams, N_rg, Nlag), dtype=simdtype)
        outnoise = sp.zeros((Ntime, Nbeams, NNs-Pulselen+1, Nlag), dtype=simdtype)
        pulses = sp.zeros((Ntime, Nbeams))
        pulsesN = sp.zeros((Ntime, Nbeams))
        timemat = sp.zeros((Ntime, 2))
        Ksysvec = self.sensdict['Ksys']
        # set up arrays that hold the location of pulses that are to be processed together
        infoname = self.datadir / 'INFO.h5'
        # Just going to assume that the info file is in the directory
        infodict = h52dict(str(infoname))
        flist = infodict['Files']
        file_list = [str(self.datadir/i) for i in flist]
        pulsen_list = infodict['Pulses']
        beamn_list = infodict['Beams']
        time_list = infodict['Time']
        file_loclist = [ifn*sp.ones(len(ifl)) for ifn, ifl in enumerate(beamn_list)]

        pulsen = sp.hstack(pulsen_list).astype(int)# pulse number
        beamn = sp.hstack(beamn_list).astype(int)# beam numbers
        ptimevec = sp.hstack(time_list).astype(float)# time of each pulse
        file_loc = sp.hstack(file_loclist).astype(int)# location in the file

        # run the time loop
        print("Forming ACF estimates")

        # For each time go through and read only the necisary files
        for itn, it in enumerate(timevec):
            update_progress(float(itn)/Ntime, "Time {0:d} of {1:d}".format(itn, Ntime))
            # do the book keeping to determine locations of data within the files
            cur_tlim = (it, it+inttime)
            curcases = sp.logical_and(ptimevec >= cur_tlim[0], ptimevec < cur_tlim[1])
            if  not sp.any(curcases):
                prog_str = "No pulses for time {0:d} of {1:d}, lagdata adjusted accordinly"
                update_progress(float(itn)/Ntime, prog_str.format(itn, Ntime))
                outdata = outdata[:itn]
                outnoise = outnoise[:itn]
                pulses = pulses[:itn]
                pulsesN = pulsesN[:itn]
                timemat = timemat[:itn]
                continue
            pulseset = set(pulsen[curcases])
            poslist = [sp.where(pulsen == item)[0] for item in pulseset]
            pos_all = sp.hstack(poslist)
            try:
                pos_all = sp.hstack(poslist)
                curfileloc = file_loc[pos_all]
            except:
                ipdb.set_trace()
            # Find the needed files and beam numbers
            curfiles = set(curfileloc)
            beamlocs = beamn[pos_all]
            timemat[itn, 0] = ptimevec[pos_all].min()
            timemat[itn, 1] = ptimevec[pos_all].max()
            # cur data pulls out all data from all of the beams and posisions
            curdata = sp.zeros((len(pos_all), N_samps), dtype=simdtype)
            curaddednoise = sp.zeros((len(pos_all), N_samps), dtype=simdtype)
            curnoise = sp.zeros((len(pos_all), NNs), dtype=simdtype)
            # Open files and get required data
            # XXX come up with way to get open up new files not have to reread in data that is already in memory
            for ifn in curfiles:
                curfileit = [sp.where(pulsen_list[ifn] == item)[0] for item in pulseset]
                curfileitvec = sp.hstack(curfileit)
                ifile = file_list[ifn]
                curh5data = h52dict(ifile)
                file_arlocs = sp.where(curfileloc == ifn)[0]
                curdata[file_arlocs] = curh5data['RawData'][curfileitvec]

                curaddednoise[file_arlocs] = curh5data['AddedNoise'].astype(simdtype)[curfileitvec]
                # Read in noise data when you have don't have ACFs

                curnoise[file_arlocs] = curh5data['NoiseData'].astype(simdtype)[curfileitvec]

            # differentiate between phased arrays and dish antennas
            if self.sensdict['Name'].lower() in ['risr', 'pfisr', 'risr-n']:
                # After data is read in form lags for each beam
                for ibeam in range(Nbeams):
                    prog_num = float(itn)/Ntime+float(ibeam)/Ntime/Nbeams
                    update_progress(prog_num, "Beam {0:d} of {1:d}".format(ibeam, Nbeams))
                    beamlocstmp = sp.where(beamlocs == ibeam)[0]
                    pulses[itn, ibeam] = len(beamlocstmp)

                    outdata[itn, ibeam] = lagfunc(curdata[beamlocstmp].copy(),
                                                  numtype=self.simparams['dtype'],
                                                  pulse=pulse, lagtype=self.simparams['lagtype'])

                    pulsesN[itn, ibeam] = len(beamlocstmp)
                    outnoise[itn, ibeam] = lagfunc(curnoise[beamlocstmp].copy(),
                                                   numtype=self.simparams['dtype'],
                                                   pulse=pulse, lagtype=self.simparams['lagtype'])
                    outaddednoise[itn, ibeam] = lagfunc(curaddednoise[beamlocstmp].copy(),
                                                        numtype=self.simparams['dtype'],
                                                        pulse=pulse, lagtype=self.simparams['lagtype'])
            else:
                for ibeam, ibeamlist in enumerate(self.simparams['outangles']):
                    prog_num = float(itn)/Ntime+float(ibeam)/Ntime/Nbeams
                    update_progress(prog_num, "Beam {0:d} of {1:d}".format(ibeam, Nbeams))
                    beamlocstmp = sp.where(sp.in1d(beamlocs, ibeamlist))[0]
                    inputdata = curdata[beamlocstmp].copy()
                    noisedata = curnoise[beamlocstmp].copy()
                    noisedataadd = curaddednoise[beamlocstmp].copy()
                    pulses[itn, ibeam] = len(beamlocstmp)
                    pulsesN[itn, ibeam] = len(beamlocstmp)
                    outdata[itn, ibeam] = lagfunc(inputdata, numtype=self.simparams['dtype'],
                                                  pulse=pulse, lagtype=self.simparams['lagtype'])
                    outnoise[itn, ibeam] = lagfunc(noisedata, numtype=self.simparams['dtype'],
                                                   pulse=pulse, lagtype=self.simparams['lagtype'])
                    outaddednoise[itn, ibeam] = lagfunc(noisedataadd, numtype=self.simparams['dtype'],
                                                        pulse=pulse, lagtype=self.simparams['lagtype'])
        # Create output dictionaries and output data
        DataLags = {'ACF':outdata, 'Pow':outdata[:, :, :, 0].real, 'Pulses':pulses,
                    'Time':timemat, 'AddedNoiseACF':outaddednoise}
        NoiseLags = {'ACF':outnoise, 'Pow':outnoise[:, :, :, 0].real, 'Pulses':pulsesN,
                     'Time':timemat}
        return(DataLags, NoiseLags)

#%% Make Lag dict to an iono container
def lagdict2ionocont(DataLags, NoiseLags, sensdict, simparams, time_vec):
    """This function will take the data and noise lags and create an instance of the
    Ionocontanier class. This function will also apply the summation rule to the lags.
    Inputs
    DataLags - A dictionary """
    # Pull in Location Data
    angles = simparams['angles']
    ang_data = sp.array([[iout[0], iout[1]] for iout in angles])
    rng_vec = simparams['Rangegates']
    n_samps = len(rng_vec)
    # pull in other data
    pulse = simparams['Pulse']
    p_samps = len(pulse)
    pulsewidth = p_samps*sensdict['t_s']
    txpower = sensdict['Pt']
    if sensdict['Name'].lower() in ['risr', 'pfisr', 'risr-n']:
        Ksysvec = sensdict['Ksys']
    else:

        beamlistlist = sp.array(simparams['outangles']).astype(int)
        inplist = sp.array([i[0] for i in beamlistlist])
        Ksysvec = sensdict['Ksys'][inplist]
        ang_data_temp = ang_data.copy()
        ang_data = sp.array([ang_data_temp[i].mean(axis=0) for i in beamlistlist])

    sumrule = simparams['SUMRULE']
    rng_vec2 = simparams['Rangegatesfinal']
    Nrng2 = len(rng_vec2)
    minrg = p_samps-1+sumrule[0].min()
    maxrg = Nrng2+minrg

    # Copy the lags
    lagsData = DataLags['ACF'].copy()
    # Set up the constants for the lags so they are now
    # in terms of density fluxtuations.
    angtile = sp.tile(ang_data, (Nrng2, 1))
    rng_rep = sp.repeat(rng_vec2, ang_data.shape[0], axis=0)
    coordlist = sp.zeros((len(rng_rep), 3))
    [coordlist[:, 0], coordlist[:, 1:]] = [rng_rep, angtile]
    (Nt, Nbeams, Nrng, Nlags) = lagsData.shape

    # make a range average to equalize out the conntributions from each gate

    plen2 = int(sp.floor(float(p_samps-1)/2))
    samps = sp.arange(0, p_samps, dtype=int)

    rng_ave = sp.zeros((Nrng, p_samps))

    for isamp in range(plen2, Nrng+plen2):
        for ilag in range(p_samps):
            toplag = int(sp.floor(float(ilag)/2))
            blag = int(sp.ceil(float(ilag)/2))
            if toplag == 0:
                sampsred = samps[blag:]
            else:
                sampsred = samps[blag:-toplag]
            cursamps = isamp-sampsred

            keepsamps = sp.logical_and(cursamps >= 0, cursamps < Nrng)
            cursamps = cursamps[keepsamps]
            rng_samps = rng_vec[cursamps]**2*1e6
            keepsamps2 = rng_samps > 0
            if keepsamps2.sum() == 0:
                continue
            rng_samps = rng_samps[keepsamps2]
            rng_ave[isamp-plen2, ilag] = 1./(sp.mean(1./(rng_samps)))
    rng_ave_temp = rng_ave.copy()
    if simparams['Pulsetype'].lower() is 'barker':
        rng_ave_temp = rng_ave[:,0][:,sp.newaxis]
    # rng_ave = rng_ave[int(sp.floor(plen2)):-int(sp.ceil(plen2))]
    # rng_ave = rng_ave[minrg:maxrg]
    rng3d = sp.tile(rng_ave_temp[sp.newaxis, sp.newaxis, :, :], (Nt, Nbeams, 1, 1))
    ksys3d = sp.tile(Ksysvec[sp.newaxis,:, sp.newaxis, sp.newaxis], (Nt, 1, Nrng, Nlags))

    # rng3d = sp.tile(rng_ave[:, sp.newaxis, sp.newaxis, sp.newaxis], (1, Nlags, Nt, Nbeams))
    # ksys3d = sp.tile(Ksysvec[sp.newaxis, sp.newaxis, sp.newaxis, :], (Nrng2, Nlags, Nt, 1))
    radar2acfmult = rng3d/(pulsewidth*txpower*ksys3d)
    pulses = sp.tile(DataLags['Pulses'][:, :, sp.newaxis, sp.newaxis], (1, 1, Nrng, Nlags))
    time_vec = time_vec[:Nt]
    # Divid lags by number of pulses
    lagsData = lagsData/pulses
    # Set up the noise lags and divid out the noise.
    lagsNoise = NoiseLags['ACF'].copy()
    lagsNoise = sp.mean(lagsNoise, axis=2)
    pulsesnoise = sp.tile(NoiseLags['Pulses'][:, :, sp.newaxis], (1, 1, Nlags))
    lagsNoise = lagsNoise/pulsesnoise
    lagsNoise = sp.tile(lagsNoise[:, :, sp.newaxis, :], (1, 1, Nrng, 1))

    # multiply the data and the sigma by inverse of the scaling from the radar
    lagsData = lagsData*radar2acfmult
    lagsNoise = lagsNoise*radar2acfmult
    # Apply summation rule
    # lags transposed from (time,beams,range,lag)to (range,lag,time,beams)
    lagsData = sp.transpose(lagsData, axes=(2, 3, 0, 1))
    lagsNoise = sp.transpose(lagsNoise, axes=(2, 3, 0, 1))
    lagsDatasum = sp.zeros((Nrng2, Nlags, Nt, Nbeams), dtype=lagsData.dtype)
    lagsNoisesum = sp.zeros((Nrng2, Nlags, Nt, Nbeams), dtype=lagsNoise.dtype)

    for irngnew, irng in enumerate(sp.arange(minrg, maxrg)):
        for ilag in range(Nlags):
            lsumtemp = lagsData[irng+sumrule[0, ilag]:irng+sumrule[1, ilag]+1, ilag].sum(axis=0)
            lagsDatasum[irngnew, ilag] = lsumtemp
            nsumtemp = lagsNoise[irng+sumrule[0, ilag]:irng+sumrule[1, ilag]+1, ilag].sum(axis=0)
            lagsNoisesum[irngnew, ilag] = nsumtemp
    # subtract out noise lags
    lagsDatasum = lagsDatasum-lagsNoisesum

    # Put everything in a parameter list
    Paramdata = sp.zeros((Nbeams*Nrng2, Nt, Nlags), dtype=lagsData.dtype)
    # Put everything in a parameter list
    # transpose from (range,lag,time,beams) to (beams,range,time,lag)
    # lagsDatasum = lagsDatasum*radar2acfmult
    # lagsNoisesum = lagsNoisesum*radar2acfmult
    lagsDatasum = sp.transpose(lagsDatasum, axes=(3, 0, 2, 1))
    lagsNoisesum = sp.transpose(lagsNoisesum, axes=(3, 0, 2, 1))

    # multiply the data and the sigma by inverse of the scaling from the radar
    # lagsDatasum = lagsDatasum*radar2acfmult
    # lagsNoisesum = lagsNoisesum*radar2acfmult

    # Calculate a variance using equation 2 from Hysell's 2008 paper. Done use full covariance matrix because assuming nearly diagonal.
    # Get the covariance matrix
    pulses_s = sp.transpose(pulses, axes=(1, 2 ,0, 3))[:, :Nrng2]
    Cttout = makeCovmat(lagsDatasum, lagsNoisesum, pulses_s, Nlags)

    Paramdatasig = sp.zeros((Nbeams*Nrng2 ,Nt, Nlags, Nlags), dtype=Cttout.dtype)

    curloc = 0
    for irng in range(Nrng2):
        for ibeam in range(Nbeams):
            Paramdata[curloc] = lagsDatasum[ibeam, irng].copy()
            Paramdatasig[curloc] = Cttout[ibeam, irng].copy()
            curloc += 1
    ionodata = IonoContainer(coordlist, Paramdata, times=time_vec, ver=1,
                             paramnames=sp.arange(Nlags)*sensdict['t_s'])
    ionosigs = IonoContainer(coordlist, Paramdatasig, times=time_vec, ver=1,
                             paramnames=sp.arange(Nlags**2).reshape(Nlags, Nlags)*sensdict['t_s'])
    return (ionodata, ionosigs)

def makeCovmat(lagsDatasum, lagsNoisesum, pulses_s, Nlags):
    """
        Makes the covariance matrix for the lags given the noise acf and number
        of pulses.
    """
    axvec = sp.roll(sp.arange(lagsDatasum.ndim), 1)
    # Get the covariance matrix
    R = sp.transpose(lagsDatasum/sp.sqrt(2.*pulses_s), axes=axvec)
    Rw = sp.transpose(lagsNoisesum/sp.sqrt(2.*pulses_s), axes=axvec)
    l = sp.arange(Nlags)
    T1, T2 = sp.meshgrid(l, l)
    R0 = R[sp.zeros_like(T1)]
    Rw0 = Rw[sp.zeros_like(T1)]
    Td = sp.absolute(T1-T2)
    Tl = T1>T2
    R12 = R[Td]
    R12[Tl] = sp.conjugate(R12[Tl])
    Rw12 = Rw[Td]
    Rw12[Tl] = sp.conjugate(Rw12[Tl])
    Ctt = R0*R12+R[T1]*sp.conjugate(R[T2])+Rw0*Rw12+Rw[T1]*sp.conjugate(Rw[T2])
    avecback = sp.roll(sp.arange(Ctt.ndim), -2)
    Cttout = sp.transpose(Ctt, avecback)
    return Cttout


"""Implementation of target attack.
  Three models (base_incep, adv_incep, res_incep) are used.
  The models are ensembled using mean norm gradient methrod.
  The inferstructure code is based on
  https://github.com/tensorflow/cleverhans/cleverhans/tree/master/
  examples/nips17_adversarial_competition/sample_targeted_attacks/iter_target_class
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import time

import numpy as np
import tensorflow as tf
from tensorflow.contrib import slim

from lib_adv import utils, attack

tf.flags.DEFINE_string(
    'output_dir',
    '../../intermediate_results/targeted_attacks_output/target_mng',
    'Output directory with images.')

tf.flags.DEFINE_string(
    'input_dir', '../../dataset/images_4/', 'Input directory with images.')

tf.flags.DEFINE_string(
    'checkpoint_path', 'mul_inception_v1.ckpt',
    'Path to checkpoint for inception network.')

tf.flags.DEFINE_string(
    'model_names',
    'base_inception,base_incpt_resnet,adv_inception,adv_incpt_resnet',
    'names of the multiple models, no space!')

tf.flags.DEFINE_string(
    'master', '', 'The address of the TensorFlow master to use.')

tf.flags.DEFINE_integer('image_width', 299, 'Width of each input images.')
tf.flags.DEFINE_integer('image_height', 299, 'Height of each input images.')
tf.flags.DEFINE_integer('image_depth', 3, 'Depth of each input images.')
tf.flags.DEFINE_integer('num_classes', 1001, 'Number of classes.')
tf.flags.DEFINE_integer('num_iter', 20, 'Number of iterations.')

tf.flags.DEFINE_integer('batch_size', 4, '#images process at one time.')
tf.flags.DEFINE_float('max_epsilon', 16.0, 'Max size of perturbation.')
tf.flags.DEFINE_integer('norm_ord', 1, 'order of norm to use')
tf.flags.DEFINE_string('labels_file', 'target_class.csv', 'target classes')

class ModelAttacker():
  """A target attacker model, multi-model ensembled."""
  def __init__(self, num_classes, image_pixels, model_names, max_epsilon,
               norm_ord, batch_shape):
    """Constructs a `ModelAttacker`.
        Set the parameters and build the base and adversarial graph"""
    prev = time.time()

    # image_factor is model dependent
    # Images for inception classifier are normalized to be in [-1, 1] interval
    # scale from [0, 255], range 255 to [-1, 1), range 2
    image_factor = 2.0 / 255.0

    list_mnames = model_names.split(',')
    eps, alpha = attack.parameters(max_epsilon, image_factor,
                                   image_pixels, norm_ord)

    utils.logger.debug('norm_ord = %d, alpha = %.3f, models: %s' % (
        norm_ord, alpha, list_mnames))

    # Prepare graph
    self._batch_shape = batch_shape
    self._x_input = tf.placeholder(tf.float32, shape=batch_shape)
    self._x_ref = tf.placeholder(tf.float32, shape=batch_shape)
    self._target_labels = tf.placeholder(tf.int32, shape=[batch_shape[0]])

    _, pred_labels, logits_w = attack.graph_base(self._x_input,
                                                 num_classes,
                                                 list_mnames)
    self._x_adv, _ = attack.graph_adv(self._x_input, num_classes,
                                      self._x_ref, self._target_labels,
                                      pred_labels, logits_w,
                                      alpha, eps, norm_ord)

    utils.logger.debug('graph init: %.3f seconds' % (time.time() - prev))

  def run(self, checkpoint_path, master, input_dir, labels_file, output_dir,
          num_iter):
    """Run the computation and generate adversarial images"""
    saver = tf.train.Saver(slim.get_model_variables())
    session_creator = tf.train.ChiefSessionCreator(
        scaffold=tf.train.Scaffold(saver=saver),
        checkpoint_filename_with_path=checkpoint_path,
        master=master)

    with tf.train.MonitoredSession(session_creator=session_creator) as sess:
      for filenames, images, labels in utils.load_images(input_dir,
                                                         self._batch_shape,
                                                         labels_file):
        adv_images = np.copy(images)
        for _ in range(num_iter):
          adv_images = sess.run(
              self._x_adv,
              feed_dict={
                  self._x_input: adv_images,
                  self._x_ref: images,
                  self._target_labels: labels})
        utils.save_images(adv_images, filenames, output_dir)

def main(_):
  tf.logging.set_verbosity(tf.logging.INFO)
  prev = time.time()

  FLG = tf.flags.FLAGS
  batch_shape = [FLG.batch_size, FLG.image_height, FLG.image_width,
                 FLG.image_depth]
  image_pixels = FLG.image_height * FLG.image_width * FLG.image_depth

  with tf.Graph().as_default():
    attacker = ModelAttacker(FLG.num_classes, image_pixels, FLG.model_names,
                             FLG.max_epsilon, FLG.norm_ord, batch_shape)
    attacker.run(FLG.checkpoint_path, FLG.master, FLG.input_dir,
                 FLG.labels_file, FLG.output_dir, FLG.num_iter)

  utils.logger.debug('%.3f seconds' % (time.time() - prev))

if __name__ == '__main__':
  tf.app.run()


#!/usr/bin/env python
#
# Original filename: loci.py
#
# Author: Tim Brandt
# Email: tbrandt@astro.princeton.edu
# Date: July 2011
#
# Summary:  Run LOCI on the input frame sequence.
#

import sys, re, gc
import numpy as np
from scipy import interpolate
import multiprocessing
from progressbar import ProgressBar
from locitools import *
from parallel import *

def loci(flux, pa, locipar, mem, mode='LOCI', fluxref=None,
         pca_arr=None, rmin=None, r_ex=None, niter=0,
         corr=None, method='matrix', do_partial_sub=True,
         sub_dir='.'):

    """

    Function loci performs Locally Optimized Combination of Images
    on an input dataset (argument 1).  It takes at least 4 arguments:

    1.  A three-dimensional flux array.  The first dimension should run
        over frames.
    2.  The position angle of the image rotator (in radians) in each frame
    3.  The LOCI configuration parameters.  This is an object defined in
        adiparam.lociparam
    4.  The available memory (in bytes)

    Optional arguments:

    5.  'mode': The mode in which to run the routine.  Currently
        accepted values are 'LOCI' (normal LOCI) and 'refine', in
        which case the LOCI coefficients minimize the difference
        between a frame and and a linear combination of the other
        frames with fluxref (argument 6) added.  Default 'LOCI'
    6.  'fluxref': A reference array, to be computed using PCA, to fit
        the PSF locally in each frame.  This is treated as an extra
        set of reference frames with infinite angular displacement.
        Ignored if None.  Default None.
    7.  'pca_arr': The flux array used to compute the LOCI subtraction
        coefficients.  This can be different from the array (argument
        1) on which those subtraction coefficients are used.  Ignored
        if None, mandatory if mode=='refine'.  Default None.
    8.  'rmin': The minimum radius at which to perform LOCI.  If None,
        this is the smaller of the saturation radius and minimum
        radius at which there are always frames to satisfy LOCI's
        angular displacement criterion.  Default None.
    9.  'r_ex': Exclusion/saturation radius.  Ignored if None; should
        be calculated using the centroiding algorithm.  Default None.
    10. 'niter': Number of steps in which to split execution to limit
        memory usage.  If 0, this is calculated within the routine.
        Default 0.
    11. 'corr': Cross-correlation coefficients between frames, to use
        only good matches in LOCI.  Possibly useful for very large
        datasets, but it can increase run time significantly.  Ignored
        if None.  Default None.
    12. 'method': Method for solving the LOCI matrix equation.
        Default is 'matrix', in which case the problem matrix is
        formed and solved using LU- decomposition.  Other choices:
        'lstsq' (SVD), 'eqcoef' (all coefficients are equal and sum to
        1--similar to mean subtraction).  Default: 'matrix'.  'lstsq'
        is more robust but will be significantly slower in most cases.
    13. 'do_partial_sub': Compute the fractional flux suppression
        using the method described in Brandt+2012.  This requires
        argument 'method' to be 'matrix', as it uses the LU
        decomposition of the problem matrix, and it requires argument
        'mode' to be 'LOCI'.  Default True

    Function loci returns the flux suppression map if argument
    'do_partial_sub' is True; otherwise, it returns None.

    """

    np.seterr(all='ignore')

    ######################################################################
    # Check arrays and sizes.
    ######################################################################

    if mode == 'refine' and fluxref is None:
        print "Cannot apply feedback to LOCI without reference output."
        sys.exit(1)
    if fluxref is not None:
        assert fluxref.shape == flux.shape, \
               "Flux and reference arrays must have the same shape."
    assert len(flux.shape) == 3, "Flux array must have three dimensions."
    assert niter >= 0 and niter < flux.shape[0], \
           "Argument niter to loci must be at least 0 and less than nframes"
    assert mode == 'refine' or mode == 'LOCI', \
           "Argument mode = " + mode + " to routine loci is not understood." \
           + "\nPlease use either 'refine' or 'LOCI'."

    ######################################################################
    # Create and populate array fluxsmooth to calculate LOCI coefficients
    ######################################################################

    nframes = flux.shape[0]
    #ngroup = 1 + int((nframes - 1) / locipar.max_n)
    ngroup = 1
    oldshape = flux.shape
    fluxsmooth = np.ndarray(flux.shape, np.float32)

    if mode == 'LOCI' and fluxref is not None:
        fluxsmooth[:] = fluxref[i]
    elif mode == 'refine' and fluxref is not None:
        fluxsmooth[:] = flux + fluxref
    elif mode == 'refine':
        print "Cannot call loci in 'refine' mode without a reference array."
        sys.exit()
    else:
        fluxsmooth[:] = flux

    fluxsmooth = np.reshape(fluxsmooth, (nframes, -1))
    flux = np.reshape(flux, (nframes, -1))
    if pca_arr is not None:
        pca_arr = np.reshape(pca_arr, (pca_arr.shape[0], -1))

    ######################################################################
    # Create arrays sorted by radius to facilitate division of frames
    # into optimization and subtraction regions.
    ######################################################################

    print "\nSorting radius array to compute indices at annuli."
    rindex, rsort, thetasort = sorted_arrays(oldshape)

    optreg = int(np.pi * locipar.fwhm**2 / 4 * locipar.npsf)
    minsep = locipar.nfwhm * locipar.fwhm
    if rmin is None:
        rmin = int(minsep * 2 / (pa.max() - pa.min())) + 1
        rmin = max(r_ex - r_ex % 5, rmin - rmin % 5) + 5
    if r_ex is None:
        r_ex = rmin

    ######################################################################
    # Calculate the total number of subtraction regions and the
    # (r, theta) boundaries of those regions.
    ######################################################################

    totn, r1, r2, theta1, theta2 = calczones(rmin, locipar.rmax, locipar.dr0,
                                             optreg)

    ######################################################################
    # Coefficients for flux suppression by a displaced source
    # They are pre-computed for an azimuthally averaged HiCIAO PSF.
    ######################################################################

    if do_partial_sub:
        try:
            sub_coefs = np.zeros((totn, nframes), np.float32)
            x, y = np.loadtxt(sub_dir + '/partial_sub.dat').T
            partial_sub = interpolate.UnivariateSpline(x, y, k=3, s=0)
        except:
            print "Unable to read flux loss data from file partial_sub.dat."
            print "Cannot calculate a fractional flux loss map in loci."
            do_partial_sub = False

    ######################################################################
    # Calculate where to add mock sources for flux loss analysis.
    # We will add a positive source (one at a time) at the center of
    # each subtraction region, and negative sources azimuthally
    # displaced from it.  See Appendix of Brandt+2012 for explanation.
    ######################################################################

    r_avg = 0.5 * (r1 + r2)
    phi = 0.5 * (theta1 + theta2)
    dphi = minsep * 2. / r_avg
    if do_partial_sub:
        np.putmask(dphi, dphi < np.std(pa), np.std(pa))

        ######################################################################
        # Size of PSF regions will be 2 * n + 1.  We will compute them
        # later; just store r**2 for now.
        ######################################################################

        n = 7
        x, y = np.meshgrid(np.arange(2 * n + 1) - n, np.arange(2 * n + 1) - n)
        r_sq_psf = (x**2 + y**2 + 0.) / (locipar.fwhm / 2)**2
        dy, dx = [oldshape[1] // 2 + 0.5, oldshape[2] // 2 + 0.5]

        xp0 = (-r_avg * np.cos(phi) + dx).astype(int) - n
        yp0 = (-r_avg * np.sin(phi) + dy).astype(int) - n
        xp1 = (-r_avg * np.cos(phi - dphi) + dx).astype(int) - n
        yp1 = (-r_avg * np.sin(phi - dphi) + dy).astype(int) - n
        xp2 = (-r_avg * np.cos(phi + dphi) + dx).astype(int) - n
        yp2 = (-r_avg * np.sin(phi + dphi) + dy).astype(int) - n
        xp3 = (-r_avg * np.cos(phi - 2 * dphi) + dx).astype(int) - n
        yp3 = (-r_avg * np.sin(phi - 2 * dphi) + dy).astype(int) - n
        xp4 = (-r_avg * np.cos(phi + 2 * dphi) + dx).astype(int) - n
        yp4 = (-r_avg * np.sin(phi + 2 * dphi) + dy).astype(int) - n

        ######################################################################
        # Two references to the same array, so that we can modify it
        # with two indices and access it with one
        ######################################################################

        sub_arr_full = np.zeros((oldshape[1], oldshape[2]))
        sub_arr = np.reshape(sub_arr_full, -1)

    p = ProgressBar('green', width=30, block='=', lastblock='>', empty=' ')

    if mode == 'LOCI':
        print "Running LOCI on {0} regions in {1} frames out to {2} pixels in radius.".format(totn, pa.size, int(locipar.rmax))
    elif mode == 'refine':
        print "Refining LOCI on {0} regions in {1} frames out to {2} pixels in radius.".format(totn, pa.size, int(locipar.rmax))
    else:
        print "Argument mode=" + mode + " to loci.py not recognized."
        sys.exit()

    ######################################################################
    # Run LOCI on each optimization section.  Do the sections in parallel,
    # with all of the frames processed in each job.  Split into multiple
    # execution blocks due to the large memory demands of the index
    # arrays.
    ######################################################################

    ######################################################################
    # Automatically set niter to use < 10% of the total memory.
    # In the 'refine' mode, this will be doubled, and incomplete freeing
    # seems to add a further factor of 2.
    ######################################################################

    if niter == 0:
        mem_iter = optreg * 4. * nframes / mem
        if mode == 'refine':
            niter = max(int(totn * mem_iter * 60), 1)
        else:
            niter = max(int(totn * mem_iter * 30), 1)

    for i in range(niter):

        tasks = multiprocessing.Queue()
        results = multiprocessing.Queue()
        ncpus = multiprocessing.cpu_count()
        consumers = [ Consumer(tasks, results)
                      for j in range(ncpus) ]
        for w in consumers:
            w.start()

        i1 = totn // niter * i
        i2 = totn // niter * (i + 1)
        if i == niter - 1:
            i2 = totn

        ##################################################################
        # Compute the indices in the original flux array corresponding
        # to the optimization and subtraction regions, using the
        # sorted radius and angle arrays to make initial guesses.
        ##################################################################

        iopt, isub, nopt, nsub = annulus_indices(
            r_ex, r1[i1:i2], r2[i1:i2], theta1[i1:i2], theta2[i1:i2],
            optreg, rindex, rsort, thetasort, p, i1, totn,
            locipar.fwhm, innerfrac=locipar.innerfrac)

        for j in range(i1, i2):
            k = j - i1

            ##################################################################
            # Add pseudo-test sources to compute the fractional flux loss.
            # We need to do this independently for each region, so we will
            # remove these same sources after submitting the jobs.
            ##################################################################

            if do_partial_sub:
                n = r_sq_psf.shape[0]

                sig = np.sort([2.5, (dphi[j] * r_avg[j] / 3), 4])[1]
                sig /= locipar.fwhm / 2.
                psf = np.exp(-r_sq_psf) + 0.05 * np.exp(-r_sq_psf / 3.)
                psf1 = np.exp(-r_sq_psf / (2. * sig**2))
                psf2 = np.exp(-r_sq_psf / (2. * 2 * sig**2))
                psf1 *= np.sum(psf) / np.sum(psf1) / 3.
                psf2 *= np.sum(psf) / np.sum(psf2) / 6.

                sub_arr_full[yp0[j]:yp0[j] + n, xp0[j]:xp0[j] + n] += psf
                sub_arr_full[yp1[j]:yp1[j] + n, xp1[j]:xp1[j] + n] -= psf1
                sub_arr_full[yp2[j]:yp2[j] + n, xp2[j]:xp2[j] + n] -= psf1
                sub_arr_full[yp3[j]:yp3[j] + n, xp3[j]:xp3[j] + n] -= psf2
                sub_arr_full[yp4[j]:yp4[j] + n, xp4[j]:xp4[j] + n] -= psf2

                sub = sub_arr[iopt[k, :nopt[k]]].copy()

                sub_arr_full[yp0[j]:yp0[j] + n, xp0[j]:xp0[j] + n] = 0
                sub_arr_full[yp1[j]:yp1[j] + n, xp1[j]:xp1[j] + n] = 0
                sub_arr_full[yp2[j]:yp2[j] + n, xp2[j]:xp2[j] + n] = 0
                sub_arr_full[yp3[j]:yp3[j] + n, xp3[j]:xp3[j] + n] = 0
                sub_arr_full[yp4[j]:yp4[j] + n, xp4[j]:xp4[j] + n] = 0
            else:
                sub = None

            if corr is not None:
                corr_rad = corr[int(r1[j])]
            else:
                corr_rad = None
            if pca_arr is not None:
                pcaopt = pca_arr[:, iopt[k, 0:nopt[k]]]
                pcasub = pca_arr[:, isub[k, 0:nsub[k]]]
            else:
                pcaopt = None
                pcasub = None

            tasks.put(Task(j, loci_sub, (fluxsmooth[:, iopt[k, 0:nopt[k]]],
                                         flux[:, isub[k, 0:nsub[k]]],
                                         pa, minsep, r_avg[j], pcaopt, pcasub,
                                         partial_sub, ngroup,  method,
                                         corr_rad, sub)))

        ##################################################################
        # "Poison Pills" to break execution
        ##################################################################

        for j in range(ncpus):
            tasks.put(None)

        for j in range(i1, i2):
            p.render([i2 * 100 / totn, (j + 1) * 100 / totn],
                     ['Calculating Indices for Region {0}'.format(i2),
                      'Processing Region {0} at Radius {1}'.format(j + 1, int(r1[j]))])
            index, result = results.get()
            if mode == 'LOCI':
                flux[:, isub[index - i1, 0:nsub[index - i1]]], \
                        sub_coefs[index] = result
            elif mode == 'refine':
                flux[:, isub[index - i1, 0:nsub[index - i1]]] = result

        del iopt, isub, nopt, nsub
        del tasks, results, consumers
        gc.collect()

    flux = np.reshape(flux, oldshape)
    if pca_arr is not None:
        pca_arr = np.reshape(pca_arr, (pca_arr.shape[0], oldshape[1], oldshape[2]))

    ##################################################################
    # Go from fractional flux loss points to a smoothed map
    ##################################################################

    if do_partial_sub:
        print "Calculating map of fractional flux loss."
        return partial_sub_map(r1, r2, theta1, theta2, sub_coefs, pa)
    else:
        return None


"""This module implements an ORIGEN v2.2 transmutation solver.
"""
from __future__ import print_function, division

import os
import subprocess
import tempfile
from collections import Mapping
from warnings import warn
from pyne.utils import QAWarning

import numpy as np

from pyne import data
from pyne import rxname
from pyne import nucname
from pyne import nuc_data
from pyne import origen22
from pyne import utils
from pyne.material import Material, from_atom_frac
from pyne.xs.data_source import NullDataSource, SimpleDataSource, EAFDataSource
from pyne.xs.cache import XSCache

warn(__name__ + " is not yet QA compliant.", QAWarning)

class Transmuter(object):
    """A class for transmuting materials using ORIGEN v2.2."""

    def __init__(self, t=0.0, phi=0.0, temp=300.0, tol=1e-7, cwd='',
                 base_tape9=origen22.BASE_TAPE9, xscache=None,
                 o2exe='o2_therm_linux.exe', *args, **kwargs):
        """Parameters
        ----------
        t : float
            Transmutations time [sec].
        phi : float or array of floats
            Neutron flux vector [n/cm^2/sec].  Currently this must either be
            a scalar or match the group structure of EAF.
        temp : float, optional
            Temperature [K] of material, defaults to 300.0.
        tol : float
            Tolerance level for chain truncation.
        cwd : str, optional
            Current working directory for origen runs. Defaults to this dir.
        base_tape9 : str or dict, optional
            A base TAPE9.INP file.  If this is a str it is interpreted as a path
            to a file, which is then read in and parsed.  If this is a dict, it is
            assumed to be in the format described in the main origen22 module.
        xscache : XSCache, optional
            A cross section cache to generate cross sections with.
        o2exe : str, optional
            Name or path to ORIGEN 2.2 executable.
        args : tuple, optional
            Other arguments ignored for compatibility with other Transmuters.
        kwargs : dict, optional
            Other keyword arguments ignored for compatibility with other Transmuters.
        """
        if not isinstance(base_tape9, Mapping):
            base_tape9 = origen22.parse_tape9(tape9=base_tape9)
        self.base_tape9 = base_tape9

        if xscache is None:
            eafds = EAFDataSource()
            eafds.load(temp=temp)
            gs = np.array([eafds.src_group_struct[0], eafds.src_group_struct[-1]])
            eafds.dst_group_struct = gs
            xscache = XSCache(group_struct=gs, data_source_classes=[SimpleDataSource,
                                                                    NullDataSource])
            xscache.load(temp=temp)
            xscache.data_sources.insert(0, eafds)
        self.xscache = xscache

        self.t = t
        self._phi = None
        self.phi = phi
        self.temp = temp
        self.tol = tol
        self.cwd = os.path.abspath(cwd)
        self.o2exe = o2exe

    @property
    def phi(self):
        return self._phi

    @phi.setter
    def phi(self, flux):
        """Ensures that the flux is correctly formatted."""
        flux = np.asarray(flux)
        if flux.ndim == 0:
            _ = np.empty(175, float)
            _.fill(flux / 175.0)
            flux = _
        elif flux.ndim == 1 and flux.shape[0] != 175:
            raise ValueError("Group structure must match EAF.")
        elif flux.ndim > 1:
            raise ValueError("The flux vector must be 0- or 1-dimensional.")
        if not np.all(flux >= 0.0):
            raise ValueError("Flux entries must be non-negative.")
        for ds in self.xscache.data_sources:
            ds.src_phi_g = flux
        self.xscache['phi_g'] = np.array([flux.sum()])
        self._phi = flux

    def transmute(self, x, t=None, phi=None, tol=None, cwd=None, xscache=None,
                  o2exe=None, *args, **kwargs):
        """Transmutes a material into its daughters.

        Parameters
        ----------
        x : Material or similar
            Input material for transmutation.
        t : float
            Transmutations time [sec].
        phi : float or array of floats
            Neutron flux vector [n/cm^2/sec].  Currently this must either be
            a scalar or match the group structure of EAF.
        tol : float
            Tolerance level for chain truncation.
        cwd : str, optional
            Current working directory for origen runs. Defaults to this dir.
        xscache : XSCache, optional
            A cross section cache to generate cross sections with.
        o2exe : str, optional
            Name or path to ORIGEN 2.2 executable.

        Returns
        -------
        y : Material
            The output material post-transmutation.

        """
        if not isinstance(x, Material):
            x = Material(x)
        if t is not None:
            self.t = t
        if phi is not None:
            self.phi = phi
        if tol is not None:
            self.tol = tol
        if cwd is not None:
            self.cwd = os.path.abspath(cwd)
        if xscache is not None:
            self.xscache = xscache
        if o2exe is not None:
            self.o2exe = o2exe

        # prepare new tape9
        nucs = set(x.comp.keys())
        base_tape9 = self.base_tape9
        decay_nlb, xsfpy_nlb = origen22.nlbs(base_tape9)
        new_tape9 = origen22.xslibs(nucs=nucs, xscache=self.xscache, nlb=xsfpy_nlb)
        t9 = origen22.merge_tape9([new_tape9, base_tape9])

        # write out files
        origen22.write_tape4(x, outfile=os.path.join(self.cwd, 'TAPE4.INP'))
        origen22.write_tape5_irradiation('IRF', self.t/86400.0, self.xscache['phi_g'][0],
            outfile=os.path.join(self.cwd, 'TAPE5.INP'), decay_nlb=decay_nlb,
            xsfpy_nlb=xsfpy_nlb, cut_off=self.tol)
        origen22.write_tape9(t9, outfile=os.path.join(self.cwd, 'TAPE9.INP'))

        # run origen & get results
        f = tempfile.NamedTemporaryFile()
        try:
            subprocess.check_call([self.o2exe], cwd=self.cwd, stdout=f, stderr=f)
        except subprocess.CalledProcessError:
            f.seek(0)
            print("ORIGEN output:\n\n{0}".format(f.read()))
            raise
        finally:
            f.close()
        t6 = origen22.parse_tape6(tape6=os.path.join(self.cwd, 'TAPE6.OUT'))
        y = t6['materials'][-1]
        return y


from jax.config import config; config.update("jax_enable_x64", True)

import jax
import numpy as np

from fe import topology

from timemachine.lib import potentials, custom_ops
from timemachine.lib import LangevinIntegrator

from ff.handlers import openmm_deserializer

def get_romol_conf(mol):
    """Coordinates of mol's 0th conformer, in nanometers"""
    conformer = mol.GetConformer(0)
    guest_conf = np.array(conformer.GetPositions(), dtype=np.float64)
    return guest_conf/10 # from angstroms to nm

class BaseFreeEnergy():

    @staticmethod
    def _get_integrator(combined_masses):
        """
        Get a integrator. The resulting impl must be bound to a python handle
        whose lifetime is concurrent with that of the context.
        """
        seed = np.random.randint(np.iinfo(np.int32).max)

        return LangevinIntegrator(
            300.0,
            1.5e-3,
            1.0,
            combined_masses,
            seed
        )

    # this will be eventually gRPC'd out to a worker
    @staticmethod
    def _simulate(lamb, box, x0, v0, final_potentials, integrator, equil_steps, prod_steps):
        all_impls = []
        bonded_impls = []
        nonbonded_impls = []

        # set up observables for du_dps here as well.

        du_dp_obs = []

        for bps in final_potentials:
            obs_list = []

            for bp in bps:
                impl = bp.bound_impl(np.float32)

                if isinstance(bp, potentials.InterpolatedPotential) or isinstance(bp, potentials.LambdaPotential):
                    bp = bp.get_u_fn()

                if isinstance(bp, potentials.Nonbonded):
                    nonbonded_impls.append(impl)
                else:
                    bonded_impls.append(impl)

                all_impls.append(impl)
                obs_list.append(custom_ops.AvgPartialUPartialParam(impl, 5))

            du_dp_obs.append(obs_list)

        intg_impl = integrator.impl()
        # context components: positions, velocities, box, integrator, energy fxns
        ctxt = custom_ops.Context(
            x0,
            v0,
            box,
            intg_impl,
            all_impls
        )

        # equilibration
        for step in range(equil_steps):
            ctxt.step(lamb)

        bonded_du_dl_obs = custom_ops.FullPartialUPartialLambda(bonded_impls, 5)
        nonbonded_du_dl_obs = custom_ops.FullPartialUPartialLambda(nonbonded_impls, 5)

        # add observable
        ctxt.add_observable(bonded_du_dl_obs)
        ctxt.add_observable(nonbonded_du_dl_obs)

        for obs_list in du_dp_obs:
            for obs in obs_list:
                ctxt.add_observable(obs)

        for _ in range(prod_steps):
            ctxt.step(lamb)

        bonded_full_du_dls = bonded_du_dl_obs.full_du_dl()
        nonbonded_full_du_dls = nonbonded_du_dl_obs.full_du_dl()

        bonded_mean, bonded_std = np.mean(bonded_full_du_dls), np.std(bonded_full_du_dls)
        nonbonded_mean, nonbonded_std = np.mean(nonbonded_full_du_dls), np.std(nonbonded_full_du_dls)

        # keep the structure of grads the same as that of final_potentials so we can properly
        # form their vjps.
        grads = []
        for obs_list in du_dp_obs:
            grad_list = []
            for obs in obs_list:
                grad_list.append(obs.avg_du_dp())
            grads.append(grad_list)

        return (bonded_mean, bonded_std), (nonbonded_mean, nonbonded_std), grads

# this class is serializable.
class AbsoluteFreeEnergy(BaseFreeEnergy):

    def __init__(self, mol, ff):
        """
        Compute the absolute free energy of a molecule via 4D decoupling.

        Parameters
        ----------
        mol: rdkit mol
            Ligand to be decoupled

        ff: ff.Forcefield
            Ligand forcefield

        """
        self.mol = mol
        self.ff = ff
        self.top = topology.BaseTopology(mol, ff)

    # this can be used for both the solvent leg and the complex leg
    def host_edge(self, lamb, host_system, host_coords, box, equil_steps=10000, prod_steps=100000):
        """
        Run equilibrium decoupling simulation at a given value of lambda in a host environment.

        Parameters
        ----------
        lamb: float [0, 1]
            0 is the fully interacting system, and 1 is the non-interacting system

        host_system: openmm.System
            OpenMM System object to be deserialized. The host can be simply a box of water, or a fully
            solvated protein

        host_coords: np.array of shape [..., 3]
            Host coordinates, in nanometers. It should be properly minimized and not have clashes
            with the ligand coordinates.

        box: np.array [3,3]
            Periodic boundary conditions, in nanometers.

        equil_steps: float
            Number of steps to run equilibration. Statistics are not gathered.

        prod_steps: float
            Number of steps to run production. Statistics are gathered.

        Returns
        -------
        float, float
            Returns a pair of average du_dl values for bonded and nonbonded terms.

        """

        ligand_masses = [a.GetMass() for a in self.mol.GetAtoms()]
        ligand_coords = get_romol_conf(self.mol)

        host_bps, host_masses = openmm_deserializer.deserialize_system(host_system, cutoff=1.2)
        num_host_atoms = host_coords.shape[0]

        final_potentials = []
        final_vjp_and_handles = []

        for bp in host_bps:
            if isinstance(bp, potentials.Nonbonded):
                host_p = bp
            else:
                final_potentials.append([bp])
                final_vjp_and_handles.append(None)

        hgt = topology.HostGuestTopology(host_p, self.top)

        # setup the parameter handlers for the ligand
        bonded_tuples = [
            [hgt.parameterize_harmonic_bond, self.ff.hb_handle],
            [hgt.parameterize_harmonic_angle, self.ff.ha_handle],
            [hgt.parameterize_proper_torsion, self.ff.pt_handle],
            [hgt.parameterize_improper_torsion, self.ff.it_handle]
        ]

        # instantiate the vjps while parameterizing (forward pass)
        for fn, handle in bonded_tuples:
            params, vjp_fn, potential = jax.vjp(fn, handle.params, has_aux=True)
            final_potentials.append([potential.bind(params)])
            final_vjp_and_handles.append((vjp_fn, handle))

        nb_params, vjp_fn, nb_potential = jax.vjp(hgt.parameterize_nonbonded, self.ff.q_handle.params, self.ff.lj_handle.params, has_aux=True)
        final_potentials.append([nb_potential.bind(nb_params)])
        final_vjp_and_handles.append([vjp_fn])

        combined_masses = np.concatenate([host_masses, ligand_masses])
        combined_coords = np.concatenate([host_coords, ligand_coords])

        return self._simulate(
            lamb,
            box,
            combined_coords,
            np.zeros_like(combined_coords),
            final_potentials,
            self._get_integrator(combined_masses),
            equil_steps,
            prod_steps
        )

# this class is serializable.
class RelativeFreeEnergy(BaseFreeEnergy):

    def __init__(self, mol_a, mol_b, core, ff):
        self.mol_a = mol_a
        self.mol_b = mol_b
        self.core = core
        self.ff = ff
        self.top = topology.SingleTopology(mol_a, mol_b, core, ff)

    def _get_integrator(self, combined_masses):
        """
        Get a integrator. The resulting impl must be bound to a python handle
        whose lifetime is concurrent with that of the context.
        """
        seed = np.random.randint(np.iinfo(np.int32).max)

        return LangevinIntegrator(
            300.0,
            1.5e-3,
            1.0,
            combined_masses,
            seed
        )

    def vacuum_edge(self, lamb, equil_steps=10000, prod_steps=100000):
        """
        Run a vacuum decoupling simulation at a given value of lambda.

        Parameters
        ----------
        lamb: float [0, 1]
            0 is the fully interacting system, and 1 is the non-interacting system

        equil_steps: float
            Number of steps to run equilibration. Statistics are not gathered.

        prod_steps: float
            Number of steps to run production. Statistics are gathered.

        Returns
        -------
        float, float
            Returns a pair of average du_dl values for bonded and nonbonded terms.

        """
        final_potentials = []
        final_vjp_and_handles = []

        ligand_masses_a = [a.GetMass() for a in self.mol_a.GetAtoms()]
        ligand_masses_b = [b.GetMass() for b in self.mol_b.GetAtoms()]

        ligand_coords_a = get_romol_conf(self.mol_a)
        ligand_coords_b = get_romol_conf(self.mol_b)

        bonded_tuples = [
            [self.top.parameterize_harmonic_bond, self.ff.hb_handle],
            [self.top.parameterize_harmonic_angle, self.ff.ha_handle],
            [self.top.parameterize_proper_torsion, self.ff.pt_handle],
            [self.top.parameterize_improper_torsion, self.ff.it_handle]
        ]

        # instantiate the vjps while parameterizing (forward pass)
        for fn, handle in bonded_tuples:
            (src_params, dst_params, uni_params), vjp_fn, (src_potential, dst_potential, uni_potential) = jax.vjp(fn, handle.params, has_aux=True)
            final_potentials.append([src_potential.bind(src_params), dst_potential.bind(dst_params), uni_potential.bind(uni_params)])
            final_vjp_and_handles.append((vjp_fn, handle))

        nb_params, vjp_fn, nb_potential = jax.vjp(self.top.parameterize_nonbonded, self.ff.q_handle.params, self.ff.lj_handle.params, has_aux=True)
        final_potentials.append([nb_potential.bind(nb_params)])
        final_vjp_and_handles.append([vjp_fn])

        combined_masses = np.mean(self.top.interpolate_params(ligand_masses_a, ligand_masses_b), axis=0)

        src_conf, dst_conf = self.top.interpolate_params(ligand_coords_a, ligand_coords_b)
        combined_coords = np.mean(self.top.interpolate_params(ligand_coords_a, ligand_coords_b), axis=0)

        box = np.eye(3) * 100.0

        return self._simulate(
            lamb,
            box,
            combined_coords,
            np.zeros_like(combined_coords),
            final_potentials,
            self._get_integrator(combined_masses),
            equil_steps,
            prod_steps
        )

    def host_edge(self, lamb, host_system, host_coords, box, equil_steps=10000, prod_steps=100000):
        """
        Run equilibrium decoupling simulation at a given value of lambda in a host environment.

        Parameters
        ----------
        lamb: float [0, 1]
            0 is the fully interacting system, and 1 is the non-interacting system

        host_system: openmm.System
            OpenMM System object to be deserialized. The host can be simply a box of water, or a fully
            solvated protein

        host_coords: np.array of shape [..., 3]
            Host coordinates, in nanometers. It should be properly minimized and not have clashes
            with the ligand coordinates.

        box: np.array [3,3]
            Periodic boundary conditions, in nanometers.

        equil_steps: float
            Number of steps to run equilibration. Statistics are not gathered.

        prod_steps: float
            Number of steps to run production. Statistics are gathered.

        Returns
        -------
        float, float
            Returns a pair of average du_dl values for bonded and nonbonded terms.

        """

        ligand_masses_a = [a.GetMass() for a in self.mol_a.GetAtoms()]
        ligand_masses_b = [b.GetMass() for b in self.mol_b.GetAtoms()]

        # extract the 0th conformer
        ligand_coords_a = get_romol_conf(self.mol_a)
        ligand_coords_b = get_romol_conf(self.mol_b)

        host_bps, host_masses = openmm_deserializer.deserialize_system(host_system, cutoff=1.2)
        num_host_atoms = host_coords.shape[0]

        final_potentials = []
        final_vjp_and_handles = []

        # keep the bonded terms in the host the same.
        # but we keep the nonbonded term for a subsequent modification
        for bp in host_bps:
            if isinstance(bp, potentials.Nonbonded):
                host_p = bp
            else:
                final_potentials.append([bp])
                # (ytz): no protein ff support for now, so we skip their vjps
                final_vjp_and_handles.append(None)

        hgt = topology.HostGuestTopology(host_p, self.top)

        # setup the parameter handlers for the ligand
        bonded_tuples = [
            [hgt.parameterize_harmonic_bond, self.ff.hb_handle],
            [hgt.parameterize_harmonic_angle, self.ff.ha_handle],
            [hgt.parameterize_proper_torsion, self.ff.pt_handle],
            [hgt.parameterize_improper_torsion, self.ff.it_handle]
        ]

        # instantiate the vjps while parameterizing (forward pass)
        for fn, handle in bonded_tuples:
            (src_params, dst_params, uni_params), vjp_fn, (src_potential, dst_potential, uni_potential) = jax.vjp(fn, handle.params, has_aux=True)
            final_potentials.append([src_potential.bind(src_params), dst_potential.bind(dst_params), uni_potential.bind(uni_params)])
            final_vjp_and_handles.append((vjp_fn, handle))

        nb_params, vjp_fn, nb_potential = jax.vjp(hgt.parameterize_nonbonded, self.ff.q_handle.params, self.ff.lj_handle.params, has_aux=True)
        final_potentials.append([nb_potential.bind(nb_params)])
        final_vjp_and_handles.append([vjp_fn, (self.ff.q_handle, self.ff.lj_handle)]) # (ytz): note the handlers are a tuple, this is checked later

        combined_masses = np.concatenate([host_masses, np.mean(self.top.interpolate_params(ligand_masses_a, ligand_masses_b), axis=0)])

        src_conf, dst_conf = self.top.interpolate_params(ligand_coords_a, ligand_coords_b)
        combined_coords = np.concatenate([host_coords, np.mean(self.top.interpolate_params(ligand_coords_a, ligand_coords_b), axis=0)])

        # (ytz): us is short form for mean and std dev.
        bonded_us, nonbonded_us, grads = self._simulate(
            lamb,
            box,
            combined_coords,
            np.zeros_like(combined_coords),
            final_potentials,
            self._get_integrator(combined_masses),
            equil_steps,
            prod_steps
        )

        grads_and_handles = []

        for du_dqs, vjps_and_handles in zip(grads, final_vjp_and_handles):
            if vjps_and_handles is not None:
                vjp_fn = vjps_and_handles[0]
                handles = vjps_and_handles[1]

                # we need to get the shapes correct (eg. nonbonded vjp emits an ndarray, not a list.)

                # (ytz): so far nonbonded grads is the only term that map back out to two
                # vjp handlers (charge and lj). the vjp also expects an nd.array, not a list. So we kill
                # two birds with one stone here, but this is quite brittle and should be refactored later on.
                if type(handles) == tuple:
                    # handle nonbonded terms
                    du_dps = vjp_fn(du_dqs[0])
                    for du_dp, handler in zip(du_dps, handles):
                        grads_and_handles.append((du_dp, type(handler)))
                else:
                    du_dp = vjp_fn(du_dqs)
                    # bonded terms return a list, so we need to flatten it here
                    grads_and_handles.append((du_dp[0], type(handles)))

        return bonded_us, nonbonded_us, grads_and_handles