#!/usr/bin/env python
## top-level script for generating probability distributions for component scores as part of CMS 2.0.
## last updated: 06.24.2017 	vitti@broadinstitute.org

import matplotlib as mp
mp.use('agg')
from dists.freqbins_func import run_traj, get_bin_strings, get_bins, check_create_dir, check_create_file, write_bin_paramfile, execute, get_concat_files, get_info_from_tped_name
from dists.scores_func import calc_ihs, calc_delihh, calc_xpehh, calc_fst_deldaf, read_neut_normfile, norm_neut_ihs, norm_sel_ihs, norm_neut_xpehh, norm_sel_xpehh, get_sim_compscore_files, get_scores_from_files, get_compscores_from_files_flatten
from dists.likes_func import plot_pdf_comparison_from_scores, get_plot_pdf_params, quick_load_likes, quick_cl_from_likes, quick_clr_from_likes, get_likes_filenames, write_master_likesfile, get_master_likefiles, get_likes_savestrings
import numpy as np
import argparse
import sys, os, subprocess
import matplotlib.pyplot as plt

#############################
## DEFINE ARGUMENT PARSER ###
#############################
def full_parser_likes_from_model():
	parser=argparse.ArgumentParser(description="This script contains command-line utilities for generating probability distributions for component scores from pre-specified demographic model(s).")
	subparsers = parser.add_subparsers(help="sub-commands")
	generate_sel_bins_parser = subparsers.add_parser('generate_sel_bins', help="Pre-processing step: generate directories with parameter files divided by selDAF, according to specified bins") #nix or clean
	generate_sel_bins_parser.add_argument('outputDir', type=str, action='store', help='location to write cosi output')

	#####################################
	## SIMULATE SELECTION TRAJECTORIES ##
	#####################################
	get_sel_traj_parser = subparsers.add_parser('get_sel_traj', help='Run forward simulations of selection trajectories to populate selscenarios by final allele frequency before running coalescent simulations for entire sample.')
	get_sel_traj_parser.add_argument('traj_outputname', type=str, action='store', help="write to file")
	get_sel_traj_parser.add_argument('--maxAttempts', type=int, action='store', help='maximum number of attempts to generate a trajectory before re-sampling selection coefficient / start time', default=100)

	##########################
	## RUN FULL SIMULATIONS ##
	##########################
	run_neut_sim_parser = subparsers.add_parser('run_neut_sim', help='run neutral simulations')
	run_sel_sim_parser = subparsers.add_parser('run_sel_sim', help='run sims with selection')
	run_sel_sim_parser.add_argument('traj_infilename', type=str, action='store', help="selection trajectory")

	####################################
	## GENERATE SCORES FROM SIMULATES ##
	####################################
	run_repscores_parser = subparsers.add_parser('run_repscores', help="run composite score calculations for simulated tped")
	run_repscores_parser.add_argument('score_basedir', type=str, action='store', help="parent directory in which to generate/populate folders for each composite score")
	run_repscores_parser.add_argument('inputTpedFile', type=str, action='store', help="TPED file with simulate data for putative selected population")
	run_repscores_parser.add_argument('--simRecomFile', type=str, action="store", help="location of input recom file", default="/n/home08/jvitti/params/test_recom.recom")
	get_neut_norm_params_parser = subparsers.add_parser('get_neut_norm_params', help="jointly normalize scores from neutral replicates to get parameters with which to normalize all replicates")
	norm_from_neut_params_parser = subparsers.add_parser('norm_from_neut_params', help="normalize component scores according to neutral distribution")
	norm_from_neut_params_parser.add_argument('--selbin', action='store', help="e.g. 0.10 -- if excluded, normalize neutral simulates")
	norm_from_neut_params_parser.add_argument('--scenario_dir', action='store', help="neut/selsim dir", default="neut")

	##################################################
	## GATHER SCORES AND CALCULATE LIKELIHOOD TABLE ##
	##################################################
	likes_from_scores_parser = subparsers.add_parser('likes_from_scores', help='Collate scores from simulated data in order to generate component test probability distributions.')
	likes_from_scores_parser.add_argument('--nrep_neut', type=int, action="store", help="number of neutral replicates", default=1000)
	likes_from_scores_parser.add_argument('--nrep_sel', type=int, action="store", help="number of replicates per selection scenario bin", default=500)
	likes_from_scores_parser.add_argument('--binMerge', action="store_true", help="if included, generate higher-order groupings of selscenario frequency bins")
	likes_from_scores_parser.add_argument('--foldDist', action="store_true", help="if included, take absolute value of score values (fold distribution over y-axis)", default=False)
	likes_from_scores_parser.add_argument('--save_suffix', type=str, action="store", help="optional string to add to filenames (e.g. to avoid overwrite)", default=None)
	likes_from_scores_parser.add_argument('--save_dir', type=str, action="store", help="if included, specify alternate outpot location", default=None)
	likes_from_scores_parser.add_argument('--save_likelihoods', action="store_true", help="in addition to plotting, save data", default=False)
	if True:
		likes_from_scores_parser.add_argument('--ihs', action="store_true", help="visualize likelihoods for iHS", default=False)
		likes_from_scores_parser.add_argument('--delihh', action="store_true", help="visualize likelihoods for delIHH", default=False)
		likes_from_scores_parser.add_argument('--nsl', action="store_true", help="visualize likelihoods for nSL", default=False)
		likes_from_scores_parser.add_argument('--xpehh', action="store_true", help="visualize likelihoods for XP-EHH", default=False)
		likes_from_scores_parser.add_argument('--deldaf', action="store_true", help="visualize likelihoods for delDAF", default=False)
		likes_from_scores_parser.add_argument('--fst', action="store_true", help="visualize likelihoods for Fst", default=False)

	########################################################
	## ANALYZE AND BUNDLE SCORE PROBABILITY DISTRIBUTIONS ##
	########################################################
	plot_likes_vs_scores_parser = subparsers.add_parser('plot_likes_vs_scores', help='useful QC step; visualize the function that converts score values into likelihood contributions towards composite scores')
	plot_likes_vs_scores_parser.add_argument('inputLikesPrefix', action="store", help="filename minus causal/linked/neutral.txt")
	plot_likes_vs_scores_parser.add_argument('--default_nregion_snps', type=int, action="store", help="presumptive number of SNPs in region (determines prior distributions for within-region CMS calculations)", default=5000)
	write_master_likes_parser = subparsers.add_parser('write_master_likes', help='specify and bundle together input distributions to pass to CMS. This allows the user to specify models, frequency of SNPs used to generated p(score|sel), etc.')
	write_master_likes_parser.add_argument('--likes_masterDir', type=str, default="/n/regal/sabeti_lab/jvitti/", help="location of likelihood tables, defined")
	#write_master_likes_parser.add_argument('--likes_nonSel', type=str, default="vsNeut", help='do we use completely neutral, or linked neutral SNPs for our non-causal distributions? by default, uses strict neutral (CMSgw)')
	#write_master_likes_parser.add_argument('--likes_freqSuffix', type=str, default="allFreqs", help='for causal SNPs, include suffix to specify which selbins to include')
	#write_master_likes_parser.add_argument('--likes_savestring', type=str, default="allFreqs", help='default savestring for individual likes files to cull')
	write_master_likes_parser.add_argument('--model', type=str, default="nulldefault", help='...')
	write_master_likes_parser.add_argument('--likes_inDir', type=str, default="/n/regal/sabeti_lab/jvitti/", help='...')

	#################
	## SHARED ARGS ##
	#################
	for common_parser in [run_neut_sim_parser, run_sel_sim_parser, run_repscores_parser, get_neut_norm_params_parser, norm_from_neut_params_parser]:
		common_parser.add_argument('--checkOverwrite', action="store_true", default=True)
	for run_sim_parser in [run_neut_sim_parser, run_sel_sim_parser]:
		run_sim_parser.add_argument('writeBase', type=str, action='store', help="write prefix")
		run_sim_parser.add_argument('--dropSings', type=float, action='store',  help='randomly thin global singletons from output dataset to model ascertainment bias', default=.25)
	for cosi_parser in [get_sel_traj_parser, run_neut_sim_parser, run_sel_sim_parser]:
		cosi_parser.add_argument('--cosiBuild', type=str, action='store', help='which version of cosi to run', default="coalescent")
		cosi_parser.add_argument('inputParamFile', type=str, action='store', help='file with model specifications for input')
		#cosi_parser.add_argument('--genmapRandomRegions', action='store_true', help='cosi option to sub-sample genetic map randomly from input')
	for cms_preconda_parser in [run_repscores_parser, get_neut_norm_params_parser, norm_from_neut_params_parser]: #TEMPORARY; conda will obviate
		cms_preconda_parser.add_argument('--cmsdir', type=str, action='store', help="location of CMS scripts (TEMPORARY; conda will obviate)", default="/n/home08/jvitti/cms/cms/")
	for interior_score_parser in [get_neut_norm_params_parser, likes_from_scores_parser]:
		interior_score_parser.add_argument('--edge', type=int, action="store", help="use interior of replicates; define per-end bp. (e.g. 1.5Mb -> 1Mb: 250000)", default=250000)
		interior_score_parser.add_argument('--chromlen', type=int, action="store", help="per bp (1.5mb = 1500000)", default=1500000)
	for norm_parser in [get_neut_norm_params_parser, norm_from_neut_params_parser, write_master_likes_parser]:
		norm_parser.add_argument('--score', type=str, action='store', default='ihs')
		norm_parser.add_argument('--simpop', type=int, action='store', default=1)
		norm_parser.add_argument('--altpop', type=int, action='store', default=2)
		norm_parser.add_argument('--nrep', type=int, action='store', default=100)
	for norm_sims_parser in [get_neut_norm_params_parser, norm_from_neut_params_parser, likes_from_scores_parser]:
		norm_sims_parser.add_argument('modeldir', type=str, action="store", help="location of component score folders for demographic scenario")
	for selbin_parser in [generate_sel_bins_parser, get_sel_traj_parser, likes_from_scores_parser]:
		selbin_parser.add_argument('--freqRange', type=str, help="range of final selected allele frequencies to simulate, e.g. .05-.95", default='.05-.95')
		selbin_parser.add_argument('--nBins', type=int, help="number of frequency bins", default=9)
	return parser

############################
## DEFINE EXEC FUNCTIONS ###
############################
### Run simuates from specified demographic
### model under various scenarios
def execute_generate_sel_bins(args):
	''' pre-processing step for sel_trajs(->sel_sim) ''' #hmmm nix this? or else, fix write_bin_paramfile: merge with quick_write_inclusive_sweep_paramfiles.py
	freqRange = args.freqRange
	nBins = args.nBins
	runDir = args.outputDir
	neutParamfile = args.inputParamFile
	fullrange, bin_starts, bin_ends, bin_medians, bin_medians_str = get_bins(freqRange, nBins)
	for ibin in range(nBins):
		populateDir = runDir + "sel_" + bin_medians_str[ibin]
		binDir = check_create_dir(populateDir)
		bounds = bin_starts[ibin], bin_ends[ibin]
		paramfilename = populateDir + "/params"
		write_bin_paramfile(neutParamfile, paramfilename, bounds)
		print('wrote to: ' + paramfilename)
	return
def execute_get_sel_traj(args):
	'''generate forward trajectories of simulated allele frequencies for demographic scenarios with selection'''
	traj_outputname = args.traj_outputname
	cosibuild = args.cosiBuild
	paramfilename = args.inputParamFile
	maxattempts = args.maxAttempts
	run_traj(traj_outputname, cosibuild, paramfilename, maxattempts)
	return
def execute_run_neut_sim(args):
	''' generates tped data for one neutral replicate from a demographic model parameter file using coalescent simulator cosi '''
	cosibuild = args.cosiBuild
	outbase = args.writeBase
	paramfilename = args.inputParamFile
	dropSing = args.dropSings
	cmd = cosibuild
	argstring = "-p " + paramfilename + " --genmapRandomRegions --drop-singletons " + str(dropSing) + " --tped " + outbase + " --output-gen-map"
	#print(cmd + " " + argstring)
	cosicreatefilename = outbase + "_0_1.tped"
	cosi_movedfilename = outbase + "_1.tped"
	proceed = check_create_file(cosicreatefilename, args.checkOverwrite)
	if proceed:
	#if not os.path.isfile(cosicreatefilename):
		fullCmd = cmd + " " + argstring
		print(fullCmd)
		execute(fullCmd)
		for ipop in [1, 2, 3, 4]:
			torenamefile = outbase + "_0_" + str(ipop) + ".tped"
			#print(torenamefile)
			assert os.path.isfile(torenamefile)
			renamed = outbase +"_" + str(ipop) + ".tped"
			renamecmd = "mv "  + torenamefile + " " + renamed
			execute(renamecmd)
	print("wrote simulates to e.g. " + renamed)
	return
def execute_run_sel_sim(args):
	''' generates tped data for one replicate from a demographic model parameter file and selection trajectory file using coalescent simulator cosi '''
	trajectory = args.traj_infilename
	cosibuild = args.cosiBuild
	outbase = args.writeBase
	paramfilename = args.inputParamFile
	dropSing = args.dropSings
	cmd = "env COSI_NEWSIM=1 env COSI_LOAD_TRAJ=" + trajectory + " " + cosibuild
	argstring = "-p " + paramfilename + " --genmapRandomRegions --drop-singletons " + str(dropSing) + " --tped " + outbase + " --output-gen-map"
	#print(cmd + " " + argstring)
	cosicreatefilename = outbase + "_0_1.tped"
	cosi_movedfilename = outbase + "_1.tped"
	proceed = check_create_file(cosi_movedfilename, args.checkOverwrite)
	if proceed:
		fullCmd = cmd + " " + argstring
		print(fullCmd)
		execute(fullCmd)
		for ipop in [1, 2, 3, 4]:
			torenamefile = outbase + "_0_" + str(ipop) + ".tped"
			#print(torenamefile)
			assert os.path.isfile(torenamefile)
			renamed = outbase +"_" + str(ipop) + ".tped"
			renamecmd = "mv "  + torenamefile + " " + renamed
			execute(renamecmd)
	print("wrote simulates to e.g. " + renamed)
	return

### Calculate component scores from
### simulated data
def execute_run_repscores(args):
	''' for one simulated replicate (agnostic wrt neut/sel), generate all component scores '''
	basedir = args.score_basedir
	cmsdir = args.cmsdir
	tped = args.inputTpedFile
	repNum, pop, tpeddir = get_info_from_tped_name(tped)
	checkOverwrite = args.checkOverwrite
	simRecomFile = args.simRecomFile
	assert os.path.isfile(tped)
	for scorefiledir in ['ihs', 'delihh', 'nsl', 'xpehh', 'freqs']:
		check_create_dir(basedir + scorefiledir)

	####### Calculate per-population
	####### scores: iHS, delIHH, nSL
	ihs_commandstring = "python " + cmsdir + "scans.py selscan_ihs"
	ihs_outfileprefix = basedir + "ihs/rep" + str(repNum) + "_" + str(pop)
	ihs_unnormedfile = ihs_outfileprefix + ".ihs.out"
	ihs_argstring = tped + " " + ihs_outfileprefix + " --threads 7 "
	ihs_fullcmd = ihs_commandstring + " " + ihs_argstring
	ihs_normedfile = ihs_unnormedfile + ".norm"
	proceed = check_create_file(ihs_unnormedfile, args.checkOverwrite)
	if proceed:
		print(ihs_fullcmd)
		execute(ihs_fullcmd)
	delihh_commandstring = "python " + cmsdir + "composite.py delihh_from_ihs"
	delihh_unnormedfile =  basedir + "delihh/rep" + str(repNum) + "_" + str(pop) + ".txt"
	delihh_argstring = ihs_unnormedfile + " "+ delihh_unnormedfile
	delihh_fullcmd = delihh_commandstring + " " + delihh_argstring
	delihh_normedfile = delihh_unnormedfile + ".norm"
	proceed = check_create_file(delihh_unnormedfile, args.checkOverwrite)
	if proceed:
		print(delihh_fullcmd)
		execute(delihh_fullcmd)
	nsl_commandstring = "python " + cmsdir + "scans.py selscan_nsl"
	nsl_unnormedfileprefix = basedir + "nsl/rep" + str(repNum) + "_" + str(pop)
	nsl_argstring = tped + " " + nsl_unnormedfileprefix
	nsl_fullcmd = nsl_commandstring + " " + nsl_argstring
	nsl_unnormedfilename = nsl_unnormedfileprefix + ".nsl.out"
	proceed = check_create_file(nsl_unnormedfilename, args.checkOverwrite)
	if proceed:
		print(nsl_fullcmd)
		execute(nsl_fullcmd)

	####### Calculate per-population-pair
	####### scores: XP-EHH, Fst, delDAF
	pops = [1, 2, 3, 4]
	altpops = pops[:]
	altpops.remove(int(pop))
	for altpop in altpops:
		xpehh_commandstring = "python " + cmsdir + "scans.py selscan_xpehh --threads 7"
		tped2 = tpeddir + "rep" + str(repNum) + "_" + str(altpop) + ".tped"
		xpehh_outfileprefix = basedir + "xpehh/rep" + str(repNum) + "_" + str(pop) + "_" + str(altpop)
		xpehh_unnormedfile = basedir + "xpehh/rep" + str(repNum) + "_" + str(pop) + "_" + str(altpop) + ".xpehh.out"
		xpehh_argumentstring = tped + " " + xpehh_outfileprefix + " " + tped2
		xpehh_fullcmd = xpehh_commandstring + " " + xpehh_argumentstring
		proceed = check_create_file(xpehh_unnormedfile, args.checkOverwrite)
		if proceed:
			print(xpehh_fullcmd)
			execute(xpehh_fullcmd)

		fstdeldaf_commandstring = "python " + cmsdir + "composite.py freqscores"
		fstdeldaf_outfilename = basedir + "freqs/rep" + str(repNum) + "_" + str(pop) + "_" + str(altpop)
		fstdeldaf_argumentstring = tped + " " + tped2 + " " + simRecomFile + " " + fstdeldaf_outfilename
		fstdeldaf_fullcmd = fstdeldaf_commandstring + " " + fstdeldaf_argumentstring
		proceed = check_create_file(fstdeldaf_outfilename, args.checkOverwrite)
		if proceed:
			print(fstdeldaf_fullcmd)
			execute(fstdeldaf_fullcmd)
	return
def execute_get_neut_norm_params(args):
	''' creates a concatenated file and saves bin output if run for the first time'''
	pop = args.simpop
	numReps = args.nrep
	basedir = args.modeldir
	if basedir[-1] != "/":
		basedir += "/"
	cmsdir = args.cmsdir
	score = args.score
	altpop = args.altpop
	chrlen, edge = args.chromlen, args.edge
	startbound, endbound = int(edge), chrlen - int(edge) #define replicate interior (conservative strategy to avoid edge effects)

	if score in ['ihs', 'delihh', 'nsl']:
		concatfilebase = basedir + "neut/concat_" + str(pop) + "_"
	elif score in ['xpehh', 'fst']:
		altpop = args.altpop
		concatfilebase = basedir + "neut/concat_" + str(pop) + "_" + str(altpop) + "_"
	else:
		print('must call per composite score')
		sys.exit(0)

	concatfilename = concatfilebase + score + ".txt"
	binfilename = concatfilebase + score + ".bins"

	if not os.path.isfile(binfilename):
		if not os.path.isfile(concatfilename):
			repfiles = []
			for irep in range(1, numReps+1):
				if score == 'ihs':
					unnormedfile = basedir + "neut/ihs/rep" + str(irep) + "_" + str(pop) + ".ihs.out"
					physpos_ind = 1
				elif score == "delihh":
					unnormedfile = basedir + "neut/delihh/rep" + str(irep) + "_" + str(pop) + ".txt"
					physpos_ind = 1
				elif score == "nsl":
					unnormedfile = basedir + "neut/nsl/rep" + str(irep) + "_" + str(pop) + ".nsl.out"
					physpos_ind = 1
				elif score == "xpehh":
					unnormedfile = basedir + "neut/xpehh/rep" + str(irep) + "_" + str(pop) + "_" + str(altpop) + ".xpehh.out"
					physpos_ind = 1
				elif score == "fst":
					unnormedfile = basedir + "neut/freqs/rep" + str(irep) + "_" + str(pop) + "_" + str(altpop)
					physpos_ind = 0
				if os.path.isfile(unnormedfile):
					repfiles.append(unnormedfile)
				else:
					print('missing: ' + unnormedfile)
			concatfile = open(concatfilename, 'w')
			for irepfile in range(len(repfiles)):
				repfile = repfiles[irepfile]
				readfile = open(repfile, 'r')
				firstline = readfile.readline()
				if score in ['xpehh', 'fst']: #header
					if irepfile == 0:
						concatfile.write(firstline)
					else:
						pass
				for line in readfile:
					entries = line.split()
					physpos = int(entries[physpos_ind])
					if physpos >= startbound and physpos <= endbound:
						concatfile.write(line)
				readfile.close()
			concatfile.close()
			print('wrote to: ' + concatfilename)

		#already have concatfile
		infilename = concatfilename
		outfilename = concatfilename + ".norm" #or just write to /tmp ?
		argstring = infilename #+ " " + outfilename
		if score in ['ihs', 'delihh']:
			cmd = "python " + cmsdir + "scans.py selscan_norm_ihs"
		elif score in ['nsl']:
			cmd = "python " + cmsdir + "scans.py selscan_norm_nsl"
		elif score in ['xpehh']:
			cmd = "python " + cmsdir + "scans.py selscan_norm_xpehh"
		else:
			cmd = ""
		fullcmd = cmd + " " + argstring

		alreadyExists = False
		if args.checkOverwrite:
			if not os.path.isfile(binfilename): #check for overwrite
				alreadyExists = False
			else:
				alreadyExists = True
		if alreadyExists == False:
			print(fullcmd)
			#execute(fullcmd)
			with open(binfilename, 'w') as outfile:
				subprocess.check_output( fullcmd.split(), stderr=outfile)
			outfile.close()
			print('wrote to: ' + binfilename)
	return
def execute_norm_from_neut_params(args):
	''' using parameters from neutral simulates, normalizes component scores '''
	pop = args.simpop
	numReps = args.nrep
	basedir = args.modeldir
	if basedir[-1] != "/":
		basedir += "/"
	cmsdir = args.cmsdir
	score = args.score
	altpop = args.altpop
	scenario_dir = args.scenario_dir
	#if args.selbin is None:
	#	scenario_dir = "neut/"
	#else:
	#	scenario_dir = "sel" + str(pop) + "/sel_" + str(args.selbin) + "/"
	concatfilename, binfilename = get_concat_files(pop, score, altpop, basedir=basedir) #I NEED TO MAKE THIS HANDLE SEL SITUATION
	print('loading normalization parameters from ' + binfilename + "...")
	###############
	## NORMALIZE ##
	###############
	for irep in range(1, numReps+1):
		if score in ['ihs']:
			unnormedfile = basedir + scenario_dir + "ihs/rep" + str(irep) + "_" + str(pop)  + ".ihs.out"
			norm_sel_ihs(unnormedfile, binfilename)
		elif score in ['delihh']:
			unnormedfile = basedir + scenario_dir + "delihh/rep" + str(irep) + "_" + str(pop) + ".txt"
			norm_sel_ihs(unnormedfile, binfilename)
		elif score in ['nsl']:
			unnormedfile = basedir + scenario_dir + "nsl/rep" + str(irep) + "_" + str(pop)+ ".nsl.out"
			norm_sel_ihs(unnormedfile, binfilename)
		elif score in ['xpehh']:
			unnormedfile = basedir + scenario_dir + "xpehh/rep" + str(irep) + "_" + str(pop) + "_" + str(altpop) + ".xpehh.out"
			norm_sel_xpehh(unnormedfile, binfilename)
		else: #if score in ['fst']:
			print('currently handling this manually: rewrite_fst_bins.py')
			pass
		if irep%100 == 0:
			print("currently rep: " + str(irep))
	return

### Define component score likelihood
### functions as histograms of simulated scores
def execute_likes_from_scores(args):
	''' define likelihood function for component score based on histograms of score values from simulated data.
	each run is per-score, per-scenario. does all pops / poppairs  '''
	modeldir, pops = args.modeldir, [1, 2, 3, 4]
	if modeldir[-1] != "/":
		modeldir += "/"
	modeldir_entries = modeldir.split('/')
	model = modeldir_entries[-2]
	nPerBin_Neut, nPerBin_Sel = args.nrep_neut, args.nrep_sel
	freqRange, nBins = args.freqRange, args.nBins #define selscenario bins
	fullrange, bin_starts, bin_ends, bin_medians, bin_medians_str = get_bins(freqRange, nBins)
	chrlen, edge = args.chromlen, args.edge #define replicate interior (conservative strategy to avoid edge effects)
	startbound, endbound = int(edge), chrlen - int(edge)
	#print(str([args.ihs, args.delihh, args.nsl, args.xpehh, args.deldaf, args.fst]))
	if [args.ihs, args.delihh, args.nsl, args.xpehh, args.deldaf, args.fst].count(True) != 1:
		print('must call one score at a time')
		sys.exit(0)
	#####################
	## LOAD NEUT FILES ##
	#####################
	neutdir = modeldir + "neut/"
	all_completed_neut = []
	for pop in pops:
		completed_neut = []
		for irep in range(1, nPerBin_Neut+1):
			loaded_neut_files = get_sim_compscore_files(pop, irep, neutdir)
			all_present = [os.path.isfile(item) for item in loaded_neut_files]
			if (all_present.count(True)) == 9: #replicate done
				completed_neut.append(loaded_neut_files)
		all_completed_neut.append(completed_neut)
	print("loaded " + str(sum([len(item) for item in all_completed_neut])) + " neutral replicates with complete component scores from " + neutdir)
	####################
	## LOAD SEL FILES ##
	####################
	binlabels = bin_medians_str
	all_completed_sel, selcounter = [], 0
	for pop in pops:
		seldir = modeldir + "sel" + str(pop) + "/"
		completed_sel = []
		## generate likelihood distributions for selection simulates
		## from multiple frequency bins
		if args.binMerge:
			#need: input bins and name of cluster
			selbin_clusters = [["0.10", "0.20", "0.30", "0.40", "0.50", "0.60", "0.70", "0.80", "0.90"],["0.70", "0.80", "0.90",], ["0.40", "0.50", "0.60",], ["0.10", "0.20", "0.30"],]
			cluster_names = ["allfreq", "hi", "mid", "low"]
			for icluster in range(len(selbin_clusters)):
				cluster, clustername = selbin_clusters[icluster], cluster_names[icluster]
				completed_cluster = []
				for selbin in cluster:
					bindir = seldir + "sel_" + str(selbin) + "/"
					for irep in range(1, nPerBin_Sel+1):
						loaded_sel_files = get_sim_compscore_files(pop, irep, bindir)
						all_present = [os.path.isfile(item) for item in loaded_sel_files]
						if (all_present.count(True)) == 9: #replicate done
							completed_cluster.append(loaded_sel_files)
							selcounter +=1
				completed_sel.append(completed_cluster)
			all_completed_sel.append(completed_sel) #[ipop][iCLUSTER][irep][iscore]
			nChunks, chunk_labels = len(cluster_names), cluster_names

		## fine-grained (ie per-sel-bin) generation of
		## score probability density functions
		else:
			for selbin in binlabels:
				bindir = seldir + "sel_" + str(selbin) + "/"
				completed_bin = []
				for irep in range(1, nPerBin_Sel+1):
					loaded_sel_files = get_sim_compscore_files(pop, irep, bindir)
					all_present = [os.path.isfile(item) for item in loaded_sel_files]
					if (all_present.count(True)) == 9: #replicate done
						completed_bin.append(loaded_sel_files)
						selcounter +=1
				completed_sel.append(completed_bin)
			all_completed_sel.append(completed_sel) #[ipop][ibin][irep][iscore]
			nChunks, chunk_labels = len(binlabels), bin_medians_str
	print("loaded " + str(selcounter) + " selection replicates with complete component scores")
	#################################
	## SORT SCORES INTO HISTOGRAMS ##
	#################################
	for ichunk in range(nChunks): #iterate over selfreqs
		chunkstring = chunk_labels[ichunk]
		####################
		## PER POP SCORES ##
		####################
		if (args.ihs or args.delihh or args.nsl): #per pop
			###########
			### SORT ##
			###########
			if args.ihs:
				score = "ihs"
				print('binning iHS scores...')
				values = get_scores_from_files(all_completed_neut, all_completed_sel, 0, ichunk, startbound, endbound, foldDists = args.foldDist)
			if args.delihh:
				score = "delihh"
				print('binning delIHH scores...')
				values = get_scores_from_files(all_completed_neut, all_completed_sel, 1, ichunk, startbound, endbound, foldDists = args.foldDist)
			if args.nsl:
				score = "nsl"
				print('binning nSL scores...')
				values = get_scores_from_files(all_completed_neut, all_completed_sel, 2, ichunk, startbound, endbound, foldDists = args.foldDist)
			if True:
				pop1vals, pop2vals, pop3vals, pop4vals = values
				neut_1, causal_1, linked_1 = pop1vals
				neut_2, causal_2, linked_2 = pop2vals
				neut_3, causal_3, linked_3 = pop3vals
				neut_4, causal_4, linked_4 = pop4vals
			##########
			## PLOT ##
			##########
			minVal, maxVal, nProbBins, annotate = get_plot_pdf_params(score, args.foldDist)
			plot_title = "PDF for " + score + ", sel_" + str(chunkstring)
			output_dir = ""
			if args.save_dir is not None:
				output_dir = args.save_dir + model + "_"
			else:
				output_dir = modeldir + "likes/"

			savefilebase = output_dir + score + "_sel_" + chunkstring
			if args.save_suffix is not None:
				savefilebase += "_" + args.save_suffix

			savefilename = savefilebase + ".png"
			#savefilebase = modeldir + "likes/" + score + "_sel_" + chunkstring #ensure consistency with get_likes_filenames
			likes_savebase_1 = output_dir + score + "_sel1_" + chunkstring + "_"
			likes_savebase_2 = output_dir + score + "_sel2_" + chunkstring + "_" #+ args.save_suffix + "_"
			likes_savebase_3 = output_dir + score + "_sel3_" + chunkstring + "_" #+ args.save_suffix + "_"
			likes_savebase_4 = output_dir + score + "_sel4_" + chunkstring + "_" #+ args.save_suffix + "_"
			if args.save_suffix is not None:
				likes_savebase_1 += args.save_suffix + "_"
				likes_savebase_2 += args.save_suffix + "_"
				likes_savebase_3 += args.save_suffix + "_"
				likes_savebase_4 += args.save_suffix + "_"

			f1, (ax1, ax2, ax3, ax4)  = plt.subplots(4, sharex=True, sharey=True)
			plt.suptitle(plot_title)
			plot_pdf_comparison_from_scores(ax1, neut_1, causal_1, linked_1, minVal, maxVal, nProbBins, "1 (AFR)", likes_savebase_1, saveFiles = args.save_likelihoods, annotate = annotate)
			plot_pdf_comparison_from_scores(ax2, neut_2, causal_2, linked_2, minVal, maxVal, nProbBins, "2 (EUR)", likes_savebase_2, saveFiles = args.save_likelihoods, annotate = annotate)
			plot_pdf_comparison_from_scores(ax3, neut_3, causal_3, linked_3, minVal, maxVal, nProbBins, "3 (EAS)", likes_savebase_3, annotate = annotate, saveFiles = args.save_likelihoods)
			plot_pdf_comparison_from_scores(ax4, neut_4, causal_4, linked_4, minVal, maxVal, nProbBins, "4 (SAS)", likes_savebase_4, annotate = annotate, saveFiles = args.save_likelihoods)
			ax4.set_xlabel('score value')
			#f1.subplots_adjust(hspace=0)
			#f1.ylabel("p(score)")
			plt.savefig(savefilename)
			plt.close()
			print('saved to ' + savefilename)

		##################
		## PER POP COMP ##
		##################
		if (args.xpehh or args.deldaf or args.fst): #per pop-comp
			###########
			### SORT ##
			###########
			if args.xpehh:
				score = "xpehh"
				print('binning XP-EHH scores...')
				values = get_compscores_from_files_flatten(all_completed_neut, all_completed_sel, "xpehh", ichunk, startbound, endbound, foldDists = args.foldDist)
			if args.deldaf:
				score = "deldaf"
				print('binning delDAF scores...')
				values = get_compscores_from_files_flatten(all_completed_neut, all_completed_sel, "deldaf", ichunk, startbound, endbound, foldDists = args.foldDist)
			if args.fst:
				score = "fst"
				print('binning Fst scores...')
				values = get_compscores_from_files_flatten(all_completed_neut, all_completed_sel, "fst", ichunk, startbound, endbound, foldDists = args.foldDist)
			if True:
				pop1vals, pop2vals, pop3vals, pop4vals = values
				neutvals1, causalvals1, linkedvals1 = pop1vals
				neutvals2, causalvals2, linkedvals2 = pop2vals
				neutvals3, causalvals3, linkedvals3 = pop3vals
				neutvals4, causalvals4, linkedvals4 = pop4vals

			##########
			## PLOT ##
			##########
			minVal, maxVal, nProbBins, annotate = get_plot_pdf_params(score, args.foldDist)
			output_dir = ""
			if args.save_dir is not None:
				output_dir = args.save_dir + model + "_"
			else:
				output_dir = modeldir + "likes/"

			savefilebase = output_dir + score + "_sel_" + str(chunkstring)
			if args.save_suffix is not None:
				savefilebase += "_" + args.save_suffix
			savefilename = savefilebase + ".png"

			plot_title = "PDF for " + score + ", " + str(chunkstring)

			likes_savebase_1 = output_dir + score + "_sel1_" + chunkstring + "_"
			likes_savebase_2 = output_dir + score + "_sel2_" + chunkstring + "_" #+ args.save_suffix + "_"
			likes_savebase_3 = output_dir + score + "_sel3_" + chunkstring + "_" #+ args.save_suffix + "_"
			likes_savebase_4 = output_dir + score + "_sel4_" + chunkstring + "_" #+ args.save_suffix + "_"
			if args.save_suffix is not None:
				likes_savebase_1 += args.save_suffix + "_"
				likes_savebase_2 += args.save_suffix + "_"
				likes_savebase_3 += args.save_suffix + "_"
				likes_savebase_4 += args.save_suffix + "_"
			f1, (ax1, ax2, ax3, ax4)  = plt.subplots(4, sharex=True, sharey=True)
			plt.suptitle(plot_title)
			plot_pdf_comparison_from_scores(ax1, neutvals1, causalvals1, linkedvals1, minVal, maxVal, nProbBins, "1 (AFR)", likes_savebase_1, saveFiles = args.save_likelihoods, annotate = annotate)
			plot_pdf_comparison_from_scores(ax2, neutvals2, causalvals2, linkedvals2, minVal, maxVal, nProbBins, "2 (EUR)", likes_savebase_2, saveFiles = args.save_likelihoods, annotate = annotate)
			plot_pdf_comparison_from_scores(ax3, neutvals3, causalvals3, linkedvals3, minVal, maxVal, nProbBins, "3 (EAS)", likes_savebase_3, annotate = annotate, saveFiles = args.save_likelihoods)
			plot_pdf_comparison_from_scores(ax4, neutvals4, causalvals4, linkedvals4, minVal, maxVal, nProbBins, "4 (SAS)", likes_savebase_4, annotate = annotate, saveFiles = args.save_likelihoods)
			ax4.set_xlabel('score value')
			#plot_pdf_comparison_from_scores(ax1, neuta, causala, linkeda, minVal, maxVal, nProbBins, ax_ylabela, savefilename_a, annotate = annotate, saveFiles = args.save_likelihoods)
			#ax3.set_xlabel('score value')
			#f1.subplots_adjust(hspace=0)
			plt.savefig(savefilename)
			plt.close()
			print('saved to ' + savefilename)
	return
def execute_plot_likes_vs_scores(args): ##REVISIT
	''' useful QC step; visualize the function that converts score values into likelihood contributions towards composite scores '''
	print("Does this still apply with the combine_amend?")
	inputprefix = args.inputLikesPrefix
	nregionsnps = args.default_nregion_snps
	causal_likes_filename, linked_likes_filename, neut_likes_filename = inputprefix + "_causal.txt", inputprefix + "_linked.txt", inputprefix + "_neut.txt"

	causal_likes = quick_load_likes(causal_likes_filename)
	linked_likes = quick_load_likes(linked_likes_filename)
	neut_likes = quick_load_likes(neut_likes_filename)

	comp_likes = quick_cl_from_likes(causal_likes, linked_likes, nSnp = nregionsnps, takeLn = True)
	comp_like_ratios = quick_clr_from_likes(causal_likes, neut_likes, takeLn = True)

	savefilename = inputprefix + "_clr_likes_vs_likesscores.png"
	fig, ax = plt.subplots()
	ax.scatter(np.arange(len(comp_like_ratios)), comp_like_ratios) #I want to make this use the actual score ranges.
	plt.savefig(savefilename)
	print('saved to: ' + savefilename)
	plt.close()

	savefilename = inputprefix + "_cl_likes_vs_likesscores.png"
	fig, ax = plt.subplots()
	ax.scatter(np.arange(len(comp_likes)), comp_likes)
	plt.savefig(savefilename)
	print('saved to: ' + savefilename)
	plt.close()
	return
def execute_write_master_likes(args):
	""" given granular output from likes_from_scores, bundle together into groups of distributions with which to composite """
	score = args.score
	writeloc = args.likes_masterDir
	basedir = args.likes_inDir
	if basedir[-1] != "/":
		basedir += "/"
	if writeloc[-1] != "/":
		writeloc += "/"
	pop = args.simpop
	model = args.model
	#like_savestring = args.likes_savestring
	score_gw_like_savestring, score_local_like_savestring = get_likes_savestrings(score, basedir) #includes _folded

	###############################
	## CMS_GW : FOLDED, VS. NEUT ##
	################## #############

	neut_filename, linked_filename, hit_hi_filename, hit_mid_filename, hit_low_filename, hit_allfreqs_filename = get_likes_filenames(basedir, model, score, pop, like_savestring = score_gw_like_savestring)
	likesFreqs_master_writefilename_global = writeloc + model + "_" + score + "_sel" + str(pop) + "_vsNeut_likesFreqs.master.txt"
	allFreqs_master_writefilename_global = writeloc + model + "_" + score + "_sel" + str(pop) + "_vsNeut_allFreqs.master.txt"
	write_master_likesfile(likesFreqs_master_writefilename_global, neut_filename, hit_hi_filename, hit_mid_filename, hit_low_filename)
	write_master_likesfile(allFreqs_master_writefilename_global, neut_filename, hit_allfreqs_filename, hit_allfreqs_filename, hit_allfreqs_filename)

	######################################
	## CMS_LOCAL : UNFOLDED, VS. LINKED ##
	######################################
	neut_filename, linked_filename, hit_hi_filename, hit_mid_filename, hit_low_filename, hit_allfreqs_filename = get_likes_filenames(basedir, model, score, pop, like_savestring = score_local_like_savestring)
	likesFreqs_master_writefilename_local = writeloc + model + "_" + score + "_sel" + str(pop) + "_vsLinked_likesFreqs.master.txt"
	allFreqs_master_writefilename_local = writeloc + model + "_" + score + "_sel" + str(pop) + "_vsLinked_allFreqs.master.txt"
	write_master_likesfile(likesFreqs_master_writefilename_local, linked_filename, hit_hi_filename, hit_mid_filename, hit_low_filename)
	write_master_likesfile(allFreqs_master_writefilename_local, linked_filename, hit_allfreqs_filename, hit_allfreqs_filename, hit_allfreqs_filename)
	return

##########
## MAIN ##
##########
if __name__ == '__main__':
	runparser = full_parser_likes_from_model()
	args = runparser.parse_args()

	# if called with no arguments, print help
	if len(sys.argv)==1:
		runparser.parse_args(['--help'])

	subcommand = sys.argv[1]
	function_name = 'execute_' + subcommand + "(args)"
	eval(function_name) #points to functions defined above, which wrap other programs in the pipeline


# -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2013-2018 GEM Foundation
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake. If not, see <http://www.gnu.org/licenses/>.

"""
Module exports :class:`CampbellBozorgnia2008`, and
:class:'CampbellBozorgnia2008Arbitrary'
"""
import numpy as np
from math import log, exp
from openquake.hazardlib.gsim.base import GMPE, CoeffsTable
from openquake.hazardlib import const
from openquake.hazardlib.imt import PGA, PGV, PGD, CAV, SA

class CampbellBozorgnia2008(GMPE):
    """
    Implements GMPE developed by Kenneth W. Campbell and Yousef Bozorgnia,
    published as "NGA Ground Motion Model for the Geometric Mean Horizontal
    Component of PGA, PGV, PGD and 5 % Damped Linear Elastic Response Spectra
    for Periods Ranging from 0.01 to 10s" (2008, Earthquake Spectra,
    Volume 24, Number 1, pages 139 - 171).
    This class implements the model for the Geometric Mean of the elastic
    spectra.
    Included in the coefficient set are the coefficients for the
    Campbell & Bozorgnia (2010) GMPE for predicting Cumulative Absolute
    Velocity (CAV), published as "A Ground Motion Prediction Equation for
    the Horizontal Component of Cumulative Absolute Velocity (CSV) Based on
    the PEER-NGA Strong Motion Database" (2010, Earthquake Spectra, Volume 26,
    Number 3, 635 - 650).
    """
    #: Supported tectonic region type is active shallow crust
    DEFINED_FOR_TECTONIC_REGION_TYPE = const.TRT.ACTIVE_SHALLOW_CRUST

    #: Supported intensity measure types are spectral acceleration, peak
    #: ground velocity, peak ground displacement and peak ground acceleration
    #: Additional model for cumulative absolute velocity defined in
    #: Campbell & Bozorgnia (2010)
    DEFINED_FOR_INTENSITY_MEASURE_TYPES = set([
        PGA,
        PGV,
        PGD,
        CAV,
        SA
    ])

    #: Supported intensity measure component is orientation-independent
    #: average horizontal :attr:`~openquake.hazardlib.const.IMC.GMRotI50`
    DEFINED_FOR_INTENSITY_MEASURE_COMPONENT = const.IMC.GMRotI50

    #: Supported standard deviation types are inter-event, intra-event
    #: and total, see section "Aleatory Uncertainty Model", page 147.
    DEFINED_FOR_STANDARD_DEVIATION_TYPES = set([
        const.StdDev.TOTAL,
        const.StdDev.INTER_EVENT,
        const.StdDev.INTRA_EVENT
    ])

    #: Required site parameters are Vs30, Vs30 type (measured or inferred),
    #: and depth (km) to the 2.5 km/s shear wave velocity layer (z2pt5)
    REQUIRES_SITES_PARAMETERS = set(('vs30', 'z2pt5'))

    #: Required rupture parameters are magnitude, rake, dip, ztor
    REQUIRES_RUPTURE_PARAMETERS = set(('mag', 'rake', 'dip', 'ztor'))

    #: Required distance measures are Rrup and Rjb.
    REQUIRES_DISTANCES = set(('rrup', 'rjb'))

    def get_mean_and_stddevs(self, sites, rup, dists, imt, stddev_types):
        """
        See :meth:`superclass method
        <.base.GroundShakingIntensityModel.get_mean_and_stddevs>`
        for spec of input and result values.
        """
        # extract dictionaries of coefficients specific to required
        # intensity measure type and for PGA
        C = self.COEFFS[imt]
        C_PGA = self.COEFFS[PGA()]

        # compute median pga on rock (vs30=1100), needed for site response
        # term calculation
        # For spectral accelerations at periods between 0.0 and 0.25 s, Sa (T)
        # cannot be less than PGA on soil, therefore if the IMT is in this
        # period range it is necessary to calculate PGA on soil
        if isinstance(imt, SA) and (imt.period > 0.0) and (imt.period < 0.25):
            get_pga_site = True
        else:
            get_pga_site = False
        pga1100, pga_site = self._compute_imt1100(C_PGA,
                                                  sites,
                                                  rup,
                                                  dists,
                                                  get_pga_site)

        # Get the median ground motion
        mean = (self._compute_magnitude_term(C, rup.mag) +
                self._compute_distance_term(C, rup, dists) +
                self._compute_style_of_faulting_term(C, rup) +
                self._compute_hanging_wall_term(C, rup, dists) +
                self._compute_shallow_site_response(C, sites, pga1100) +
                self._compute_basin_response_term(C, sites.z2pt5))

        # If it is necessary to ensure that Sa(T) >= PGA (see previous comment)
        if get_pga_site:
            idx = mean < np.log(pga_site)
            mean[idx] = np.log(pga_site[idx])

        stddevs = self._get_stddevs(C,
                                    sites,
                                    pga1100,
                                    C_PGA['s_lny'],
                                    stddev_types)
        return mean, stddevs

    def _compute_imt1100(self, C, sites, rup, dists, get_pga_site=False):
        """
        Computes the PGA on reference (Vs30 = 1100 m/s) rock.
        """
        # Calculates simple site response term assuming all sites 1100 m/s
        fsite = (C['c10'] + (C['k2'] * C['n'])) * log(1100. / C['k1'])
        # Calculates the PGA on rock
        pga1100 = np.exp(self._compute_magnitude_term(C, rup.mag) +
                         self._compute_distance_term(C, rup, dists) +
                         self._compute_style_of_faulting_term(C, rup) +
                         self._compute_hanging_wall_term(C, rup, dists) +
                         self._compute_basin_response_term(C, sites.z2pt5) +
                         fsite)
        # If PGA at the site is needed then remove factor for rock and
        # re-calculate on correct site condition
        if get_pga_site:
            pga_site = np.exp(np.log(pga1100) - fsite)
            fsite = self._compute_shallow_site_response(C, sites, pga1100)
            pga_site = np.exp(np.log(pga_site) + fsite)
        else:
            pga_site = None
        return pga1100, pga_site

    def _compute_magnitude_term(self, C, mag):
        """
        Returns the magnitude scaling factor (equation (2), page 144)
        """
        fmag = C['c0'] + C['c1'] * mag
        if mag <= 5.5:
            return fmag
        elif mag > 6.5:
            return fmag + (C['c2'] * (mag - 5.5)) + (C['c3'] * (mag - 6.5))
        else:
            return fmag + (C['c2'] * (mag - 5.5))

    def _compute_distance_term(self, C, rup, dists):
        """
        Returns the distance scaling factor (equation (3), page 145)
        """
        return (C['c4'] + C['c5'] * rup.mag) * \
            np.log(np.sqrt(dists.rrup ** 2. + C['c6'] ** 2.))

    def _compute_style_of_faulting_term(self, C, rup):
        """
        Returns the style of faulting factor, depending on the mechanism (rake)
        and top of rupture depth (equations (4) and (5), pages 145 - 146)
        """
        frv, fnm = self._get_fault_type_dummy_variables(rup.rake)

        if frv > 0.:
            # Top of rupture depth term only applies to reverse faults
            if rup.ztor < 1.:
                ffltz = rup.ztor
            else:
                ffltz = 1.
        else:
            ffltz = 0.
        return (C['c7'] * frv * ffltz) + (C['c8'] * fnm)

    def _get_fault_type_dummy_variables(self, rake):
        """
        Returns the coefficients FRV and FNM, describing if the rupture is
        reverse (FRV = 1.0, FNM = 0.0), normal (FRV = 0.0, FNM = 1.0) or
        strike-slip/oblique-slip (FRV = 0.0, FNM = 0.0). Reverse faults are
        classified as those with a rake in the range 30 to 150 degrees. Normal
        faults are classified as having a rake in the range -150 to -30 degrees
        :returns:
            FRV, FNM
        """
        if (rake > 30.0) and (rake < 150.):
            return 1., 0.
        elif (rake > -150.0) and (rake < -30.0):
            return 0., 1.
        else:
            return 0., 0.

    def _compute_hanging_wall_term(self, C, rup, dists):
        """
        Returns the hanging wall scaling term, the product of the scaling
        coefficient and four separate scaling terms for distance, magnitude,
        rupture depth and dip (equations 6 - 10, page 146). Individual
        scaling terms defined in separate functions
        """
        return (C['c9'] *
                self._get_hanging_wall_distance_term(dists, rup.ztor) *
                self._get_hanging_wall_magnitude_term(rup.mag) *
                self._get_hanging_wall_depth_term(rup.ztor) *
                self._get_hanging_wall_dip_term(rup.dip))

    def _get_hanging_wall_distance_term(self, dists, ztor):
        """
        Returns the hanging wall distance scaling term (equation 7, page 146)
        """
        fhngr = np.ones_like(dists.rjb, dtype=float)
        idx = dists.rjb > 0.
        if ztor < 1.:
            temp_rjb = np.sqrt(dists.rjb[idx] ** 2. + 1.)
            r_max = np.max(np.column_stack([dists.rrup[idx], temp_rjb]),
                           axis=1)
            fhngr[idx] = (r_max - dists.rjb[idx]) / r_max
        else:
            fhngr[idx] = (dists.rrup[idx] - dists.rjb[idx]) / dists.rrup[idx]
        return fhngr

    def _get_hanging_wall_magnitude_term(self, mag):
        """
        Returns the hanging wall magnitude scaling term (equation 8, page 146)
        """
        if mag <= 6.0:
            return 0.
        elif mag >= 6.5:
            return 1.
        else:
            return 2. * (mag - 6.0)

    def _get_hanging_wall_depth_term(self, ztor):
        """
        Returns the hanging wall depth scaling term (equation 9, page 146)
        """
        if ztor >= 20.0:
            return 0.
        else:
            return (20. - ztor) / 20.0

    def _get_hanging_wall_dip_term(self, dip):
        """
        Returns the hanging wall dip scaling term (equation 10, page 146)
        """
        if dip > 70.0:
            return (90.0 - dip) / 20.0
        else:
            return 1.0

    def _compute_shallow_site_response(self, C, sites, pga1100):
        """
        Returns the shallow site response term (equation 11, page 146)
        """
        stiff_factor = C['c10'] + (C['k2'] * C['n'])
        # Initially default all sites to intermediate rock value
        fsite = stiff_factor * np.log(sites.vs30 / C['k1'])
        # Check for soft soil sites
        idx = sites.vs30 < C['k1']
        if np.any(idx):
            pga_scale = np.log(pga1100[idx] +
                               (C['c'] * ((sites.vs30[idx] / C['k1']) **
                                C['n']))) - np.log(pga1100[idx] + C['c'])
            fsite[idx] = C['c10'] * np.log(sites.vs30[idx] / C['k1']) + \
                (C['k2'] * pga_scale)
        # Any very hard rock sites are rendered to the constant amplification
        # factor
        idx = sites.vs30 >= 1100.
        if np.any(idx):
            fsite[idx] = stiff_factor * log(1100. / C['k1'])

        return fsite

    def _compute_basin_response_term(self, C, z2pt5):
        """
        Returns the basin response term (equation 12, page 146)
        """
        fsed = np.zeros_like(z2pt5, dtype=float)
        idx = z2pt5 < 1.0
        if np.any(idx):
            fsed[idx] = C['c11'] * (z2pt5[idx] - 1.0)

        idx = z2pt5 > 3.0
        if np.any(idx):
            fsed[idx] = (C['c12'] * C['k3'] * exp(-0.75)) *\
                (1.0 - np.exp(-0.25 * (z2pt5[idx] - 3.0)))
        return fsed

    def _get_stddevs(self, C, sites, pga1100, sigma_pga, stddev_types):
        """
        Returns the standard deviations as described in the "ALEATORY
        UNCERTAINTY MODEL" section of the paper. Equations 13 to 19, pages 147
        to 151
        """
        std_intra = self._compute_intra_event_std(C,
                                                  sites.vs30,
                                                  pga1100,
                                                  sigma_pga)

        std_inter = C['t_lny'] * np.ones_like(sites.vs30)
        stddevs = []
        for stddev_type in stddev_types:
            assert stddev_type in self.DEFINED_FOR_STANDARD_DEVIATION_TYPES
            if stddev_type == const.StdDev.TOTAL:
                stddevs.append(self._get_total_sigma(C, std_intra, std_inter))
            elif stddev_type == const.StdDev.INTRA_EVENT:
                stddevs.append(std_intra)
            elif stddev_type == const.StdDev.INTER_EVENT:
                stddevs.append(std_inter)
        return stddevs

    def _compute_intra_event_std(self, C, vs30, pga1100, sigma_pga):
        """
        Returns the intra-event standard deviation at the site, as defined in
        equation 15, page 147
        """
        # Get intra-event standard deviation at the base of the site profile
        sig_lnyb = np.sqrt(C['s_lny'] ** 2. - C['s_lnAF'] ** 2.)
        sig_lnab = np.sqrt(sigma_pga ** 2. - C['s_lnAF'] ** 2.)
        # Get linearised relationship between f_site and ln PGA
        alpha = self._compute_intra_event_alpha(C, vs30, pga1100)

        return np.sqrt(
            (sig_lnyb ** 2.) +
            (C['s_lnAF'] ** 2.) +
            ((alpha ** 2.) * (sig_lnab ** 2.)) +
            (2.0 * alpha * C['rho'] * sig_lnyb * sig_lnab))

    def _compute_intra_event_alpha(self, C, vs30, pga1100):
        """
        Returns the linearised functional relationship between fsite and
        pga1100, determined from the partial derivative defined on equation 17
        on page 148
        """
        alpha = np.zeros_like(vs30, dtype=float)
        idx = vs30 < C['k1']
        if np.any(idx):
            temp1 = (pga1100[idx] +
                     C['c'] * (vs30[idx] / C['k1']) ** C['n']) ** -1.
            temp1 = temp1 - ((pga1100[idx] + C['c']) ** -1.)
            alpha[idx] = C['k2'] * pga1100[idx] * temp1

        return alpha

    def _get_total_sigma(self, C, std_intra, std_inter):
        """
        Returns the total sigma term as defined by equation 16, page 147
        This method is defined here as the Campbell & Bozorgnia (2008) model
        can also be applied to the "arbitrary" horizontal component
        definition, in which case the total sigma is modified.
        """
        return np.sqrt(std_intra ** 2. + std_inter ** 2.)

    COEFFS = CoeffsTable(sa_damping=5, table="""\
      imt      c0     c1      c2      c3      c4    c5    c6     c7      c8     c9     c10    c11   c12    k1      k2     k3     c     n  s_lny  t_lny s_lnAF  c_lny    rho
      cav  -4.354  0.942  -0.178  -0.346  -1.309 0.087  7.24  0.111  -0.108  0.362   2.549  0.090  1.277  400  -2.690  1.000  1.88  1.18  0.371  0.196  0.300  0.089  0.735
      pgd  -5.270  1.600  -0.070   0.000  -2.000  0.17  4.00  0.000   0.000  0.000  -0.820  0.300  1.000  400   0.000  2.744  1.88  1.18  0.667  0.485  0.300  0.290  0.174
      pgv   0.954  0.696  -0.309  -0.019  -2.016  0.17  4.00  0.245   0.000  0.358   1.694  0.092  1.000  400  -1.955  1.929  1.88  1.18  0.484  0.203  0.300  0.190  0.691
      pga  -1.715  0.500  -0.530  -0.262  -2.118  0.17  5.60  0.280  -0.120  0.490   1.058  0.040  0.610  865  -1.186  1.839  1.88  1.18  0.478  0.219  0.300  0.166  1.000
    0.010  -1.715  0.500  -0.530  -0.262  -2.118  0.17  5.60  0.280  -0.120  0.490   1.058  0.040  0.610  865  -1.186  1.839  1.88  1.18  0.478  0.219  0.300  0.166  1.000
    0.020  -1.680  0.500  -0.530  -0.262  -2.123  0.17  5.60  0.280  -0.120  0.490   1.102  0.040  0.610  865  -1.219  1.840  1.88  1.18  0.480  0.219  0.300  0.166  0.999
    0.030  -1.552  0.500  -0.530  -0.262  -2.145  0.17  5.60  0.280  -0.120  0.490   1.174  0.040  0.610  908  -1.273  1.841  1.88  1.18  0.489  0.235  0.300  0.165  0.989
    0.050  -1.209  0.500  -0.530  -0.267  -2.199  0.17  5.74  0.280  -0.120  0.490   1.272  0.040  0.610 1054  -1.346  1.843  1.88  1.18  0.510  0.258  0.300  0.162  0.963
    0.075  -0.657  0.500  -0.530  -0.302  -2.277  0.17  7.09  0.280  -0.120  0.490   1.438  0.040  0.610 1086  -1.471  1.845  1.88  1.18  0.520  0.292  0.300  0.158  0.922
    0.100  -0.314  0.500  -0.530  -0.324  -2.318  0.17  8.05  0.280  -0.099  0.490   1.604  0.040  0.610 1032  -1.624  1.847  1.88  1.18  0.531  0.286  0.300  0.170  0.898
    0.150  -0.133  0.500  -0.530  -0.339  -2.309  0.17  8.79  0.280  -0.048  0.490   1.928  0.040  0.610  878  -1.931  1.852  1.88  1.18  0.532  0.280  0.300  0.180  0.890
    0.200  -0.486  0.500  -0.446  -0.398  -2.220  0.17  7.60  0.280  -0.012  0.490   2.194  0.040  0.610  748  -2.188  1.856  1.88  1.18  0.534  0.249  0.300  0.186  0.871
    0.250  -0.890  0.500  -0.362  -0.458  -2.146  0.17  6.58  0.280   0.000  0.490   2.351  0.040  0.700  654  -2.381  1.861  1.88  1.18  0.534  0.240  0.300  0.191  0.852
    0.300  -1.171  0.500  -0.294  -0.511  -2.095  0.17  6.04  0.280   0.000  0.490   2.460  0.040  0.750  587  -2.518  1.865  1.88  1.18  0.544  0.215  0.300  0.198  0.831
    0.400  -1.466  0.500  -0.186  -0.592  -2.066  0.17  5.30  0.280   0.000  0.490   2.587  0.040  0.850  503  -2.657  1.874  1.88  1.18  0.541  0.217  0.300  0.206  0.785
    0.500  -2.569  0.656  -0.304  -0.536  -2.041  0.17  4.73  0.280   0.000  0.490   2.544  0.040  0.883  457  -2.669  1.883  1.88  1.18  0.550  0.214  0.300  0.208  0.735
    0.750  -4.844  0.972  -0.578  -0.406  -2.000  0.17  4.00  0.280   0.000  0.490   2.133  0.077  1.000  410  -2.401  1.906  1.88  1.18  0.568  0.227  0.300  0.221  0.628
    1.000  -6.406  1.196  -0.772  -0.314  -2.000  0.17  4.00  0.255   0.000  0.490   1.571  0.150  1.000  400  -1.955  1.929  1.88  1.18  0.568  0.255  0.300  0.225  0.534
    1.500  -8.692  1.513  -1.046  -0.185  -2.000  0.17  4.00  0.161   0.000  0.490   0.406  0.253  1.000  400  -1.025  1.974  1.88  1.18  0.564  0.296  0.300  0.222  0.411
    2.000  -9.701  1.600  -0.978  -0.236  -2.000  0.17  4.00  0.094   0.000  0.371  -0.456  0.300  1.000  400  -0.299  2.019  1.88  1.18  0.571  0.296  0.300  0.226  0.331
    3.000 -10.556  1.600  -0.638  -0.491  -2.000  0.17  4.00  0.000   0.000  0.154  -0.820  0.300  1.000  400   0.000  2.110  1.88  1.18  0.558  0.326  0.300  0.229  0.289
    4.000 -11.212  1.600  -0.316  -0.770  -2.000  0.17  4.00  0.000   0.000  0.000  -0.820  0.300  1.000  400   0.000  2.200  1.88  1.18  0.576  0.297  0.300  0.237  0.261
    5.000 -11.684  1.600  -0.070  -0.986  -2.000  0.17  4.00  0.000   0.000  0.000  -0.820  0.300  1.000  400   0.000  2.291  1.88  1.18  0.601  0.359  0.300  0.237  0.200
    7.500 -12.505  1.600  -0.070  -0.656  -2.000  0.17  4.00  0.000   0.000  0.000  -0.820  0.300  1.000  400   0.000  2.517  1.88  1.18  0.628  0.428  0.300  0.271  0.174
    10.00 -13.087  1.600  -0.070  -0.422  -2.000  0.17  4.00  0.000   0.000  0.000  -0.820  0.300  1.000  400   0.000  2.744  1.88  1.18  0.667  0.485  0.300  0.290  0.174
    """)

class CampbellBozorgnia2008Arbitrary(CampbellBozorgnia2008):
    """
    Implements the Campbell & Bozorgnia (2008) GMPE as modified to represent
    the arbitrary horizontal component of ground motion, instead of the
    Rotationally Independent Geometric Mean (GMRotI) originally defined in
    the paper.
    """

    #: Supported intensity measure component is arbitrary horizontal
    #: :attr:`~openquake.hazardlib.const.IMC.HORIZONTAL`,
    DEFINED_FOR_INTENSITY_MEASURE_COMPONENT = const.IMC.HORIZONTAL

    def _get_total_sigma(self, C, std_intra, std_inter):
        """
        Returns the total sigma term for the arbitrary horizontal component of
        ground motion defined by equation 18, page 150
        """
        return np.sqrt(std_intra ** 2. + std_inter ** 2. + C['c_lny'] ** 2.)


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

__author__ = "Luocheng Huang"
__copyright__ = "Copyright 2018"
__credits__ = ["Luocheng Huang"]
__license__ = "GPL"
__version__ = "2.0.2"
__maintainer__ = "Luocheng Huang"
__email__ = "luocheng@uw.edu"
__status__ = "Production"

import pandas as pd
import os
import numpy as np
import math
import time as time
from multiprocessing import Pool
import matplotlib.pyplot as plt
import S4

#CWD = os.getcwd()+'/'

def simulate_one(config, field=1):
    """a wrapper for pickling
    simulates one instance
    """
    global FIELD
    FIELD = field
    new = RCWA.Simulation()
    new.load_input(config)
    return new.run()

class RCWA:
    """
    An rcwa simulation ojbect
    """

    def __init__(self, input_list=None, cores=10, field = None):
        self.input = input_list
        self.cores = cores
        self.field = field
        global FIELD
        """field is the amount of field in the center measured. 1 is 100% .5 is 25%
            returns the mean value if field==None
            else returns the center field."""
        FIELD = field
        pass

    def timeit(method):
        """a decor for timing methods, should be used as a decorator
        """
        def timed(*args, **kw):
            ts = time.time()
            result = method(*args, **kw)
            te = time.time()
            if 'log_time' in kw:
                name = kw.get('log_name', method.__name__.upper())
                kw['log_time'][name] = int((te - ts) * 1000)
            else:
                print '%r  %2.2f s' % \
                      (method.__name__, (te - ts) )
            return result
        return timed

    @timeit
    def simulate_and_graph(self):
        """
        inputs the attribute list,
        :return: simulate results
        """

        p = Pool(self.cores)
        df = pd.DataFrame()
        FIELD = self.field
        solution = p.map(simulate_one, self.input)  # create a pool of workers and map it onto different cores
        p.close() # prevent memory leakage
        p.join() # synchronization point

        t_list = []
        p_list = []
        for s in solution:
            t_list.append(s[0])
            p_list.append(s[1])
        # make phase continuous
        # new_p = unwrap(p_list)
        new_p = [x[0] / (2*math.pi) for x in p_list]

        config = self.input[0]

        # plot data
        plt.figure(0)
        plt.figure(figsize=(10, 5), dpi=200)
        label1 = 'phase'
        label2 = 'amplitude'
        plt.plot(new_p, label=label1)
        plt.plot(t_list, label=label2)
        plt.xlabel('Grating post width')
        plt.ylabel('amp & phase')
        plt.title('Wavelength = {} nm'.format('unknown'))
        plt.legend()

        plt.savefig("{}.png".format(('unknown')))
        plt.show();
        plt.clf();

        # save to csv
        a = (new_p, t_list)

        np.savetxt("data_{}.csv".format(('unknown')), a, delimiter=",")

        pass

    @timeit
    def simulate(self):
        """
        inputs the attribute list,
        :return: simulate results
        """

        p = Pool(self.cores)
        df = pd.DataFrame(self.input)
        solution = p.map(simulate_one, self.input)  # create a pool of workers and map it onto different cores
        solutions = pd.DataFrame(solution)
        p.close() # prevent memory leakage
        p.join() # synchronization point

        result = pd.concat([df, solutions], axis=1)
        result.columns = ['CSCS', 'e_field']

        # make phase continuous
        # new_p = unwrap(p_list)

        return result

    class Simulation:
        """A simulation object"""

        def __init__(self):
            self.layers = []
            self.basis = ((1,0),(0,1))
            self.NumBasis=150
            self.layers = None
            self.layer_material = None
            self.layer_thickness = None
            self.layer_pattern = None
            self.transmission = None
            self.wavelength = None
            self.buffer = 0.1
            self.Nxy = 52

        def load_input(self,input_instance):
            """loades the canonical id into local variables"""

            # the first step is to assign values to the variables from the input_instance
            parts = input_instance.split('/')
            first_arg = parts[0].split('-')
            if first_arg[0] == 't':
                self.transmission = True
            elif first_arg[0] == 'r':
                self.transmission = False

            self.wavelength = float(first_arg[1])
            self.basis = eval(first_arg[2])

            self.layers = parts[1:]
            layers_split = [layer.split('=') for layer in self.layers]
            self.layer_material = [layer_split[0] for layer_split in layers_split]
            split_2 = [layer_split[1] for layer_split in layers_split]
            self.layer_thickness = []
            self.layer_pattern = []
            for x in split_2:
                x_split = x.split(':')
                if len(x_split) == 1:
                    self.layer_thickness.append(float(x_split[0]))
                    self.layer_pattern.append(None)
                else:
                    self.layer_thickness.append(float(x_split[0]))
                    self.layer_pattern.append(x_split[1].split('-'))

        def run(self):
            """create layers and return a S4 object
            :type attr: dict
            return S
            """
            def querry_n(material, wavelength):
                """
                given the material in string, wavelength in µm
                output (n,k)
                """
                mat = pd.read_csv('../rcwa/n_data/{}_n.csv'.format(material))

                w = mat['Wavelength, µm']
                n_list = mat['n']
                k_list = mat['k']

                n = np.interp(wavelength, w, n_list) # interpolate the refractive index value
                k = np.interp(wavelength, w, k_list) # interpolate the k value

                return n + k*1j

            # define the S object
            S = S4.New(Lattice=self.basis,
                       NumBasis=self.NumBasis)  # orthogonal basis, numbasis scales with memory O(N^2)

            # define materials
            for material in set(self.layer_material):
                n = querry_n(material,self.wavelength)
                # defines permittivity in terms of refractive index
                eps = n**2
                S.SetMaterial(Name=material, Epsilon=eps)
            S.SetMaterial(Name='Vacuum', Epsilon=1.0)

            # define layers
            S.AddLayer(Name='air_above', Thickness=self.buffer, Material='Vacuum')
            index = 0
            for layer in zip(self.layer_material, self.layer_thickness, self.layer_pattern):
                index += 1
                if layer[2] == None:
                    S.AddLayer(Name=str(index), Thickness=layer[1], Material=layer[0])
                else:
                    S.AddLayer(Name=str(index), Thickness=layer[1], Material='Vacuum')
                    if '+' in layer[2]:
                        patterns = layer[2].split('+')
                    else:
                        patterns = layer[2]
                    for pattern in patterns:
                        shape = pattern[0]
                        coor = eval(pattern[1:])
                        if shape == 'C':
                            # create circle features
                            if coor[2] != 0:
                                S.SetRegionCircle(
                                    Layer = str(index),
                                    Material = layer[0],
                                    Center = (coor[0],coor[1]),
                                    Radius = coor[2]
                                )
                        elif shape == 'S':
                        # create circle features
                            if coor[2] != 0:
                                S.SetRegionSquare(
                                    Layer = str(index),
                                    Material = layer[0],
                                    Center = (coor[0],coor[1]),
                                    Angle = 0,
                                    Halfwidths = (coor[2], coor[2])

                                )

            S.AddLayer(Name='below', Thickness=self.buffer, Material='Vacuum')

            # simulation settings
            S.SetExcitationPlanewave(IncidenceAngles=(0.0, 0.0), sAmplitude=0.0,
                                     pAmplitude=1.0, Order=0)

            S.SetFrequency(1/self.wavelength)

            # For higher accuracy
            S.SetOptions(SubpixelSmoothing=True)

            # e is array of dimension Nx by Ny, each element is a tuple of length 3

            if  not self.transmission:
                z_pos = self.buffer/(2.0)
            else:
                z_pos = self.buffer*1.5+sum(self.layer_thickness)

            e, h = S.GetFieldsOnGrid(z=z_pos, NumSamples=(self.Nxy, self.Nxy),
                                     Format='Array')
            e = np.array(e)
            if not self.transmission:
                e = e - np.exp(1j*2*np.pi/self.wavelength*self.buffer/2)

            if FIELD==None:
                e_field = np.mean(e[:,:,0])  # adding 1 removes incident wave
            else:
                (a,b,_) = np.shape(e)
                edge = (1-FIELD)/2.0
                e_field = e[int(a*edge):int(a*(1-edge)),int(b*edge):int(b*(1-edge)),0]

            e_field = np.round(e_field, 3)

            return pd.Series(str(e_field.tolist()))


"""Model of somatostatin brain extracellular release, diffusion, and loss.
"""

from neuron import h, rxd
from neuron.units import nM, uM, cm, s, M, nm, um
import numpy as np

# Avogadro's Number from scipy
from scipy.constants import N_A
from .basemodel import ModelBase

class Model(ModelBase):
    """A model of somatostatin release, diffusion, and loss in the brain extracellular space.

    This model simulates the photorelease of the neuropeptide somatostatin from
    gold coated nanovesicles and its subsequent diffusion in the brain
    extracellular space. An initial cylindrical bolus of somatostatin is placed in
    the center of the extracellular domain to approximate the condition after
    photorelease from gold coated nanovesicles with release stimulated by a
    tornado scan of radius 30 micron with a multi-photon microscope. Somatostatin
    can then diffuse away from the release site and can be lost from the
    extracellular space through a first order decay reaction.

    Model components (accessible as model attributes):

        Compartments:
            ecs -> rxd.Extracellular: The extracellular space domain and its
                parameters like volume fraction and tortuosity.
        Species:
            sst -> rxd.Species: Neuropeptide somatostatin (sst) that is
                photoreleased in the brain extracellular space.
        Parameters:
            loss_rate -> rxd.Parameter: Defines the kinetic rate for first order
                loss of sst from the extracellular space.
        Reactions:
            sst_loss -> rxd.Rate: The reaction corresponding to first order loss
                of sst from the extracellular space.
                Forward reaction: sst --loss_rate--> None

    """

    def __init__(
        self,
        volume_fraction: float = 0.2,
        tortuosity: float = 1.0,
        dx: float = 5,
        xlo: float = -200.0,
        xhi: float = 200.0,
        ylo: float = -200.0,
        yhi: float = 200.0,
        zlo: float = -200.0,
        zhi: float = 200.0,
    ):
        """Defines and constructs all the model components.

        Args:
            volume_fraction: The volume fraction of extracellular space
                (unitless). DEFAULT: 0.2
            tortuosity: The tortuosity for diffusion in the extracellular space
                (unitless). DEFAULT: 1.
            dx: The discrettization size for volume voxels in the extracellular
                space in microns. DEFAUT: 5
            xlo: The position of the lower edge of the extracellular simulation
                domain along the x-direction in microns. DEFAULT: -200.
            xhi: The position of the upper edge of the extracellular simulation
                domain along the x-direction in microns. DEFAULT: 200.
            ylo: The position of the lower edge of the extracellular simulation
                domain along the y-direction in microns. DEFAULT: -200.
            yhi: The position of the upper edge of the extracellular simulation
                domain along the y-direction in microns. DEFAULT: 200.
            zlo: The position of the lower edge of the extracellular simulation
                domain along the z-direction in microns. DEFAULT: -200.
            zhi: The position of the upper edge of the extracellular simulation
                domain along the z-direction in microns. DEFAULT: 200.
        """
        # Where? -- specify the regions
        # For extracellular reaction-diffusion we just have one
        # Extracellular region.
        self.ecs = rxd.Extracellular(
            xlo=xlo,
            ylo=ylo,
            zlo=zlo,
            xhi=xhi,
            yhi=yhi,
            zhi=zhi,
            dx=dx,
            volume_fraction=volume_fraction,
            tortuosity=tortuosity,
        )
        # Who? -- define all the species
        # Calcein
        # The effective diffusion coefficient for fluorescently labelled SST
        # in acute brain slices from integrative optical imaging (IOI) is
        # 8.9+-1.3 x10^-7 cm^2/s. Xiong et al. 2021 bioRxiv https://doi.org/10.1101/2021.09.10.459853
        d_sst = (
            8.9e-7 * cm ** 2 / s
        )  # effective diffusion coefficient (factoring in tortuosity)
        # Approximate SST photorelease from gold coated nanovesicles as in
        # Xiong et al. 2021 bioRxiv https://doi.org/10.1101/2021.09.10.459853
        # Use an initial bolus of SST with a uniform concentration corresponding
        # to the estimated 1.2x10^8 released molecs inside a disc that
        # approximates the two-photon tornado scan area used for nanovesicle
        # release.
        r_stim = 30 * um  # radius of the tornado scan area
        Qsst = 1.2e8  # Number of SST molecules released during photostimulation
        focal_disc_z = 12  # z-height of the 2-photon focal plane
        disc_vol = (
            np.pi * r_stim ** 2 * focal_disc_z
        )  # volume of the tornado scan area.
        # SST concentration accounting for the disc volume and volume faction.
        sst_0 = ((Qsst / N_A) / disc_vol) / volume_fraction * 1e15 * M
        self.sst = rxd.Species(
            self.ecs,
            name="somatostatin",
            d=d_sst,
            charge=0,
            initial=lambda nd: sst_0
            if (nd.x3d ** 2 + nd.y3d ** 2 < r_stim ** 2)
            and (np.abs(nd.z3d) < focal_disc_z / 2)
            else 0,
            ecs_boundary_conditions=0.0,
        )
        # What? -- specify all the reactions
        # 1st-order loss of SST
        # Loss rate of SST estimated in range 0.023-0.048 per second.
        # Xiong et al. 2021 bioRxiv https://doi.org/10.1101/2021.09.10.459853
        kf = 3.6e-2 / s  # the reaction rate
        self.loss_rate = rxd.Parameter(self.ecs, value=kf)
        self.sst_loss = rxd.Rate(self.sst, -self.loss_rate * self.sst)

        # Initialize private variables -- required for all models
        self._times = None  # We'll store the time points in our simulation trajectory.
        self._observables = None  # We'll assign our observables to this variable.

    def simulate(self, time_step: float, n_steps: int, output_frequency: int = 1):
        """Run the simulation and set the desired trajectory ouptuts.

        Args:
            time_step: The time interval between each step of the simulation
                trajectory in milliseconds. Times not in milliseconds can
                be converted using unit conversions from the neuron.units
                module.
            n_steps: The total number of simulation steps to run.
            output_frequency: Set the frequency for computing and storing
                any observables during the simulation. DEFAULT: 1

        """
        # Set the time step and initialize the simulator.
        h.dt = time_step
        h.finitialize()
        # define any observables that you want to store
        times = list()
        sst_zproject_mean = list()
        # Run each step.
        for f in range(n_steps + 1):
            if (f == 0) or ((f % output_frequency) == 0):
                # simulation time
                times.append(f * time_step)
                # Get the mean value z-projection of the Calcein concentration.
                sst_zproject_mean.append(self.sst[self.ecs].states3d.mean(2))
            h.fadvance()
        self._times = np.array(times)
        self._observables = {"zproject_mean_sst": sst_zproject_mean}
        return

model = Model()


import os
import sys
from functools import lru_cache, wraps

import astropy.units as astropy_units
import numpy as np
import six
from astropy.io import fits

from astromodels.functions.function import Function1D, FunctionMeta
from astromodels.utils import configuration
from astromodels.utils import _get_data_file_path
import gc

try:

    import pyatomdb

    has_atomdb = True

except:

    has_atomdb = False

if has_atomdb:
    # APEC class

    class APEC(Function1D, metaclass=FunctionMeta):
        r"""
        description :
            The Astrophysical Plasma Emission Code (APEC, Smith et al. 2001)
            contributed by Dominique Eckert
        parameters :
            K :
                desc : Normalization in units of 1e-14/(4*pi*(1+z)^2*dA*2)*EM
                initial value : 1.0
                is_normalization : True
                transformation : log10
                min : 1e-30
                max : 1e3
                delta : 0.1
            kT :
                desc : Plasma temperature
                initial value : 1.0
                min : 0.08
                max : 64
                delta : 0.1
            abund :
                desc : Metal abundance
                initial value : 1
                min : 0.0
                max : 5.0
                delta : 0.01
                fix : yes
            redshift :
                desc : Source redshift
                initial value : 0.1
                min : 0.0
                max : 10.0
                delta : 1e-3
                fix : yes

        """

        def _set_units(self, x_unit, y_unit):
            self.kT.unit = astropy_units.keV

            self.abund.unit = astropy_units.dimensionless_unscaled

            self.redshift.unit = astropy_units.dimensionless_unscaled

            self.K.unit = y_unit

        def init_session(self, abund_table="AG89"):
            # initialize PyAtomDB session
            self.session = pyatomdb.spectrum.CIESession(abundset=abund_table)

        def clean(self):
            """
            Clean the current APEC session to avoid having too many open files
            :returns:
            """

            self.session = None
            del self.session
            gc.collect()

        def evaluate(self, x, K, kT, abund, redshift):
            assert self.session is not None, "please run init_session(abund)"

            sess = self.session

            nval = len(x)

            xz = x * (1.0 + redshift)

            ebplus = (np.roll(xz, -1) + xz)[: nval - 1] / 2.0

            ebounds = np.empty(nval + 1)

            ebounds[1:nval] = ebplus

            ebounds[0] = xz[0] - (ebplus[0] - xz[0])

            ebounds[nval] = xz[nval - 1] + (xz[nval - 1] - ebplus[nval - 2])

            binsize = (np.roll(ebounds, -1) - ebounds)[:nval]

            sess.set_response(ebounds, raw=True)

            sess.set_abund(
                [
                    6,
                    7,
                    8,
                    9,
                    10,
                    11,
                    12,
                    13,
                    14,
                    16,
                    17,
                    18,
                    19,
                    20,
                    21,
                    22,
                    23,
                    24,
                    25,
                    26,
                    27,
                    28,
                    29,
                    30,
                ],
                abund,
            )

            spec = sess.return_spectrum(kT) / binsize / 1e-14

            return K * spec

    # VAPEC class

    class VAPEC(Function1D, metaclass=FunctionMeta):
        r"""
        description :
            The Astrophysical Plasma Emission Code (APEC, Smith et al. 2001), variable abundances for individual elements
            contributed by Dominique Eckert
        parameters :
            K :
                desc : Normalization in units of 1e-14/(4*pi*(1+z)^2*dA*2)*EM
                initial value : 1.0
                is_normalization : True
                transformation : log10
                min : 1e-30
                max : 1e3
                delta : 0.1
            kT :
                desc : Plasma temperature
                initial value : 1.0
                min : 0.08
                max : 64
                delta : 0.1
            Fe :
                desc : Fe abundance
                initial value : 1
                min : 0.0
                max : 5.0
                delta : 0.01
                fix : yes
            C :
                desc : C abundance
                initial value : 1
                min : 0.0
                max : 5.0
                delta : 0.01
                fix : yes
            N :
                desc : N abundance
                initial value : 1
                min : 0.0
                max : 5.0
                delta : 0.01
                fix : yes
            O :
                desc : O abundance
                initial value : 1
                min : 0.0
                max : 5.0
                delta : 0.01
                fix : yes
            Ne :
                desc : Ne abundance
                initial value : 1
                min : 0.0
                max : 5.0
                delta : 0.01
                fix : yes
            Mg :
                desc : Mg abundance
                initial value : 1
                min : 0.0
                max : 5.0
                delta : 0.01
                fix : yes
            Al :
                desc : Al abundance
                initial value : 1
                min : 0.0
                max : 5.0
                delta : 0.01
                fix : yes
            Si :
                desc : Si abundance
                initial value : 1
                min : 0.0
                max : 5.0
                delta : 0.01
                fix : yes
            S :
                desc : S abundance
                initial value : 1
                min : 0.0
                max : 5.0
                delta : 0.01
                fix : yes
            Ar :
                desc : Ar abundance
                initial value : 1
                min : 0.0
                max : 5.0
                delta : 0.01
                fix : yes
            Ca :
                desc : Ca abundance
                initial value : 1
                min : 0.0
                max : 5.0
                delta : 0.01
                fix : yes
            Ni :
                desc : Ni abundance
                initial value : 1
                min : 0.0
                max : 5.0
                delta : 0.01
                fix : yes
            redshift :
                desc : Source redshift
                initial value : 0.1
                min : 0.0
                max : 10.0
                delta : 1e-3
                fix : yes

        """

        def _set_units(self, x_unit, y_unit):
            self.kT.unit = astropy_units.keV

            self.Fe.unit = astropy_units.dimensionless_unscaled

            self.C.unit = astropy_units.dimensionless_unscaled

            self.N.unit = astropy_units.dimensionless_unscaled

            self.O.unit = astropy_units.dimensionless_unscaled

            self.Ne.unit = astropy_units.dimensionless_unscaled

            self.Mg.unit = astropy_units.dimensionless_unscaled

            self.Al.unit = astropy_units.dimensionless_unscaled

            self.Si.unit = astropy_units.dimensionless_unscaled

            self.Ar.unit = astropy_units.dimensionless_unscaled

            self.Ca.unit = astropy_units.dimensionless_unscaled

            self.Ni.unit = astropy_units.dimensionless_unscaled

            self.redshift.unit = astropy_units.dimensionless_unscaled

            self.K.unit = y_unit

        def init_session(self, abund_table="AG89"):
            # initialize PyAtomDB session
            self.session = pyatomdb.spectrum.CIESession(abundset=abund_table)

        def clean(self):
            """
            Clean the current APEC session to avoid having too many open files
            :returns:
            """

            self.session = None
            del self.session
            gc.collect()

        def evaluate(
            self, x, K, kT, Fe, C, N, O, Ne, Mg, Al, Si, S, Ar, Ca, Ni, redshift
        ):
            assert self.session is not None, "please run init_session(abund)"

            sess = self.session

            nval = len(x)

            xz = x * (1.0 + redshift)

            ebplus = (np.roll(xz, -1) + xz)[: nval - 1] / 2.0

            ebounds = np.empty(nval + 1)

            ebounds[1:nval] = ebplus

            ebounds[0] = xz[0] - (ebplus[0] - xz[0])

            ebounds[nval] = xz[nval - 1] + (xz[nval - 1] - ebplus[nval - 2])

            binsize = (np.roll(ebounds, -1) - ebounds)[:nval]

            sess.set_response(ebounds, raw=True)

            sess.set_abund(
                [6, ], C,
            )

            sess.set_abund(
                [7, ], N,
            )

            sess.set_abund(
                [8, ], O,
            )

            sess.set_abund(
                [10, ], Ne,
            )

            sess.set_abund(
                [12, ], Mg,
            )

            sess.set_abund(
                [13, ], Al,
            )

            sess.set_abund(
                [14, ], Si,
            )

            sess.set_abund(
                [16, ], S,
            )

            sess.set_abund(
                [18, ], Ar,
            )

            sess.set_abund(
                [20, ], Ca,
            )

            sess.set_abund(
                [26, ], Fe,
            )

            sess.set_abund(
                [28, ], Ni,
            )

            sess.set_abund(
                [9, 11, 15, 17, 19, 21, 22, 23, 24, 25, 27, 29, 30], Fe
            )  # Remaining elements are set to Fe

            spec = sess.return_spectrum(kT) / binsize / 1e-14

            return K * spec

import numpy,import astropy,from astropy