# ISMachinery.py
# Author: Juliette Zerick (jnzerick@math.ucdavis.edu)
# EAD 221 + PHY 250 (Spring 2009)
# Purpose: This is the machinery behind the immune system algorithm.
# 	It does the heavy lifting--mostly just running a single generation
# 	which is complicated enough to warrant an entire file full of
#	helper functions.

from UtilsNConstants import *
from Antibody import *
from Antigen import *
from PluginFuncs import *
from Trajectory import *

import pylab
import numpy

# a counter used to uniquely identify generated Antibodies--handy for
# keeping track of the heritage of Antibodies (over the generations of
# mutations)
AB_ID_COUNT = 1

# prints the whole population; str() neesd to be defined for the object
def print_population(population):
	for A in population:
		print A

# helper function that checks whether the Antigen is in the "cloud"
# of the Antibody--if the distance is, at least
def Ag_in_Ab_cloud(Ab,Ag):
	d = Ab.AgAb_aff(Ag)
	if (1.0/d) < AB_CLOUD_RAD:
		return True
	return False

# calculates the SUM affinities between an Antibody and every Antigen
# within the "cloud" radius, AB_CLOUD_RAD
def Ab_vs_allAg(Ab,Ag_pop):
	dij = 0

	for Ag in Ag_pop:
		if Ag_in_Ab_cloud(Ab,Ag):
			dij+= Ab.AgAb_aff(Ag)

	return dij

# a cmp function for sorting a list of Antibodies (or Antigens) by
# affinity; note that A1 and A2 better have aff objects defined
def aff_cmp(A1,A2):
	if A1.aff == A2.aff:
		return 0
	elif A1.aff < A2.aff:
		return -1
	
	return 1 

# helper function that strips aff off whatever population of objects is
# passed in (Antigens or Antibodies, in this case)
def remove_aff(population):
	for A in population:
		del A.aff

# calculate the affinities between an Antigen and each Antibody in
# the given population; the result is stored in an aff object attached to
# each Antibody object in Ab_pop, which later is stripped off with the
# function above
def calc_AbAg_aff(Ab_pop,Ag):
	for A in Ab_pop:
		A.aff = A.AgAb_aff(Ag)

# whatever objects are marked for removal in the pool get removed; to be
# flagged for removal, the object must have an attribute rem stuck to it,
# and its value must be 1 (True)
def cull_pool(pool):
	culling = []

	for i in xrange(len(pool)):
		if pool[i].rem == 1:
			culling.append(pool[i])

	for A in culling:
		pool.remove(A)

# remove objects at random from the pool; this is to get the size of the pool
# down to NUM_ANTIBODIES size--trimming excess Antibodies
def cull_at_random(pool):
	numtocull = len(pool) - NUM_ANTIBODIES

	N = len(pool)

	for i in xrange(numtocull):
		sel = Ncointoss(N-1)

		while pool[sel].rem == 1:
			sel = Ncointoss(N-1)

		pool[sel].rem = 1

	cull_pool(pool)

# this function compares the old population of Antibodies with the new
# generation; if each Antibody in the old population has (mutated) descendents
# in the new generation, it is not considered "culled"; those without surviving
# kin have "gone to the grave."
def find_culled(oldpop,newpop):
	culled = []

	for A in oldpop:
		for i in xrange(len(newpop)):
			if A.ID == newpop[i].ID:
				break

		if A.ID == newpop[i].ID:
			continue

		culled.append(A)

	print len(culled),"go to the grave"

	return culled

# runs one generation or iteration of the immune system algorithm
def Ab_generation(Ab_pop,Ag_pop,gen_num):
	pool = []

	# for each antigen...
	for Ag in Ag_pop:
		H = []

		# find the highest-affinity antibodies
		calc_AbAg_aff(Ab_pop,Ag)
		Ab_pop.sort(aff_cmp)

		# grab the highest-affinity antibodies
		for i in xrange(H_CULL):
			H.append(Ab_pop[-i])

		# clone the high-affinity antibodies proportionate to their
		# affinities
		for A in H:
			C = []
			for i in xrange(int(math.floor(0.25/A.aff))):
				C.append(A.copy())
			Cstar = C

			# and mutate inversely proportional to affinity
			prob_mut = A.aff/20
			for c in Cstar:
				c.mutate(prob_mut)
			pool += Cstar

	print "Size of the pool is:",len(pool)

	# WARNING: DO NOT USE THIS BIT OF CODE (I left it here as a warning).
	# If the unique Antibodies are sorted out of the pool, the effects
	# of competition (and adaptation) go down the drain. Competitive
	# advantage is translated into greater representation in the pool,
	# which then means a higher likelihood of selection for reproduction.
	# So if the pool is reduced in size, this effect is lost. The result
	# is a bunch of Antibodies that don't cluster.

	#cream = get_uniques(pool)
	#print "%d unique Antibodies in the pool" % (len(cream))

	# calculate "cloud affinities"--the sum of the affinities between the
	# Antibody and the Antigens within a fixed radius.
	for Ab in pool:
		Ab.aff = Ab_vs_allAg(Ab,Ag_pop)

	# sort the pool by affinity
	pool.sort(aff_cmp)

	# remove the antibodies that are too far out			
	for Ab in pool:
		if Ab.aff < SIGMA_D:
			Ab.rem = 1
		else:
			Ab.rem = 0
	cull_pool(pool)

	print "After first culling, pool size is:",len(pool)

	# remove antibodies that are too close
	i = len(pool)-1
	while i >= 0:
		if pool[i].rem == 1:
			i -= 1
			continue
		for A in pool:
			if pool[i].Ab_aff(A) < SIGMA_S:
				if pool[i] != A:
					A.rem = 1
		i -= 1
	cull_pool(pool)
 
	print "After second culling, pool size is:",len(pool)

	if len(pool) > NUM_ANTIBODIES:
		cull_at_random(pool)

	print "after third culling, pool size is",len(pool)

	# strip off the aff data
	remove_aff(pool)

	# try removing the rem data--if it exists
	for A in pool:
		try: 
			del A.rem
		except AttributeError:
			continue

	# if the size of the pool is less than NUM_ANTIBODIES, fill the pool
	# with randomly-generated Antibodies
	B_diversity = []
	if len(pool) < NUM_ANTIBODIES:
		for i in xrange(NUM_ANTIBODIES-len(pool)):
			B_diversity.append(spawn_Ab(gen_num))

	pool += B_diversity

	# optional--see where the Antibodies that didn't reproduce went
	#culled = find_culled(Ab_pop,pool)
	#plot_Ab(culled)

	return pool

# generate a new chromosome (bit string)
def gen_chromosome():
	return gen_binstring(CHROM_LENGTH)

# spawn a random Antibody
def spawn_Ab(gen_num):
	# get a random chromosome
	C = gen_chromosome()
	
	# set up its history
	T = Trajectory(plugin_2Dplot_conn,[])
	T.feed_point(chrom2coord(C))

	global AB_ID_COUNT # it already was declared globally, but this'll
		# keep the compiler happy

	# append a newly-constructed Antibody
	A = Antibody(plugin_mutate_Ab,plugin_AbAg_aff,plugin_Ab_aff,
		plugin_printAb,plugin_Ab_tostr,C,gen_num,T,AB_ID_COUNT)

	AB_ID_COUNT += 1

	del C
	return A

# spawn a fresh population of random Antibodies
def spawn_Ab_pop():
	population = []

	# first generation
	init_gen = 0

	# generate POP_SIZE Antibodies
	for i in xrange(NUM_ANTIBODIES):
		population.append(spawn_Ab(init_gen))

	return population

# generate a "clump point"--a point within a small, fixed radius of the
# given point P
def gen_clump_point(P):
	C = []

	# get len(P) random numbers in [0,1] and map them appropriately
	for i in xrange(len(P)):
		p_i = P[i]
		r = random.random()
		C.append(mapUI2ab(r,p_i-CLUMP_RADIUS,p_i+CLUMP_RADIUS))

	return C

# spawn a random Antigen
def spawn_Ag(chromosome,gen_num):
	A = Antigen(plugin_mutate_Ag,plugin_AbAg_aff,plugin_Ag_aff,
		plugin_printAg,plugin_Ag_tostr,chromosome,gen_num)

	return A

# this calculates the minimum distance between all the cluster centers (points)
def calc_min_dist(centers):
	D = []

	# just find all the distances...
	for i in xrange(len(centers)):
		for j in xrange(len(centers)):
			if i != j:
				D.append(dist_metric(centers[i],centers[j]))

	# and return the smallest one
	return min(D)
	
# set up an initial population of Antigens
def spawn_Ag_pop():
	population = []
	init_gen = 0

	centers = []

	for i in xrange(NUM_AG_CLUMPS):
		# get a center by mapping random points in [0,1] to [0,ALPHA]
		C = []
		for d in xrange(DIMS):
			C.append(ALPHA*random.random())

		centers.append(C)

		# for each center, spawn a clump of points around it
		for j in xrange(NUM_ANTIGENS_PER_CLUMP):
			chromosome = gen_clump_point(C)
			population.append(spawn_Ag(chromosome,init_gen))

	global SIGMA_S # keeps the compiler happy

	# reset SIGMA_S to be the smallest distance between the centers of the
	# clusters, divided by two
	SIGMA_S = 0.5*calc_min_dist(centers)

	print "SIGMA_S =",SIGMA_S

	return population

# copy a population; provides a deeper copy than copylist
def copy_pop(population):
	L = []
	for I in population:
		L.append(I.copy())

	return L

# this takes an Antibody and returns the coordinates in its chromosome
def Ab2point(Ab):
	# some calculations required
	return chrom2coord(Ab.chromosome)

# this takes an Antigen and returns the coordinates in its chromosome
def Ag2point(Ag):
	return Ag.chromosome

# plot a population of Antibodies using their histories (Trajectory objects)
def plot_Ab(Ab_pop):
	for Ab in Ab_pop:
		Ab.history.plot_traj() 

# plot a population of Antigens
def plot_Ag(Ag_pop):
	# Antigens do not have Trajectory objects as data, but the 
	# Trajectory class is handy enough to use for plotting
	Ag_T = Trajectory(plugin_2Dplot_dots,[])

	for Ag in Ag_pop:
		Ag_T.feed_point(Ag2point(Ag))

	Ag_T.plot_traj()

# this "nudges" an Antigen by a small, random distance
def nudge_Ag(Ag):
	r1 = (2*random.random()-1)*CLUMP_RADIUS*0.1
	r2 = (2*random.random()-1)*CLUMP_RADIUS*0.1

	Ag.chromosome[0] += r1
	Ag.chromosome[1] += r2

# this shifts or "nudges" an entire population of Antigens
def shift_Ag_pop(Ag_pop):
	for Ag in Ag_pop:
		nudge_Ag(Ag)