#! /usr/bin/python
import math
import os
import string
import struct
import sys
import time
import csv
import numpy

#functions including
#1. simple classes to store all the information about each species and reaction.
#2. Functions to read in the species and reaction file and check for sanity
#3. Functions to write out files necessary for UCLCHEM

##########################################################################################
#1. simple classes to store all the information about each species and reaction.
#largely just to make the other functions more readable.
##########################################################################################
class Species:
	def __init__(self, inputRow):
		self.name = inputRow[0]
		self.mass = inputRow[1]
		self.bindener = float(inputRow[2])
		self.solidFraction = float(inputRow[3])
		self.monoFraction = float(inputRow[4])
		self.volcFraction = float(inputRow[5])
		self.enthalpy = float(inputRow[6])

	def is_grain_species(self):
		return self.name[0] == '#'

	def is_ion(self):
		return (self.name[-1] == "+" or self.name[-1] == "-")

class Reaction:
	def __init__(self, inputRow):
		self.reactants = [inputRow[0], inputRow[1], self.NANCheck(inputRow[2])]
		self.products = [inputRow[3], self.NANCheck(
			inputRow[4]), self.NANCheck(inputRow[5]), self.NANCheck(inputRow[6])]
		self.alpha = float(inputRow[7])
		self.beta = float(inputRow[8])
		self.gamma = float(inputRow[9])
		self.templow = inputRow[10]
		self.temphigh = inputRow[11]

	def NANCheck(self, a):
		aa = a if a else 'NAN'
		return aa

##########################################################################################
#2. Functions to read in the species and reaction file and check for sanity
##########################################################################################

# Read the entries in the specified species file
def read_species_file(fileName):
	speciesList = []
	f = open(fileName, 'rb')
	reader = csv.reader(f, delimiter=',', quotechar='|')
	for row in reader:
		if row[0] != "NAME" and "!" not in row[0]:
			speciesList.append(Species(row))
	nSpecies = len(speciesList)
	return nSpecies, speciesList

# Read the entries in the specified reaction file and keep the reactions that involve the species in our species list

def read_reaction_file(fileName, speciesList, ftype):
	reactions = []
	keepList = []
	# keeplist includes the elements that ALL the reactions should be formed from
	keepList.extend(['', 'NAN', '#', 'E-', 'e-', 'ELECTR', 'PHOTON', 'CRP', 'CRPHOT', 'FREEZE', 'CRH',
                  'PHOTD', 'THERM', 'XRAY', 'XRSEC', 'XRLYA', 'XRPHOT', 'DESOH2', 'DESCR', 'DEUVCR', "CHEMDES", "DIFF"])
	for species in speciesList:
		keepList.append(species.name)
	if ftype == 'UMIST':  # if it is a umist database file
		f = open(fileName, 'rb')
		reader = csv.reader(f, delimiter=':', quotechar='|')
		for row in reader:
			# if all the reaction elements belong to the keeplist
			if all(x in keepList for x in [row[2], row[3], row[4], row[5], row[6], row[7]]):
				#umist file doesn't have third reactant so add space and has a note for how reactions there are so remove that
				reactions.append(Reaction(row[2:4]+['']+row[4:8]+row[9:]))
	if ftype == 'UCL':  # if it is a ucl made (grain?) reaction file
		f = open(fileName, 'rb')
		reader = csv.reader(f, delimiter=',', quotechar='|')
		for row in reader:
			# if all the reaction elements belong to the keeplist
			if all(x in keepList for x in row[0:7]):
				row[10] = 0.0
				row[11] = 10000.0
				reactions.append(Reaction(row))

	nReactions = len(reactions)
	return nReactions, reactions

def remove_duplicate_species(speciesList):
		#check for duplicate species
	duplicates = 0
	duplicate_list = []
	for i in range(0, len(speciesList)):
		for j in range(0, len(speciesList)):
			if speciesList[i].name == speciesList[j].name:
				if (j != i) and speciesList[i].name not in duplicate_list:
					print "\t {0} appears twice in input species list".format(speciesList[i].name)
					duplicate_list.append(speciesList[i].name)

	for duplicate in duplicate_list:
		removed = False
		i = 0
		while not removed:
			if speciesList[i].name == duplicate:
				del speciesList[i]
				print "\tOne entry of {0} removed from list".format(duplicate)
				removed = True
			else:
				i += 1
	return speciesList

#Look for possibly incorrect parts of species list

def filter_species(speciesList, reactionList):
	#check for species not involved in any reactions
	lostSpecies = []
	for species in speciesList:
		keepFlag = False
		for reaction in reactionList:
			if species.name in reaction.reactants or species.name in reaction.products:
				keepFlag = True
		if not keepFlag:
			lostSpecies.append(species.name)
			speciesList.remove(species)

	print '\tSpecies in input list that do not appear in final list:'
	print '\t', lostSpecies
	print '\n'
	return speciesList

#All species should freeze out at least as themselves and all grain species should desorb according to their binding energy
#This function adds those reactions automatically to slim down the grain file

def add_desorb_reactions(speciesList, reactionList):
	for species in speciesList:
		if species.is_grain_species():
			for reacType in ['DESOH2', "DESCR", "DEUVCR"]:
				newReaction = Reaction([species.name, reacType, 'NAN', species.name[1:],
                                    'NAN', 'NAN', 'NAN', 1, 0, species.bindener, 0.0, 10000.0])
				reactionList.append(newReaction)
	return reactionList

#check reactions to alert user of potential issues including repeat reactions
#and multiple freeze out routes

def reaction_check(speciesList, reactionList, freeze_check=True):

	#first check for multiple freeze outs so user knows to do alphas
	print "\tSpecies with multiple freeze outs, check alphas:"
	for spec in speciesList:
		freezes = 0
		for reaction in reactionList:
			if (spec.name in reaction.reactants and 'FREEZE' in reaction.reactants):
				freezes += 1
		if (freezes > 1):
			print "\t{0} freezes out through {1} routes".format(spec.name, freezes)
		if freezes < 1 and not spec.is_grain_species() and freeze_check:
			print "\t{0} does not freeze out".format(spec.name, freezes)
	#now check for duplicate reactions
	duplicate_list = []
	print "\n\tPossible duplicate reactions for manual removal:"
	duplicates = 0
	for i, reaction1 in enumerate(reactionList):
		if i not in duplicate_list:
			for j, reaction2 in enumerate(reactionList):
				if i != j:
					if set(reaction1.reactants) == set(reaction2.reactants):
						if set(reaction1.products) == set(reaction2.products):
							print "\tReactions {0} and {1} are possible duplicates".format(i+1, j+1)
							print "\t", str(i+1), reaction1.reactants, "-->", reaction1.products
							print "\t", str(j+1), reaction1.reactants, "-->", reaction2.products
							duplicates += 1
							duplicate_list.append(i)
							duplicate_list.append(j)

	if (duplicates == 0):
		print "\tNone"

#capitalize files

def make_capitals(fileName):
	a = open(fileName).read()
	output = open(fileName, mode='w')
	output.write(a.upper())
	output.close()

def find_constituents(speciesList):
	elementList = ['H', 'D', 'HE', 'C', 'N', 'O', 'F',
                'P', 'S', 'CL', 'LI', 'NA', 'MG', 'SI', 'PAH', '15N']
	elementMass = [1, 2, 4, 12, 14, 16, 19, 31, 32, 35, 3, 23, 24, 28, 420, 15]
	symbols = ['#', '+', '-', '(', ')']

	for species in speciesList:
		speciesName = species.name
		i = 0
		atoms = []
		bracket = False
		bracketContent = []
		#loop over characters in species name to work out what it is made of
		while i < len(speciesName):
			#if character isn't a #,+ or - then check it otherwise move on
			if speciesName[i] not in symbols:
				if i+1 < len(speciesName):
					#if next two characters are (eg) 'MG' then atom is Mg not M and G
					if speciesName[i:i+3] in elementList:
						j = i+3
					elif speciesName[i:i+2] in elementList:
						j = i+2
					#otherwise work out which element it is
					elif speciesName[i] in elementList:
						j = i+1

				#if there aren't two characters left just try next one
				elif speciesName[i] in elementList:
					j = i+1
				#if we've found a new element check for numbers otherwise print error
				if j > i:
					if bracket:
						bracketContent.append(speciesName[i:j])
					else:
						atoms.append(speciesName[i:j])  # add element to list
					if j < len(speciesName):
						if is_number(speciesName[j]):
							if int(speciesName[j]) > 1:
								for k in range(1, int(speciesName[j])):
									if bracket:
										bracketContent.append(speciesName[i:j])
									else:
										atoms.append(speciesName[i:j])
								i = j+1
							else:
								i = j
						else:
							i = j
					else:
						i = j
				else:
					print speciesName[i]
					print"\t{0} contains elements not in element list:".format(speciesName)
					print elementList
			else:
				#if symbol is start of a bracketed part of molecule, keep track
				if (speciesName[i] == "("):
					bracket = True
					bracketContent = []
					i += 1
				#if it's the end then add bracket contents to list
				elif speciesName[i] == ")":
					if is_number(speciesName[i+1]):
						for k in range(0, int(speciesName[i+1])):
							atoms.extend(bracketContent)
						i += 2
					else:
						atoms.extend(bracketContent)
						i += 1
				#otherwise move on
				else:
					i += 1

		species.n_atoms = len(atoms)
		mass = 0
		for atom in atoms:
			mass += elementMass[elementList.index(atom)]
		if mass != float(species.mass):
			#print "\tcalculated mass of {0} does not match input mass".format(speciesName)
			#print "\tcalculated mass: {0} \t input mass: {1}\n".format(mass,species.mass)
			#print "\tkeeping calculated mass\n"
			species.mass = str(mass)
	return speciesList

def is_number(s):
    try:
        float(s)
        return True
    except ValueError:
        return False
##########################################################################################
#3. Functions to write out files necessary for UCLCHEM
##########################################################################################

# Write the species file in the desired format

def write_species(fileName, speciesList):
	f = open(fileName, 'wb')
	writer = csv.writer(f, delimiter=',', quotechar='|',
	                    quoting=csv.QUOTE_MINIMAL, lineterminator='\n')
	nSpecies = len(speciesList)
	for species in speciesList:
		writer.writerow([species.name, species.mass, species.n_atoms])

# Write the reaction file in the desired format

def write_reactions(fileName, reactionList):
	f = open(fileName, 'wb')
	writer = csv.writer(f, delimiter=',', quotechar='|',
	                    quoting=csv.QUOTE_MINIMAL, lineterminator='\n')
	nReactions = len(reactionList)
	for reaction in reactionList:
		#if statement changes beta for ion freeze out to 1. This is how ucl_chem recognises ions when calculating freeze out rate
		if ('FREEZE' in reaction.reactants and reaction.reactants[0][-1] == '+'):
			reaction.beta = 1
		writer.writerow(reaction.reactants+reaction.products +
		                [reaction.alpha, reaction.beta, reaction.gamma, reaction.templow, reaction.temphigh])

def write_odes_f90(fileName, speciesList, reactionList):
	output = open(fileName, mode='w')
	#go through every species and build two strings, one with eq for all destruction routes and one for all formation
	ydotString = build_ode_string(speciesList, reactionList)
	output.write(ydotString)
	output.close()

def build_ode_string(speciesList, reactionList):
	odeString = ""
	nSpecies = len(speciesList)
	nReactions = len(reactionList)

	for n, species in enumerate(speciesList):
		lossString = ''
		formString = ''
		#go through entire reaction list
		for i, reaction in enumerate(reactionList):

			#if species appear in reactants, reaction is a destruction route
			if species.name in reaction.reactants:
				bodyCount = 0  # two or more bodies in a reaction mean we multiply rate by density so need to keep track
				#easy for h2 formation
				if is_H2_formation(reaction.reactants, reaction.products):
					lossString += '-2*RATE('+str(i+1)+')*D'
					continue
				#multiply string by number of time species appears in reaction. multiple() defined below
				#so far reaction string is rate(reaction_index) indexs are all +1 for fortran array indexing
				lossString += '-' + \
					multiple(reaction.reactants.count(species.name))+'RATE('+str(i+1)+')'

				#now add *Y(species_index) to string for every reactant
				for reactant in set(reaction.reactants):
					n_appearances = reaction.reactants.count(reactant)
					#every species appears at least once in its loss reaction
					#so we multiply entire loss string by Y(species_index) at end
					#thus need one less Y(species_index) per reaction
					if reactant == species.name:
						bodyCount += n_appearances
						if n_appearances > 1:
							for j, possibleReactants in enumerate(speciesList):
								if reactant == possibleReactants.name:
									for appearance in range(1, n_appearances):
										lossString += '*Y('+str(j+1)+')'
									continue
					else:
						#look through species list and find reactant
						for j, possibleReactants in enumerate(speciesList):
							if reactant == possibleReactants.name:
								for appearance in range(n_appearances):
									lossString += '*Y('+str(j+1)+')'
									bodyCount += 1
								continue
				#now string is rate(reac_index)*Y(species_index1)*Y(species_index2) may need *D if total rate is
				#proportional to density
				if reaction.reactants.count('FREEZE') > 0 or reaction.reactants.count('DESOH2') > 0:
					lossString += '*D'
				for body in range(1, bodyCount):
					lossString += "*D"

			#same process as above but rate is positive for reactions where species is positive
			if species.name in reaction.products:
				bodyCount = 0  # two or more bodies in a reaction mean we multiply rate by density so need to keep track

				if is_H2_formation(reaction.reactants, reaction.products):
					#honestly H should be index 1 but lets check
					H_index = speciesList.index(
						next((x for x in speciesList if x.name == 'H')))
					formString += '+RATE('+str(i+1)+')*Y('+str(H_index+1)+')*D'
					continue

				#multiply string by number of time species appears in reaction. multiple() defined below
				#so far reaction string is rate(reaction_index) indexs are all +1 for fortran array indexing
				formString += '+' + \
					multiple(reaction.products.count(species.name))+'RATE('+str(i+1)+')'

				#now add *Y(species_index) to string for every reactant
				for reactant in set(reaction.reactants):
					n_appearances = reaction.reactants.count(reactant)
					for j, possibleReactants in enumerate(speciesList):
						if reactant == possibleReactants.name:
							for appearance in range(n_appearances):
								formString += '*Y('+str(j+1)+')'
								bodyCount += 1
							continue

				#now string is rate(reac_index)*Y(species_index1)*Y(species_index2) may need *D if total rate is
				#proportional to density
				if reaction.reactants.count('FREEZE') > 0 or reaction.reactants.count('DESOH2') > 0:
					formString += '*D'
				for body in range(1, bodyCount):
					formString += "*D"

		if lossString != '':
			lossString = '      LOSS = '+lossString+'\n'
			lossString = truncate_line(lossString)
			odeString += lossString
		if formString != '':
			formString = '      PROD = '+formString+'\n'
			formString = truncate_line(formString)
			odeString += formString

		#start with empty string and add production and loss terms if they exists
		ydotString = ''
		if formString != '':
			ydotString += 'PROD'
			if lossString != '':
				ydotString += '+'
		if lossString != '':
			ydotString += 'Y('+str(n+1)+')*LOSS'

		#if we have prod and/or loss add ydotstring to odes
		if ydotString != '':
			ydotString = '      YDOT('+str(n+1)+') = '+ydotString+"\n"
			ydotString = truncate_line(ydotString)
			odeString += ydotString
		else:
			ydotString = '      YDOT('+str(n+1)+') = 0.0\n'
			ydotString = truncate_line(ydotString)
			odeString += ydotString
	return odeString

#create a file containing length of each list of moleculetypes and then the two lists (gas and grain) of species in each type
#as  well as fraction that evaporated in each type of event

def write_evap_lists(openFile, speciesList):
	grainlist = []
	mgrainlist = []
	solidList = []
	monoList = []
	volcList = []
	bindEnergyList = []
	enthalpyList = []

	for i, species in enumerate(speciesList):
		if species.name[0] == '#':
			#find gas phase version of grain species. For #CO it looks for first species in list with just CO and then finds the index of that
			try:
				j = speciesList.index(
					next((x for x in speciesList if x.name == species.name[1:])))
			except:
				print "\n**************************************\nWARNING\n**************************************"
				print "{0} has no gas phase equivalent in network. Every species should at least freeze out and desorb.".format(species.name)
				print "ensure {0} is in the species list, and at least one reaction involving it exists and try again".format(species.name[1:])
				print "Alternatively, provide the name of the gas phase species you would like {0} to evaporate as".format(species.name)
				input = raw_input(
					"type x to quit Makerates or any species name to continue\n")
				if input.lower() == "x":
					exit()
				else:
					j = speciesList.index(
						next((x for x in speciesList if x.name == input.upper())))

			#plus ones as fortran and python label arrays differently
			mgrainlist.append(i+1)
			grainlist.append(j+1)
			solidList.append(species.solidFraction)
			monoList.append(species.monoFraction)
			volcList.append(species.volcFraction)
			bindEnergyList.append(species.bindener)
			enthalpyList.append(species.enthalpy)

	openFile.write(array_to_string("gasGrainList", grainlist, type="int"))
	openFile.write(array_to_string("grainList", mgrainlist, type="int"))
	openFile.write(array_to_string("solidFractions", solidList, type="float"))
	openFile.write(array_to_string("monoFractions", monoList, type="float"))
	openFile.write(array_to_string("volcanicFractions", volcList, type="float"))
	openFile.write(array_to_string(
		"bindingEnergy", bindEnergyList, type="float", parameter=False))
	openFile.write(array_to_string(
		"formationEnthalpy", enthalpyList, type="float"))

	return len(grainlist)

# Create the appropriate multiplication string for a given number

def multiple(number):
    if number == 1:
    	return ''
    else:
    	return str(number)+'*'

def truncate_line(input, codeFormat='F90', continuationCode=None):
	lineLength = 72
	maxlines = 300
	lines = 0
	result = ''
	i = 0
	j = 0
	while i+j < len(input):
		j += 1
		if j > lineLength:
			#important not to break entries so split lines at ,s
			try:
				k = input[i+j-16:i+j].index(",")
			except:
				try:
					k = input[i+j-16:i+j].index("*")
				except:
					k = input[i+j-16:i+j].index(")")
			j = j-16+k
			result += input[i:i+j]+"&\n    &"
			i = i+j
			j = 0
	result += input[i:i+j]
	return result

def is_H2_formation(reactants, products):
    nReactants = len([species for species in reactants if species != ''])
    nProducts = len([species for species in products if species != ''])
    if nReactants == 2 and nProducts == 1:
        if reactants[0] == 'H' and reactants[1] == 'H' and products[0] == 'H2':
        	return True
    if nReactants == 3 and nProducts == 2:
        if reactants[0] == 'H' and reactants[1] == 'H' and reactants[2] == '#' and products[0] == 'H2' and products[1] == '#':
        	return True
    return False

def write_network_file(fileName, speciesList, reactionList):
	openFile = open(fileName, "wb")
	openFile.write("integer, parameter :: nSpec={0}, nReac={1}\n".format(
		len(speciesList), len(reactionList)))

	#write arrays of all species stuff
	names = []
	atoms = []
	masses = []
	for species in speciesList:
		names.append(species.name)
		masses.append(float(species.mass))
		atoms.append(species.n_atoms)
	openFile.write(array_to_string("specname", names, type="string"))
	openFile.write(array_to_string("mass", masses, type="float"))
	openFile.write(array_to_string("atomCounts", atoms, type="int"))

	#then write evaporation stuff
	nGrain = write_evap_lists(openFile, speciesList)

	#finally all reactions
	reactant1 = []
	reactant2 = []
	reactant3 = []
	prod1 = []
	prod2 = []
	prod3 = []
	prod4 = []
	alpha = []
	beta = []
	gama = []
	for reaction in reactionList:
		reactant1.append(reaction.reactants[0])
		reactant2.append(reaction.reactants[1])
		reactant3.append(reaction.reactants[2])
		prod1.append(reaction.products[0])
		prod2.append(reaction.products[1])
		prod3.append(reaction.products[2])
		prod4.append(reaction.products[3])
		alpha.append(reaction.alpha)
		beta.append(reaction.beta)
		gama.append(reaction.gamma)
	openFile.write(array_to_string("re1", reactant1, type="string"))
	openFile.write(array_to_string("re2", reactant2, type="string"))
	openFile.write(array_to_string("re3", reactant3, type="string"))
	openFile.write(array_to_string("p1", prod1, type="string"))
	openFile.write(array_to_string("p2", prod2, type="string"))
	openFile.write(array_to_string("p3", prod3, type="string"))
	openFile.write(array_to_string("p4", prod4, type="string"))
	openFile.write(array_to_string("alpha", alpha, type="float", parameter=False))
	openFile.write(array_to_string("beta", beta, type="float", parameter=False))
	openFile.write(array_to_string("gama", gama, type="float", parameter=False))

def array_to_string(name, array, type="int", parameter=True):
	if parameter:
		outString = ", parameter :: "+name+" ({0})=(/".format(len(array))
	else:
		outString = " :: "+name+" ({0})=(/".format(len(array))
	if type == "int":
		outString = "integer"+outString
		for value in array:
			outString += "{0},".format(value)
	elif type == "float":
		outString = "double precision"+outString
		for value in array:
			outString += "{0:.4e},".format(value)
	elif type == "string":
		outString = "character(Len=10)"+outString
		for value in array:
			outString += "\""+value.ljust(10)+"\","
	else:
		print "Not a valid type for array to string"
	outString = outString[:-1]+"/)\n"
	outString = truncate_line(outString)
	return outString


from collections.abc import Mapping
from collections import defaultdict

import numba
from numba import generated_jit, types
import numpy as np

from africanus.util.patterns import Multiton
from africanus.experimental.rime.fused.arguments import ArgumentDependencies
from africanus.experimental.rime.fused.intrinsics import IntrinsicFactory
from africanus.experimental.rime.fused.specification import RimeSpecification

DATASET_TYPES = []

try:
    from daskms.dataset import Dataset as dmsds
except ImportError:
    pass
else:
    DATASET_TYPES.append(dmsds)

try:
    from xarray import Dataset as xrds
except ImportError:
    pass
else:
    DATASET_TYPES.append(xrds)

def rime_impl_factory(terms, transformers, ncorr):
    @generated_jit(nopython=True, nogil=True, cache=True)
    def rime(names, *inargs):
        if len(inargs) != 1 or not isinstance(inargs[0], types.BaseTuple):
            raise TypeError(f"{inargs[0]} must be be a Tuple")

        if not isinstance(names, types.BaseTuple):
            raise TypeError(f"{names} must be a Tuple of strings")

        if len(names) != len(inargs[0]):
            raise ValueError(f"len(names): {len(names)} "
                             f"!= {len(inargs[0])}")

        if not all(isinstance(n, types.Literal) for n in names):
            raise TypeError(f"{names} must be a Tuple of strings")

        if not all(n.literal_type is types.unicode_type for n in names):
            raise TypeError(f"{names} must be a Tuple of strings")

        # Get literal argument names
        names = tuple(n.literal_value for n in names)

        # Generate intrinsics
        argdeps = ArgumentDependencies(names, terms, transformers)
        factory = IntrinsicFactory(argdeps)
        out_names, pack_opts_indices = factory.pack_optionals_and_indices_fn()
        out_names, pack_transformed = factory.pack_transformed_fn(out_names)
        term_state = factory.term_state_fn(out_names)
        term_sampler = factory.term_sampler_fn()

        try:
            lm_i = out_names.index("lm")
            uvw_i = out_names.index("uvw")
            chan_freq_i = out_names.index("chan_freq")
        except ValueError as e:
            raise ValueError(f"{str(e)} is required")

        def impl(names, *inargs):
            args_opt_idx = pack_opts_indices(inargs)
            args = pack_transformed(args_opt_idx)
            state = term_state(args)

            nsrc, _ = args[lm_i].shape
            nrow, _ = args[uvw_i].shape
            nchan, = args[chan_freq_i].shape

            vis = np.zeros((nrow, nchan, ncorr), np.complex128)
            # Kahan summation compensation
            compensation = np.zeros_like(vis)

            for s in range(nsrc):
                for r in range(nrow):
                    t = state.time_index[r]
                    a1 = state.antenna1[r]
                    a2 = state.antenna2[r]
                    f1 = state.feed1[r]
                    f2 = state.feed2[r]

                    for ch in range(nchan):
                        X = term_sampler(state, s, r, t, f1, f2, a1, a2, ch)

                        for c, value in enumerate(numba.literal_unroll(X)):
                            # Kahan summation
                            y = value - compensation[r, ch, c]
                            current = vis[r, ch, c]
                            x = current + y
                            compensation[r, ch, c] = (x - current) - y
                            vis[r, ch, c] = x

            return vis

        return impl

    return rime

class RimeFactory(metaclass=Multiton):
    REQUIRED_ARGS = ArgumentDependencies.REQUIRED_ARGS
    REQUIRED_ARGS_LITERAL = tuple(types.literal(n) for n in REQUIRED_ARGS)
    DEFAULT_SPEC = "(Kpq, Bpq): [I, Q, U, V] -> [XX, XY, YX, YY]"

    def __reduce__(self):
        return (RimeFactory, (self.rime_spec,))

    def __hash__(self):
        return hash(self.rime_spec)

    def __eq__(self, rhs):
        return (isinstance(rhs, RimeFactory) and
                self.rime_spec == rhs.rime_spec)

    def __init__(self, rime_spec=DEFAULT_SPEC):
        if isinstance(rime_spec, RimeSpecification):
            pass
        elif isinstance(rime_spec, (list, tuple)):
            rime_spec = RimeSpecification(*rime_spec)
        elif isinstance(rime_spec, str):
            rime_spec = RimeSpecification(rime_spec)

        self.rime_spec = rime_spec
        self.impl = rime_impl_factory(
            rime_spec.terms,
            rime_spec.transformers,
            len(rime_spec.corrs))

    def dask_blockwise_args(self, **kwargs):
        """ Get the dask schema """
        argdeps = ArgumentDependencies(
            tuple(kwargs.keys()),
            self.rime_spec.terms,
            self.rime_spec.transformers)
        # Holds kwargs + any dummy outputs from transformations
        dummy_kw = kwargs.copy()

        dask_schema = defaultdict(list)
        for a in argdeps.REQUIRED_ARGS:
            dask_schema[a].append(("internal", ("row",)))

        POISON = object()

        for transformer in argdeps.can_create.values():
            kw = {}

            for a in transformer.ARGS:
                v = dummy_kw.get(a, None if a in argdeps.KEY_ARGS else POISON)
                kw[a] = v

            for a, d in transformer.KWARGS.items():
                kw[a] = dummy_kw.get(a, d)

            inputs, outputs = transformer.dask_schema(**kw)

            for k, schema in inputs.items():
                dask_schema[k].append((transformer, schema))

            dummy_kw.update(outputs)

        for term in self.rime_spec.terms:
            kw = {a: dummy_kw[a] for a in term.ALL_ARGS if a in dummy_kw}

            for k, v in term.dask_schema(**kw).items():
                dask_schema[k].append((term, v))

        merged_schema = {}

        for a, candidates in dask_schema.items():
            dims = set(pair[1] for pair in candidates)
            if len(dims) != 1:
                raise ValueError(
                    f"Multiple candidates provided conflicting "
                    f"dimension definitions for {a}: {candidates}.")

            merged_schema[a] = dims.pop()

        names = list(sorted(argdeps.valid_inputs & set(kwargs.keys())))
        blockwise_args = [e for n in names
                          for e in (kwargs[n], merged_schema.get(n, None))]

        assert 2 * len(names) == len(blockwise_args)
        return names, blockwise_args

    def __call__(self, time, antenna1, antenna2, feed1, feed2, **kwargs):
        keys = (self.REQUIRED_ARGS_LITERAL +
                tuple(map(types.literal, kwargs.keys())))

        return self.impl(keys, time,
                         antenna1, antenna2,
                         feed1, feed2,
                         *kwargs.values())

def consolidate_args(args, kw):
    mapping = {}
    oargs = []

    for element in args:
        if isinstance(element, tuple(DATASET_TYPES)):
            mapping.update((k.lower(), v.data) for k, v in element.items())
        elif isinstance(element, Mapping):
            mapping.update(element)
        else:
            oargs.append(element)

    mapping.update(zip(oargs, RimeFactory.REQUIRED_ARGS))
    mapping.update(kw)

    return mapping

def rime(rime_spec, *args, **kw):
    """
    Evaluates the Radio Interferometer Measurement Equation (RIME), given
    the Specification of the RIME :code:`rime_spec`, as well as the
    inputs to the RIME given in :code:`*args` and :code:`**kwargs`.
    """
    mapping = consolidate_args(args, kw)
    factory = RimeFactory(rime_spec=rime_spec)
    return factory(**mapping)


# Generate primers of a specified length and annealing temperature,
# with a specified distance from other primers.
# Accept a random seed and starting position, to guarantee the results can be repeated.
#import primer3
import ConfigParser
import sys
import os
import math
import numpy as np
import pandas as pd
import copy
import primer3

DNA = ['A','C','G','T']
DNA_C = {'A':'T', 'C':'G', 'G':'C', 'T':'A'}
N_PROBS = [0.25, 0.25, 0.25, 0.25]
MAX_NN_LENGTH = 60 # primer3 default
PRIMER3_METHOD = 'santalucia' # primer3 default
MAX_LOOP = 30 # primer3 default

def rc(seq):
  return(''.join([DNA_C[q] for q in seq])[::-1])

class primer_set(object):

  def __init__(self, cfg):
    params = dict(cfg.items('Params'))
    self.toe_anneal = float(params['toe_anneal'])
    self.toe_tolerance = float(params['toe_tolerance'])
    self.full_anneal = float(params['full_anneal'])
    self.full_tolerance = float(params['full_tolerance'])
    self.misprime_tolerance = float(params['misprime_tolerance'])
    self.evolve_iterations = int(params['evolve_iterations']) # number of iterations allowed for one call to 'evolve_new_seq' to get one new seq
    self.seqs = {'Toehold':[], 'Full':[]}
    self.seq_base = params['SEQ']
    self.base_probs = {'N': N_PROBS} # base probabilities for the toehold: include e.g. BsaI sites this way
    for q in self.seq_base:
      if not q in DNA and not q in self.base_probs.keys():
        x = [float(q) for q in params[q].strip().split(':')]
        self.base_probs[q] = x/sum(x)
    self.len_full_add = int(params['len_full_add'])# length of sequence to add to toehold to get something that anneals at self.full_anneal
    self.add_to_left = params['add_to_left'] == 'True'
    self.mv_conc = float(params['mv_conc']) # 50
    self.dv_conc = float(params['dv_conc']) # 0
    self.dntp_conc = float(params['dntp_conc']) # 0.8
    self.dna_conc = float(params['dna_conc']) # 50
    np.random.seed(int(params['random_seed']))

  # Get the annealing temperature of two sequences
  def melting_temp(self, seq):
    return( primer3.calcTm(seq, mv_conc=self.mv_conc, dv_conc=self.dv_conc, dntp_conc=self.dntp_conc,
                      dna_conc=self.dna_conc, max_nn_length=MAX_NN_LENGTH, tm_method=PRIMER3_METHOD, salt_corrections_method=PRIMER3_METHOD) )

  # Find whether a concerning heterodimer exists between two sequences
  def get_misprime_tm(self, seq1, seq2, temp):
    tm_f1 = self.get_misprime_tm_inner(seq1, seq2, temp)
    tm_r1 = self.get_misprime_tm_inner(seq1, rc(seq2), temp)
    tm_f2 = self.get_misprime_tm_inner(seq2, seq1, temp)
    tm_r2 = self.get_misprime_tm_inner(seq2, rc(seq1), temp)
    return(max(tm_f1, tm_r1, tm_f2, tm_r2))

  def get_misprime_tm_inner(self, seq1, seq2, temp):
    heterodimer = primer3.calcHeterodimer(seq1, seq2,
                    mv_conc=self.mv_conc, dv_conc=self.dv_conc, dntp_conc=self.dntp_conc, dna_conc=self.dna_conc,
                    temp_c=temp, max_loop=MAX_LOOP)
    if not heterodimer.structure_found:
      return -100.
    return heterodimer.tm

  def true_if_no_misprime(self, seq1, seq2, temp):
    struct_tm = self.get_misprime_tm(seq1, seq2, temp)
    #if(struct_tm >= (temp - self.misprime_tolerance)):
    #  raise Exception('oops!' + str(seq1) + ' ' + str(seq2) + ' ' + str(struct_tm))

    return struct_tm < (temp - self.misprime_tolerance) # e.g. if temp = 45, self.misprime_tolerance = 5, OK if t_anneal < 40; will be melted by 45

  # placeholder, for maybe rejecting e.g. hairpins
  def primer_passed_filter(self, seq):
    return True

  # generate a random sequence
  def get_random_seq(self, parent = None):
    if parent == None: # this is Python's fault
      parent = self.seq_base
    ans = []
    for b in parent:
      if b in DNA:
        ans.append(b)
      else:
        ans.append( np.random.choice(DNA, None, self.base_probs[b]) )
    return(''.join(ans))

  # get all variants of a random sequence, and find the one with Tm closest to some target
  # Compare to a template to see which bases aren't allowed to be changed
  def evolve_one_step(self, seq, template, target):
    assert(len(seq) == len(template))
    ans = seq
    err = abs( self.melting_temp(seq) - target )
    for i in range(len(seq)):
      if template[i] not in DNA:
        for b in DNA:
          if seq[i] != b:
            tmp = [q for q in seq]
            tmp[i] = b
            tmp = ''.join(tmp)
            new_err = abs(self.melting_temp(tmp) - target)
            if new_err < err and self.primer_passed_filter(tmp):
              err = new_err
              ans = tmp
    return(ans)

  def true_if_no_misprimes_to_seqs(self, seq):
    if len(self.seqs['Full']) == 0:
      return True
    return(all([self.true_if_no_misprime(seq, q, self.toe_anneal) for q in self.seqs['Full']]))

  # evolve a valid new toe first, then extend
  # 'self.evolve_iterations' available to get both tasks done.
  def evolve_new_seq_inner(self):
    seq = self.get_random_seq()
    is_full = False
    template = self.seq_base
    toe = None; full = None
    for i in range(self.evolve_iterations):
      if is_full: # get temperature targets based on what part of the process this is
        target = self.full_anneal
        target_low = target - self.full_tolerance; target_high = target + self.full_tolerance
      else:
        target = self.toe_anneal
        target_low = target - self.toe_tolerance; target_high = target + self.toe_tolerance
      tm = self.melting_temp(seq)
      if tm > target_low and tm < target_high:
        if not is_full:
          is_full = True; toe = seq
          new_pad = ''.join(['N']*self.len_full_add)
          if self.add_to_left:
            template = new_pad + seq
            seq = self.get_random_seq(template)
          else:
            template = seq + new_pad
            seq = self.get_random_seq(template)
        else:
          full = seq
          return(toe, full)
      seq = self.evolve_one_step(seq, template, target)
    return(None, None)

  def evolve_new_seq(self):
    toe, full = self.evolve_new_seq_inner()
    if full != None:
      if self.true_if_no_misprimes_to_seqs(full):
        self.seqs['Toehold'].append(toe)
        self.seqs['Full'].append(full)

def get_seqs(cfg):
  num_seqs = int(cfg.get('Params','num_seqs'))
  num_iters = int(cfg.get('Params','num_iters'))
  primers = primer_set(cfg)
  for i in range(num_iters):
    primers.evolve_new_seq()
    if len(primers.seqs['Full']) >= num_seqs:
      return(primers)
  if len(primers.seqs['Full']) < num_seqs:
    raise Exception(str(num_seqs) + ' primers desired, but ' + str(len(primers.seqs)) + ' found.')
  return(primers)

if __name__ == '__main__':
  cfg = ConfigParser.RawConfigParser(allow_no_value=True); cfg.optionxform=str
  cfg.read(sys.argv[1])
  primers = get_seqs(cfg)
  for (toe, full) in zip(primers.seqs['Toehold'], primers.seqs['Full']):
    print('Full: ' + str(full))
    print('Toe: ' + str(toe))
    print('T_melt (toe): ' + str(primers.melting_temp(toe)))
    print('T_melt (full): ' + str(primers.melting_temp(full)))
    t_ann = float(cfg.get('Params','toe_anneal'))
    worst_seq = ''
    worst_t = -1000.
    for q in primers.seqs['Full']:
      if q != full:
        new_t = primers.get_misprime_tm(full, q, t_ann)
        if new_t > worst_t:
          worst_t = new_t; worst_seq = q

    print('T_misprime (max): ' + str(worst_t))
    print('Worst misprime sequence: ' + str(worst_seq))
  print('Sequences generated: ' + str(len(primers.seqs['Full'])))
  df_out = pd.DataFrame(primers.seqs)
  df_out.to_csv(os.path.expanduser(cfg.get('Files','out_fn')), index = False)


# coding: utf-8
#
#    Project: Azimuthal integration
#             https://github.com/silx-kit/pyFAI
#
#    Copyright (C) 2015-2021 European Synchrotron Radiation Facility, Grenoble, France
#    Copyright (C)      2016 Synchrotron SOLEIL - L'Orme des Merisiers Saint-Aubin
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.

"""Module for treating simultaneously multiple detector configuration
within a single integration"""

__author__ = "Jérôme Kieffer"
__contact__ = "Jerome.Kieffer@ESRF.eu"
__license__ = "MIT"
__copyright__ = "European Synchrotron Radiation Facility, Grenoble, France"
__date__ = "19/01/2021"
__status__ = "stable"
__docformat__ = 'restructuredtext'

import collections
import logging
logger = logging.getLogger(__name__)
from .azimuthalIntegrator import AzimuthalIntegrator
from .containers import Integrate1dResult
from .containers import Integrate2dResult
from . import units
import threading
import numpy
from .method_registry import IntegrationMethod
error = None

class MultiGeometry(object):
    """This is an Azimuthal integrator containing multiple geometries,
    for example when the detector is on a goniometer arm
    """

    def __init__(self, ais, unit="2th_deg",
                 radial_range=(0, 180), azimuth_range=(-180, 180),
                 wavelength=None, empty=0.0, chi_disc=180):
        """
        Constructor of the multi-geometry integrator
        :param ais: list of azimuthal integrators
        :param radial_range: common range for integration
        :param azimuthal_range: common range for integration
        :param empty: value for empty pixels
        :param chi_disc: if 0, set the chi_discontinuity at
        """
        self._sem = threading.Semaphore()
        self.abolute_solid_angle = None
        self.ais = [ai if isinstance(ai, AzimuthalIntegrator)
                    else AzimuthalIntegrator.sload(ai)
                    for ai in ais]
        self.wavelength = None
        if wavelength:
            self.set_wavelength(wavelength)
        self.radial_range = tuple(radial_range[:2])
        self.azimuth_range = tuple(azimuth_range[:2])
        self.unit = units.to_unit(unit)
        self.abolute_solid_angle = None
        self.empty = empty
        if chi_disc == 0:
            for ai in self.ais:
                ai.setChiDiscAtZero()
        elif chi_disc == 180:
            for ai in self.ais:
                ai.setChiDiscAtPi()
        else:
            logger.warning("Unable to set the Chi discontinuity at %s", chi_disc)

    def __repr__(self, *args, **kwargs):
        return "MultiGeometry integrator with %s geometries on %s radial range (%s) and %s azimuthal range (deg)" % \
            (len(self.ais), self.radial_range, self.unit, self.azimuth_range)

    def integrate1d(self, lst_data, npt=1800,
                    correctSolidAngle=True,
                    lst_variance=None, error_model=None,
                    polarization_factor=None,
                    normalization_factor=None,
                    lst_mask=None, lst_flat=None,
                    method="splitpixel"):
        """Perform 1D azimuthal integration

        :param lst_data: list of numpy array
        :param npt: number of points int the integration
        :param correctSolidAngle: correct for solid angle (all processing are then done in absolute solid angle !)
        :param lst_variance: list of array containing the variance of the data. If not available, no error propagation is done
        :type lst_variance: list of ndarray
        :param error_model: When the variance is unknown, an error model can be given: "poisson" (variance = I), "azimuthal" (variance = (I-<I>)^2)
        :type error_model: str
        :param polarization_factor: Apply polarization correction ? is None: not applies. Else provide a value from -1 to +1
        :param normalization_factor: normalization monitors value (list of floats)
        :param all: return a dict with all information in it (deprecated, please refer to the documentation of Integrate1dResult).
        :param lst_mask: numpy.Array or list of numpy.array which mask the lst_data.
        :param lst_flat: numpy.Array or list of numpy.array which flat the lst_data.
        :param method: integration method, a string or a registered method
        :return: 2th/I or a dict with everything depending on "all"
        :rtype: Integrate1dResult, dict
        """
        method = IntegrationMethod.select_one_available(method, dim=1)

        if len(lst_data) == 0:
            raise RuntimeError("List of images cannot be empty")
        if normalization_factor is None:
            normalization_factor = [1.0] * len(self.ais)
        elif not isinstance(normalization_factor, collections.Iterable):
            normalization_factor = [normalization_factor] * len(self.ais)
        if lst_variance is None:
            lst_variance = [None] * len(self.ais)
        if lst_mask is None:
            lst_mask = [None] * len(self.ais)
        elif isinstance(lst_mask, numpy.ndarray):
            lst_mask = [lst_mask] * len(self.ais)
        if lst_flat is None:
            lst_flat = [None] * len(self.ais)
        elif isinstance(lst_flat, numpy.ndarray):
            lst_flat = [lst_flat] * len(self.ais)
        signal = numpy.zeros(npt, dtype=numpy.float64)
        normalization = numpy.zeros_like(signal)
        count = numpy.zeros_like(signal)
        variance = None
        for ai, data, monitor, var, mask, flat in zip(self.ais, lst_data, normalization_factor, lst_variance, lst_mask, lst_flat):
            res = ai.integrate1d_ng(data, npt=npt,
                                    correctSolidAngle=correctSolidAngle,
                                    variance=var, error_model=error_model,
                                    polarization_factor=polarization_factor,
                                    radial_range=self.radial_range,
                                    azimuth_range=self.azimuth_range,
                                    method=method, unit=self.unit, safe=True,
                                    mask=mask, flat=flat, normalization_factor=monitor)
            sac = (ai.pixel1 * ai.pixel2 / ai.dist ** 2) if correctSolidAngle else 1.0
            count += res.count
            normalization += res.sum_normalization * sac
            signal += res.sum_signal
            if res.sigma is not None:
                if variance is None:
                    variance = res.sum_variance.astype(dtype=numpy.float64)  # explicit copy
                else:
                    variance += res.sum_variance

        tiny = numpy.finfo("float32").tiny
        norm = numpy.maximum(normalization, tiny)
        invalid = count <= 0.0
        I = signal / norm
        I[invalid] = self.empty

        if variance is not None:
            sigma = numpy.sqrt(variance) / norm
            sigma[invalid] = self.empty
            result = Integrate1dResult(res.radial, I, sigma)
        else:
            result = Integrate1dResult(res.radial, I)
        result._set_compute_engine(res.compute_engine)
        result._set_unit(self.unit)
        result._set_sum_signal(signal)
        result._set_sum_normalization(normalization)
        result._set_sum_variance(variance)
        result._set_count(count)
        return result

    def integrate2d(self, lst_data, npt_rad=1800, npt_azim=3600,
                    correctSolidAngle=True,
                    lst_variance=None, error_model=None,
                    polarization_factor=None,
                    normalization_factor=None, lst_mask=None,
                    lst_flat=None, method="splitpixel"):
        """Performs 2D azimuthal integration of multiples frames, one for each geometry

        :param lst_data: list of numpy array
        :param npt: number of points int the integration
        :param correctSolidAngle: correct for solid angle (all processing are then done in absolute solid angle !)
        :param lst_variance: list of array containing the variance of the data. If not available, no error propagation is done
        :type lst_variance: list of ndarray
        :param error_model: When the variance is unknown, an error model can be given: "poisson" (variance = I), "azimuthal" (variance = (I-<I>)^2)
        :type error_model: str
        :param polarization_factor: Apply polarization correction ? is None: not applies. Else provide a value from -1 to +1
        :param normalization_factor: normalization monitors value (list of floats)
        :param all: return a dict with all information in it (deprecated, please refer to the documentation of Integrate2dResult).
        :param lst_mask: numpy.Array or list of numpy.array which mask the lst_data.
        :param lst_flat: numpy.Array or list of numpy.array which flat the lst_data.
        :param method: integration method (or its name)
        :return: I/2th/chi or a dict with everything depending on "all"
        :rtype: Integrate2dResult, dict
        """
        if len(lst_data) == 0:
            raise RuntimeError("List of images cannot be empty")
        if normalization_factor is None:
            normalization_factor = [1.0] * len(self.ais)
        elif not isinstance(normalization_factor, collections.Iterable):
            normalization_factor = [normalization_factor] * len(self.ais)
        if lst_variance is None:
            lst_variance = [None] * len(self.ais)
        if lst_mask is None:
            lst_mask = [None] * len(self.ais)
        elif isinstance(lst_mask, numpy.ndarray):
            lst_mask = [lst_mask] * len(self.ais)
        if lst_flat is None:
            lst_flat = [None] * len(self.ais)
        elif isinstance(lst_flat, numpy.ndarray):
            lst_flat = [lst_flat] * len(self.ais)

        method = IntegrationMethod.select_one_available(method, dim=2)

        signal = numpy.zeros((npt_azim, npt_rad), dtype=numpy.float64)
        count = numpy.zeros_like(signal)
        normalization = numpy.zeros_like(signal)
        variance = None

        for ai, data, monitor, var, mask, flat in zip(self.ais, lst_data, normalization_factor, lst_variance, lst_mask, lst_flat):
            res = ai.integrate2d_ng(data, npt_rad=npt_rad, npt_azim=npt_azim,
                                    correctSolidAngle=correctSolidAngle,
                                    variance=var, error_model=error_model,
                                    polarization_factor=polarization_factor,
                                    radial_range=self.radial_range,
                                    azimuth_range=self.azimuth_range,
                                    method=method, unit=self.unit, safe=True,
                                    mask=mask, flat=flat, normalization_factor=monitor)
            sac = (ai.pixel1 * ai.pixel2 / ai.dist ** 2) if correctSolidAngle else 1.0
            count += res.count
            signal += res.sum_signal
            normalization += res.sum_normalization * sac
            if var is not None:
                if variance is None:
                    variance = res.sum_variance.astype(numpy.float64)  # explicit copy !
                else:
                    variance += res.sum_variance

        tiny = numpy.finfo("float32").tiny
        norm = numpy.maximum(normalization, tiny)
        invalid = count <= 0
        I = signal / norm
        I[invalid] = self.empty

        if variance is not None:
            sigma = numpy.sqrt(variance) / norm
            sigma[invalid] = self.empty
            result = Integrate2dResult(I, res.radial, res.azimuthal, sigma)
        else:
            result = Integrate2dResult(I, res.radial, res.azimuthal)
        result._set_sum(signal)
        result._set_compute_engine(res.compute_engine)
        result._set_unit(self.unit)
        result._set_sum_signal(signal)
        result._set_sum_normalization(normalization)
        result._set_sum_variance(variance)
        result._set_count(count)
        return result

    def set_wavelength(self, value):
        """
        Changes the wavelength of a group of azimuthal integrators
        """
        self.wavelength = float(value)
        for ai in self.ais:
            ai.set_wavelength(self.wavelength)


"""
Implementation of BIMODAL to generate SMILES
"""

import numpy as np
import torch
import torch.nn as nn
from bidir_lstm import BiDirLSTM

torch.manual_seed(1)
np.random.seed(5)

class BIMODAL():

    def __init__(self, molecule_size=7, encoding_dim=55, lr=.01, hidden_units=128):

        self._molecule_size = molecule_size
        self._input_dim = encoding_dim
        self._output_dim = encoding_dim
        self._layer = 2
        self._hidden_units = hidden_units

        # Learning rate
        self._lr = lr

        # Build new model
        self._lstm = BiDirLSTM(self._input_dim, self._hidden_units, self._layer)

        # Check availability of GPUs
        self._gpu = torch.cuda.is_available()
        self._device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        if torch.cuda.is_available():
            self._lstm = self._lstm.cuda()
            print('GPU available')

        # Adam optimizer
        self._optimizer = torch.optim.Adam(self._lstm.parameters(), lr=self._lr, betas=(0.9, 0.999))
        # Cross entropy loss
        self._loss = nn.CrossEntropyLoss(reduction='mean')

    def build(self, name=None):
        """Build new model or load model by name
        :param name:    model name
        """

        if (name is None):
            self._lstm = BiDirLSTM(self._input_dim, self._hidden_units, self._layer)

        else:
            self._lstm = torch.load(name + '.dat', map_location=self._device)

        if torch.cuda.is_available():
            self._lstm = self._lstm.cuda()

        self._optimizer = torch.optim.Adam(self._lstm.parameters(), lr=self._lr, betas=(0.9, 0.999))

    def print_model(self):
        '''Print name and shape of all tensors'''
        for name, p in self._lstm.state_dict().items():
            print(name)
            print(p.shape)

    def train(self, data, label, epochs=1, batch_size=1):
        '''Train the model
        :param  data:   data array (n_samples, molecule_size, encoding_length)
        :param  label:  label array (n_samples, molecule_size)
        :param  epochs: number of epochs for the training
        :param  batch_size: batch size for the training
        :return statistic:  array storing computed losses (epochs, batch size)
        '''

        # Number of samples
        n_samples = data.shape[0]

        # Change axes from (n_samples, molecule_size, encoding_dim) to (molecule_size, n_samples, encoding_dim)
        data = np.swapaxes(data, 0, 1)

        # Create tensor from label
        label = torch.from_numpy(label).to(self._device, dtype=torch.long)

        # Calculate number of batches per epoch
        if (n_samples % batch_size) is 0:
            n_iter = n_samples // batch_size
        else:
            n_iter = n_samples // batch_size + 1

        # To store losses
        statistic = np.zeros((epochs, n_iter))

        # Prepare model for training
        self._lstm.train()

        # Iteration over epochs
        for i in range(epochs):

            # Iteration over batches
            for n in range(n_iter):

                # Set gradient to zero for batch
                self._optimizer.zero_grad()

                # Store losses in list
                losses = []

                # Compute indices used as batch
                batch_start = n * batch_size
                batch_end = min((n + 1) * batch_size, n_samples)

                # Reset model with correct batch size
                self._lstm.new_sequence(batch_end - batch_start, self._device)

                # Current batch
                batch_data = torch.from_numpy(data[:, batch_start:batch_end, :].astype('float32')).to(self._device)

                # Initialize start and end position of sequence read by the model
                start = self._molecule_size // 2
                end = start + 1

                for j in range(self._molecule_size - 1):
                    self._lstm.new_sequence(batch_end - batch_start, self._device)

                    # Select direction for next prediction
                    if j % 2 == 0:
                        dir = 'right'
                    else:
                        dir = 'left'

                    # Predict next token
                    pred = self._lstm(batch_data[start:end], dir, self._device)

                    # Compute loss and extend sequence read by the model
                    if j % 2 == 0:
                        loss = self._loss(pred, label[batch_start:batch_end, end])
                        end += 1

                    else:
                        loss = self._loss(pred, label[batch_start:batch_end, start - 1])
                        start -= 1

                    # Append loss of current position
                    losses.append(loss.item())

                    # Accumulate gradients
                    # (NOTE: This is more memory-efficient than summing the loss and computing the final gradient for the sum)
                    loss.backward()

                # Store statistics: loss per token (middle token not included)
                statistic[i, n] = np.sum(losses) / (self._molecule_size - 1)

                # Perform optimization step
                self._optimizer.step()

        return statistic

    def validate(self, data, label, batch_size=128):
        ''' Validation of model and compute error
        :param data:    test data (n_samples, molecule_size, encoding_size)
        :param label:   label data (n_samples_molecules_size)
        :param batch_size:  batch size for validation
        :return:            mean loss over test data
        '''

        # Use train mode to get loss consistent with training
        self._lstm.train()

        # Gradient is not computed to reduce memory requirements
        with torch.no_grad():
            # Compute tensor of labels
            label = torch.from_numpy(label).to(self._device, dtype=torch.long)

            # Number of samples
            n_samples = data.shape[0]

            # Change axes from (n_samples, molecule_size, encoding_dim) to (molecule_size , n_samples, encoding_dim)
            data = np.swapaxes(data, 0, 1).astype('float32')

            # Initialize loss for complete validation set
            tot_loss = 0

            # Calculate number of batches per epoch
            if (n_samples % batch_size) is 0:
                n_iter = n_samples // batch_size
            else:
                n_iter = n_samples // batch_size + 1

            for n in range(n_iter):

                # Compute indices used as batch
                batch_start = n * batch_size
                batch_end = min((n + 1) * batch_size, n_samples)

                # Data used in this batch
                batch_data = torch.from_numpy(data[:, batch_start:batch_end, :].astype('float32')).to(self._device)

                # Initialize loss for molecule
                molecule_loss = 0

                # Reset model with correct batch size and device
                self._lstm.new_sequence(batch_end - batch_start, self._device)

                start = self._molecule_size // 2
                end = start + 1

                for j in range(self._molecule_size - 1):
                    self._lstm.new_sequence(batch_end - batch_start, self._device)

                    # Select direction for next prediction
                    if j % 2 == 0:
                        dir = 'right'
                    if j % 2 == 1:
                        dir = 'left'

                    # Predict next token
                    pred = self._lstm(batch_data[start:end], dir, self._device)

                    # Extend reading of the sequence
                    if j % 2 == 0:
                        loss = self._loss(pred, label[batch_start:batch_end, end])
                        end += 1

                    if j % 2 == 1:
                        loss = self._loss(pred, label[batch_start:batch_end, start - 1])
                        start -= 1
                    # Sum loss over molecule
                    molecule_loss += loss.item()

                # Add loss per token to total loss (start token and end token not counted)
                tot_loss += molecule_loss / (self._molecule_size - 1)

            return tot_loss / n_iter

    def sample(self, middle_token, T=1):
        '''Generate new molecule
        :param middle_token:    starting sequence
        :param T:               sampling temperature
        :return molecule:       newly generated molecule (molecule_length, encoding_length)
        '''

        # Prepare model
        self._lstm.eval()

        # Gradient is not compute to reduce memory requirements
        with torch.no_grad():
            # Output array with merged forward and backward directions

            # New sequence
            seq = np.zeros((self._molecule_size, 1, self._output_dim))
            seq[self._molecule_size // 2, 0] = middle_token

            # Create tensor for data and select correct device
            seq = torch.from_numpy(seq.astype('float32')).to(self._device)

            # Define start/end values for reading
            start = self._molecule_size // 2
            end = start + 1

            for j in range(self._molecule_size - 1):
                self._lstm.new_sequence(1, self._device)

                # Select direction for next prediction
                if j % 2 == 0:
                    dir = 'right'
                if j % 2 == 1:
                    dir = 'left'

                pred = self._lstm(seq[start:end], dir, self._device)

                # Compute new token
                token = self.sample_token(np.squeeze(pred.cpu().detach().numpy()), T)

                # Set new token within sequence
                if j % 2 == 0:
                    seq[end, 0, token] = 1.0
                    end += 1

                if j % 2 == 1:
                    seq[start - 1, 0, token] = 1.0
                    start -= 1

        return seq.cpu().numpy().reshape(1, self._molecule_size, self._output_dim)

    def sample_token(self, out, T=1.0):
        ''' Sample token
        :param out: output values from model
        :param T:   sampling temperature
        :return:    index of predicted token
        '''
        # Explicit conversion to float64 avoiding truncation errors
        out = out.astype('float64')

        # Compute probabilities with specific temperature
        out_T = out / T
        p = np.exp(out_T) / np.sum(np.exp(out_T))

        # Generate new token at random
        char = np.random.multinomial(1, p, size=1)
        return np.argmax(char)

    def save(self, name='test_model'):
        torch.save(self._lstm, name + '.dat')


# Copyright (c) 2009-2019 The Regents of the University of Michigan
# This file is part of the HOOMD-blue project, released under the BSD 3-Clause License.

from hoomd import _hoomd
#from hoomd.jit import _jit
import hoomd
import tempfile
import shutil
import subprocess
import os

import numpy as np
from hoomd.swapmc import _swapmc as _plugin_patch
#from hoomd.jit import patch

#This python file faciltates the creation of relevant patch energies
#we'll do shifted lennard jones and polydisperse's system for now
#let's make different classes for different pair potentials. Of course, we'll let the user defined path energy to be available as well
class lj(object):
    R''' Define the patch energy of a lennard jones particle.

    Args:
        scaledr_cut (float): the scale of cutoff radius relative to particle center to center, beyond which all pair interactions are assumed 0.
        eps (float): uniform energetics
        llvm_ir_fname (str): File name of the llvm IR file to load.
        clang_exec (str): The Clang executable to use

    Patch energies define energetic interactions between pairs of shapes in :py:mod:`hpmc <hoomd.hpmc>` integrators.
    Shapes within a cutoff distance of *r_cut* are potentially interacting and the energy of interaction is a function
    the type and orientation of the particles and the vector pointing from the *i* particle to the *j* particle center.

    The :py:class:`user` patch energy takes C++ code, JIT compiles it at run time and executes the code natively
    in the MC loop at with full performance. It enables researchers to quickly and easily implement custom energetic
    interactions without the need to modify and recompile HOOMD.

    .. rubric:: C++ code

    Supply C++ code to the *code* argument and :py:class:`user` will compile the code and call it to evaluate
    patch energies. Compilation assumes that a recent ``clang`` installation is on your PATH. This is convenient
    when the energy evaluation is simple or needs to be modified in python. More complex code (i.e. code that
    requires auxiliary functions or initialization of static data arrays) should be compiled outside of HOOMD
    and provided via the *llvm_ir_file* input (see below).

    The text provided in *code* is the body of a function with the following signature:

    .. code::

        float eval(const vec3<float>& r_ij,
                   unsigned int type_i,
                   const quat<float>& q_i,
                   float d_i,
                   float charge_i,
                   unsigned int type_j,
                   const quat<float>& q_j,
                   float d_j,
                   float charge_j)

    * ``vec3`` and ``quat`` are defined in HOOMDMath.h.
    * *r_ij* is a vector pointing from the center of particle *i* to the center of particle *j*.
    * *type_i* is the integer type of particle *i*
    * *q_i* is the quaternion orientation of particle *i*
    * *d_i* is the diameter of particle *i*
    * *charge_i* is the charge of particle *i*
    * *type_j* is the integer type of particle *j*
    * *q_j* is the quaternion orientation of particle *j*
    * *d_j* is the diameter of particle *j*
    * *charge_j* is the charge of particle *j*
    * Your code *must* return a value.
    * When \|r_ij\| is greater than *r_cut*, the energy *must* be 0. This *r_cut* is applied between
      the centers of the two particles: compute it accordingly based on the maximum range of the anisotropic
      interaction that you implement.

    Example:

    .. code-block:: python

        square_well = """float rsq = dot(r_ij, r_ij);
                            if (rsq < 1.21f)
                                return -1.0f;
                            else
                                return 0.0f;
                      """
        patch = hoomd.jit.patch.user(mc=mc, r_cut=1.1, code=square_well)

    .. rubric:: LLVM IR code

    You can compile outside of HOOMD and provide a direct link
    to the LLVM IR file in *llvm_ir_file*. A compatible file contains an extern "C" eval function with this signature:

    .. code::

        float eval(const vec3<float>& r_ij,
                   unsigned int type_i,
                   const quat<float>& q_i,
                   float d_i,
                   float charge_i,
                   unsigned int type_j,
                   const quat<float>& q_j,
                   float d_j,
                   float charge_j)

    ``vec3`` and ``quat`` are defined in HOOMDMath.h.

    Compile the file with clang: ``clang -O3 --std=c++11 -DHOOMD_LLVMJIT_BUILD -I /path/to/hoomd/include -S -emit-llvm code.cc`` to produce
    the LLVM IR in ``code.ll``.

    .. versionadded:: 2.3
    '''
    def __init__(self, mc, kT, scaledr_cut, eps, mode='shifted', llvm_ir_file=None, clang_exec=None):
        hoomd.util.print_status_line();
        #there is no need to write your own code. By definition, we're
        if (mode == 'truncated'):
            code = """
                            float rsq = dot(r_ij, r_ij);
                            float sigma  = 0.5*( d_i+d_j);
                            float rcut  = {}*sigma;
                            if (rsq <= rcut*rcut)
                               {{
                               float eps   = {};
                               float sigmasq = sigma*sigma;
                               float rsqinv = sigmasq / rsq;
                               float r6inv = rsqinv*rsqinv*rsqinv;
                               return 4.0f*eps*r6inv*(r6inv-1.0f);
                               }}
                            else
                               {{
                               return 0.0f;
                               }}
                          """.format(scaledr_cut,eps/kT);
        elif (mode == 'shifted'):
            code = """
                            float rsq = dot(r_ij, r_ij);
                            float sigma  = 0.5*( d_i+d_j);
                            float rcut  = {}*sigma;
                            if (rsq <= rcut*rcut)
                               {{
                               float eps   = {};
                               float sigmasq = sigma*sigma;
                               float rsqinv = sigmasq / rsq;
                               float r6inv = rsqinv*rsqinv*rsqinv;

                               float rsqinv_cut = sigmasq /(rcut*rcut);
                               float r6inv_cut = rsqinv_cut*rsqinv_cut*rsqinv_cut;
                               return 4.0f*eps*(r6inv*(r6inv-1.0f)-r6inv_cut*(r6inv_cut-1.0f));
                               }}
                            else
                               {{
                               return 0.0f;
                               }}
                          """.format(scaledr_cut,eps/kT);
        elif (mode == 'force-shift'):
            code = """
                            float rsq = dot(r_ij, r_ij);
                            float sigma  = 0.5*( d_i+d_j);
                            float rcut  = {}*sigma;
                            if (rsq <= rcut*rcut)
                               {{
                               float eps   = {};
                               float sigmasq = sigma*sigma;
                               float rsqinv = sigmasq / rsq;
                               float r6inv = rsqinv*rsqinv*rsqinv;

                               float rsqinv_cut = sigmasq /(rcut*rcut);
                               float r6inv_cut = rsqinv_cut*rsqinv_cut*rsqinv_cut;
                               return 4.0f*eps*( r6inv*(r6inv-1.0f)-r6inv_cut*(r6inv_cut-1.0f) +6*(sqrtf(rsq)/rcut-1.0f)*r6inv_cut*(2*r6inv_cut-1.0f) );
                               }}
                            else
                               {{
                               return 0.0f;
                               }}
                          """.format(scaledr_cut,eps/kT);
        # check if initialization has occurred
        if hoomd.context.exec_conf is None:
            raise RuntimeError('Error creating Lennard-Jones patch energy, call context.initialize() first');

        # raise an error if this run is on the GPU
        if hoomd.context.exec_conf.isCUDAEnabled():
            hoomd.context.msg.error("Patch energies are not supported on the GPU\n");
            raise RuntimeError("Error initializing patch energy");

        # Find a clang executable if none is provided
        if clang_exec is not None:
            clang = clang_exec;
        else:
            clang = 'clang'

        if code is not None and llvm_ir_file is None:
            llvm_ir = self.compile_user(code, clang)
        else:
            # IR is a text file
            with open(llvm_ir_file,'r') as f:
                llvm_ir = f.read()

        self.compute_name = "patch"
        self.cpp_evaluator = _plugin_patch.PatchEnergyJITCustom(hoomd.context.exec_conf, llvm_ir, scaledr_cut);
        mc.set_PatchEnergyEvaluator(self);

        self.mc = mc
        self.enabled = True
        self.log = False

    def compile_user(self, code, clang_exec, fn=None):
        R'''Helper function to compile the provided code into an executable

        Args:
            code (str): C++ code to compile
            clang_exec (str): The Clang executable to use
            fn (str): If provided, the code will be written to a file.

        .. versionadded:: 2.3
        '''
        cpp_function = """
#include "hoomd/HOOMDMath.h"
#include "hoomd/VectorMath.h"

extern "C"
{
float eval(const vec3<float>& r_ij,
    unsigned int type_i,
    const quat<float>& q_i,
    float d_i,
    float charge_i,
    unsigned int type_j,
    const quat<float>& q_j,
    float d_j,
    float charge_j)
    {
"""
        cpp_function += code
        cpp_function += """
    }
}
"""

        include_path = os.path.dirname(hoomd.__file__) + '/include';
        include_path_source = hoomd._hoomd.__hoomd_source_dir__;

        if clang_exec is not None:
            clang = clang_exec;
        else:
            clang = 'clang';

        if fn is not None:
            cmd = [clang, '-O3', '--std=c++11', '-DHOOMD_LLVMJIT_BUILD', '-I', include_path, '-I', include_path_source, '-S', '-emit-llvm','-x','c++', '-o',fn,'-']
        else:
            cmd = [clang, '-O3', '--std=c++11', '-DHOOMD_LLVMJIT_BUILD', '-I', include_path, '-I', include_path_source, '-S', '-emit-llvm','-x','c++', '-o','-','-']
        p = subprocess.Popen(cmd,stdin=subprocess.PIPE,stdout=subprocess.PIPE,stderr=subprocess.PIPE)

        # pass C++ function to stdin
        output = p.communicate(cpp_function.encode('utf-8'))
        llvm_ir = output[0].decode()

        f = open("llvm_ir_file",'w')
        f.write(llvm_ir)
        f.close()

        if p.returncode != 0:
            hoomd.context.msg.error("Error compiling provided code\n");
            hoomd.context.msg.error("Command "+' '.join(cmd)+"\n");
            hoomd.context.msg.error(output[1].decode()+"\n");
            raise RuntimeError("Error initializing patch energy");

        return llvm_ir

    R''' Disable the patch energy and optionally enable it only for logging

    Args:
        log (bool): If true, only use patch energy as a log quantity

    '''
    def disable(self,log=None):
        hoomd.util.print_status_line();

        if log:
            # enable only for logging purposes
            self.mc.cpp_integrator.disablePatchEnergyLogOnly(log)
            self.log = True
        else:
            # disable completely
            self.mc.cpp_integrator.setPatchEnergy(None);
            self.log = False

        self.enabled = False

    R''' (Re-)Enable the patch energy

    '''
    def enable(self):
        hoomd.util.print_status_line()
        self.mc.cpp_integrator.setPatchEnergy(self.cpp_evaluator);

class softrepulsive(object):
    R''' Define the patch energy of a soft-repulsive particle.
    '''
    def __init__(self, mc, kT, scaledr_cut, eps, mode='shifted', llvm_ir_file=None, clang_exec=None):
        hoomd.util.print_status_line();
        #there is no need to write your own code. By definition, we're
        if (mode == 'truncated'):
            code = """
                            float rsq = dot(r_ij, r_ij);
                            float sigma  = 0.5*( d_i+d_j);
                            float rcut  = {}*sigma;
                            if (rsq <= rcut*rcut)
                               {{
                               float eps   = {};
                               float sigmasq = sigma*sigma;
                               float rsqinv = sigmasq / rsq;
                               float r6inv = rsqinv*rsqinv*rsqinv;
                               return 4.0f*eps*r6inv*r6inv;
                               }}
                            else
                               {{
                               return 0.0f;
                               }}
                          """.format(scaledr_cut,eps/kT);
        elif (mode == 'shifted'):
            code = """
                            float rsq = dot(r_ij, r_ij);
                            float sigma  = 0.5*( d_i+d_j);
                            float rcut  = {}*sigma;
                            if (rsq <= rcut*rcut)
                               {{
                               float eps   = {};
                               float sigmasq = sigma*sigma;
                               float rsqinv = sigmasq / rsq;
                               float r6inv = rsqinv*rsqinv*rsqinv;

                               float rsqinv_cut = sigmasq /(rcut*rcut);
                               float r6inv_cut = rsqinv_cut*rsqinv_cut*rsqinv_cut;
                               return 4.0f*eps*(r6inv*r6inv-r6inv_cut*r6inv_cut);
                               }}
                            else
                               {{
                               return 0.0f;
                               }}
                          """.format(scaledr_cut,eps/kT);
        elif (mode == 'force-shift'):
            code = """
                            float rsq = dot(r_ij, r_ij);
                            float sigma  = 0.5*( d_i+d_j);
                            float rcut  = {}*sigma;
                            if (rsq <= rcut*rcut)
                               {{
                               float eps   = {};
                               float sigmasq = sigma*sigma;
                               float rsqinv = sigmasq / rsq;
                               float r6inv = rsqinv*rsqinv*rsqinv;

                               float rsqinv_cut = sigmasq /(rcut*rcut);
                               float r6inv_cut = rsqinv_cut*rsqinv_cut*rsqinv_cut;
                               return 4.0f*eps*( r6inv*(r6inv)-r6inv_cut*(r6inv_cut) +6*(sqrtf(rsq)/rcut-1.0f)*r6inv_cut*(2*r6inv_cut) );
                               }}
                            else
                               {{
                               return 0.0f;
                               }}
                          """.format(scaledr_cut,eps/kT);
        # check if initialization has occurred
        if hoomd.context.exec_conf is None:
            raise RuntimeError('Error creating Lennard-Jones patch energy, call context.initialize() first');

        # raise an error if this run is on the GPU
        if hoomd.context.exec_conf.isCUDAEnabled():
            hoomd.context.msg.error("Patch energies are not supported on the GPU\n");
            raise RuntimeError("Error initializing patch energy");

        # Find a clang executable if none is provided
        if clang_exec is not None:
            clang = clang_exec;
        else:
            clang = 'clang'

        if code is not None and llvm_ir_file is None:
            llvm_ir = self.compile_user(code, clang)
        else:
            # IR is a text file
            with open(llvm_ir_file,'r') as f:
                llvm_ir = f.read()

        self.compute_name = "patch"
        self.cpp_evaluator = _plugin_patch.PatchEnergyJITCustom(hoomd.context.exec_conf, llvm_ir, scaledr_cut);
        mc.set_PatchEnergyEvaluator(self);

        self.mc = mc
        self.enabled = True
        self.log = False

    def compile_user(self, code, clang_exec, fn=None):
        R'''Helper function to compile the provided code into an executable

        Args:
            code (str): C++ code to compile
            clang_exec (str): The Clang executable to use
            fn (str): If provided, the code will be written to a file.

        .. versionadded:: 2.3
        '''
        cpp_function = """
#include "hoomd/HOOMDMath.h"
#include "hoomd/VectorMath.h"

extern "C"
{
float eval(const vec3<float>& r_ij,
    unsigned int type_i,
    const quat<float>& q_i,
    float d_i,
    float charge_i,
    unsigned int type_j,
    const quat<float>& q_j,
    float d_j,
    float charge_j)
    {
"""
        cpp_function += code
        cpp_function += """
    }
}
"""

        include_path = os.path.dirname(hoomd.__file__) + '/include';
        include_path_source = hoomd._hoomd.__hoomd_source_dir__;

        if clang_exec is not None:
            clang = clang_exec;
        else:
            clang = 'clang';

        if fn is not None:
            cmd = [clang, '-O3', '--std=c++11', '-DHOOMD_LLVMJIT_BUILD', '-I', include_path, '-I', include_path_source, '-S', '-emit-llvm','-x','c++', '-o',fn,'-']
        else:
            cmd = [clang, '-O3', '--std=c++11', '-DHOOMD_LLVMJIT_BUILD', '-I', include_path, '-I', include_path_source, '-S', '-emit-llvm','-x','c++', '-o','-','-']
        p = subprocess.Popen(cmd,stdin=subprocess.PIPE,stdout=subprocess.PIPE,stderr=subprocess.PIPE)

        # pass C++ function to stdin
        output = p.communicate(cpp_function.encode('utf-8'))
        llvm_ir = output[0].decode()

        f = open("llvm_ir_file",'w')
        f.write(llvm_ir)
        f.close()

        if p.returncode != 0:
            hoomd.context.msg.error("Error compiling provided code\n");
            hoomd.context.msg.error("Command "+' '.join(cmd)+"\n");
            hoomd.context.msg.error(output[1].decode()+"\n");
            raise RuntimeError("Error initializing patch energy");

        return llvm_ir

    R''' Disable the patch energy and optionally enable it only for logging

    Args:
        log (bool): If true, only use patch energy as a log quantity

    '''
    def disable(self,log=None):
        hoomd.util.print_status_line();

        if log:
            # enable only for logging purposes
            self.mc.cpp_integrator.disablePatchEnergyLogOnly(log)
            self.log = True
        else:
            # disable completely
            self.mc.cpp_integrator.setPatchEnergy(None);
            self.log = False

        self.enabled = False

    R''' (Re-)Enable the patch energy

    '''
    def enable(self):
        hoomd.util.print_status_line()
        self.mc.cpp_integrator.setPatchEnergy(self.cpp_evaluator);

class polydisperse(object):
    R''' Define the patch energy of a polydisperse soft-repulsive
    '''
    def __init__(self, mc, kT, scaledr_cut, v0, eps, model, llvm_ir_file=None, clang_exec=None,kappa=3.0):
        hoomd.util.print_status_line();

        if (model == "polydisperse12"):
            a = scaledr_cut
            A = np.array([  [1,a**2,a**4],
                            [0,2*a,4*a**3],
                            [0,2,12*a**2]
                            ])
            b = np.array([-v0/a**12,12*v0/a**13,-156*v0/a**14])
            c = np.linalg.solve(A,b)
            code = """
                            float rsq = dot(r_ij, r_ij);
                            float sigma  = 0.5*( d_i+d_j)*(1-{}*fabs(d_i-d_j) );
                            float rcut  = {}*(sigma);
                            if (rsq <= rcut*rcut)
                               {{
                               float v0   = {};
                               float c0   = {};
                               float c1   = {};
                               float c2   = {};
                               float sigmasq = sigma*sigma;
                               float rsqinv = sigmasq / rsq;
                               float _rsq = rsq / sigmasq;
                               float r6inv = rsqinv*rsqinv*rsqinv;
                               return v0*r6inv*r6inv+c0+c1*_rsq+c2*_rsq*_rsq;
                               }}
                            else
                               {{
                               return 0.0f;
                               }}
                          """.format(eps,scaledr_cut,v0/kT,c[0]/kT,c[1]/kT,c[2]/kT);
        elif (model == "lennardjones"):
            a = scaledr_cut
            A = np.array([  [1,a**2,a**4],
                            [0,2*a,4*a**3],
                            [0,2,12*a**2]
                            ])
            b = np.array([-v0*(1/a**12-1/a**6),v0*(12*v0/a**13-6/a**7),-v0*(156/a**14-42/a**8)])
            c = np.linalg.solve(A,b)
            code = """
                            float rsq = dot(r_ij, r_ij);
                            float sigma  = 0.5*( d_i+d_j)*(1-{}*fabs(d_i-d_j) );
                            float rcut  = {}*(sigma);
                            if (rsq <= rcut*rcut)
                               {{
                               float v0   = {};
                               float c0   = {};
                               float c1   = {};
                               float c2   = {};
                               float sigmasq = sigma*sigma;
                               float rsqinv = sigmasq / rsq;
                               float _rsq = rsq / sigmasq;
                               float r6inv = rsqinv*rsqinv*rsqinv;
                               return v0*(r6inv*r6inv-r6inv)+c0+c1*_rsq+c2*_rsq*_rsq;
                               }}
                            else
                               {{
                               return 0.0f;
                               }}
                          """.format(eps,scaledr_cut,v0/kT,c[0]/kT,c[1]/kT,c[2]/kT);
        elif (model == "polydisperse18"):
            a = scaledr_cut
            A = np.array([  [1,a**2,a**4],
                            [0,2*a,4*a**3],
                            [0,2,12*a**2]
                            ])
            b = np.array([-v0*(1/a**18),(18*v0/a**19),-v0*(342/a**20)])
            c = np.linalg.solve(A,b)
            code = """
                            float rsq = dot(r_ij, r_ij);
                            float sigma  = 0.5*( d_i+d_j)*(1-{}*fabs(d_i-d_j) );
                            float rcut  = {}*(sigma);
                            if (rsq <= rcut*rcut)
                               {{
                               float v0   = {};
                               float c0   = {};
                               float c1   = {};
                               float c2   = {};
                               float sigmasq = sigma*sigma;
                               float rsqinv = sigmasq / rsq;
                               float _rsq = rsq / sigmasq;
                               float r6inv = rsqinv*rsqinv*rsqinv;
                               return v0*(r6inv*r6inv*r6inv)+c0+c1*_rsq+c2*_rsq*_rsq;
                               }}
                            else
                               {{
                               return 0.0f;
                               }}
                          """.format(eps,scaledr_cut,v0/kT,c[0]/kT,c[1]/kT,c[2]/kT);
        elif (model == "polydisperse10"):
            a = scaledr_cut
            c0 =  -(56.0)*v0/(a**10);
            c1 =  (140.0)*v0/(a**12);
            c2 =  -(120.0)*v0/(a**14);
            c3 =  (35.0)*v0/(a**16);
            code = """
                            float rsq = dot(r_ij, r_ij);
                            float sigma  = 0.5*( d_i+d_j)*(1-{}*fabs(d_i-d_j) );
                            float rcut  = {}*(sigma);
                            if (rsq <= rcut*rcut)
                               {{
                               float v0   = {};
                               float c0   = {};
                               float c1   = {};
                               float c2   = {};
                               float c3   = {};
                               float sigmasq = sigma*sigma;
                               float rsqinv = sigmasq / rsq;
                               float _rsq = rsq / sigmasq;
                               float r10inv = rsqinv*rsqinv*rsqinv*rsqinv*rsqinv;
                               return v0*r10inv+c0+c1*_rsq+c2*_rsq*_rsq+c3*_rsq*_rsq*_rsq;
                               }}
                            else
                               {{
                               return 0.0f;
                               }}
                          """.format(eps,scaledr_cut,v0/kT,c0/kT,c1/kT,c2/kT,c3/kT);
        elif (model == "polydisperse106"):
            a = scaledr_cut
            c0 = (-21 + 10*a**4)*v0/a**10
            c1 =  -((5*(-7 + 3*a**4)*v0)/a**12)
            c2 = (3*(-5 + 2*a**4)*v0)/a**14
            code = """
                            float rsq = dot(r_ij, r_ij);
                            float sigma  = 0.5*( d_i+d_j)*(1-{}*fabs(d_i-d_j) );
                            float rcut  = {}*(sigma);
                            if (rsq <= rcut*rcut)
                               {{
                               float v0   = {};
                               float c0   = {};
                               float c1   = {};
                               float c2   = {};
                               float sigmasq = sigma*sigma;
                               float rsqinv = sigmasq / rsq;
                               float _rsq = rsq / sigmasq;
                               float r10inv = rsqinv*rsqinv*rsqinv*rsqinv*rsqinv;
                               return v0*(r10inv-rsqinv*rsqinv*rsqinv)+c0+c1*_rsq + c2*_rsq*_rsq;
                               }}
                            else
                               {{
                               return 0.0f;
                               }}
                          """.format(eps,scaledr_cut,v0/kT,c0/kT,c1/kT,c2/kT);
        elif (model == "polydisperseyukawa"):
            a = scaledr_cut
            c0 = -((np.exp(-kappa*a)*(15 + 7*kappa*a + kappa**2*a**2)*v0)/(8*a))
            c1 = (np.exp(-kappa*a)*(5 + 5*kappa*a + kappa**2*a**2)*v0)/(4*a**3)
            c2 = -((np.exp(-kappa*a)*(3 + 3*kappa*a + kappa**2*a**2)*v0)/(8*a**5))
            code = """
                            float rsq = dot(r_ij, r_ij);
                            float sigma  = 0.5*( d_i+d_j)*(1-{}*fabs(d_i-d_j) );
                            float rcut  = {}*(sigma);
                            if (rsq <= rcut*rcut)
                               {{
                               float v0   = {};
                               float kappa = {};
                               float c0   = {};
                               float c1   = {};
                               float c2   = {};
                               float sigmasq = sigma*sigma;
                               float rsqinv = sigmasq / rsq;
                               float rinv = sqrtf(rsqinv);
                               float _rsq = rsq / sigmasq;
                               float _r = sqrtf(_rsq);
                               float r10inv = rsqinv*rsqinv*rsqinv*rsqinv*rsqinv;
                               return v0*exp(-_r*kappa)*rinv+c0+c1*_rsq + c2*_rsq*_rsq;
                               }}
                            else
                               {{
                               return 0.0f;
                               }}
                          """.format(eps,scaledr_cut,v0/kT,kappa,c0/kT,c1/kT,c2/kT);
        else:
            raise RuntimeError('Error creating Polydisperse patch energy. Not one of the available models. Perhaps theres a typo?');

        # check if initialization has occurred
        if hoomd.context.exec_conf is None:
            raise RuntimeError('Error creating Lennard-Jones patch energy, call context.initialize() first');

        # raise an error if this run is on the GPU
        if hoomd.context.exec_conf.isCUDAEnabled():
            hoomd.context.msg.error("Patch energies are not supported on the GPU\n");
            raise RuntimeError("Error initializing patch energy");

        # Find a clang executable if none is provided
        if clang_exec is not None:
            clang = clang_exec;
        else:
            clang = 'clang'

        if code is not None and llvm_ir_file is None:
            llvm_ir = self.compile_user(code, clang)
        else:
            # IR is a text file
            with open(llvm_ir_file,'r') as f:
                llvm_ir = f.read()

        self.compute_name = "patch"
        self.cpp_evaluator = _plugin_patch.PatchEnergyJITCustom(hoomd.context.exec_conf, llvm_ir, scaledr_cut);
        mc.set_PatchEnergyEvaluator(self);

        self.mc = mc
        self.enabled = True
        self.log = False

    def compile_user(self, code, clang_exec, fn=None):
        R'''Helper function to compile the provided code into an executable

        Args:
            code (str): C++ code to compile
            clang_exec (str): The Clang executable to use
            fn (str): If provided, the code will be written to a file.

        .. versionadded:: 2.3
        '''
        cpp_function = """
#include "hoomd/HOOMDMath.h"
#include "hoomd/VectorMath.h"

extern "C"
{
float eval(const vec3<float>& r_ij,
    unsigned int type_i,
    const quat<float>& q_i,
    float d_i,
    float charge_i,
    unsigned int type_j,
    const quat<float>& q_j,
    float d_j,
    float charge_j)
    {
"""
        cpp_function += code
        cpp_function += """
    }
}
"""

        include_path = os.path.dirname(hoomd.__file__) + '/include';
        include_path_source = hoomd._hoomd.__hoomd_source_dir__;

        if clang_exec is not None:
            clang = clang_exec;
        else:
            clang = 'clang';

        if fn is not None:
            cmd = [clang, '-O3', '--std=c++11', '-DHOOMD_LLVMJIT_BUILD', '-I', include_path, '-I', include_path_source, '-S', '-emit-llvm','-x','c++', '-o',fn,'-']
        else:
            cmd = [clang, '-O3', '--std=c++11', '-DHOOMD_LLVMJIT_BUILD', '-I', include_path, '-I', include_path_source, '-S', '-emit-llvm','-x','c++', '-o','-','-']
        p = subprocess.Popen(cmd,stdin=subprocess.PIPE,stdout=subprocess.PIPE,stderr=subprocess.PIPE)

        # pass C++ function to stdin
        output = p.communicate(cpp_function.encode('utf-8'))
        llvm_ir = output[0].decode()

        f = open("llvm_ir_file",'w')
        f.write(llvm_ir)
        f.close()

        if p.returncode != 0:
            hoomd.context.msg.error("Error compiling provided code\n");
            hoomd.context.msg.error("Command "+' '.join(cmd)+"\n");
            hoomd.context.msg.error(output[1].decode()+"\n");
            raise RuntimeError("Error initializing patch energy");

        return llvm_ir

    R''' Disable the patch energy and optionally enable it only for logging

    Args:
        log (bool): If true, only use patch energy as a log quantity

    '''
    def disable(self,log=None):
        hoomd.util.print_status_line();

        if log:
            # enable only for logging purposes
            self.mc.cpp_integrator.disablePatchEnergyLogOnly(log)
            self.log = True
        else:
            # disable completely
            self.mc.cpp_integrator.setPatchEnergy(None);
            self.log = False

        self.enabled = False

    R''' (Re-)Enable the patch energy

    '''
    def enable(self):
        hoomd.util.print_status_line()
        self.mc.cpp_integrator.setPatchEnergy(self.cpp_evaluator);