import os
import sys
import time
import datetime
try:
	import knockadapt
except ImportError:
	# This is super hacky, will no longer be necessary after development ends
	file_directory = os.path.dirname(os.path.abspath(__file__))
	parent_directory = os.path.split(file_directory)[0]
	knockoff_directory = parent_directory + '/knockadapt'
	sys.stdout.write(f'Knockoff dir is {knockoff_directory}\n')
	sys.path.insert(0, os.path.abspath(knockoff_directory))
	import knockadapt
from knockadapt import knockoff_stats as kstats
from knockadapt import utilities
from knockadapt.knockoff_filter import KnockoffFilter

import warnings
import numpy as np
import pandas as pd

import itertools
from multiprocessing import Pool
from functools import partial

# Global: the set of antisymmetric functions we use
PAIR_AGGS = ['cd', 'sm']
# q-values to evaluate for
q_values = 0.05*np.arange(1, 21)
# Degrees of freedom for t-distributions
DEFAULT_DF_T = knockadapt.graphs.DEFAULT_DF_T

def fetch_kwarg(kwargs, key, default=None):
	""" Utility function for parsing """
	if key in kwargs:
		return kwargs.pop(key)
	else:
		return default

def str2bool(v):
	""" Helper function, converts strings to boolean vals"""
	if isinstance(v, bool):
		return v
	elif isinstance(v, str):
		if v.lower() in ['yes', 'true', 't', 'y', '1']:
			return True
		elif v.lower() in ['no', 'false', 'f', 'n', '0']:
			return False
	return v

def obj2int(v):
	try:
		v = float(v)
		if v.is_integer():
			v = int(v)
		return v
	except:
		return str2bool(v)

def val2list(val):
	""" Turns discrete values into lists, otherwise returns """
	if not isinstance(val, list):
		if isinstance(val, np.ndarray):
			if len(val.shape) == 1:
				return val
		return [val]
	return val

def dict2keyproduct(dictionary):
	""" Takes a dictionary mapping to lists
	and returns the sorted list of keys and
	the cartesian product of each list."""
	keys = sorted([key for key in dictionary])
	components = []
	for key in keys:
		dictionary[key] = val2list(dictionary[key])
		components.append(dictionary[key])
	product = list(itertools.product(*components))
	return keys, list(product)

def apply_pool(func, all_inputs, num_processes):
	""" Utility function"""

	# Don't use the pool object if num_processes=1
	if num_processes == 1:
		all_outputs = []
		for inp in all_inputs:
			all_outputs.append(func(inp))
	else:
		with Pool(num_processes) as thepool:
			all_outputs = thepool.map(func, all_inputs)

	return all_outputs

def fetch_competitor_S(
	Sigma,
	groups,
	time0,
	S_curve=False,
	verbose=False,
	**kwargs
):
	"""
	Creates various different S-matrices.
	"""
	# If S_curve is true, just do gamma * I for many gammas
	p = Sigma.shape[0]
	fetch_kwarg(kwargs, 'sdp_tol', None) # Prevent double arg errors
	if S_curve:
		mineig = np.linalg.eigh(Sigma)[0].min()
		gammas = 2*mineig*np.arange(0, 11, 1)/10
		gammas[0] += 0.001
		gammas[-1] -= 0.001
		S_matrices = {
			f'S{np.around(gamma, 3)}':gamma*np.eye(p) for gamma in gammas
		}
		S_matrices['mvr'] = knockadapt.knockoffs.compute_S_matrix(
			Sigma=Sigma,
			groups=groups,
			method='mvr',
			solver='cd',
		)
		S_matrices['maxent'] = knockadapt.knockoffs.compute_S_matrix(
			Sigma=Sigma,
			groups=groups,
			method='maxent',
			solver='cd'
		)
		S_matrices['sdp'] = knockadapt.knockoffs.compute_S_matrix(
			Sigma=Sigma,
			groups=groups,
			method='sdp',
			tol=1e-10,
		)
		S_matrices['ci'] = knockadapt.knockoffs.compute_S_matrix(
			Sigma=Sigma,
			groups=groups,
			method='ciknock',
		)
		# Create perturbed option when G_SDP is low rank
		if Sigma[0,1] >= 0.5:
			S_matrices['sdp_perturbed'] = 0.99*S_matrices['sdp']
		return S_matrices

	### Special case: detect if Sigma is equicorrelated,
	# in which case we can calculate the solution analytically
	# for ungrouped knockoffs.
	if Sigma is not None:
		if np.unique(groups).shape[0] == p:
			rho = Sigma[0, 1]
			if rho >= 0:
				equicorr = rho*np.ones((p, p)) + (1-rho)*np.eye(p)
				if np.all(np.abs(Sigma-equicorr) < 1e-10):
					print(f"Sigma is equicorr (rho={rho}), using analytical solution")
					S_SDP = min(1, 2-2*rho)*np.eye(p)
					S_matrices = {'sdp':S_SDP}
					# Compute exact solution for smaller p
					if p < 1000:
						S_matrices['mvr'] = knockadapt.knockoffs.compute_S_matrix(
							Sigma=Sigma,
							groups=groups,
							method='mvr',
							solver='cd',
						)
						S_matrices['mmi'] = knockadapt.knockoffs.compute_S_matrix(
							Sigma=Sigma,
							groups=groups,
							method='maxent',
							solver='cd'
						)
						S_matrices['ci'] = knockadapt.knockoffs.compute_S_matrix(
							Sigma=Sigma,
							method='ciknock',
						)
					# Else, asymptotically these methods are the same
					else:
						S_matrices['mvr'] = (1-rho)*np.eye(p)
					if rho < 0.5:
						return S_matrices
					else:
						S_matrices['sdp_perturbed'] = S_SDP*(0.99)
						return S_matrices

	### Calculate (A)SDP S-matrix
	if time0 is None:
		time0 = time.time()

	# Compute SDP matrix
	S_matrices = {}
	if verbose:
		(f'Now computing SDP S matrices, time is {time.time() - time0}')
	S_SDP = knockadapt.knockoffs.compute_S_matrix(
		Sigma=Sigma,
		groups=groups,
		max_block=501,
		sdp_tol=1e-5,
		method='sdp' if p <= 500 else 'asdp',
		**kwargs
	)
	S_matrices['sdp'] = S_SDP
	if verbose:
		print(f'Finished computing SDP matrices, time is {time.time() - time0}')

	### Calculate MCR matrices
	rej_rate = fetch_kwarg(kwargs, 'rej_rate', 0)
	if rej_rate != 0:
		kwargs['rej_rate'] = rej_rate
	for new_method in ['mvr','maxent']:

		# Choose between projected gradient decent and coord descent
		# solvers
		if new_method == 'maxent' and rej_rate != 0:
			solver = 'psgd'
		else:
			solver = 'cd'
		S_MRC = knockadapt.knockoffs.compute_S_matrix(
			Sigma=Sigma,
			groups=groups,
			method=new_method,
			solver=solver,
			**kwargs
		)
		S_matrices[new_method] = S_MRC
	S_matrices['ci'] = knockadapt.knockoffs.compute_S_matrix(
		Sigma=Sigma,
		method='ciknock',
	)
	if verbose:
		print(f'Finished computing MRC matrices, time is {time.time() - time0}')

	return S_matrices

def Z2selections(Z, groups, q, **kwargs):

	# Calculate W statistics
	W = kstats.combine_Z_stats(Z, groups, **kwargs)

	# Calculate selections
	T = kstats.data_dependent_threshhold(W=W, fdr=q)
	selected_flags = (W >= T).astype("float32")
	return selected_flags, W

def selection2power(selections, group_nonnulls):
	# Calculate fdp, power
	fdp = np.sum(selections*(1-group_nonnulls))
	power = np.sum(selections*group_nonnulls)

	# Normalize
	fdp = fdp/max(1, np.sum(selections))
	power = power/max(1, np.sum(group_nonnulls))

	return [power, fdp]

def single_dataset_power_fdr(
	seed,
	Sigma,
	beta,
	groups,
	normalize=True,
	sample_kwargs={},
	filter_kwargs={
		'feature_stat_kwargs':{},
		'knockoff_kwargs':{},
	},
	S_matrices=None,
	time0=None,
	S_curve=False,
):
	"""
	Knockoff kwargs should be included in filter_kwargs
	"""

	# Fetch groups, infer q
	if Sigma is not None:
		p = Sigma.shape[0]
	else:
		p = sample_kwargs['p']
	if groups is None:
		groups = np.arange(1, p+1, 1)

	# Prevent global changes to kwargs
	filter_kwargs = filter_kwargs.copy()
	sample_kwargs = sample_kwargs.copy()

	# Sample data, record time
	localtime = time.time()
	np.random.seed(seed)
	X, y, beta, _, Sigma = knockadapt.graphs.sample_data(
		corr_matrix=Sigma,
		beta=beta,
		**sample_kwargs
	)

	# Parse nonnulls (this allows for group knockoffs although
	# we do not actually use this in our experiments)
	group_nonnulls = knockadapt.utilities.fetch_group_nonnulls(
		beta, groups
	)

	# Some updates for fixedX knockoffs
	# and MX knockoffs when Sigma must be inferred
	fixedX = fetch_kwarg(filter_kwargs, 'fixedx', default=False)
	infer_sigma = fetch_kwarg(filter_kwargs, 'infer_sigma', default=False)

	# For the metro sampler, we can compute better S matrices if we have a
	# guess of the rejection rate. We do not use this in our experiments
	# by default.
	rej_rate = fetch_kwarg(filter_kwargs, 'rej_rate', default=0)

	# We calculate S-matrices if we do not already know them.
	if 'knockoff_kwargs' in filter_kwargs:
		kwargs = filter_kwargs['knockoff_kwargs']
	else:
		kwargs = {}
	# Case 1: We have to estimate Sigma from the data
	if infer_sigma:
		shrinkage = fetch_kwarg(filter_kwargs, 'shrinkage', default='ledoitwolf')
		Sigma, _ = knockadapt.utilities.estimate_covariance(X, shrinkage=shrinkage)
		Sigma = utilities.cov2corr(Sigma)
		invSigma = utilities.chol2inv(Sigma)
	# Case 2: We are running FX knockoffs
	if fixedX:
		Sigma = np.dot(X.T, X)
		invSigma = None
	# Calculate S-matrices from Sigma
	if infer_sigma or fixedX or S_matrices is None:
		verbose = fetch_kwarg(kwargs, 'verbose', default=False)
		S_matrices = fetch_competitor_S(
			Sigma=Sigma,
			groups=groups,
			time0=time0,
			rej_rate=rej_rate,
			verbose=verbose,
			S_curve=S_curve,
			**kwargs
		)

	# Now we loop through the S matrices
	degen_flag = 'sdp_perturbed' in S_matrices
	output = []
	for S_method in S_matrices:

		# Pull S matrix
		S = S_matrices[S_method]

		# If the method produces fully degenerate knockoffs,
		# signal this as part of the filter kwargs.
		_sdp_degen = (degen_flag and S_method == 'sdp')

		# Do NOT run OLS or debiased lasso for degenerate case,
		# because this will lead to linear algebra errors.
		if _sdp_degen and 'feature_stat' in filter_kwargs:
			if filter_kwargs['feature_stat'] == 'dlasso':
				for pair_agg in PAIR_AGGS:
					for q in q_values:
						output.append([
							S_method,
							0,
							0,
							q,
							0,
							"NULL",
							pair_agg,
							np.zeros(p),
							np.zeros(p),
							np.zeros(p),
							np.zeros(p),
							seed
						])
				continue

		# Create knockoff_kwargs
		if 'knockoff_kwargs' not in filter_kwargs:
			filter_kwargs['knockoff_kwargs'] = {}
		filter_kwargs['knockoff_kwargs']['S'] = S
		filter_kwargs['knockoff_kwargs']['method'] = S_method
		filter_kwargs['knockoff_kwargs']['_sdp_degen'] = _sdp_degen
		if 'verbose' not in filter_kwargs['knockoff_kwargs']:
			filter_kwargs['knockoff_kwargs']['verbose'] = False

		# Possibly pass a few parameters to metro sampler
		if 'x_dist' in sample_kwargs:
			# For Ising
			if 'gibbs_graph' in sample_kwargs:
				filter_kwargs['knockoff_kwargs']['gibbs_graph'] = sample_kwargs['gibbs_graph']
			# For t-distributions
			if str(sample_kwargs['x_dist']).lower() in ['ar1t', 'blockt']:
				if 'df_t' in sample_kwargs:
					filter_kwargs['knockoff_kwargs']['df_t'] = sample_kwargs['df_t']
				else:
					filter_kwargs['knockoff_kwargs']['df_t'] = DEFAULT_DF_T # This matters

		# Run MX knockoff filter to obtain
		# Z statistics
		knockoff_filter = KnockoffFilter(fixedX=fixedX)
		knockoff_filter.forward(
			X=X,
			y=y,
			mu=np.zeros(p),
			Sigma=Sigma,
			groups=groups,
			**filter_kwargs
		)
		Z = knockoff_filter.Z
		score = knockoff_filter.score
		score_type = knockoff_filter.score_type

		# Calculate power/fdp/score for a variety of
		# antisymmetric functions
		for pair_agg in PAIR_AGGS:
			for q in q_values:
				# Start by creating selections
				selections, W = Z2selections(
					Z=Z,
					groups=groups,
					q=q,
					pair_agg=pair_agg
				)
				# Then create power/fdp
				power, fdp = selection2power(
					selections, group_nonnulls
				)
				output.append([
					S_method,
					power,
					fdp,
					q,
					score,
					score_type,
					pair_agg,
					np.around(W, 4),
					np.around(Z[0:p], 4),
					np.around(Z[p:], 4),
					selections,
					seed,
				])

	# Possibly log progress
	try:
		if seed % 10 == 0 and sample_kwargs['n'] == MAXIMUM_N:
			overall_cost = time.time() - time0
			local_cost = time.time() - localtime
			print(f"Finished one seed {seed}, took {local_cost} per seed, {overall_cost} total")
	except:
		# In notebooks this will error
		pass

	# Output: list of
	# [S_method, power, fdp, q, score, score_type, pair_agg, W, Z, tildeZ, selection, seed]
	return output

def calc_power_and_fdr(
	time0,
	Sigma,
	beta,
	groups=None,
	q=0.2,
	reps=100,
	num_processes=1,
	sample_kwargs={},
	filter_kwargs={},
	seed_start=0,
	S_matrices={'sdp':None, 'mvr':None},
	S_curve=False,
):

	# Fetch nonnulls
	if Sigma is not None:
		p = Sigma.shape[0]
	else:
		p = sample_kwargs['p']

	# Sample data reps times and calc power/fdp
	partial_sd_power_fdr = partial(
		single_dataset_power_fdr,
		Sigma=Sigma,
		beta=beta,
		groups=groups,
		sample_kwargs=sample_kwargs,
		filter_kwargs=filter_kwargs,
		S_matrices=S_matrices,
		time0=time0,
		S_curve=S_curve,
	)
	# (Possibly) apply multiprocessing
	all_inputs = list(range(seed_start, seed_start+reps))
	num_processes = min(len(all_inputs), num_processes)
	all_outputs = apply_pool(
		func = partial_sd_power_fdr,
		all_inputs = all_inputs,
		num_processes=num_processes
	)

	# Extract output and return
	final_out = []
	for output in all_outputs:
		final_out.extend(output)

	# Final out: list of arrays:
	# S_method, power, fdp, q, score, score_type, pair_agg, W, Z, tildeZ, selection
	return final_out

def analyze_degen_solns(
	Sigma,
	beta,
	dgp_number,
	groups=None,
	sample_kwargs={},
	filter_kwargs={},
	fstat_kwargs={},
	reps=50,
	num_processes=5,
	seed_start=0,
	time0=None,
	S_curve=False,
	storew=True,
	):
	"""
	:param dgp_number: A number corresponding to which dgp
	each row corresponds to.
	:param sample_kwargs:
	A dictionary. Each key is a sample parameter name, with the value
	being either a single value or a list or array of values.
	:param filter_kwargs:
	A dictionary. Each key is a mx_filter parameter name, with the value
	being either a single value or a list or array of values.
	"""
	global MAXIMUM_N # A hack to allow for better logging

	# Infer p and set n defaults
	if Sigma is not None:
		p = Sigma.shape[0]
	else:
		p = sample_kwargs['p']
		if isinstance(p, list):
			p = p[0]
	sample_kwargs['p'] = [p]
	if 'n' not in sample_kwargs:
		sample_kwargs['n'] = [
			int(p/4), int(p/2), int(p/1.5), int(p), int(2*p), int(4*p)
		]
	if not isinstance(sample_kwargs['n'], list):
		sample_kwargs['n'] = [sample_kwargs['n']]
	MAXIMUM_N = max(sample_kwargs['n']) # Helpful for logging
	if groups is None:
		groups = np.arange(1, p+1, 1)

	# Construct/iterate cartesian product of sample, filter, fstat kwargs
	sample_keys, sample_product = dict2keyproduct(sample_kwargs)
	filter_keys, filter_product = dict2keyproduct(filter_kwargs)
	fstat_keys, fstat_product = dict2keyproduct(fstat_kwargs)

	# Initialize final output
	counter = 0
	columns = ['power', 'fdp', 'q', 'S_method', 'antisym', 'score', 'score_type']
	if storew:
		columns += [f'W{i}' for i in range(1, p+1)]
		columns += [f'Z{i}' for i in range(1, p+1)]
		columns += [f'tildeZ{i}' for i in range(1, p+1)]
		columns += [f'selection{i}' for i in range(1, p+1)]
	columns += ['dgp_number']
	columns += [x for x in sample_keys if x != 'gibbs_graph']
	columns += filter_keys + fstat_keys + ['seed']
	result_df = pd.DataFrame(columns=columns)

	# Check if we are going to ever fit MX knockoffs on the
	# "ground truth" covariance matrix. If so, we'll memoize
	# the SDP/MVR/MMI results.
	if 'fixedx' in filter_kwargs:
		fixedX_vals = filter_kwargs['fixedx']
		if False in fixedX_vals:
			MX_flag = True
		else:
			MX_flag = False
	else:
		MX_flag = True
	if 'infer_sigma' in filter_kwargs:
		infer_sigma_vals = filter_kwargs['infer_sigma']
		if False in infer_sigma_vals:
			ground_truth = True
		else:
			ground_truth = False
	else:
		ground_truth = True
	# Possibly save S-matrices
	if ground_truth and MX_flag and Sigma is not None:
		rej_rate = fetch_kwarg(filter_kwargs, 'rej_rate', default=[0])[0]
		print(f"Storing SDP/MVR/MMI results with rej_rate={rej_rate}")
		S_matrices = fetch_competitor_S(
			Sigma=Sigma,
			groups=groups,
			time0=time0,
			S_curve=S_curve,
			rej_rate=rej_rate,
			verbose=True,
		)
	else:
		print(f"Not storing SDP/MVR results")
		S_matrices = None

	### Calculate power of knockoffs for the different methods
	for filter_vals in filter_product:
		filter_vals = list(filter_vals)
		new_filter_kwargs = {
			key:val for key, val in zip(filter_keys, filter_vals)
		}

		# In high dimensional cases or binomial cases,
		# don't fit OLS.
		if 'feature_stat' in new_filter_kwargs:
			fstat = new_filter_kwargs['feature_stat']
		else:
			fstat = 'lasso'

		for sample_vals in sample_product:
			sample_vals = list(sample_vals)
			new_sample_kwargs = {
				key:val for key, val in zip(sample_keys, sample_vals)
			}

			# Don't run OLS in certain cases
			if fstat == 'ols':
				if new_sample_kwargs['n'] < 2*new_sample_kwargs['p']:
					continue
				if 'y_dist' in new_sample_kwargs:
					if new_sample_kwargs['y_dist'] == 'binomial':
						continue
			if fstat == 'dlasso':
				if 'y_dist' in new_sample_kwargs:
					if new_sample_kwargs['y_dist'] == 'binomial':
						continue

			for fstat_vals in fstat_product:
				# Extract feature-statistic kwargs
				# and place them properly (as a dictionary value
				# of the filter_kwargs)
				fstat_vals = list(fstat_vals)
				new_fstat_kwargs = {
					key:val for key, val in zip(fstat_keys, fstat_vals)
				}
				if 'feature_stat_kwargs' in new_filter_kwargs:
					new_filter_kwargs['feature_stat_kwargs'] = dict(
						new_filter_kwargs['feature_stat_kwargs'],
						**new_fstat_kwargs
					)
				else:
					new_filter_kwargs['feature_stat_kwargs'] = new_fstat_kwargs

				# Power/FDP for the S methods
				out = calc_power_and_fdr(
					Sigma=Sigma,
					beta=beta,
					groups=groups,
					reps=reps,
					num_processes=num_processes,
					sample_kwargs=new_sample_kwargs,
					filter_kwargs=new_filter_kwargs,
					seed_start=seed_start,
					time0=time0,
					S_matrices=S_matrices,
					S_curve=S_curve,
				)

				# Loop through antisymmetric functions and S matrices
				for vals in out:
					S_method, power, fdp, q, score, score_type, agg, W, Z, tildeZ, selections, seed = vals
					row = [power, fdp, q, S_method, agg, score, score_type]
					if storew:
						row.extend(W.tolist())
						row.extend(Z.tolist())
						row.extend(tildeZ.tolist())
						row.extend(selections.astype(np.int32).tolist())
					row += [dgp_number]
					sample_vals_to_add = [
						sample_vals[i] for i in range(len(sample_vals)) if sample_keys[i] != 'gibbs_graph'
					]
					row += sample_vals_to_add + filter_vals + fstat_vals + [seed]
					result_df.loc[counter] = row
					counter += 1

	return result_df, S_matrices

def parse_args(args):
	""" Homemade argument parser
	Usage: --argname_{dgp/sample/filter} value
	Value should be one of:
		- integer/float
		- arbitrary string
		- Alternatively, string of the form
		"start{num}end{num}numvals{num}"
		which indicates that the parameter ought
		to be varied along a linear interpolation
		from the specified start/end with the specified
		number of values.
		- Alternatively, a string like
		"[str1, str2, str3]" which will be parsed as
		["str1", "str2", "str3"]

	Note that "reps" and "num_processes" are special kwargs.
	They get placed as sample_kwargs by default but will be
	extracted.

	The --description argument must be the last argument, as
	all arguments after the --description argument will be ignored.
	"""

	# Get rid of script name
	args = args[1:]

	# Initialize kwargs constructor
	special_keys = [
		'reps', 'num_processes', 'seed_start', 'description',
		's_curve', 'resample_beta', 'resample_sigma', 'storew',
	]
	key_types = ['dgp', 'sample', 'filter', 'fstat']
	all_kwargs = {ktype:{} for ktype in key_types}
	key = None
	key_type = None # One of dgp, sample, filter
	description_index = None # At the end, can write description

	#Parse
	for i, arg in enumerate(args):
		arg = str(arg).lower()

		# Even placements should be keyword
		if i % 2 == 0:
			if str(arg)[0:2] != '--':
				raise ValueError(
					f'{i}th arg ({arg}) should be a kwarg name preceded by "--"'
				)
			key = arg[2:]

			# Check what type of keyword
			if key in special_keys:
				key_type = 'sample'
			else:
				key_split = key.split('_')
				key_type = key_split[-1]
				key = ('_').join(key_split[0:-1])

			# Raise error if unexpected key type
			if key_type not in key_types:
				raise ValueError(
					f'{i}th arg ({arg}) has key_type {key_type}, must be one of {key_types}'
				)

			# Friendly reminder
			if key == 'feature_stat_fn':
				raise ValueError("feature_stat_fn is depreciated, use feature_stat")

			# Description
			if key == 'description':
				description_index = i
				break

		# Parse values
		if i % 2 == 1:

			# Check if it's a list written out
			if arg[0] == '[':
				# Strip brackets, whitspace, and split by commas
				value = arg.replace('[', '').replace(']', '')
				value = value.replace(' ', '').split(',')
				# Try to process
				value = [obj2int(v) for v in value]

			# Check if it's start{num}end{num}numvals{num}
			elif arg[0:5] == 'start':
				start = float(arg.replace('start', '').split('end')[0])
				end = float(arg.split('end')[-1].split('numvals')[0])
				numvals = int(arg.split('numvals')[-1])
				value = np.linspace(
					start, end, numvals
				)

			# Apply obj2int (preserves strings, infers floats, bools, ints)
			else:
				value = obj2int(arg)

			all_kwargs[key_type][key] = val2list(value)

	# Parse description
	description = ''
	if description_index is not None:
		description += (' ').join(args[description_index+1:])
	description += f'\n \n Other arguments were: {args[0:description_index]}'
	all_kwargs['sample']['description'] = description

	out = (all_kwargs[key_type] for key_type in key_types)
	return out

def main(args):
	""" Layers of keyword arguments are as follows.
	- dgp_kwargs: kwargs needed to create Sigma (e.g. p, method)
	- sample_kwargs: kwargs needed to sample data (e.g. n)
	- filter_kwargs: kwargs for the knockoff filter. (e.g. recycling)
	- fstat_kwargs: kwargs for the feature statistic. (e.g. pair_agg)

	For each of these, keys mapping to iterables like lists or
	numpy arrays will be varied. See parse_args comments.

	The MAIN constraint is that the same sample_kwargs must be used
	for all dgp_kwargs. E.g., it will be difficult to vary both
	a for an AR1 covariance matrix and delta for an ErdosRenyi covariance
	matrix. The same goes for filter_kwargs abd fstat_kwargs.
	"""

	# Create kwargs
	dgp_kwargs, sample_kwargs, filter_kwargs, fstat_kwargs = parse_args(args)
	print(f"Args were {args}")

	# Make sure pair_aggs is not being duplicated
	if 'pair_agg' in fstat_kwargs:
		raise ValueError("Many pair_aggs will be analyzed anyway. Do not add this as a fstat_kwarg.")
	# Some common errors I make
	if 'y_dist' in dgp_kwargs:
		raise ValueError("y_dist ought to be in sample_kwargs")

	# Parse some special non-graph kwargs
	reps = fetch_kwarg(sample_kwargs, 'reps', default=[1])[0]
	num_processes = fetch_kwarg(sample_kwargs, 'num_processes', default=[12])[0]
	seed_start = fetch_kwarg(sample_kwargs, 'seed_start', default=[0])[0]
	description = fetch_kwarg(sample_kwargs, 'description', default='')
	S_curve = fetch_kwarg(sample_kwargs, 's_curve', default=[False])[0]
	resample_beta = fetch_kwarg(sample_kwargs, 'resample_beta', default=[True])[0]
	resample_sigma = fetch_kwarg(sample_kwargs, 'resample_sigma', default=[True])[0]
	storew = fetch_kwarg(sample_kwargs, 'storew', default=[False])[0]
	print(f"S_curve flag = {S_curve}")
	print(f"DGP kwargs are {dgp_kwargs}")
	print(f"Sample kwargs are {sample_kwargs}")
	print(f"Filter kwargs are {filter_kwargs}")
	print(f"Ftat kwargs are {fstat_kwargs}")
	print(f"Description is {description.split('Other arguments were:')[0]}")

	# Create output paths
	today = str(datetime.date.today())
	hour = str(datetime.datetime.today().time())
	hour = hour.replace(':','-').split('.')[0]
	output_path = f'data/mrcfinal/{today}/{hour}'
	all_key_types = ['dgp', 'sample', 'filter', 'fstat']
	all_kwargs = [dgp_kwargs,sample_kwargs, filter_kwargs, fstat_kwargs]

	for key_type,kwargs in zip(all_key_types, all_kwargs):
		output_path += f'/{key_type}_'
		keys = sorted([k for k in kwargs])
		for k in keys:
			path_val = ('').join(str(kwargs[k]).split(' '))
			output_path += f'{k}{path_val}_'

	# Put it all together and ensure directory exists
	output_path += f'seedstart{seed_start}_reps{reps}_results.csv'
	beta_path = output_path.split('.csv')[0] + '_betas.csv'
	S_path = output_path.split('.csv')[0] + '_S.csv'
	description_path = f'data/mrcfinal/{today}/{hour}/' + 'description.txt'
	output_dir = os.path.dirname(output_path)
	if not os.path.exists(output_dir):
		os.makedirs(output_dir)
	print(f"Output path is {output_path}")

	# Save description
	with open(description_path, 'w') as thefile:
		thefile.write(description)

	# Initialize final final output
	all_results = pd.DataFrame()

	# Loop through DGP parameters
	dgp_keys = sorted([key for key in dgp_kwargs])
	dgp_components = []
	for key in dgp_keys:
		dgp_kwargs[key] = val2list(dgp_kwargs[key])
		dgp_components.append(dgp_kwargs[key])
	dgp_product = itertools.product(*dgp_components)

	# p sadly does need to have 1 value for a variety of reasons
	if 'p' in dgp_kwargs:
		if len(dgp_kwargs['p']) > 1:
			raise ValueError(f"Must have only one value of p, not {dgp_kwarys[p]}")
		else:
			p = dgp_kwargs['p'][0]
	else:
		p = 100

	# Initialize ways to save beta, dgp
	dgp_number = 0
	beta_df = pd.DataFrame(columns=np.arange(1, p+1, 1))
	beta_df.index.name = 'dgp_number'
	# Stores diagonal of S
	S_diags_df = pd.DataFrame(
		columns = ['dgp_number', 'S_method'] + [i for i in range(1, p+1)]
	)
	S_counter = 0

	for dgp_vals in dgp_product:

		# Create DGP using dgp kwargs
		dgp_vals = list(dgp_vals)
		new_dgp_kwargs = {key:val for key,val in zip(dgp_keys, dgp_vals)}
		print(f"DGP kwargs are now: {new_dgp_kwargs}")
		np.random.seed(110)
		_, _, beta, invSigma, Sigma = knockadapt.graphs.sample_data(
			**new_dgp_kwargs
		)

		# Use precomputed Sigma for ising model
		if 'x_dist' in new_dgp_kwargs:
			if new_dgp_kwargs['x_dist'] == 'gibbs':
				if 'method' not in new_dgp_kwargs:
					raise ValueError(f"Method must be supplied for x_dist == gibbs")
				if new_dgp_kwargs['method'] == 'ising':
					p = new_dgp_kwargs['p']
					file_dir = os.path.dirname(os.path.abspath(__file__))
					v_file = f'{file_dir}/qcache/vout{p}.txt'
					if os.path.exists(v_file):
						print(f"Loading custom Sigma for gibbs ising model")
						Sigma = np.loadtxt(f'{file_dir}/qcache/vout{p}.txt')
					else:
						print(f"No custom Sigma available for gibbs ising model: using default")
			# Save gibbs_graph
			sample_kwargs['gibbs_graph'] = [invSigma]

		# Cache beta
		if not resample_beta:
			beta_df.loc[dgp_number] = beta
			beta_df.to_csv(beta_path)

		# Push sample_kwargs to full output
		if resample_beta or resample_sigma:
			for key in new_dgp_kwargs:
				val = new_dgp_kwargs[key]
				if key in sample_kwargs:
					old_val = sample_kwargs[key]
					if len(old_val) != 1:
						raise ValueError(f"DGP / sample keys {key} conflict when resampling Sigma / beta")
					old_val = old_val[0]
					if old_val not in dgp_kwargs[key]:
						raise ValueError(f"DGP / sample keys {key} conflict when resampling Sigma / beta")
				sample_kwargs[key] = [val]

		# Create results
		result, S_matrices = analyze_degen_solns(
			Sigma=Sigma if not resample_sigma else None,
			beta = beta if not resample_beta else None,
			dgp_number=dgp_number,
			groups=None,
			sample_kwargs=sample_kwargs,
			filter_kwargs=filter_kwargs,
			fstat_kwargs=fstat_kwargs,
			reps=reps,
			num_processes=num_processes,
			seed_start=seed_start,
			time0=time0,
			S_curve=S_curve,
			storew=storew,
		)
		for key in dgp_keys:
			if key in sample_kwargs:
				continue
			result[key] = new_dgp_kwargs[key]
		all_results = all_results.append(
			result,
			ignore_index = True
		)
		all_results.to_csv(output_path)

		# Cache S outputs
		if S_matrices is not None:
			for S_method in S_matrices:
				S = S_matrices[S_method]
				if S is None:
					continue
				S_diag = np.diag(S)
				S_diags_df.loc[S_counter] = [dgp_number, S_method] + S_diag.tolist()
				S_counter += 1
			S_diags_df.to_csv(S_path)

		# Increment dgp number
		dgp_number += 1

	return all_results

if __name__ == '__main__':

	time0 = time.time()
	MAXIMUM_N = 0
	main(sys.argv)


'''
author: HAK
time  : 10:00 PM, 28/10/2017
'''

from watchdog.events import FileSystemEventHandler
from watchdog.observers import Observer
import time
from INFO_PROVIDER.fileInfo import FILE_INFO
from INFO_PROVIDER.directoryInfo import PATH_INFO_PROVIDER
import subprocess
from colorama import Fore
import requests
import sys
import json
import emailWatcher
import datetime as dt

'''
An event handler that use to listen triggered events from FileSystemEvent
This class onsist of one function on_created which provide two basic event logging generators
1. When a file is created
2. When a directory is created

Returns:
-------FILE--------
1. Path
2. Name
3. Extension
4. Time of creation
5. Size (BYTES)
-------------------

-----DIRECTORY-----
1. Path
2. Name
3. Time of creation
4. Size (BYTES)
-------------------
'''

class Handler(FileSystemEventHandler):
    data = None

    def objectInfo(self, getInfo):
        self.data = getInfo

    def on_modified(self, event):
        print("Modified " + event.src_path)

    def on_created(self, event):
        PATH = event.src_path
        # Check if a event is generated for directory or not
        if event.is_directory:
            DIR = PATH_INFO_PROVIDER(PATH)
            print(Fore.LIGHTCYAN_EX,
                  '----------------------------- NEW DIRECTORY CREATED -----------------------------')
            print(Fore.LIGHTCYAN_EX, "PATH  : ", DIR.DIRBASIC()[0])
            print(Fore.LIGHTCYAN_EX, "NAME  : ", DIR.DIRBASIC()[1])
            print(Fore.LIGHTCYAN_EX, "CTIME : ", DIR.DIRBASIC()[2])
            print(Fore.LIGHTCYAN_EX, "SIZE  : ", DIR.DIRBASIC()[3], "BYTES")
            print(Fore.LIGHTCYAN_EX,
                  '---------------------------------------------------------------------------------', Fore.RESET)
        else:
            FILE = FILE_INFO(event.src_path)
            print(Fore.LIGHTCYAN_EX)
            print('-------------------------------- NEW FILE CREATED -------------------------------')
            print("PATH  : ", FILE.FILEBASIC()[0])
            print("NAME  : ", FILE.FILEBASIC()[1])
            print("EXT   : ", FILE.FILEBASIC()[2])
            print("CTIME : ", FILE.FILEBASIC()[3])
            print("SIZE  : ", FILE.FILEBASIC()[4], "BYTES")
            print('---------------------------------------------------------------------------------', Fore.RESET)
            CHG_PATH = '\"' + event.src_path + '\"' if ' ' in event.src_path else event.src_path
            #self.data['email'].sendMailForInconsistentFiles()
            # may call call instead of Popen
            self.data['fileReceived'] = CHG_PATH
            subprocess.Popen(['python3', self.data['pyfile'], str(self.data)])

'''
class: watcher
used to watch for events create and modified
'''

class Watcher:
    PYFILE = None
    DBINFO = None
    DIRECTORY = None
    SITENAME = None
    MAIL = None

    def __init__(self, DIRECTORY=None, PYFILE=None, DBINFO=None, SITEID=None):
        self.PYFILE = PYFILE
        self.DBINFO = DBINFO
        self.DIRECTORY = DIRECTORY
        self.SITEID = SITEID

    def getInfo(self):
        return {
            'pyfile': self.PYFILE,
            'db': self.DBINFO,
            'fileReceived': None
        }

    def run(self):
        try:
            event_handler = Handler()
            event_handler.objectInfo(self.getInfo())
            #emailService = requests.get('http://0.0.0.0:5000/api/emailapi/emailconfig?siteId=' + self.SITEID)
            emailService = requests.get('https://x45k5kd3hj.execute-api.us-east-2.amazonaws.com/dev/emailconfig?siteId='+self.SITEID)
            emailConfig = json.loads(emailService.content.decode('utf-8'))
            email = emailWatcher.EmailHandler(path=self.DIRECTORY,
                                              timeThreshold=emailConfig['timeDiffer'],
                                              siteId=self.SITEID,
                                              inconsistentFiles=emailConfig['checkInconsistentFiles'],
                                              inconsistenTime=emailConfig['checkTimeInconsistencies'])
            observer = Observer()
            observer.schedule(event_handler, self.DIRECTORY, recursive=True)
            observer.start()
        except:
            observer.stop()
            #requests.put('http://0.0.0.0:5000/api/emailapi/emailstatusupdate?siteId='+str(self.SITEID)+'&serviceStatus='+str(0))
            print("Observer Stopped Unexpectedly.")
            sys.exit(1)
        try:
            while True:
                time.sleep(emailConfig['monitoringPing'])
                email.update()
        except:
            observer.stop()
            print("Observer Stopped Manually.")
            #requests.put('http://0.0.0.0:5000/api/emailapi/emailstatusupdate?siteId=' + str(self.SITEID) + '&serviceStatus=' + str(0))
            #requests.put('http://0.0.0.0:5000/api/emailapi/sendmailtogroupid?siteId=' + str(self.SITEID) +
            #             '&isMaintainance=' + str(1) +
            #             '&emailType=' + str('INTERNAL_ISSUE') +
            #             '&datetime=' + str(dt.datetime.now()))
        observer.join()


import time

import numpy
import torch

RESO_H = 768
RESO_W = RESO_H // 2
GRID_H = 16
GRID_W = RESO_W // 5

class LineModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.feature_extractor = torch.nn.Sequential(
            # 3 x 768 x 384
            Conv(3, 12, 5, padding=2),
            # 12 x 768 x 384
            Residual(12),
            ResidualReduce(12, 24),
            # 24 x 384 x 192
            Residual(24),
            ResidualReduce(24, 48),
            # 48 x 192 x 96
            Residual(48),
            ResidualReduce(48, 96),
            # 96 x 96 x 48
            Residual(96),
            ResidualReduce(96, 192),
            # 192 x 48 x 24
        )
        self.line_detector = torch.nn.Sequential(
            # 192 x 48 x 24
            Conv(192, 216, (1, 3), padding=0, groups=3),
            # 240 x 48 x 22
            Conv(216, 240, (1, 3), padding=0, groups=3),
            # 240 x 48 x 20
            torch.nn.Conv2d(240, 240, (1, RESO_H // 32 - 4), groups=3),
            # 240 x 48 x 1
            torch.nn.Conv2d(240, 3, 1, groups=3),
            # 3 x 48 x 1
        )

    def forward(self, inpt):
        features = self.feature_extractor(inpt)

        lines = self.line_detector(features).squeeze(dim=3)

        if self.training:
            confidence = lines[:, 0, :]
        else:
            confidence = torch.sigmoid(lines[:, 0, :])
        offset = torch.nn.functional.softsign(lines[:, 1, :])
        scaling = torch.nn.functional.softplus(lines[:, 2, :])

        return torch.stack([confidence, offset, scaling], dim=2)

class GridModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.feature_extractor = torch.nn.Sequential(
            # 3 x 768 x 384
            Conv(3, 15, 5, padding=2),
            # 15 x 768 x 384
            Residual(15),
            ResidualOne(15),
            Residual(15),
            ResidualOne(15),
            ResidualReduce(15, 30),
            # 30 x 384 x 192
            Residual(30),
            ResidualOne(30),
            Residual(30),
            ResidualOne(30),
            ResidualReduce(30, 60),
            # 60 x 192 x 96
            Residual(60),
            ResidualOne(60),
            Residual(60),
            ResidualOne(60),
            ResidualReduce(60, 120),
            # 120 x 96 x 48
            Residual(120),
            ResidualOne(120),
            Residual(120),
            ResidualOne(120),
            ResidualReduce(120, 240),
            # 240 x 48 x 24
        )
        self.grid_detector = torch.nn.Sequential(
            # 240 x 48 x 24
            Conv(240, 240, (1, 3), padding=0, groups=5),
            Conv(240, 240, 1, padding=0, groups=5),
            # 240 x 48 x 22
            Conv(240, 240, (1, 3), padding=0, groups=5),
            Conv(240, 240, 1, padding=0, groups=5),
            # 240 x 48 x 20
            torch.nn.Conv2d(240, 240, (1, 4), stride=(1, 4), groups=5),
            # 240 x 48 x 5
            torch.nn.Conv2d(240, 5, 1, groups=5),
            # 5 x 48 x 5
        )

    def forward(self, input):
        features = self.feature_extractor(input)

        grids = self.grid_detector(features)

        if self.training:
            confidence = grids[:, 0, :, :]
        else:
            confidence = torch.sigmoid(grids[:, 0, :, :])
        offset_x = grids[:, 1, :, :]
        offset_y = grids[:, 2, :, :]
        scale_x = torch.nn.functional.softplus(grids[:, 3, :, :])
        scale_y = torch.nn.functional.softplus(grids[:, 4, :, :])

        return torch.stack([confidence, offset_x, offset_y, scale_x, scale_y], dim=1)

class Residual(torch.nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.conv = torch.nn.Sequential(
            Conv(channels, channels, 3, padding=1), Conv(channels, channels, 1, padding=0)
        )

    def forward(self, inpt):
        oupt = self.conv(inpt)
        return oupt + inpt

class ResidualOne(torch.nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.conv = torch.nn.Sequential(
            Conv(channels, channels, 1, padding=0), Conv(channels, channels, 1, padding=0)
        )

    def forward(self, inpt):
        oupt = self.conv(inpt)
        return oupt + inpt

class ResidualReduce(torch.nn.Module):
    def __init__(self, in_chan, out_chan):
        super().__init__()
        self.branch_0 = torch.nn.Sequential(
            Conv(in_chan, out_chan, 3, padding=1, stride=2), Conv(out_chan, out_chan, 1, padding=0)
        )
        self.branch_1 = torch.nn.Sequential(
            torch.nn.MaxPool2d(2), Conv(in_chan, out_chan, 1, padding=0)
        )

    def forward(self, inpt):
        oupt_0 = self.branch_0(inpt)
        oupt_1 = self.branch_1(inpt)
        return oupt_0 + oupt_1

class Conv(torch.nn.Module):
    def __init__(self, in_chan, out_chan, kernel_size, padding, stride=1, groups=1):
        super().__init__()
        self.conv = torch.nn.Sequential(
            torch.nn.Conv2d(in_chan, out_chan, kernel_size, stride, padding, groups=groups),
            torch.nn.BatchNorm2d(out_chan),
            torch.nn.LeakyReLU(inplace=True),
        )

    def forward(self, inpt):
        return self.conv(inpt)

if __name__ == "__main__":
    numpy.set_printoptions(precision=2, suppress=True)

    model = LineModel()
    model.train()

    data = torch.randn(2, 3, RESO_H, RESO_W)

    start = time.time()
    pred = model(data)

    print("T:", time.time() - start)
    print(pred.size())
    print(pred.max().item(), pred.min().item())


import copy

import numpy as np

from PySide6.QtCore import QItemSelectionModel, QObject, Signal
from PySide6.QtWidgets import (
    QComboBox, QLineEdit, QSizePolicy, QTableWidgetItem
)

from hexrd.constants import chargestate
from hexrd.material import Material

from hexrdgui.periodic_table_dialog import PeriodicTableDialog
from hexrdgui.scientificspinbox import ScientificDoubleSpinBox
from hexrdgui.thermal_factor_editor import ThermalFactorEditor
from hexrdgui.ui_loader import UiLoader
from hexrdgui.utils import block_signals

COLUMNS = {
    'symbol': 0,
    'charge': 1,
    'occupancy': 2,
    'thermal_factor': 3
}

DEFAULT_CHARGE = '0'
DEFAULT_U = Material.DFLT_U[0]

OCCUPATION_MIN = 0
OCCUPATION_MAX = 1

THERMAL_FACTOR_MIN = -1.e7
THERMAL_FACTOR_MAX = 1.e7

U_TO_B = 8 * np.pi ** 2
B_TO_U = 1 / U_TO_B

class MaterialSiteEditor(QObject):

    site_modified = Signal()

    def __init__(self, site, parent=None):
        super().__init__(parent)

        loader = UiLoader()
        self.ui = loader.load_file('material_site_editor.ui', parent)

        self._site = site

        self.charge_comboboxes = []
        self.occupancy_spinboxes = []
        self.thermal_factor_spinboxes = []

        self.update_gui()

        self.setup_connections()

    def setup_connections(self):
        self.ui.select_atom_types.pressed.connect(self.select_atom_types)
        self.ui.thermal_factor_type.currentIndexChanged.connect(
            self.thermal_factor_type_changed)
        for w in self.site_settings_widgets:
            w.valueChanged.connect(self.update_config)
        self.ui.table.selectionModel().selectionChanged.connect(
            self.selection_changed)
        self.ui.remove_atom_type.pressed.connect(self.remove_selected_atom)
        self.ui.convert_u_to_tensors.toggled.connect(self.convert_u_to_tensors)

    def select_atom_types(self):
        dialog = PeriodicTableDialog(self.atom_types, self.ui)
        if not dialog.exec():
            return

        self.atom_types = dialog.selected_atoms

    @property
    def site(self):
        return self._site

    @site.setter
    def site(self, v):
        self._site = v
        self.update_gui()

    @property
    def atoms(self):
        return self.site['atoms']

    @property
    def total_occupancy(self):
        return sum(x['occupancy'] for x in self.atoms)

    @property
    def fractional_coords(self):
        return self.site['fractional_coords']

    @property
    def thermal_factor_type(self):
        return self.ui.thermal_factor_type.currentText()

    def U(self, val):
        # Take a thermal factor from a spin box and convert it to U
        type = self.thermal_factor_type
        if type == 'U':
            multiplier = 1
        elif type == 'B':
            multiplier = B_TO_U
        else:
            raise Exception(f'Unknown type: {type}')

        return val * multiplier

    def B(self, val):
        # Take a thermal factor from a spin box and convert it to B
        type = self.thermal_factor_type
        if type == 'U':
            multiplier = U_TO_B
        elif type == 'B':
            multiplier = 1
        else:
            raise Exception(f'Unknown type: {type}')

        return val * multiplier

    def thermal_factor(self, atom):
        # Given an atom, return the thermal factor in either B or U
        type = self.thermal_factor_type
        if type == 'U':
            multiplier = 1
        elif type == 'B':
            multiplier = U_TO_B
        else:
            raise Exception(f'Unknown type: {type}')

        return atom['U'] * multiplier

    @property
    def atom_types(self):
        return [x['symbol'] for x in self.site['atoms']]

    @atom_types.setter
    def atom_types(self, v):
        if v == self.atom_types:
            # No changes needed...
            return

        # Reset all the occupancies
        atoms = self.atoms
        previous_u_values = {x['symbol']: x['U'] for x in atoms}
        previous_charges = {x['symbol']: x['charge'] for x in atoms}
        atoms.clear()

        for symbol in v:
            # Use previous values if available. Otherwise, use the defaults.
            atom = {
                'symbol': symbol,
                'U': previous_u_values.get(symbol, DEFAULT_U),
                'charge': previous_charges.get(symbol, DEFAULT_CHARGE),
            }
            atoms.append(atom)

        self.reset_occupancies()
        self.update_table()
        self.emit_site_modified_if_valid()

    @property
    def num_rows(self):
        return self.ui.table.rowCount()

    @property
    def selected_row(self):
        selected = self.ui.table.selectionModel().selectedRows()
        return selected[0].row() if selected else None

    def select_row(self, i):
        if i is None or i >= self.num_rows:
            # Out of range. Don't do anything.
            return

        # Select the row
        selection_model = self.ui.table.selectionModel()
        selection_model.clearSelection()

        model_index = selection_model.model().index(i, 0)
        command = QItemSelectionModel.Select | QItemSelectionModel.Rows
        selection_model.select(model_index, command)

    def selection_changed(self):
        self.update_enable_states()

    def update_enable_states(self):
        enable_remove = self.num_rows > 1 and self.selected_row is not None
        self.ui.remove_atom_type.setEnabled(enable_remove)

    def remove_selected_atom(self):
        if self.selected_row is None:
            return

        atom_types = self.atom_types
        del atom_types[self.selected_row]
        self.atom_types = atom_types

    def create_symbol_label(self, v):
        w = QTableWidgetItem(v)
        return w

    def create_charge_combobox(self, charge, symbol):
        cb = QComboBox(self.ui.table)

        if charge not in chargestate[symbol]:
            raise Exception(f'Invalid charge {charge} for {symbol}')

        cb.addItems(chargestate[symbol])
        cb.setCurrentText(charge)
        cb.currentIndexChanged.connect(self.update_config)

        size_policy = QSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
        cb.setSizePolicy(size_policy)

        self.charge_comboboxes.append(cb)
        return cb

    def create_occupancy_spinbox(self, v):
        sb = ScientificDoubleSpinBox(self.ui.table)
        sb.setKeyboardTracking(False)
        sb.setMinimum(OCCUPATION_MIN)
        sb.setMaximum(OCCUPATION_MAX)
        sb.setValue(v)
        sb.valueChanged.connect(self.update_config)

        size_policy = QSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
        sb.setSizePolicy(size_policy)

        self.occupancy_spinboxes.append(sb)
        return sb

    def create_thermal_factor_spinbox(self, v):
        sb = ThermalFactorSpinBox(self.ui.table)
        sb.setKeyboardTracking(False)
        sb.setMinimum(THERMAL_FACTOR_MIN)
        sb.setMaximum(THERMAL_FACTOR_MAX)
        sb.setValue(v)
        sb.valueChanged.connect(self.update_config)
        sb.setToolTip('Double-click to open tensor editor')

        size_policy = QSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
        sb.setSizePolicy(size_policy)

        self.thermal_factor_spinboxes.append(sb)
        return sb

    def clear_table(self):
        self.charge_comboboxes.clear()
        self.occupancy_spinboxes.clear()
        self.thermal_factor_spinboxes.clear()
        self.ui.table.clearContents()

    def update_gui(self):
        with block_signals(*self.site_settings_widgets):
            for i, w in enumerate(self.fractional_coords_widgets):
                w.setValue(self.fractional_coords[i])

            self.update_total_occupancy()
            self.update_table()
            self.reset_scalar_tensor_toggle()

    def reset_scalar_tensor_toggle(self):
        any_scalars = any(not isinstance(w.value(), np.ndarray)
                          for w in self.thermal_factor_spinboxes)

        with block_signals(self.ui.convert_u_to_tensors):
            self.ui.convert_u_to_tensors.setChecked(not any_scalars)

    def update_table(self):
        prev_selected = self.selected_row

        block_list = [
            self.ui.table,
            self.ui.table.selectionModel()
        ]
        with block_signals(*block_list):
            atoms = self.site['atoms']
            self.clear_table()
            self.ui.table.setRowCount(len(atoms))
            for i, atom in enumerate(atoms):
                w = self.create_symbol_label(atom['symbol'])
                self.ui.table.setItem(i, COLUMNS['symbol'], w)

                w = self.create_charge_combobox(atom['charge'], atom['symbol'])
                self.ui.table.setCellWidget(i, COLUMNS['charge'], w)

                w = self.create_occupancy_spinbox(atom['occupancy'])
                self.ui.table.setCellWidget(i, COLUMNS['occupancy'], w)

                v = self.thermal_factor(atom)
                w = self.create_thermal_factor_spinbox(v)
                self.ui.table.setCellWidget(i, COLUMNS['thermal_factor'], w)

            self.update_occupancy_validity()

            if prev_selected is not None:
                select_row = (prev_selected if prev_selected < self.num_rows
                              else self.num_rows - 1)
                self.select_row(select_row)

            # Just in case the selection actually changed...
            self.selection_changed()

    def thermal_factor_type_changed(self):
        self.update_thermal_factor_header()
        self.update_table()

        # Update the text for the tensor toggle as well
        text = f'Convert {self.thermal_factor_type} to tensors'
        self.ui.convert_u_to_tensors.setText(text)

    def update_thermal_factor_header(self):
        w = self.ui.table.horizontalHeaderItem(COLUMNS['thermal_factor'])
        w.setText(self.thermal_factor_type)

    def update_config(self):
        for i, w in enumerate(self.fractional_coords_widgets):
            self.fractional_coords[i] = w.value()

        for atom, combobox in zip(self.atoms, self.charge_comboboxes):
            atom['charge'] = combobox.currentText()

        for atom, spinbox in zip(self.atoms, self.occupancy_spinboxes):
            atom['occupancy'] = spinbox.value()

        for atom, spinbox in zip(self.atoms, self.thermal_factor_spinboxes):
            atom['U'] = self.U(spinbox.value())

        self.update_total_occupancy()
        self.update_occupancy_validity()

        self.emit_site_modified_if_valid()

    def update_total_occupancy(self):
        self.ui.total_occupancy.setValue(self.total_occupancy)

    def reset_occupancies(self):
        total = 1.0
        atoms = self.atoms
        num_atoms = len(atoms)
        for atom in atoms:
            atom['occupancy'] = total / num_atoms

        self.update_total_occupancy()
        self.update_occupancy_validity()

    @property
    def site_valid(self):
        return self.occupancies_valid

    @property
    def occupancies_valid(self):
        return self.total_occupancy <= 1.0

    def update_occupancy_validity(self):
        valid = self.occupancies_valid
        color = 'white' if valid else 'red'
        msg = '' if valid else 'Sum of occupancies must be <= 1'

        self.ui.total_occupancy.setStyleSheet(f'background-color: {color}')
        self.ui.total_occupancy.setToolTip(msg)

    def emit_site_modified_if_valid(self):
        if not self.site_valid:
            return

        self.site_modified.emit()

    @property
    def fractional_coords_widgets(self):
        return [
            self.ui.coords_x,
            self.ui.coords_y,
            self.ui.coords_z
        ]

    @property
    def site_settings_widgets(self):
        return self.fractional_coords_widgets

    def convert_u_to_tensors(self, b):

        def scalar_to_tensor(spinbox):
            if isinstance(spinbox.value(), np.ndarray):
                # Already a tensor
                return

            tensor = np.zeros(6, dtype=np.float64)
            tensor[:3] = spinbox.value()
            spinbox.setValue(tensor)

        def tensor_to_scalar(spinbox):
            value = spinbox.value()
            if not isinstance(value, np.ndarray):
                # Already a scalar
                return

            # Use the previous spinbox value if available
            scalar = spinbox.editor.ui.scalar_value.value()
            if (np.isclose(scalar, 0) and np.allclose(value[:3], value[0]) and
                    np.allclose(value[3:], 0)):
                # If the previous value is zero, and the tensor is diagonal,
                # use the diagonal value
                scalar = value[0]

            spinbox.setValue(scalar)

        f = scalar_to_tensor if b else tensor_to_scalar

        for spinbox in self.thermal_factor_spinboxes:
            f(spinbox)

class ThermalFactorSpinBox(ScientificDoubleSpinBox):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.editor = ThermalFactorEditor(0, parent)
        self.setLineEdit(ThermalFactorLineEdit(self, self))
        self.valueChanged.connect(self.update_editor_value)

    def value(self):
        return self.editor.value

    def setValue(self, v):
        self.editor.value = v
        if self.editor.is_tensor:
            # Force an update
            super().setValue(super().value())
            self.valueChanged.emit(super().value())
            self.setReadOnly(True)
        else:
            super().setValue(v)
            self.valueChanged.emit(v)
            self.setReadOnly(False)

    def update_editor_value(self):
        if not self.editor.is_tensor:
            self.editor.value = super().value()

    def textFromValue(self, value):
        if not hasattr(self, 'editor') or not self.editor.is_tensor:
            return super().textFromValue(value)

        return 'Tensor'

    def open_editor(self):
        original = copy.deepcopy(self.editor.value)
        if not self.editor.exec():
            self.editor.value = original
            return

        self.setValue(self.editor.value)

class ThermalFactorLineEdit(QLineEdit):
    def __init__(self, spinbox, parent=None):
        super().__init__(parent)

        self.spinbox = spinbox

    def mousePressEvent(self, event):
        if self.isReadOnly():
            self.open_editor()
            return

        super().mousePressEvent(event)

    def mouseDoubleClickEvent(self, event):
        self.open_editor()

    def open_editor(self):
        self.spinbox.open_editor()


from django.urls import re_path

from user.views import LoginViewSet, UserViewSet

user_detail = UserViewSet.as_view({
    'post': 'post',
    'get': 'list'
})

urlpatterns = [
    re_path(r'^/?$', user_detail, name='user-detail'),
    re_path(r'^/(?P<pk>\d+)/?$', UserViewSet.as_view({'get': 'get'}), name='get'),
    re_path(r'^/me/?$', UserViewSet.as_view({'delete': 'delete'}), name='delete'),
    re_path(r'^/login/?$', LoginViewSet.as_view({'post': 'login'}), name='login'),
]


import toml
import os

import streamlit as st

"""Load environment variables from .streamlit/secrets.toml to os.environ"""
def load_secrets():
    if st.secrets:
        os.environ.update(st.secrets)
        return

    try:
        with open(".streamlit/secrets.toml") as f:
            secrets = toml.load(f)
        for k, v in secrets.items():
            os.environ[k] = v
    except FileNotFoundError:
        print("No secrets.toml file found")

load_secrets()


#!/usr/bin/env python3

import json
import argparse
from jinja2 import Template

"""
make sure to install jinja2 in your python environment:

    `pip install jinja2`

example:

    ./kctf-init.py --template ./terraform/custom_templates/lambda_function.j2 --module_dir ./terraform/modules/lambda_function --output_file ./testing/tf/kctf-init-lambda-function.tf --params ./terraform/tools/sample_lambda_params.json
"""

parser = argparse.ArgumentParser()
parser.add_argument('--template', required=True, help="Jinja Template location")
parser.add_argument('--module_dir', required=True, help="Terraform module location")
parser.add_argument('--output_file', required=True, help="Generated terraform file full path")
parser.add_argument('--params', required=False, help="Json file with dynamic params/configuration")
args = parser.parse_args()

with open(args.template, "r") as template_file:
    try:
        template_file_contents = template_file.read()
        if template_file_contents:
            resource_file = open(args.output_file, 'a')
            params = {}
            if args.params:
                with open(args.params, "r") as params_file:
                    params = json.load(params_file)
            params['module_source'] = args.module_dir
            jinja_template = Template(template_file_contents)
            rendered_jinja_template = jinja_template.render(params)
            resource_file.write(rendered_jinja_template)
            resource_file.close()
            print(f'saved: {args.output_file}')
    except Exception as e:
        print(e)


# -*- coding: utf-8 -*-
import scrapy
import re

import datetime
from urllib.parse import urljoin
from NBANews.utils.common import get_md5
from NBANews.items import NewsItem, NBAnewsItemLoader
from scrapy_splash import SplashRequest
from scrapy.loader import ItemLoader

class NanewsSpider(scrapy.Spider):
    name = 'nbaNews'
    allowed_domains = ['sports.qq.com']
    start_urls = ['http://sports.qq.com/l/basket/nba/list20181018164449.htm']

    def start_requests(self):
        for url in self.start_urls:
            yield SplashRequest(url, self.parse, args={'wait': 0.5})

    def parse(self, response):
        '''
        1.解析新闻列表中的所有url  并交给scrapy 解析每个页面的数据
        2.将下一页的url交给scrapy
        :param response:
        :return:
        '''
        post_nodes = response.xpath('//*[@id="listInfo"]/li')

        for post_node in post_nodes:
            news_url = post_node.xpath('./a/@href').extract_first()
            image_url = urljoin("http://", post_node.xpath('./a/img/@src').extract_first())
            # 获取每页列表中的新闻
            yield SplashRequest(news_url, self.parse_detail, args={'wait': 0.5}, meta={"front-image-url": image_url})

        '''
        #获取下一页的的url
        next_url = response.xpath('//*[@class="next"]/@href').extract_first()
        print("next_url :"+str(next_url))
        if next_url:
            yield SplashRequest(next_url, self.parse, args={'wait': 0.5})
        '''

    def parse_detail(self, response):
        '''
        解析页面中的详细数据
        :param response:
        :return:
        '''

        #newsItem = NewsItem()
        # title = response.xpath('//div[@class="hd"]/h1/text()').extract()[0]
        # create_time = response.xpath('//span[@class="a_time"]/text()').extract()[0]
        # comment_nums = response.xpath('//*[@id="cmtNum"]/text()').extract()[0]
        # content = response.xpath('//div[@class="Cnt-Main-Article-QQ"]//p[@class="text"]/text()').extract_first()
        #image_path = [urljoin('http://', re.match(r'<img src="(.+?)"', p).group(1)) for p in
        #             response.xpath('//*[@align="center"]/img[@src]').extract()]

        '''
        newsItem['title'] = title
        newsItem['create_time'] = create_time
        newsItem['url'] = response.url
        newsItem['url_object_id'] = get_md5(response.url)
        newsItem['front_image_url'] = [front_image_url]
        newsItem['comment_nums'] = comment_nums
        newsItem['image_path'] = image_path
        newsItem['content'] = content
        '''

        # Itemloader 加载item

        item_loader = NBAnewsItemLoader(item=NewsItem(), response=response)
        front_image_url = response.meta.get("front-image-url", "")
        item_loader.add_xpath('title', '//div[@class="hd"]/h1/text()')
        item_loader.add_value('create_time', '/span[@class="a_time"]/text()')
        item_loader.add_xpath('comment_nums', '//*[@id="cmtNum"]/text()')
        item_loader.add_xpath('image_path', '//*[@align="center"]/img[@src]')
        item_loader.add_xpath('content', '//div[@class="Cnt-Main-Article-QQ"]//p[@class="text"]/text()')
        item_loader.add_value('url', response.url)
        item_loader.add_value('url_object_id', get_md5(response.url))
        item_loader.add_value('front_image_url', [front_image_url])
        newsItem = item_loader.load_item()

        yield newsItem


from pydub import AudioSegment
# from moviepy.video.io.ffmpeg_tools import ffmpeg_extract_subclip
import argparse
import sys, subprocess, math
import time
import os

def crop(times, text, audio_path, audio_save_dir, video_path=None, video_save_dir=None):
    """
    Used to crop the audio and video from the time given in lab files.
    Input:
        times: A List[Tuples] containing start and end times.
        audio_path: path to the audio file.
        audio_save_dir: path to the folder where cropped audio files have to be stored.
        video_path: path to the video clip.
        video_save_dir: path to the folder where cropped video files have to be stored.
    Output:
        None.
        Crops and saves audio and video clips.
    """
    audio = AudioSegment.from_wav(audio_path)
    for i, (t1, t2) in enumerate(times):
        id = '%02d' % i
        text_sequence = text[i]
        f = open(audio_save_dir + '/' + id + '.txt', 'w')
        f.write(text_sequence)
        f.close()

        t1_new = t1*1000
        t2_new = t2*1000
        cropped_audio = audio[t1_new:t2_new]
        cropped_audio.export(audio_save_dir + "/" + id +".wav", format="wav", parameters=["-ar", "16000"])
        if video_path is not None:
            extension = video_path.split('.')[-1]
            # ffmpeg command to crop
            command = ("ffmpeg -ss %s -i %s -t %s %s" % (t1, video_path, t2-t1, video_save_dir + "/" + id + "." + extension))
            output = subprocess.call(command, shell=True, stdout=None)

    print("Finished cropping and saving {} files".format(i+1))

def read_lab(path):
    """
    Used to read the lab file and extract time information from it.
    Input:
        path: path to the lab file
    Output:
        times: List[Tupes] containg start and end time of each utterance.
    """
    times = []
    text = []
    with open(path, 'r') as f:
        data = f.readlines()
        for d in data:
            d = d.split()
            t1, t2 = float(d[0]), float(d[1])
            text.append(' '.join(d[2:]))
            times.append((t1,t2))

    return times, text

def create_metafile(path, save_dir):
    """
    Converts the lab file into a format which can be directly read by the TTS system.
    Input:
        path: path to the lab file.
        save_dir: folder in which the metafile is stored.
    Output:
        None
        Save the metafile.txt in the specified folder.
    """
    with open(save_dir + '/' +"metadata.txt", 'w') as s:
        with open(path, 'r') as f:
            data = f.readlines()
            for i,d in enumerate(data):
                d = d.split()
                s.write(str(i)+"|"+" ".join(d[2:])+"\n")

    print("Completed writing metafile")

def main():

    parser = argparse.ArgumentParser()
    parser.add_argument('--lab_file', type=str, default='', required=True)
    parser.add_argument('--audio_path', type=str, default='', required=False)
    parser.add_argument('--audio_save_dir', type=str, default='', required=False)
    parser.add_argument('--video_path', type=str, default=None, required=False)
    parser.add_argument('--video_save_dir', type=str, default=None, required=False)
    parser.add_argument('--metafile_save_dir', type=str, default='', required=False)
    args = parser.parse_args()

    # Make dirs
    if not os.path.exists(args.video_save_dir):
        os.makedirs(args.video_save_dir)
    if not os.path.exists(args.audio_save_dir):
        os.makedirs(args.audio_save_dir)

    times, text = read_lab(args.lab_file)
    crop(times, text, args.audio_path, args.audio_save_dir, args.video_path, args.video_save_dir)
    create_metafile(args.lab_file, args.metafile_save_dir)

if __name__ == '__main__':
    main()


import plotDataPoints as pdp
import drawLine as dl
import matplotlib.pyplot as plt

def plotProgresskMeans(X, centroids, previous, idx, K, i):
    #PLOTPROGRESSKMEANS is a helper function that displays the progress of
    #k-Means as it is running. It is intended for use only with 2D data.
    #   PLOTPROGRESSKMEANS(X, centroids, previous, idx, K, i) plots the data
    #   points with colors assigned to each centroid. With the previous
    #   centroids, it also plots a line between the previous locations and
    #   current locations of the centroids.
    #

    # plt.hold(True)

    # Plot the examples
    pdp.plotDataPoints(X, idx, K)

    # Plot the centroids as black x's
    plt.scatter(centroids[:,0], centroids[:,1], marker='x', s=400, c='k', linewidth=1)

    # Plot the history of the centroids with lines
    for j in xrange(centroids.shape[0]):
        dl.drawLine(centroids[j, :], previous[j, :], c='b')

    # Title
    plt.title('Iteration number {:d}'.format(i+1))

    return