text:
"""
   Copyright 2018 Scott Almond

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

     http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.

Purpose:
This class manages the configuration file IO of the hyperspace map
and regular data access to this info.  This includes:
- Subway Map (2D)
- Hyperspace Map (3D)
Layers:
- LINEAR_DEFINITION, what you will encounter between where-you-were and where-you-are, and how that will be represented in 2D/3D
- BRANCH_DEFINITION, where-you-were, where-you-are, where-you're-going relationship, if this relationship is reversable, and what where-you-are looks like
- SEGMENT_DEFINITION, refinement of what the 3D visuals looks like
- DEBRIS_DEFINITION, refinement of what the 3D visuals looks like

Structure:
The 2D/3D maps are inter-related.  They share "nodes" where transitions occur.  Several node relationships exist:
- corner, where one linear sequence comes in and one linear sequence leaves
  The 2D map shows a 45 degree bend, typically, here with one line coming in and one leaving
  The 3D visuals does make any clear indiciation a bend has occurred
  This distinction is useful for clearly visualizing progress of the pod through hyperspace on the 2D map as opposed to have one node-node relationship consist of multiple bends in the config file
- branch, where one linear sequence comes in and two leave
  The 2D map shows a circle where thee sequences exit
  The 3D visuals show a branching structure that the pod travels through
- jump/land, where a hyperspace jump starts/ends
  The 2D map shows the pod stopping at the "jump" node, only proceeding
  once the pod has reached the "land" node.  The jump has a perpendicular line
  The land has a perpendicular line plus a triangle pointing in the direction of travel
  The 3D visuals use a different ring type to indicate a hyperspace jump
- start, where the start of the map is
  The 2D map shows a square where only one sequence exits
  The 3D visuals play an intro sequence before the first life.
  On subsequnet lives, the player pod starts here, blinking briefly to indicate a new life
- dead, a dead end
  On the 2D map, an 'X' is displayed
  On the 3D visuals, a black hole of soome kind is displayed
- end, the goal of traversing the map
  In the 2D map, this is the same as "start"
  In the 3D visuals, something analagous to an "dead" is displayed, but in a different color (white)

Between each pair of nodes, one or more segments exist
- The 2D map only tracks the total process made between two nodes
- Each segment consists of a singular 3D circular arcing path.  To make
  a "back-and-forth" motion requires multiple segments back-to-back.  Each
  segment consists of some number of rings, which in turn each contain debris

python3 -m chapters.wall.hyperspace_helper.Maze

File types:
  linear_definition
    defines how nodes are conceptually inter-related: branching, jumps, dead ends, start, finish
    is effectively the core of the maze definition
    defines how nodes are represented on a 2D screen (support for 2D graphic rendering)
  branch_definition
	defines how the straight segments defiend in linear_definition are inter-related
  segment_definition
    defines a list of ring-debris relationships for each segment (support for 3D grphi rendering)
  debris_definition
    defines one rings' worth of debris (quick-reference dictionary for segment_definition)

File columns:
	linear_definition
		node_id_prev, node_id_next, segments, colors
			note: segments is a space-separated list of numbers from segment_definition
			precon: each segment_id only appears once ever in the config file
			precon: the first node-node relationship in the file is the start of the maze
			precon: branches only contain one segment (otherwise state machine would lose sense of direction in maze)
			note: colors is a space-separated list of paranthetical underscore numerical RGB triplets: "0_255_0 0_0_255"
			    would draw a green and blue line next to each other on the subway map
	branch_definition
		node_id_prev, node_id_curr, node_id_next, is_two_way, sprite, x_pixels, y_pixels
			note: sprite is one of the following strings: start, end, corner, branch, dead, jump2, land3
				where jump2 and land3 can have any numeric after them
				where default
			note: is_two_way is a boolean: 'TRUE' or 'FALSE'
				used as a flag to allow/deny the player from going back down
				a path in the backward direction.  If false, then only the forward
				direction (as listed in teh config file) is permitted/accessible
	segment_definition
		segment_id, type, curvature_degrees,  orientation_degrees,
		ring_model, ring_angle_degrees, ring_angle_degrees_per_second,
		debris_name, debris_angle_degrees, debris_angle_degrees_per_second
			note: type can be one of the following strings: straight, dead, end
			note: ring_model is an int index per AssetLibrary.ring_filename
			note: debris_name can be 'None' string or a string to 'debris_definition'
			note: ring_angle_rate_degrees_per_second and debris_... can be the
				string 'DEFAULT' to use the built-in default constant
			note: all positions/velocities are made to be True precisely
				when the pod passes through them.  Forward/backward propagation is
				used to compute the state before/after an encounter with the player
			precon: the last ring in a branch should be a double-ring .obj model
			precon: if segment_id are used in 'linear_definition' where
				node_id_prev==node_id_next then segment_type must be type 'branch'
	debris_definition
		debris_name, debris_model, radius, x_scale, y_scale, z_scale,
		x, y, z, x_degrees, y_degrees, z_degrees,
		x_degrees_per_second, y_degrees_per_second, z_degrees_per_second
			note: debris_name is a unique string used to identify the debris construction
				if the debris_name is blank, this row is combined with the previous row
				to create a composite debris definition
			note: debris_model is an int index from AssetLibrary.asteroid_filename
			note: scale will shrink/expand the CAD model
				by the stated ratio (1.0 default), leave blank to use default
			note: radius is a distance measurement in pi3d units to help with
				rendering the cross hairs on an asteroid that can be destroyed
			note: rows with a negative radius are ignored by the targeting algorithm
				(cannot eb targeted)
			note: radius is updated (scaled) automatically in RingAssembly.addDebris
				if xyz_scale are non-unity
			note: x, y, and z are offset in pi3d distance units from the center of the ring
				x is left/right, y is up/down, z is forward/back (direction of travel)
			note: all positions/velocities are made to be True precisely
				when the pod passes through them.  Forward/backward propagation is
				used to compute the state before/after an encounter with the player
			precon: first row must have a name in the 'debris_name' field

Note: idea to separate node/segments was a poor choice since it results
  in duplicate tracking and indexing in many places which gets confusing/verbose
  A better architecture would have been to have only atomic segments
  and then define "node" markers in the map as the juncture between segments
"""

from enum import Enum
import copy #deepcopy
import numpy as np

class MAZE_CONFIG(Enum):
	LINEAR_DEFINITION={"filename":"linear_definition.csv"}
	BRANCH_DEFINITION={"filename":"branch_definition.csv"}
	SEGMENT_DEFINITION={"filename":"segment_definition.csv"}
	DEBRIS_DEFINITION={"filename":"debris_definition.csv"}

class Maze:
	#CONFIG_FILE_PATH='./chapters/wall/assets/hyperspace/configuration/'
	CONFIG_FILE_PATH='/home/pi/Documents/EscapeRoom/src/chapters/wall/assets/hyperspace/configuration/'

	def __init__(self):
		try: #resource paths depend on whether Maze is running headless for testing, or as part of application
			from chapters.wall.hyperspace_helper.RingAssembly import RingAssembly
			from util.ResourceManager import ResourceManager
			from chapters.wall.hyperspace_helper.Segment import Segment
			self.Segment=Segment
		except ImportError:
			from RingAssembly import RingAssembly #to fetch constants for ring/debris rotation rate
			from ResourceManager import ResourceManager
			#from Segment import Segment
		self.rm_library=ResourceManager
		self.ra_library=RingAssembly

	def clean(self):
		self.linear_definition=self.__load(MAZE_CONFIG.LINEAR_DEFINITION)
		self.branch_definition=self.__load(MAZE_CONFIG.BRANCH_DEFINITION)
		self.segment_definition=self.__load(MAZE_CONFIG.SEGMENT_DEFINITION)
		self.debris_definition=self.__load(MAZE_CONFIG.DEBRIS_DEFINITION)
		self.__verify_linear_and_branch_consistency()
		self.__verify_segment_definition()
		self.__verify_debris_definition()
		self.__merge_debris_into_segment_definition()

	#ensure that every ring pair in linear_definition is mated up with at least one branch in branch_definition and vice versa
	def __verify_linear_and_branch_consistency(self):
		#TODO: sanity check to make sure no definitions are missing...
		#  ex: nodes used in linear_definition/brach_definition, but not used elsewhere, same with segments
		#  TODO: update segment_definition to directly reference debris_definition row contents... "ring_list" --> debris_list=[debris_name[], debris_name[] ... ]
		unique_node_id=[]
		for is_check in [False,True]: #check that all nodes are in agreement between linear_definition and branch_definition
			for definition in [self.linear_definition,self.branch_definition]:
				for row in definition:
					if(definition is self.linear_definition):
						node_list=[row["node_id_prev"],row["node_id_next"]]
						definition_string="linear_definition"
					else:
						node_list=[row["node_id_prev"],row["node_id_curr"],row["node_id_next"]]
						definition_string="branch_definition"
					for node_id in node_list:
						if(is_check):
							if(node_id>=0 and not node_id in unique_node_id):
								#raise ValueError("Maze.__verify_linear_and_branch_consistency: node_id: ",node_prev," not defined in ",definition_string," config file")
								print("Maze.__verify_linear_and_branch_consistency: node_id: ",node_prev," not defined in ",definition_string," config file")
						else:
							if(not node_id in unique_node_id): unique_node_id.append(node_id)
		return True

	#ensure every segment defined in linear_definition has a concrete definition in segment_definition
	def __verify_segment_definition(self):
		#check that all segments in linear_definition are defined
		for linear_row in self.linear_definition:
			for linear_segment_id in linear_row["segment_list"]:
				is_found=False
				for segment_definition in self.segment_definition:
					if(segment_definition["segment_id"]==linear_segment_id):
						is_found=True
						break
				if(not is_found):
					print("Maze.__verify_segment_definition: segment_id "+str(linear_segment_id)+" not defined in linear_definition config file")
					#raise ValueError("Maze.__verify_segment_definition: segment_id "+str(linear_segment_id)+" not defined in linear_definition config file")

	#ensure that for every string debris_name called in segment_definition that there is a corresponding debris_definition
	def __verify_debris_definition(self):
		for segment_row in self.segment_definition:
			for ring_row in segment_row["ring_list"]:
				if(self.__get_debris_definition_by_name(ring_row["debris_name"]) is None):
					#raise ValueError("Maze.__verify_debris_definition: debris_name "+str(segment_row["debris_name"])+" not defined in debris_definition config file")
					print("Maze.__verify_debris_definition: debris_name "+str(ring_row["debris_name"])+" not defined in debris_definition config file")

	#merge debris_definition into segment_definition to simplify later code lookups
	def __merge_debris_into_segment_definition(self):
		for segment_row in self.segment_definition:
			for ring_row in segment_row["ring_list"]:
				ring_row["debris_list"]=self.__get_debris_definition_by_name(ring_row["debris_name"])

	def __get_debris_definition_by_name(self,debris_name):
		if(debris_name in ["","None","NONE","Null","NULL"]):
			return []
		if(not debris_name in self.debris_definition):
			return None
		return self.debris_definition[debris_name]

	#load csv file
	#linear_definition = [ [prev,next], [prev,next], ... ]
	#branch_definition = [ [prev,curr,next], [prev,curr,next], ... ]
	#segment_definition = [ {"segment_id":id,"ring_list":[ring_assy, ring_assy, ...}, {"segment_id":id,"ring_list":[ring_assy, ring_assy, ...}, ...]
	#debris_definition = { name:[debris, debris, ...], name:[debris, debris, ...], ... }
	def __load(self,maze_config):
		filename=maze_config.value["filename"]
		file_contents=self.rm_library.loadCSV(self.CONFIG_FILE_PATH+filename)
		if(maze_config==MAZE_CONFIG.DEBRIS_DEFINITION):
			output={}
		else:
			output=[]
		previous_row=[]
		previous_row_name=''
		for file_row_id in range(len(file_contents)):
			file_row=file_contents[file_row_id]
			if(maze_config==MAZE_CONFIG.LINEAR_DEFINITION):
				file_row["node_id_prev"]=int(file_row["node_id_prev"])
				file_row["node_id_next"]=int(file_row["node_id_next"])
				segment_list=file_row["segment_list"].split()
				file_row["segment_list"]=[]
				for segment in segment_list: file_row["segment_list"].append(int(segment))
				color_list=file_row["colors"].split()
				file_row["colors"]=[]
				for color_str in color_list:
					this_color=color_str.split('_')
					if(len(this_color) != 3):
						raise NotImplementedError('Maze.__load: unknown color format specified in LINEAR_DEFINITION: ',this_color)
					for idx in range(3):
						this_color[idx]=int(this_color[idx])
					file_row["colors"].append(this_color)
				output.append(file_row)
			elif(maze_config==MAZE_CONFIG.BRANCH_DEFINITION):
				file_row["node_id_prev"]=int(file_row["node_id_prev"])
				file_row["node_id_curr"]=int(file_row["node_id_curr"])
				file_row["node_id_next"]=int(file_row["node_id_next"])
				is_two_way=False
				if(file_row["is_two_way"]==1 or file_row["is_two_way"]=='True' or
				   file_row["is_two_way"]=='TRUE'):
					is_two_way=True
				file_row["is_two_way"]=is_two_way

				#TODO, handle jump/land sprite case

				output.append(file_row)
			elif(maze_config==MAZE_CONFIG.SEGMENT_DEFINITION):
				file_row["segment_id"]=int(file_row["segment_id"])
				for prev_row in output: #ease-of-use implementation: copy-paste first row contents into subsequent rows with same id number
					if(prev_row["segment_id"]==file_row["segment_id"]):
						file_row["type"]=prev_row["type"]
						file_row["curvature_degrees"]=str(prev_row["curvature_degrees"])
						file_row["orientation_degrees"]=str(prev_row["orientation_degrees"])
						break
				file_row["type"]=file_row["type"].strip()
				if(not file_row["type"] in ["straight","dead","end","branch"]):
					raise NotImplementedError('Maze.__load: unknown segment type in SEGMENT_DEFINITION: ',file_row["type"])
				for deg in ["curvature","orientation","ring_angle","debris_angle"]:
					file_row[deg+"_degrees"]=file_row[deg+"_degrees"].strip()
					if(len(file_row[deg+"_degrees"])<=0): #if no string present, then populate default
						file_row[deg+"_degrees"]=0.0
					else:
						file_row[deg+"_degrees"]=float(file_row[deg+"_degrees"])
				if(len(file_row["ring_angle_degrees_per_second"].strip())<=0):
					file_row["ring_angle_degrees_per_second"]=RingAssembly.RING_ROTATION_DEGREES_PER_SECOND
				else:
					file_row["ring_angle_degrees_per_second"]=float(file_row["ring_angle_degrees_per_second"])
				if(len(file_row["debris_angle_degrees_per_second"].strip())<=0):
					file_row["debris_angle_degrees_per_second"]=RingAssembly.DEBRIS_ROTATION_DEGREES_PER_SECOND
				else:
					file_row["debris_angle_degrees_per_second"]=float(file_row["debris_angle_degrees_per_second"])
				segment_row={"segment_id":file_row["segment_id"],"type":file_row["type"],
				    "curvature_degrees":file_row["curvature_degrees"],
				    "orientation_degrees":file_row["orientation_degrees"],"ring_list":[]}
				is_new_segment=True
				for row in output:#find first row in output where this segment_id has laready been defined
					if(row["segment_id"]==file_row["segment_id"]):
						segment_row=row
						is_new_segment=False
						break
				ring_row={}
				ring_parameters=["ring_model","debris_name",
								 "ring_angle_degrees","ring_angle_degrees_per_second",
								 "debris_angle_degrees","debris_angle_degrees_per_second"]
				for parameter in ring_parameters:
					if(parameter=="debris_name"):
						ring_row[parameter]=file_row[parameter]
					elif(parameter=="ring_model"):
						ring_row[parameter]=int(file_row[parameter])
					else:
						ring_row[parameter]=float(file_row[parameter])
				segment_row["ring_list"].append(ring_row)
				if(is_new_segment):
					output.append(segment_row)
			elif(maze_config==MAZE_CONFIG.DEBRIS_DEFINITION):
				file_row["debris_name"]=file_row["debris_name"].strip()
				this_row_name=file_row["debris_name"]
				if(len(file_row["debris_name"])<=0):
					this_row_name=previous_row_name
				file_row["radius"]=int(file_row["radius"])
				file_row["debris_model"]=int(file_row["debris_model"])
				for var in ["","_scale","_degrees","_degrees_per_second"]:
					for dim in ["x","y","z"]:
						file_row[dim+var]=file_row[dim+var].strip()
						if(len(file_row[dim+var])<=0): #defaults
							if(var==""): #0 pos
								file_row[dim+var]=0.0
							elif(var=="_scale"): #1x scale
								file_row[dim+var]=1.0
							elif(var=="_degrees"): #0 deg
								file_row[dim+var]=0.0
							elif(var=="_degrees_per_second"):
								file_row[dim+var]=0.0 #0 deg/sec
						else:
							file_row[dim+var]=float(file_row[dim+var])
				if(not this_row_name in output):
					output[this_row_name]=[]
				output[this_row_name].append(file_row)
				previous_row_name=this_row_name
			else:
				raise NotImplementedError('Maze.__load: Unable to load config file type: ',maze_config)
		return output

	def getFirstPopulatedSegment(self,asset_library,start_time_seconds):
		segment_id=1 #hard coded/precon that the first segment is 1
		is_forward=True
		is_branch=False
		start_position=np.array([0.0,0.0,0.0])
		start_rotation_matrix=np.eye(3)
		return self.getPopulatedSegment(segment_id,is_forward,is_branch,asset_library,
			start_position,start_rotation_matrix,start_time_seconds)

	#return Segment with RingAssemblies assembled/installed/added
	#  TODO: only use/load RingAssemblies if in Hyperspace.py, not Map.py ...
	#can't remember, is padding (space between rings) put BEFORE ring
	#  CAD model, or after ...? TODO
	#  It's after:  u=idx/len - so u goes from [0,1)
	def getPopulatedSegment(self,segment_id,is_forward,is_branch,asset_library,
		start_position,start_rotation_matrix,start_time_seconds):
		segment_definition=self.__getSegmentDefinition(segment_id)
		if(segment_definition is None):
			raise IndexError("Maze.getSegment: cannot find segment_definition for segment_id: ",segment_id)
		#is_branch=segment_definition["sprite"]=="branch" #provided by method call...
		curvature_degrees=segment_definition["curvature_degrees"]
		orientation_degrees=segment_definition["orientation_degrees"]
		ring_count=len(segment_definition["ring_list"])
		segment=self.Segment(asset_library,is_branch,start_position,start_rotation_matrix,
			start_time_seconds,curvature_degrees,orientation_degrees,ring_count,segment_id)
		ring_list=segment_definition["ring_list"]
		if(not is_forward): ring_list=list(reversed(ring_list))
		#if(not is_branch):
		if(True):
			for ring_index in range(len(ring_list)):
				ring_definition=ring_list[ring_index]
				u=ring_index/len(ring_list)
				curve_id=0 #always add custom RingAssemblies to first curve (assumes non-branch)
				ring_index=ring_definition["ring_model"]
				ring_rotation_degrees=ring_definition["ring_angle_degrees"]
				ring_rotation_rate=ring_definition["ring_angle_degrees_per_second"]
				debris_rotation_degrees=ring_definition["debris_angle_degrees"]
				debris_rotation_rate=ring_definition["debris_angle_degrees_per_second"]
				ring_assembly=segment.addRingAssembly(asset_library,u,curve_id,ring_index,
					ring_rotation_degrees,ring_rotation_rate,
					debris_rotation_degrees,debris_rotation_rate)
				debris_list=ring_definition["debris_list"]
				for debris_definition in debris_list:
					debris_model_index=debris_definition["debris_model"]
					location=[debris_definition["x"],debris_definition["y"],debris_definition["z"]]
					angle=[debris_definition["x_degrees"],debris_definition["y_degrees"],debris_definition["z_degrees"]]
					angular_velocity=[debris_definition["x_degrees_per_second"],
									  debris_definition["y_degrees_per_second"],
									  debris_definition["z_degrees_per_second"]]
					scale=[debris_definition["x_scale"],debris_definition["y_scale"],debris_definition["z_scale"]]
					radius=debris_definition["radius"]
					#print("Maze.getPopulatedSegment: debris_model_index: ",debris_model_index)
					this_debris=ring_assembly.addDebris(debris_model_index,location,angle,angular_velocity,scale,radius)
		return segment

	#search through file for definition (curvature, orientation, etc)
	def __getSegmentDefinition(self,segment_id):
		for row in self.segment_definition:
			if(row["segment_id"]==segment_id):#match
				return copy.deepcopy(row) #should no longer need deepcopy since debris is now copied into segment_definition at start ratehr than after this method call
		return None

	#returns list of segment_id that come after the curr_segment, prev_segment sequence
	#returns {"segment_id":X,"is_forward":bool,"is_branch":bool}
	#note: uses recursion to evaluate if segment is a branch (if next segment has mulitple segments after it)
	def getSegmentIdAfter(self,prev_segment_id,curr_segment_id,evaluate_is_branch=True):
		#consider five cases:
		#  prev_nodes==curr_nodes
		#    is between two nodes
		#    then next segment lies either 1) between these nodes, 2) at the start/end of the next nodes
		#  else prev_nodes!=curr_nodes
		#    if curr_nodes is branch, 3) return two followers coming from prev_nodes direction
		#
		print("Maze.getSegmentIdAfter: prev_segment_id: "+str(prev_segment_id)+", curr_segment_id: "+str(curr_segment_id)+", evaluate_is_branch: "+str(evaluate_is_branch))
		lin_def_pair=self.getOrderedLinearDefinition(curr_segment_id,prev_segment_id)
		prev_lin_def=lin_def_pair["prev"]
		curr_lin_def=lin_def_pair["curr"]
		#print("Maze.next_node_list: prev_lin_def[node_id_prev]: ",prev_lin_def["node_id_prev"])
		#print("Maze.next_node_list: curr_lin_def[node_id_prev]: ",curr_lin_def["node_id_prev"])
		#print("Maze.next_node_list: curr_lin_def[node_id_next]: ",curr_lin_def["node_id_next"])
		next_node_list=self.getNodeIdAfter(prev_lin_def["node_id_prev"],curr_lin_def["node_id_prev"],curr_lin_def["node_id_next"])
		out_list=[]
		#print("Maze.next_node_list: next_node_list: ",next_node_list)
		#if(prev_lin_def["node_id_prev"]==curr_lin_def["node_id_prev"] and
		#   prev_lin_def["node_id_next"]==curr_lin_def["node_id_next"]): #then operating in the same node-node region
		if(curr_lin_def["segment_id_index"]==(curr_lin_def["segment_list_len"]-1)):#then at end of this node-node region
			for next_node_id in next_node_list: #get the first segment_id from the next node-node region
				segment_list=self.getSegmentsBetweenNodes(curr_lin_def["node_id_next"],next_node_id)
				#print("Maze.getSegmentIdAfter: segment_list: ",segment_list)
				is_branch=False
				branch_list=[]#list of segment id's following this branch
				if(curr_segment_id>=0 and evaluate_is_branch): #only evaluate if is_branch if the current node has something potentially after it
					try:
						next_list=self.getSegmentIdAfter(curr_segment_id,segment_list["segment_list"][0],evaluate_is_branch=False)
						if(len(next_list)>1):
							is_branch=True #next segment is a branch only if there are two segments that come after it
							for inner_row in next_list:
								branch_list.append(inner_row["segment_id"])
					except TypeError:#attempt to get next segment and failed due to there not being any declared segments at end of queue
						is_branch=False
				out_list.append({"is_forward":segment_list["is_forward"],"segment_id":segment_list["segment_list"][0],"is_branch":is_branch,"branch_list":branch_list}) #precon: len(segment_list)>0
		else:#fetch next from current node-node region
			segment_list=curr_lin_def["segment_list"]
			curr_segment_index=curr_lin_def["segment_id_index"]
			#print("Maze.getSegmentIdAfter: curr_segment_index",curr_segment_index)
			#print("Maze.getSegmentIdAfter: curr_segment_index",curr_segment_index)
			out_list.append({"is_forward":curr_lin_def["is_forward"],"segment_id":segment_list[curr_segment_index+1],"is_branch":False,"branch_list":[]}) #get segment_id after curr_segment_id in this node-node region
		#else:#prev_ and curr_segment_id straddle a node-node relationship
		#	if(curr_lin_def["segment_id_index"]==(curr_lin_def["segment_list_len"]-1)):
		return out_list

	#reutrn a list of next_node_id (may include branches)
	#curr_node_id_from aka prev_node_id_to
	def getNodeIdAfter(self,prev_node_id_from,curr_node_id_from,curr_node_id_to):
		prev_node_id=curr_node_id_from
		if(curr_node_id_from==curr_node_id_to):#if in branch, use previous node as anchor point to determine next node
			prev_node_id=prev_node_id_from
		curr_node_id=curr_node_id_to
		out_list=[]
		for row in self.branch_definition:
			#print("Maze.getNodeIdAfter: row: ",row)
			if(row["node_id_curr"]==curr_node_id):
				if(row["node_id_prev"]==prev_node_id): #forward relationship
					out_list.append(row["node_id_next"])
				elif(row["node_id_next"]==prev_node_id and row["is_two_way"]): #reverse relationship
					out_list.append(row["node_id_prev"])
		#make sure not get stuck in a loop - remove [curr_node_id_from,curr_node_id_to] relationship, if present
		#if(curr_node_id_from==curr_node_id_to):
		#	out_list=[row for row in out_list if not row==curr_node_id_to]
		if(len(out_list)>1): #if looking to branch out
			if(not curr_node_id_from==curr_node_id_to): #and not already in a branch
				out_list=[curr_node_id_to] #then specifically go to this upcoming branch
		return out_list

	#get the linear_definition for curr_ and prev_segment_id, but ordered (is_forward) for the partiular input sequence
	#returns {"prev":{"prev_id_node":X,"next_id_node":Y,"is_forward":bool,"segment_id":Z,"segment_list":prev_curr_sorted_segment_list},
	#		  "curr":{"prev_id_node":X,"next_id_node":Y,"is_forward":bool,"segment_id":Z,"segment_list":prev_curr_sorted_segment_list}}
	#put from_ and to_ nodes in order, and include the index of the segment in the segment_list
	def getOrderedLinearDefinition(self,curr_segment_id,prev_segment_id):
		prev_lin_def=self.getLinearDefinitionOfSegment(prev_segment_id)
		curr_lin_def=self.getLinearDefinitionOfSegment(curr_segment_id)
		if(curr_lin_def["node_id_prev"]==prev_lin_def["node_id_prev"] and
		   curr_lin_def["node_id_next"]==prev_lin_def["node_id_next"]): #then both segments within same node-node range
			#if(prev_node_id["segment_id_index"]>curr_node_id["segment_id_index"]): #need to flip because headed backwards #varaible not defined error
			if(prev_lin_def["segment_id_index"]>curr_lin_def["segment_id_index"]): #need to flip because headed backwards
				prev_lin_def=self.__reverseLinearDefinition(prev_lin_def)
				curr_lin_def=self.__reverseLinearDefinition(curr_lin_def)
		else:
			#else, stradelling some form of node-node boundary, and may need to reverse one or both segment linear_definition rows
			is_reverse_prev=True
			is_reverse_curr=True
			if(prev_lin_def["node_id_prev"]==prev_lin_def["node_id_next"]):
				is_reverse_prev=False #prev is branch
			if(curr_lin_def["node_id_prev"]==curr_lin_def["node_id_next"]):
				is_reverse_curr=False #curr is branch
			if(prev_lin_def["node_id_next"] in [curr_lin_def["node_id_prev"],curr_lin_def["node_id_next"]]):
				is_reverse_prev=False #prev in correct order
			if(curr_lin_def["node_id_prev"] in [prev_lin_def["node_id_prev"],prev_lin_def["node_id_next"]]):
				is_reverse_curr=False #curr in correct order
			if(is_reverse_prev): prev_lin_def=self.__reverseLinearDefinition(prev_lin_def)
			if(is_reverse_curr): curr_lin_def=self.__reverseLinearDefinition(curr_lin_def)
		#assert prev_lin_def["node_id_next"]==curr_lin_def["node_id_prev"]
		return {"prev":prev_lin_def,"curr":curr_lin_def}

	def __reverseLinearDefinition(self,row):
		row["is_forward"]=False
		was_node_id_prev=row["node_id_prev"]
		was_node_id_next=row["node_id_next"]
		row["node_id_next"]=was_node_id_prev#flip order of nodes so that _from_ is always first, before _to_
		row["node_id_prev"]=was_node_id_next
		row["segment_id_index"]=row["segment_list_len"]-row["segment_id_index"]-1 #flip orientation of list_index
		row["segment_list"]=list(reversed(row["segment_list"]))
		return row

	#given a segment_id, find the two nodes around it
	#always returns is_forward=True order
	def getLinearDefinitionOfSegment(self,segment_id):
		for row in self.linear_definition:
			if(segment_id in row["segment_list"]):
				return {
				"node_id_prev":row["node_id_prev"],
				"node_id_next":row["node_id_next"],
				"is_forward":True,
				"segment_id_index":row["segment_list"].index(segment_id),
				"segment_id":segment_id,
				"segment_list":row["segment_list"],
				"segment_list_len":len(row["segment_list"])}
		return None

	#given two (precon: sequential) nodes (OR same node to find branch segment(s)), find the segment list
	#segment list will always be in order from prev_node to curr_node
	#returns {"is_forward":True/False,"segment_list":[segment_id,segment_id,segment_id]}
	#note: method will blindly return segments even if they are not accessible in the stated order (ex. branch_definition is_two_way=False)
	def getSegmentsBetweenNodes(self,prev_node,next_node):
		for row in self.linear_definition:
			this_prev=row["node_id_prev"]
			this_next=row["node_id_next"]
			this_segments=row["segment_list"]
			hit=False
			is_forward=True
			if(this_prev==prev_node and this_next == next_node):
				hit=True
			elif(this_next==prev_node and this_prev==next_node):
				hit=True
				is_forward=False
				this_segments=list(reversed(this_segments))
			if(hit):
				return {"is_forward":is_forward,"node_id_prev":this_prev,"node_id_next":this_next,"segment_list":this_segments}
		return None

	#given the previous and current segment id's, determine how far along the pod is between two nodes
	#precon: segment_ratio is between 0 and 1
	#precon: prev_segment_id and curr_segment_id are defined sequentially in self.linear_definition
	#returns dictionary with node id's and ratio of completion between endpoints
	# weighted for each Segment's total length (to ensure ratio increments linearly
	# with time as pod traverses between nodes)
	def getNodeStatus(self,prev_segment_id,curr_segment_id,segment_ratio):
		ordered_definition=getOrderedLinearDefinition(self,curr_segment_id,prev_segment_id)
		curr_node_definition=ordered_definition["curr"]
		curr_node_segment_list=curr_node_definition["segment_list"] #already ordered from prev to curr nodes
		progress_duration=0
		total_duration=0
		found_current_segment=False
		for this_segment_id in curr_node_segment_list:
			if(not found_current_segment):
				if(curr_segment_id==this_segment_id):
					found_current_segment=True
					progress_duration+=segment_ratio*self.__getSegmentDuration(this_segment_id)
				else:
					progress_duration+=self.__getSegmentDuration(this_segment_id)
			total_duration+=self.__getSegmentDuration(this_segment_id)
		node_progress=progress_duration/total_duration
		return {"prev_id_node":curr_node_definition["prev_id_node"],
				"curr_id_node":curr_node_definition["curr_id_node"],
				"node_progress":node_progress}

	#determine the number of ring assemblies in the given segment
	# used as a proxy to estimate time elapsed when traversing the given segment
	# info is combined in super-methods with other segment durations to
	# gauge progresss between nodes
	#note: currently "duration" is defined as a ring count, not a explict duration
	# though they are related through a constant rate of speed along all Segment curves
	def __getSegmentDuration(self,segment_id):
		for file_row in self.segment_definition:
			if(file_row["segment_id"]==segment_id):
				return len(file_row["ring_list"])

if __name__ == "__main__":
	import sys
	sys.path.append('/home/pi/Documents/EscapeRoom/src/util')
	#from ResourceManager import *
	#csv=ResourceManager.loadCSV(Maze.CONFIG_FILE_PATH+MAZE_CONFIG.LINEAR_DEFINITION.value["filename"])
	#print(csv)
	maze=Maze()
	maze.clean()
	print(maze.getSegmentsBetweenNodes(100,91))
	print(maze.getSegmentsBetweenNodes(91,100))
	print(maze.getSegmentsBetweenNodes(91,91))
	#print(maze.linear_definition)
	#print(maze.branch_definition)
	#print(maze.segment_definition)
	#print(maze.debris_definition)
	segments=maze.getSegmentIdAfter(3,4)
	#segments=maze.getSegmentIdAfter(2,3)
	print("Next segment: ",segments)

	#too involved to de-couple from embedded library imports in helper classes, so can only test on live system
	#segments=segments[0]
	#import numpy as np
	#is_branch=False
	#temp2=maze.getPopulatedSegment(segments["segment_id"],segments["is_forward"],is_branch,None,np.array([0,0,0]),np.eye(3),0)

max_stars_repo_name: scottalmond/EscapeRoom

max_issues_repo_name: scottalmond/EscapeRoom

max_forks_repo_name: scottalmond/EscapeRoom


text:
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import argparse
import os
import sys
import numpy as np
import tensorflow as tf
import random
from collections import OrderedDict
from six.moves import xrange
import json
import math
import time
import datetime
import model
import dag
#from utils import data_parallelism

_HEIGHT = 32
_WIDTH = 32
_DEPTH = 3
_NUM_DATA_FILES = 5

_WEIGHT_DECAY = 5e-4 #1e-4
_MOMENTUM = 0.9

_NUM_IMAGES = {
    'train': 45000,
    'valid': 5000,
    'test': 10000,
}

_TEST_BATCH_SIZE = 100

parser = argparse.ArgumentParser()

# Basic model parameters.
parser.add_argument('--data_dir', type=str, default='/tmp/cifar10_data',
                    help='The path to the CIFAR-10 data directory.')

parser.add_argument('--dataset', type=str, default='cifar10',
                    help='CIFAR-10, CIFAR-100.')

parser.add_argument('--model_dir', type=str, default='/tmp/cifar10_model',
                    help='The directory where the model will be stored.')

parser.add_argument('--seed', type=int, default=None,
                    help='Seed to use.')

parser.add_argument(
    '--data_format', type=str, default=None,
    choices=['channels_first', 'channels_last'],
    help='A flag to override the data format used in the model. channels_first '
         'provides a performance boost on GPU but is not always compatible '
         'with CPU. If left unspecified, the data format will be chosen '
         'automatically based on whether TensorFlow was built for CPU or GPU.')

def record_dataset(filenames, dataset, mode):
  """Returns an input pipeline Dataset from `filenames`."""
  if dataset == 'cifar10':
    record_bytes = _HEIGHT * _WIDTH * _DEPTH + 1
    return tf.data.FixedLengthRecordDataset(filenames, record_bytes)
  elif dataset == 'cifar100':
    record_bytes = _HEIGHT * _WIDTH * _DEPTH + 2
    return tf.data.FixedLengthRecordDataset(filenames, record_bytes)

def get_filenames(split, mode, data_dir, dataset):
  """Returns a list of filenames."""
  if dataset == 'cifar10':
    if not split:
      data_dir = os.path.join(data_dir, 'cifar-10-batches-bin')

    assert os.path.exists(data_dir), (
        'Run cifar10_download_and_extract.py first to download and extract the '
        'CIFAR-10 data.')

    if split:
      if mode == 'train':
        return [
          os.path.join(data_dir, 'train_batch_%d.bin' % i)
          for i in range(1, _NUM_DATA_FILES + 1)]
      elif mode == 'valid':
        return [os.path.join(data_dir, 'valid_batch.bin')]
      else:
        return [os.path.join(data_dir, 'test_batch.bin')]
    else:
      if mode == 'train':
        return [
          os.path.join(data_dir, 'data_batch_%d.bin' % i)
          for i in range(1, _NUM_DATA_FILES + 1)
    ]
      else:
        return [os.path.join(data_dir, 'test_batch.bin')]

  elif dataset == 'cifar100':
    data_dir = os.path.join(data_dir, 'cifar-100-binary')

    assert os.path.exists(data_dir)

    if mode == 'train':
      return [os.path.join(data_dir, 'train.bin')]
    else:
      return [os.path.join(data_dir, 'test.bin')]

def parse_record(raw_record, dataset):
  #Parse CIFAR image and label from a raw record.
  # Every record consists of a label followed by the image, with a fixed number
  # of bytes for each.
  if dataset == 'cifar10':
    label_bytes = 1
  elif dataset == 'cifar100':
    label_bytes = 2
  image_bytes = _HEIGHT * _WIDTH * _DEPTH
  record_bytes = label_bytes + image_bytes

  # Convert bytes to a vector of uint8 that is record_bytes long.
  record_vector = tf.decode_raw(raw_record, tf.uint8)

  # The first byte represents the label, which we convert from uint8 to int32
  # and then to one-hot.
  if dataset == 'cifar10':
    label = tf.cast(record_vector[0], tf.int32)
    label = tf.one_hot(label, 10)
  elif dataset == 'cifar100':
    label = tf.cast(record_vector[1], tf.int32)
    label = tf.one_hot(label, 100)

  # The remaining bytes after the label represent the image, which we reshape
  # from [depth * height * width] to [depth, height, width].
  depth_major = tf.reshape(
      record_vector[label_bytes:record_bytes], [_DEPTH, _HEIGHT, _WIDTH])

  # Convert from [depth, height, width] to [height, width, depth], and cast as
  # float32.
  image = tf.cast(tf.transpose(depth_major, [1, 2, 0]), tf.float32)

  return image, label

def preprocess_image(image, mode, cutout_size):
  """Preprocess a single image of layout [height, width, depth]."""
  if mode == 'train':
    # Resize the image to add four extra pixels on each side.
    image = tf.image.resize_image_with_crop_or_pad(
        image, _HEIGHT + 8, _WIDTH + 8)

    # Randomly crop a [_HEIGHT, _WIDTH] section of the image.
    image = tf.random_crop(image, [_HEIGHT, _WIDTH, _DEPTH])

    # Randomly flip the image horizontally.
    image = tf.image.random_flip_left_right(image)

  # Subtract off the mean and divide by the variance of the pixels.
  image = tf.image.per_image_standardization(image)

  if mode == 'train' and cutout_size is not None:
    mask = tf.ones([cutout_size, cutout_size], dtype=tf.int32)
    start = tf.random_uniform([2], minval=0, maxval=32, dtype=tf.int32)
    mask = tf.pad(mask, [[cutout_size + start[0], 32 - start[0]],
                        [cutout_size + start[1], 32 - start[1]]])
    mask = mask[cutout_size: cutout_size + 32,
                cutout_size: cutout_size + 32]
    mask = tf.reshape(mask, [32, 32, 1])
    mask = tf.tile(mask, [1, 1, 3])
    image = tf.where(tf.equal(mask, 0), x=image, y=tf.zeros_like(image))
  return image

def input_fn(split, mode, data_dir, dataset, batch_size, cutout_size, num_epochs=1):
  """Input_fn using the tf.data input pipeline for CIFAR-10 dataset.

  Args:
    mode: train, valid or test.
    data_dir: The directory containing the input data.
    batch_size: The number of samples per batch.
    num_epochs: The number of epochs to repeat the dataset.

  Returns:
    A tuple of images and labels.
  """
  data_set = record_dataset(get_filenames(split, mode, data_dir, dataset), dataset, mode)

  if mode == 'train':
      data_set = data_set.shuffle(buffer_size=50000)

  data_set = data_set.map(lambda x:parse_record(x, dataset), num_parallel_calls=16)
  data_set = data_set.map(
      lambda image, label: (preprocess_image(image, mode, cutout_size), label),
      num_parallel_calls=4)

  data_set = data_set.repeat(num_epochs)
  data_set = data_set.batch(batch_size)
  data_set = data_set.prefetch(10)
  iterator = data_set.make_one_shot_iterator()
  images, labels = iterator.get_next()

  return images, labels

def _log_variable_sizes(var_list, tag):
  """Log the sizes and shapes of variables, and the total size.

    Args:
      var_list: a list of varaibles
      tag: a string
  """
  name_to_var = {v.name: v for v in var_list}
  total_size = 0
  for v_name in sorted(list(name_to_var)):
    v = name_to_var[v_name]
    v_size = int(np.prod(np.array(v.shape.as_list())))
    tf.logging.info("Weight    %s\tshape    %s\tsize    %d",
      v.name[:-2].ljust(80),
      str(v.shape).ljust(20), v_size)
    total_size += v_size
  tf.logging.info("%s Total size: %d", tag, total_size)

def get_test_ops(x, y, params, reuse=False):
  with tf.device('/gpu:0'):
    inputs = tf.reshape(x, [-1, _HEIGHT, _WIDTH, _DEPTH])
    labels = y
    res = model.build_model(inputs, params, False, reuse)
    logits = res['logits']
    cross_entropy = tf.losses.softmax_cross_entropy(
      logits=logits, onehot_labels=labels)
    # Add weight decay to the loss.
    loss = cross_entropy + params['weight_decay'] * tf.add_n(
      [tf.nn.l2_loss(v) for v in tf.trainable_variables()])

    if 'aux_logits' in res:
      aux_logits = res['aux_logits']
      aux_loss = tf.losses.softmax_cross_entropy(
        logits=aux_logits, onehot_labels=labels, weights=params['aux_head_weight'])
      loss += aux_loss

    predictions = tf.argmax(logits, axis=1)
    labels = tf.argmax(y, axis=1)
    test_accuracy = tf.reduce_mean(tf.cast(tf.equal(predictions, labels), dtype=tf.float32))
    return loss, test_accuracy

def test(params):
  g = tf.Graph()
  with g.as_default(), tf.device('/cpu:0'):
    tf.set_random_seed(params['seed'])
    x_test, y_test = input_fn(False, 'test', params['data_dir'], params['dataset'], 100, None, None)
    _log_variable_sizes(tf.trainable_variables(), 'Trainable Variables')
    test_loss, test_accuracy = get_test_ops(x_test, y_test, params, True)
    tf.logging.info('Starting Session')
    config = tf.ConfigProto(allow_soft_placement=True)
    with tf.train.SingularMonitoredSession(
      config=config, checkpoint_dir=params['model_dir']) as sess:
      test_ops = [
        test_loss, test_accuracy
      ]
      test_start_time = time.time()
      test_loss_list = []
      test_accuracy_list = []
      for _ in range(_NUM_IMAGES['test'] // 100):
        test_loss_v, test_accuracy_v = sess.run(test_ops)
        test_loss_list.append(test_loss_v)
        test_accuracy_list.append(test_accuracy_v)
      test_time = time.time() - test_start_time
      log_string =  "Evaluation on test data\n"
      log_string += "loss={:<6f} ".format(np.mean(test_loss_list))
      log_string += "test_accuracy={:<8.6f} ".format(np.mean(test_accuracy_list))
      log_string += "secs={:<10.2f}".format((test_time))
      tf.logging.info(log_string)

def get_params():
  if not os.path.exists(os.path.join(FLAGS.model_dir, 'hparams.json')):
    raise ValueError('model_dir/hparams.json does not exist ')
  with open(os.path.join(FLAGS.model_dir, 'hparams.json'), 'r') as f:
    params = json.load(f)
  params.update({
    'model_dir': FLAGS.model_dir,
    'data_dir': FLAGS.data_dir,
    'data_format': FLAGS.data_format,
    'dataset': FLAGS.dataset,
  })
  return params

def main(unused_argv):
  # Using the Winograd non-fused algorithms provides a small performance boost.
  os.environ['TF_ENABLE_WINOGRAD_NONFUSED'] = '1'
  params = get_params()
  test(params)

if __name__ == '__main__':
  tf.logging.set_verbosity(tf.logging.INFO)
  FLAGS, unparsed = parser.parse_known_args()
  tf.app.run(argv=[sys.argv[0]] + unparsed)

ext: py

reason: import numpy


text:
import os, sys, math, argparse, time
import torch
import torch.optim as optim
from torch.autograd import Variable
import torch.backends.cudnn as cudnn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from copy import copy
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import matplotlib.gridspec as gridspec

from scipy.ndimage.interpolation import rotate
from itertools import product, cycle
import pickle

from resnet18 import *
from metaopt.util import *
from metaopt.util_ml import *
from metaopt.optimizer import SGD_Multi_LR, SGD_Quotient_LR
from metaopt.cifar.main import feval, meta_update, get_grad_valid, update_optimizer_hyperparams, load_cifar10, parse_args

TRAIN=0
VALID=1
TEST =2

torch.manual_seed(5)
#torch.manual_seed(0)

def main(args, trial=0, ifold=0, device='cuda', quotient=None):

    dataset = load_cifar10(args)
    print('Model Type: %s Opt Type: %s meta-lr %f lr %f l2 %f, Update Freq %d Reset Freq %d |Nvl| %d Epoch %d' \
            % (args.model_type, args.opt_type, args.mlr,  args.lr, \
                args.lambda_l2, args.update_freq, args.reset_freq, args.valid_size, args.num_epoch))

    # Model
    print('==> Building model..')
    ## Initialize Model and Optimizer
    if args.model_type == 'qrez18':
        model = QResNet18(args.lr, args.lambda_l2, quotient=args.quotient)
        optimizer = SGD_Quotient_LR(model.parameters(), lr=args.lr, weight_decay=args.lambda_l2, quotient=quotient)
    elif args.model_type == 'arez18':
        model = AResNet18(args.lr, args.lambda_l2)
        optimizer = SGD_Multi_LR(model.parameters(), lr=args.lr, weight_decay=args.lambda_l2)
    elif args.model_type == 'rez18':
        model = ResNet18(args.lr, args.lambda_l2)
        if args.opt_type == 'adam':
            optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.lambda_l2, betas=(args.beta1,args.beta2))
        else:
            optimizer = optim.SGD(model.parameters(), lr=args.lr, weight_decay=args.lambda_l2)
        #optimizer = optim.SGD(model.parameters(), lr=args.lr, weight_decay=args.lambda_l2)
    elif args.model_type == 'rez18drop':
        model = ResNet18_Drop(args.lr, args.lambda_l2)
        optimizer = optim.SGD(model.parameters(), lr=args.lr, weight_decay=args.lambda_l2)
    elif args.model_type == 'rez50':
        model = ResNet50(args.lr, args.lambda_l2)
        optimizer = optim.SGD(model.parameters(), lr=args.lr, weight_decay=args.lambda_l2)

    #Update optimizer with new eta
    optimizer = update_optimizer_hyperparams(args, model, optimizer)

    if args.is_cuda:
        model = model.to(device)
        model = torch.nn.DataParallel(model)
        cudnn.benchmark = True

    if args.sampler_type is None:
        os.makedirs('%s/exp/cifar10/ood/' % args.save_dir, exist_ok=True)
        os.makedirs('%s/exp/cifar10/ood/%s/' % (args.save_dir, args.ood_type), exist_ok=True)
        os.makedirs('%s/exp/cifar10/ood/%s/mlr%f_lr%f_l2%f/' % (args.save_dir, args.ood_type, args.mlr, args.lr, args.lambda_l2), exist_ok=True)
        fdir = '%s/exp/cifar10/ood/%s/mlr%f_lr%f_l2%f/%s_%depoch_%dvlbz_%s_%dupdatefreq_%dresetfreq_%dupdatelabmda_fold%d/' \
                % (args.save_dir, args.ood_type, args.mlr, args.lr, args.lambda_l2, args.model_type, \
                    args.num_epoch, args.batch_size_vl, args.opt_type, \
                    args.update_freq, args.reset_freq, args.update_lambda, ifold)

        if args.opt_type == 'sgd_step':
            fdir = fdir.rstrip('/') + '_stepsize%d/' % args.step_size
        elif args.opt_type == 'sgd_expstep':
            fdir = fdir.rstrip('/') + '_gamma%f/' % args.gamma

        if quotient is not None:
            fdir = fdir.rstrip('/') + '_quotient%d/' % quotient

    else:
        os.makedirs('%s/cifar10/trial%d/' % (args.save_dir, trial), exist_ok=True)
        os.makedirs('%s/cifar10/trial%d/%s/' % (args.save_dir, trial, args.sampler_type), exist_ok=True)
        os.makedirs('%s/cifar10/trial%d/%s/mlr%f_lr%f_l2%f/' % (args.save_dir, trial, args.sampler_type, args.mlr, args.lr, args.lambda_l2), exist_ok=True)
        fdir = '%s/cifar10/trial%d/%s/mlr%f_lr%f_l2%f/%s_%depoch_%dvlbz_%s_%dupdatefreq_%dresetfreq_%dupdatelabmda_fold%d' \
                % (args.save_dir, trial, args.sampler_type, args.mlr, args.lr, args.lambda_l2, args.model_type, \
                    args.num_epoch, args.batch_size_vl, args.opt_type, \
                    args.update_freq, args.reset_freq, args.update_lambda, ifold)
    os.makedirs(fdir, exist_ok=True)
    os.makedirs(fdir+'/checkpoint/', exist_ok=True)
    args.fdir = fdir
    print(args.fdir)

    ## Train
    Wn_list, l2_list, lr_list, dFdlr_list, dFdl2_list, gang_list, tr_epoch, vl_epoch, te_epoch,\
                                tr_acc_list, te_acc_list, \
                                tr_loss_list, vl_loss_list, te_loss_list \
                                = train(args, dataset, model, optimizer, is_cuda=args.is_cuda)

    if args.save:
        os.makedirs(fdir, exist_ok=True)
        np.save(fdir+'Wn', Wn_list)
        np.save(fdir+'lr', lr_list)
        np.save(fdir+'l2', l2_list)
        np.save(fdir+'gang_list', gang_list)
        np.save(fdir+'dFdlr_list', dFdlr_list)
        np.save(fdir+'dFdl2_list', dFdl2_list)
        np.save(fdir+'tr_epoch', tr_epoch)
        np.save(fdir+'vl_epoch', vl_epoch)
        np.save(fdir+'te_epoch', te_epoch)
        np.save(fdir+'tr_loss', tr_loss_list)
        np.save(fdir+'vl_loss', vl_loss_list)
        np.save(fdir+'te_loss', te_loss_list)
        np.save(fdir+'tr_acc', tr_acc_list)
        np.save(fdir+'te_acc', te_acc_list)

    print('Final test loss %f' % te_loss_list[-1])
    print(type(te_loss_list[-1]))
    return te_loss_list[-1]

def manipulate(args, data, target, epoch, num_quotient=5, ood_type=None):

    if ood_type is None: ood_type = args.ood_type

    freq = args.num_epoch / num_quotient
    if ood_type == 'target':
        if epoch > freq and epoch < freq * 2:
            target = (target + 1) % args.ydim
        elif epoch > freq * 2 and epoch < freq * 3:
            target = (target + 2) % args.ydim
        elif epoch > freq * 3 and epoch < freq * 4:
            target = (target + 3) % args.ydim

    elif ood_type == 'hflip':
        import pdb; pdb.set_trace()

    elif ood_type == 'rotate':
        N,C,D1,D2 = data.shape
        angle = 0
        if epoch > freq and epoch < freq * 2:
            angle=90
        elif epoch > freq * 2 and epoch < freq * 3:
            angle=180
        elif epoch > freq * 3 and epoch < freq * 4:
            angle=270

        if angle > 0:
            data = data.permute(0,2,3,1)
            data = np.asarray([rotate(data[i].squeeze(), angle, reshape=False) for i in range(N)])
            data = data.reshape([N,D1,D2,C]).transpose(0,3,1,2)
            data = torch.FloatTensor(data)
            #if args.is_cuda: data.cuda()

        #new_data3 = np.asarray([rotate(data[0].squeeze(), 90, reshape=False)])
        #new_data2 = np.asarray([rotate(data[0].squeeze(), 180, reshape=False)])
        #new_data1 = np.asarray([rotate(data[0].squeeze(), 270, reshape=False)])
        #plt.figure(figsize=(4, 1))
        #gs = gridspec.GridSpec(1, 12)
        #gs.update(wspace=0, hspace=0)
        #ax = plt.subplot(gs[0])
        #ax.imshow(new_data3.squeeze().transpose(1,2,0),interpolation='bicubic');
        #ax.axis('off')
        #ax = plt.subplot(gs[1])
        #ax.imshow(new_data2.squeeze().transpose(1,2,0),interpolation='bicubic');
        #ax.axis('off')
        #ax = plt.subplot(gs[2])
        #ax.imshow(new_data1.squeeze().transpose(1,2,0),interpolation='bicubic');
        #ax = plt.axis('off');
        #ax = plt.subplot(gs[3])
        #
        #ax.imshow(data[0].data.cpu().numpy().transpose(1,2,0),interpolation='bicubic');
        #ax = plt.axis('off');
        #plt.savefig('figs/cifar_rotation.png')
        #import pdb; pdb.set_trace()
    return data, target

def train(args, dataset, model, optimizer, saveF=0, is_cuda=1):

    start_time0 = time.time()
    opt_type = args.opt_type
    if 'step' in opt_type:
        lrsch_type = opt_type.split('_')[-1]
        if 'sgd_expstep' == opt_type:
            scheduler = lr_scheduler_init(optimizer, lrsch_type, gamma=args.gamma)
        elif 'sgd_step' == opt_type:
            scheduler = lr_scheduler_init(optimizer, lrsch_type, step_size=args.step_size)
        else:
            scheduler = lr_scheduler_init(optimizer, lrsch_type, N=args.num_epoch+1)
        #scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=1000, gamma=0.95)

    counter = 0
    lr_list, l2_list = [], []
    dFdlr_list, dFdl2_list, Wn_list, gang_list = [], [], [], []
    tr_epoch, tr_loss_list, tr_acc_list = [], [], []
    vl_epoch, vl_loss_list, vl_acc_list = [], [], []
    te_epoch, te_loss_list, te_acc_list = [], [], []
    for epoch in range(args.num_epoch+1):

        if epoch % 10 == 0:
            te_losses, te_accs = [], []
            for batch_idx, (data, target) in enumerate(dataset[TEST]):

                data, target = manipulate(args, data, target, epoch, num_quotient=5)
                data, target = to_torch_variable(data, target, is_cuda)
                _, loss, accuracy, _, _ = feval(data, target, model, optimizer, mode='eval', is_cuda=is_cuda)
                te_losses.append(loss)
                te_accs.append(accuracy)
            te_epoch.append(epoch)
            te_loss_list.append(np.mean(te_losses))
            te_acc_list.append(np.mean(te_accs))

            print('Valid Epoch: %d, Loss %f Acc %f' %
                (epoch, np.mean(te_losses), np.mean(te_accs)))

            if '_step' in opt_type:
                scheduler.step()
                model.module.eta = optimizer.param_groups[0]['lr']

        if 'expstep' in opt_type:
            scheduler.step()
            model.module.eta = optimizer.param_groups[0]['lr']

        start_time = time.time()
        for batch_idx, (data, target) in enumerate(dataset[TRAIN]):

            data, target = manipulate(args, data, target, epoch, num_quotient=5)
            data, target = to_torch_variable(data, target, is_cuda)
            if 'step' in opt_type:
                model, loss, accuracy, output, noise = feval(data, target, model, optimizer, \
                                is_cuda=is_cuda, mode='train', opt_type=opt_type)
            else:
                model, loss, accuracy, output, noise = feval(data, target, model, optimizer, \
                                is_cuda=is_cuda, mode='meta-train', opt_type=opt_type)
            tr_epoch.append(counter)
            tr_loss_list.append(loss)
            tr_acc_list.append(accuracy)

            if args.reset_freq > 0 and counter % args.reset_freq == 0:
                model_ = model.module if 'DataParallel' in str(type(model)) else model
                model_.reset_jacob(is_cuda)

            if counter % args.update_freq == 0 and 'step' not in opt_type and args.mlr != 0.0:
                data_vl, target_vl = next(dataset[VALID])
                data_vl, target_vl = manipulate(args, data_vl, target_vl, epoch, num_quotient=5)
                data_vl, target_vl = to_torch_variable(data_vl, target_vl, is_cuda)
                model, loss_vl, optimizer = meta_update(args, data_vl, target_vl, data, target, model, optimizer, noise, is_cuda=is_cuda)
                vl_epoch.append(counter)
                vl_loss_list.append(loss_vl.item())
            counter += 1

        end_time = time.time()
        if epoch == 0: print('Single epoch timing %f' % ((end_time-start_time) / 60))

        if epoch % args.checkpoint_freq == 0:
            os.makedirs(args.fdir+ '/checkpoint/', exist_ok=True)
            save(model, args.fdir+ '/checkpoint/epoch%d' % epoch)

        lambda_str = str(model.module.lambda_l2) if args.update_lambda else 'Fixed'
        #if args.update_lambda: lambda_str = str(model_.lambda_l2)  #else 'Fixed'
        model_ = model.module if 'DataParallel' in str(type(model)) else model
        fprint = 'Train Epoch: %d, Tr Loss %f Vl loss %f Acc %f Eta %s, L2 %s, |dFdlr| %.2f |dFdl2| %.2f |G| %.4f |G_vl| %.4f Gang %.3f |W| %.2f'
        print(fprint % (epoch, np.mean(tr_loss_list[-100:]), \
                        np.nanmean(vl_loss_list[-100:]), \
                        np.nanmean(tr_acc_list[-100:]), \
                        str(model_.eta), str(model_.lambda_l2), \
                        model_.dFdlr_norm, model_.dFdl2_norm,\
                        model_.grad_norm,  model_.grad_norm_vl, \
                        model_.grad_angle, model_.param_norm))

        Wn_list.append(model_.param_norm)
        dFdlr_list.append(model_.dFdlr_norm)
        dFdl2_list.append(model_.dFdl2_norm)
        if args.model_type == 'arez18' or args.model_type == 'qrez18':
            lr_list.append(model_.eta.copy())
            l2_list.append(model_.lambda_l2.copy())
        else:
            lr_list.append(model_.eta)
            l2_list.append(model_.lambda_l2)
        gang_list.append(model_.grad_angle)

    Wn_list = np.asarray(Wn_list)
    l2_list = np.asarray(l2_list)
    lr_list = np.asarray(lr_list)
    dFdlr_list = np.asarray(dFdlr_list)
    dFdl2_list = np.asarray(dFdl2_list)
    tr_epoch = np.asarray(tr_epoch)
    vl_epoch = np.asarray(vl_epoch)
    te_epoch = np.asarray(te_epoch)
    tr_acc_list = np.asarray(tr_acc_list)
    te_acc_list = np.asarray(te_acc_list)
    tr_loss_list = np.asarray(tr_loss_list)
    vl_loss_list = np.asarray(vl_loss_list)
    te_loss_list = np.asarray(te_loss_list)
    gang_list = np.asarray(gang_list)

    end_time0 = time.time()
    print('Total training timing %f' % ((end_time0-start_time0) / 3600))

    return Wn_list, l2_list, lr_list, dFdlr_list, dFdl2_list, gang_list, \
                tr_epoch, vl_epoch, te_epoch, \
                tr_acc_list, te_acc_list, tr_loss_list, vl_loss_list, te_loss_list

def feval(data, target, model, optimizer, mode='eval', is_cuda=0, opt_type='sgd', N=50000):

    if mode == 'eval':
        model.eval()
        with torch.no_grad():
            output = model(data)
    else:
        model.train()
        optimizer.zero_grad()
        output = model(data)

    # Compute Loss
    loss = F.nll_loss(output, target)
    #criterion = nn.CrossEntropyLoss()
    #loss = criterion(output, target)
    pred = output.argmax(dim=1, keepdim=True).flatten()  # get the index of the max log-probability
    accuracy = pred.eq(target).float().mean()

    noise = None
    if 'train' in mode:
        loss.backward()

        if opt_type == 'sgld':
            for i,param in enumerate(model.parameters()):

                noise = torch.randn(size=param.shape)
                model_ = model.module if 'DataParallel' in str(type(model)) else model
                if type(model_.eta) == type(np.array([])):
                    eps = np.sqrt(model_.eta[i]*2/ N) * noise  if model_.eta[i] > 0 else 0 * noise
                else:
                    eps = np.sqrt(model_.eta*2/ N) * noise  if model_.eta > 0 else 0 * noise
                eps = to_torch_variable(eps, is_cuda=is_cuda)
                param.grad.data = param.grad.data + eps.data

        optimizer.step()

    elif 'grad' in mode:
        loss.backward()

    return model, loss.item(), accuracy.item(), output, noise

if __name__ == '__main__':

    args = parse_args()
    is_cuda = args.is_cuda
    device = 'cuda' if is_cuda else 'cpu'

    main(args)

reason: import numpy,from scipy


text:
import numpy as np

from mygrad.operation_base import Operation
from mygrad import Tensor

__all__ = ["huber_loss"]

class HuberLoss(Operation):
    """ Returns the Huber loss (smooth L1).

    Extended Description
    --------------------
    The Huber loss is given by

    .. math::
        L_\delta(x, y) = \frac{1}{N}\sum\limits_1^N \bigl\{ \begin{array}{l l}
            \frac{(x_i - y_i)^2}{2} & |x_i - y_i| \leq \delta\\
            \delta|x_i - y_i| - \frac{\delta}{2} & |x_i - y_i| > \delta\end{array}
    """

    scalar_only = True

    def __call__(self, outputs, targets, delta=1):
        """
        Parameters
        ----------
        outputs : mygrad.Tensor, shape=(N, any)
            The output for each of the N pieces of data.

        targets : Union[mygrad.Tensor, numpy.ndarray], shape=(N, any)
            The target for each datum.

        delta : Real > 0, optional (default=1)
            The value under which to use a squared error.

        Returns
        -------
        numpy.ndarray
            The average Huber loss.
        """
        if delta <= 0:
            raise ValueError("Delta must be > 0")

        if isinstance(targets, Tensor):
            targets = targets.data

        self.variables = (outputs,)
        outs = outputs.data

        diff = outs - targets
        sign = np.sign(diff)
        np.abs(diff, out=diff)

        loss = diff.copy()
        np.multiply(loss, diff / 2, where=(diff < 1), out=loss)
        loss[diff >= 1] -= 0.5 * delta
        loss[diff >= 1] *= delta

        self.back = np.where(diff < 1, outs - targets, delta * sign) / outs.size

        return np.mean(loss)

    def backward_var(self, grad, index, **kwargs):
        return grad * self.back

def huber_loss(x, y, *, delta=1):
    """ Returns the Huber loss (smooth L1).

    Parameters
    ----------
    outputs : mygrad.Tensor, shape=(N, any)
        The output for each of the N pieces of data.

    targets : Union[mygrad.Tensor, numpy.ndarray], shape=(N, any)
        The target for each datum.

    delta : Real > 0, optional (default=1)
        The value under which to use a squared error.

    Returns
    -------
    mygrad.Tensor, shape=()
        The average Huber loss.

    Extended Description
    --------------------
    The Huber loss is given by

    .. math::
        L_\delta(x, y) = \frac{1}{N}\sum\limits_1^N \bigl\{ \begin{array}{l l}
            \frac{(x_i - y_i)^2}{2} & |x_i - y_i| \leq \delta\\
            \delta|x_i - y_i| - \frac{\delta}{2} & |x_i - y_i| > \delta\end{array}

    """
    return Tensor._op(HuberLoss, x, op_args=(y, delta))

max_stars_repo_name: davidmascharka/MyNN

max_issues_repo_name: davidmascharka/MyNN

max_forks_repo_name: davidmascharka/MyNN


text:
"""
Examples:
  To smear an already smeared spectrum with a light yield of 200 to a
  a light yield of 190 then the following lines are required::

    >>> smearer = smear.SmearEnergySmearLY()
    >>> ly = smearer.calc_smear_ly(190., cur_ly=200.)
    >>> smearer.set_resolution(ly)
    >>> smeared_spec = smearer.weighted_smear(spectrum)

.. note:: Similar methods are available in all other smearing classes.
"""
import numpy
import itertools
import copy

class Smear(object):
    """ The base class for smearing spectra.

    Args:
      name (string): The name of the smearing class.

    Attributes:
      _name (string): name of the smeaing class.
      _num_sigma (float): The width of the window in terms of number of sigma
        you wish to apply weights to.
    """

    def __init__(self, name):
        """ Initialise the Smear class by seeding the random number generator.
        """
        numpy.random.seed()
        self._name = name
        self._num_sigma = 5.

    def calc_gaussian(self, x, mean, sigma):
        """ Calculates the value of a gaussian whose integral is equal to
          one at position x with a given mean and sigma.

          Args:
            x : Position to calculate the gaussian
            mean : Mean of the gaussian
            sigma : Sigma of the gaussian

          Returns:
            float: Value of the gaussian at the given position
        """
        return (numpy.exp(-(x - mean) ** 2 / (2 * sigma ** 2)) /
                (sigma*numpy.sqrt(2 * numpy.pi)))

    def get_bin_mean(self, low, bin, width):
        """ Calculates the mean value of a bin.

        Args:
          low (float): The lower bound value of the parameter.
          bin (int): The number of the bin you wish to calculate the mean of.
          width (float): The width of the bin.

        Returns:
          float: The mean value of the bin.
        """
        return low + (bin + 0.5)*width

    def get_num_sigma(self):
        """ Returns the width of the window in terms of number of sigma
          you wish to apply weights to.

        Returns:
          float: The number of sigma.
        """
        return self._num_sigma

    def set_num_sigma(self, num_sigma):
        """ Sets the width of the window in terms of number of sigma
          you wish to apply weights to.

        Args:
          num_sigma (float): The number of sigma you wish to apply weights to.
        Raises:
          ValueError: If the number of sigma is zero or negative.
        """
        if (num_sigma > 0.):
            self._num_sigma = float(num_sigma)
        else:
            raise ValueError("%s is an invalid num_sigma. Value must be "
                             "greater than zero." % num_sigma)

    def get_bounds(self, mean, sigma):
        """ Calculates the boundaries you wish to apply the smearing
          weights to.

        Args:
          mean (float): The mean value you are smearing.
          sigma (float): The sigma of the gaussian you are using to smear.

        Returns:
          tuple: First value of the tuple is the lower bound. The second is
            the upper bound.
        """
        low = mean - self._num_sigma*sigma
        high = mean + self._num_sigma*sigma
        return low, high

class EnergySmearLY(Smear):
    """ The class which smears energy. It accepts resolution in terms of light
      yield (LY) in units of NHit per MeV.

    Args:
      poisson (bool): If True, use poisson smearing.

    Attributes:
      _light_yield (float): The light yield of the scintillator in NHits per
        MeV.
      _poisson_smear (Bool): True if poisson smearing is to be applied. False
        if gaussian smearing is to be applied.
    """

    def __init__(self, poisson=True):
        """ Initialises the class.
        """
        super(EnergySmearLY, self).__init__("energy_light_yield")
        self._poisson_smear = poisson
        self._light_yield = 200  # Nhit/MeV
        self._log_factorial = {}

    def calc_poisson_energy(self, x, lamb):
        """ Calculates the value of a poisson whose integral is equal to
        one at position x with a given lambda value.

        Args:
          x : Number of events
          lamb : Lambda of the poisson

        Returns:
          float: The value of the poisson at the given position
        """
        photons = int(x*self._light_yield)
        expected = lamb*self._light_yield
        if photons not in self._log_factorial:
            self._log_factorial[photons] = (
                numpy.sum(numpy.log(numpy.arange(1, (photons+1)))))
        log_pois = (photons*numpy.log(expected) -
                    self._log_factorial[photons] -
                    expected)
        return numpy.exp(log_pois)

    def calc_smear_ly(self, new_ly, cur_ly=None):
        """Calculates the value of light yield (ly) required to smear a
          data set which has already been smeared with a light yield of cur_ly
          to achieve a smeared data set with a new light yield of new_ly.

        Args:
          new_ly (float): The value of light yield wanted for the smeared PDF.
          cur_ly (float, optional): Current value of light yield the PDF
            has been convolved with from the true value PDF.

        Raises:
          ValueError: If new_ly is smaller than cur_sigma. Can't smear to
            higher light yields (smaller sigmas)

        Returns:
          float: The value of light yield needed to smear the current
            PDF to obtain a new light yield: new_ly.
        """
        if not cur_ly:
            cur_ly = self.get_resolution()
        if new_ly > cur_ly:
            raise ValueError("New light yield must be smaller than the"
                             "current light yield. cur_ly: %s. new_ly: %s."
                             % (cur_ly, new_ly))
        return new_ly*cur_ly/(cur_ly-new_ly)

    def get_resolution(self):
        """ Returns the light yield.

        Returns:
          float: The light yield.
        """
        return self._light_yield

    def get_sigma(self, energy):
        """ Calculates sigma at a given energy.

        Args:
          energy (float): Energy value of data point(s)

        Returns:
          float: Sigma equivalent to sqrt(energy/_light_yield)
        """
        return numpy.sqrt(energy/self._light_yield)

    def set_resolution(self, light_yield):
        """ Sets the light yield

        Args:
          light_yield (float): The value you wish to set the light yield to.

        Raises:
          ValueError: If the light yield is zero or negative.
        """
        if light_yield > 0.:
            self._light_yield = float(light_yield)
        else:
            raise ValueError("%s is an invalid light yield. Light yield "
                             "must be greater than zero.")

    def weighted_smear(self, spectrum, par="energy_mc"):
        """ Smears the energy of a :class:`spectra.Spectra` by
        calculating a Gaussian PDF for each bin. Weights are then
        applied to a window of width specified by the number of sigma
        depending on the value of the Gaussian PDF at the mean of the
        bin.

        Args:
          spectrum (:class:`spectra.Spectra`): Spectrum you wish to
            smear.
          par (string, optional): The name of the parameter you wish
            to smear. The default is energy_mc.

        Raises:
          IndexError: If par is not in the specta config.

        Returns:
          :class:`spectra.Spectra`: The smeared spectrum
        """
        if par not in spectrum.get_config().get_pars():
            raise IndexError("%s is not a parameter in the spectrum" % par)
        idx = spectrum.get_config().get_index(par)
        bins = []
        lows = []
        widths = []
        par_names = []
        for par_name in spectrum.get_config().get_pars():
            bins.append(range(spectrum.get_config().get_par(par_name)._bins))
            lows.append(spectrum.get_config().get_par(par_name)._low)
            widths.append(spectrum.get_config().get_par(par_name).get_width())
            par_names.append(par_name)
        smeared_spec = copy.copy(spectrum)
        smeared_spec._name = spectrum._name + "_ly" + str(self._light_yield)
        smeared_spec._data = numpy.zeros(spectrum._data.shape)
        for bin in itertools.product(*bins):
            entries = spectrum._data[bin]
            if entries:
                data = {}
                low = None
                high = None
                for i in range(len(bin)):
                    mean = self.get_bin_mean(lows[i], bin[i], widths[i])
                    if i == idx:
                        sigma = self.get_sigma(mean)
                        low, high = self.get_bounds(mean, sigma)
                        low = spectrum.get_config().get_par(
                            par).round(low - 0.5 * widths[i]) + 0.5 * widths[i]
                        high = spectrum.get_config().get_par(
                            par).round(high + 0.5 * widths[i]) + \
                            0.5 * widths[i]
                        if low < spectrum.get_config().get_par(par)._low:
                            low = spectrum.get_config().get_par(par)._low + \
                                0.5 * widths[i]
                        if high > spectrum.get_config().get_par(par)._high:
                            high = spectrum.get_config().get_par(par)._high - \
                                0.5 * widths[i]
                        weights = []
                        for energy in numpy.arange(low, high, widths[i]):
                            if self._poisson_smear is True:
                                weights.append(self.calc_poisson_energy(energy,
                                                                        mean))
                            else:
                                weights.append(self.calc_gaussian(energy,
                                                                  mean,
                                                                  sigma))
                    else:
                        data[par_names[i]] = mean
                total_weight = sum(weights)
                i = 0
                for energy in numpy.arange(low, high, widths[idx]):
                    data[par] = energy
                    smeared_spec.fill(weight=entries*weights[i]/total_weight,
                                      **data)
                    i += 1
        smeared_spec._raw_events = spectrum._raw_events
        return smeared_spec

    def random_smear(self, spectrum, par="energy_mc"):
        """ Smears the energy of a :class:`spectra.Spectra` by
        generating a number of random points from  Gaussian PDF generated
        from that bins mean value and the corresponding sigma. The number
        of points generated is equivalent to the number of entries in that
        bin.

        Args:
          spectrum (:class:`spectra.Spectra`): Spectrum you wish to
            smear.
          par (string, optional): The name of the parameter you wish to
            smear. The default is energy_mc.

        Raises:
          IndexError: If par is not in the specta config.

        Returns:
          :class:`spectra.Spectra`: The smeared spectrum
        """
        if par not in spectrum.get_config().get_pars():
            raise IndexError("%s is not a parameter in the spectrum" % par)
        idx = spectrum.get_config().get_index(par)
        bins = []
        lows = []
        widths = []
        par_names = []
        for par_name in spectrum.get_config().get_pars():
            bins.append(range(spectrum.get_config().get_par(par_name)._bins))
            lows.append(spectrum.get_config().get_par(par_name)._low)
            widths.append(spectrum.get_config().get_par(par_name).get_width())
            par_names.append(par_name)
        smeared_spec = copy.copy(spectrum)
        smeared_spec._name = spectrum._name + "_ly" + str(self._light_yield)
        smeared_spec._data = numpy.zeros(spectrum._data.shape)
        for bin in itertools.product(*bins):
            entries = int(spectrum._data[bin])
            if entries:
                data = {}
                for i in range(len(bin)):
                    mean = self.get_bin_mean(lows[i], bin[i], widths[i])
                    if i == idx:
                        mean_e = mean
                        sigma = self.get_sigma(mean)
                    else:
                        data[par_names[i]] = mean
                for i in range(entries):
                    if self._poisson_smear is True:
                        photons = (numpy.fabs
                                   (numpy.random.poisson(mean_e *
                                                         self._light_yield)))
                        data[par] = photons / float(self._light_yield)
                    else:
                        data[par] = (numpy.fabs
                                     (numpy.random.normal(mean_e, sigma)))
                    try:
                        smeared_spec.fill(**data)
                    except ValueError:
                        print "WARNING: Smeared energy out of bounds. Skipping"
        smeared_spec._raw_events = spectrum._raw_events
        return smeared_spec

class EnergySmearRes(Smear):
    """ Allows you to smear directly by supplied energy resolution
      (in :math:`\sqrt{MeV}`).

    Inherits from :class:`Smear`

    Args:
      poisson (bool): If True, use poisson smearing.

    Attributes:
      _energy_resolution (float): Energy resolution in :math:`\sqrt{MeV}`
        e.g. 0.05 for :math:`\sigma = 5\%/\sqrt{E[MeV]}`.
      _poisson_smear (Bool): True if poisson smearing is to be applied.
        False if gaussian smearing is to be applied.
    """

    def __init__(self, poisson=True):
        """ Initialise the class
        """
        super(EnergySmearRes, self).__init__("energy_resolution")
        self._poisson_smear = poisson
        self._light_yield = 200  # Nhit/MeV
        self._log_factorial = {}

    def calc_poisson_energy(self, x, lamb):
        """ Calculates the value of a poisson whose integral is equal to
        one at position x with a given lambda value.

        Args:
          x : Number of events
          lamb : Lambda of the poisson

        Returns:
          float: The value of the poisson at the given position
        """
        photons = int(x*self._light_yield)
        expected = lamb*self._light_yield
        if photons not in self._log_factorial:
            self._log_factorial[photons] = (
                numpy.sum(numpy.log(numpy.arange(1, (photons+1)))))
        log_pois = (photons*numpy.log(expected) -
                    self._log_factorial[photons] -
                    expected)
        return numpy.exp(log_pois)

    def calc_smear_resoluton(self, new_res, cur_res=None):
        """Calculates the value of resolution required to smear a data set
          which has already been smeared with a resolution of cur_res to
          achieve a new resolution of new_res.

        Args:
          new_res (float): The value of resolution wanted for the smeared PDF.
          cur_res (float, optional): Current value of resolution the PDF
            has been convolved with from the true value PDF.

        Raises:
          ValueError: If new_res is smaller than cur_sigma. Can't smear to
            higher resolutions (smaller sigmas)

        Returns:
          float: The value of resolution needed to smear the current
            PDF to obtain a new resolution with sigma value new_res.
        """
        if not cur_res:
            cur_res = self.get_resolution()
        if new_res < cur_res:
            raise ValueError("New resolution must be larger than the"
                             "current resolution. cur_res: %s. new_res: %s."
                             % (cur_res, new_res))
        return numpy.fabs(numpy.sqrt(new_res**2 - cur_res**2))

    def get_resolution(self):
        """ Get the energy resolution

        Returns:
          float: Energy resolution in  :math:`\sqrt{MeV}`
          e.g. 0.05 for :math:`\sigma = 5\%/\sqrt{E[MeV]}`
        """
        return self._resolution

    def get_sigma(self, energy):
        """ Calculates sigma at a given energy.

        Args:
          energy (float): Energy value of data point(s)

        Returns:
          float: Sigma (MeV) equivalent to energy_resolution *
            :math:`\sqrt{energy}`
        """
        return self._resolution * numpy.power(energy, (1. / 2.))

    def set_resolution(self, resolution):
        """ Set the energy resolution in :math:`\sqrt{MeV}`
        e.g. 0.05 for :math:`\sigma = 5\%/\sqrt{E[MeV]}`.

        Args:
          resolution (float): Energy resolution in :math:`\sqrt{MeV}`
            e.g. 0.05 for :math:`\sigma = 5\%/\sqrt{E[MeV]}`.

        Raises:
          ValueError: If the resolution is not between 0 and 1.
        """
        if (resolution > 0. and resolution < 1.):
            self._resolution = resolution
        else:
            raise ValueError("%s is an invalid energy resolution. Value "
                             "must be between 0. and 1." % resolution)

    def weighted_smear(self, spectrum, par="energy_mc"):
        """ Smears the energy of a :class:`spectra.Spectra` by
        calculating a Gaussian PDF for each bin. Weights are then applied
        to a window of width specified by the number of sigma depending on
        the value of the Gaussian PDF at the mean of the bin.

        Args:
          spectrum (:class:`spectra.Spectra`): Spectrum you wish to
            smear.
          par (string, optional): The name of the parameter you wish to
            smear. The default is energy_mc.

        Raises:
          IndexError: If par is not in the specta config.

        Returns:
          :class:`spectra.Spectra`: The smeared spectrum
        """
        if par not in spectrum.get_config().get_pars():
            raise IndexError("%s is not a parameter in the spectrum" % par)
        idx = spectrum.get_config().get_index(par)
        bins = []
        lows = []
        widths = []
        par_names = []
        for par_name in spectrum.get_config().get_pars():
            bins.append(range(spectrum.get_config().get_par(par_name)._bins))
            lows.append(spectrum.get_config().get_par(par_name)._low)
            widths.append(spectrum.get_config().get_par(par_name).get_width())
            par_names.append(par_name)
        smeared_spec = copy.copy(spectrum)
        smeared_spec._name = (spectrum._name + "_" +
                              str(100. * self._resolution) + "%")
        smeared_spec._data = numpy.zeros(spectrum._data.shape)
        for bin in itertools.product(*bins):
            entries = spectrum._data[bin]
            if entries:
                data = {}
                low = None
                high = None
                for i in range(len(bin)):
                    mean = self.get_bin_mean(lows[i], bin[i], widths[i])
                    if i == idx:
                        sigma = self.get_sigma(mean)
                        low, high = self.get_bounds(mean, sigma)
                        low = spectrum.get_config().get_par(
                            par).round(low - 0.5 * widths[i]) + 0.5 * widths[i]
                        high = spectrum.get_config().get_par(
                            par).round(high + 0.5 * widths[i]) + \
                            0.5 * widths[i]
                        if low < spectrum.get_config().get_par(par)._low:
                            low = spectrum.get_config().get_par(par)._low + \
                                0.5 * widths[i]
                        if high > spectrum.get_config().get_par(par)._high:
                            high = spectrum.get_config().get_par(par)._high - \
                                0.5 * widths[i]
                        weights = []
                        for energy in numpy.arange(low, high, widths[i]):
                            if self._poisson_smear is True:
                                weights.append(self.calc_poisson_energy(energy,
                                                                        mean))
                            else:
                                weights.append(self.calc_gaussian(energy,
                                                                  mean,
                                                                  sigma))
                    else:
                        data[par_names[i]] = mean
                total_weight = sum(weights)
                i = 0
                for energy in numpy.arange(low, high, widths[idx]):
                    data[par] = energy
                    smeared_spec.fill(weight=entries*weights[i]/total_weight,
                                      **data)
                    i += 1
        smeared_spec._raw_events = spectrum._raw_events
        return smeared_spec

    def random_smear(self, spectrum, par="energy_mc"):
        """ Smears the energy of a :class:`spectra.Spectra` by
        generating a number of random points from  Gaussian PDF
        generated from that bins mean value and the corresponding
        sigma. The number of points generated is equivalent to the
        number of entries in that bin.

        Args:
          spectrum (:class:`spectra.Spectra`): Spectrum you wish to
            smear.
          par (string, optional): The name of the parameter you wish to
            smear. The default is energy_mc.

        Raises:
          IndexError: If par is not in the specta config.

        Returns:
          :class:`spectra.Spectra`: The smeared spectrum
        """
        if par not in spectrum.get_config().get_pars():
            raise IndexError("%s is not a parameter in the spectrum" % par)
        idx = spectrum.get_config().get_index(par)
        bins = []
        lows = []
        widths = []
        par_names = []
        for par_name in spectrum.get_config().get_pars():
            bins.append(range(spectrum.get_config().get_par(par_name)._bins))
            lows.append(spectrum.get_config().get_par(par_name)._low)
            widths.append(spectrum.get_config().get_par(par_name).get_width())
            par_names.append(par_name)
        smeared_spec = copy.copy(spectrum)
        smeared_spec._name = (spectrum._name + "_" +
                              str(100. * self._resolution) + "%")
        smeared_spec._data = numpy.zeros(spectrum._data.shape)
        for bin in itertools.product(*bins):
            entries = int(spectrum._data[bin])
            if entries:
                data = {}
                for i in range(len(bin)):
                    mean = self.get_bin_mean(lows[i], bin[i], widths[i])
                    if i == idx:
                        mean_e = mean
                        sigma = self.get_sigma(mean)
                    else:
                        data[par_names[i]] = mean
                for i in range(entries):
                    if self._poisson_smear is True:
                        photons = (numpy.fabs
                                   (numpy.random.poisson(mean_e *
                                                         self._light_yield)))
                        data[par] = photons / float(self._light_yield)
                    else:
                        data[par] = (numpy.fabs
                                     (numpy.random.normal(mean_e, sigma)))
                    try:
                        smeared_spec.fill(**data)
                    except ValueError:
                        print "WARNING: Smeared energy out of bounds. Skipping"
        smeared_spec._raw_events = spectrum._raw_events
        return smeared_spec

class RadialSmear(Smear):
    """ The class which smears the radius. It accepts resolution in terms of
      sigma in units of mm.

    Args:
      poisson (bool): If True, use poisson smearing.

    Attributes:
      _resolution (float): The position resolution (mm).
      _poisson_smear (Bool): True if poisson smearing is to be applied.
        False if gaussian smearing is to be applied.
    """

    def __init__(self):
        """ Initialises the class.
        """
        super(RadialSmear, self).__init__("radial")
        self._resolution = 100.  # mm

    def calc_smear_resoluton(self, new_res, cur_res=None):
        """Calculates the value of resolution required to smear a data set
          which has already been smeared with a resolution of cur_res to
          achieve a new resolution of new_res.

        Args:
          new_res (float): The value of resolution wanted for the smeared PDF.
          cur_res (float, optional): Current value of resolution the PDF
            has been convolved with from the true value PDF.

        Raises:
          ValueError: If new_res is smaller than cur_sigma. Can't smear to
            higher resolutions (smaller sigmas)

        Returns:
          float: The value of resolution needed to smear the current
            PDF to obtain a new resolution: new_res.
        """
        if not cur_res:
            cur_res = self.get_resolution()
        if new_res < cur_res:
            raise ValueError("New resolution must be larger than the"
                             "current resolution. cur_res: %s. new_res: %s."
                             % (cur_res, new_res))
        return numpy.fabs(numpy.sqrt(new_res**2 - cur_res**2))

    def get_resolution(self):
        """Gets the position resolution.

        Returns:
          float: Position resolution.
        """
        return self._resolution

    def set_resolution(self, resolution):
        """Sets the position resolution:

        Raises:
          ValueError: If resolution is zero or less.

        Args:
          resolution (float): Position resolution in mm.
        """
        if resolution > 0:
            self._resolution = resolution
        else:
            raise ValueError("%s is an incorrect position resolutioin. Value "
                             "must be greater than zero." % resolution)

    def get_sigma(self):
        """Sigma and resolution are equivalent for radial dimensions
        currently. This function calls self.get_resolution()

        Returns:
          float: Sigma in mm equivalent to resolution
        """
        return self.get_resolution()

    def weighted_smear(self, spectrum, par="radial_mc"):
        """ Smears the radius of a :class:`spectra.Spectra` by
        calculating a Gaussian PDF for each bin. Weights are then
        applied to a window of width specified by the number of sigma
        depending on the value of the Gaussian PDF at the mean of the
        bin.

        Args:
          spectrum (:class:`spectra.Spectra`): Spectrum you wish to
            smear.
          par (string, optional): The name of the parameter you wish to
            smear. The default is radial_mc.

        Raises:
          IndexError: If par is not in the specta config.

        Returns:
          :class:`spectra.Spectra`: The smeared spectrum
        """
        if par not in spectrum.get_config().get_pars():
            raise IndexError("%s is not a parameter in the spectrum" % par)
        idx = spectrum.get_config().get_index(par)
        bins = []
        lows = []
        widths = []
        par_names = []
        for par_name in spectrum.get_config().get_pars():
            bins.append(range(spectrum.get_config().get_par(par_name)._bins))
            lows.append(spectrum.get_config().get_par(par_name)._low)
            widths.append(spectrum.get_config().get_par(par_name).get_width())
            par_names.append(par_name)
        smeared_spec = copy.copy(spectrum)
        smeared_spec._name = (spectrum._name + "_" +
                              str(100. * self._resolution) + "mm")
        smeared_spec._data = numpy.zeros(spectrum._data.shape)
        for bin in itertools.product(*bins):
            entries = spectrum._data[bin]
            if entries:
                data = {}
                low = None
                high = None
                for i in range(len(bin)):
                    mean = self.get_bin_mean(lows[i], bin[i], widths[i])
                    if i == idx:
                        sigma = self.get_sigma()
                        low, high = self.get_bounds(mean, sigma)
                        low = spectrum.get_config().get_par(
                            par).round(low - 0.5 * widths[i]) + 0.5 * widths[i]
                        high = spectrum.get_config().get_par(
                            par).round(high + 0.5 * widths[i]) + \
                            0.5 * widths[i]
                        if low < spectrum.get_config().get_par(par)._low:
                            low = spectrum.get_config().get_par(par)._low + \
                                0.5 * widths[i]
                        if high > spectrum.get_config().get_par(par)._high:
                            high = spectrum.get_config().get_par(par)._high - \
                                0.5 * widths[i]
                        weights = []
                        for radius in numpy.arange(low, high, widths[i]):
                            weights.append(self.calc_gaussian(radius,
                                                              mean,
                                                              sigma))
                    else:
                        data[par_names[i]] = mean
                total_weight = sum(weights)
                i = 0
                for radius in numpy.arange(low, high, widths[idx]):
                    data[par] = radius
                    smeared_spec.fill(weight=entries*weights[i]/total_weight,
                                      **data)
                    i += 1
        smeared_spec._raw_events = spectrum._raw_events
        return smeared_spec

    def random_smear(self, spectrum, par="radial_mc"):
        """ Smears the radius of a :class:`spectra.Spectra` by
        generating a number of random points from  Gaussian PDF
        generated from that bins mean value and the corresponding
        sigma. The number of points generated is equivalent to the
        number of entries in that bin.

        Args:
          spectrum (:class:`spectra.Spectra`): Spectrum you wish to
            smear.
          par (string, optional): The name of the parameter you wish to
            smear. The default is radial_mc.

        Raises:
          IndexError: If par is not in the specta config.

        Returns:
          :class:`spectra.Spectra`: The smeared spectrum
        """
        if par not in spectrum.get_config().get_pars():
            raise IndexError("%s is not a parameter in the spectrum" % par)
        idx = spectrum.get_config().get_index(par)
        bins = []
        lows = []
        widths = []
        par_names = []
        for par_name in spectrum.get_config().get_pars():
            bins.append(range(spectrum.get_config().get_par(par_name)._bins))
            lows.append(spectrum.get_config().get_par(par_name)._low)
            widths.append(spectrum.get_config().get_par(par_name).get_width())
            par_names.append(par_name)
        smeared_spec = copy.copy(spectrum)
        smeared_spec._name = (spectrum._name + "_" +
                              str(100. * self._resolution) + "mm")
        smeared_spec._data = numpy.zeros(spectrum._data.shape)
        for bin in itertools.product(*bins):
            entries = spectrum._data[bin]
            if entries:
                data = {}
                for i in range(len(bin)):
                    mean = self.get_bin_mean(lows[i], bin[i], widths[i])
                    if i == idx:
                        mean_r = mean
                        sigma = self.get_sigma()
                    else:
                        data[par_names[i]] = mean
                for i in range(entries):
                    data[par] = numpy.fabs(numpy.random.normal(mean_r, sigma))
                    try:
                        smeared_spec.fill(**data)
                    except ValueError:
                        print "WARNING: Smeared radius out of bounds. Skipping"
        smeared_spec._raw_events = spectrum._raw_events
        return smeared_spec

max_stars_repo_name: snoplusuk/echidna

max_issues_repo_name: snoplusuk/echidna

max_forks_repo_name: snoplusuk/echidna