//
//	Deilephila.cc
//
//	Anna Balkenius & Christian Balkenius 2002-01-04
//
//	anna.balkenius@cob.lu.se
//	christian.balkenius@lucs.lu.se
//
//	Changed 2003-07-24 - documentation added; minor changes
//
//	This code simulates the hawkmoth Deilephila elpenor in an instrumental learning situation
//	where it has to select one of five artificial flowers of different colours.
//
//	The program was used for the simulations in:
//
//		Balkenius, A., Kelber, A., Balkenius, C. (2002)
//		Simulations of Learning and Behaviour in the Hawkmoth Deilephila elpenor
//		In Hallam, et al. , From Animals to Animats 7. Cambridge, MA: MIT Press.
//
//		Balkenius, A., Kelber, A., Balkenius, C. (2003)
//		A model of selection between stimulus and place strategy in a Hawkmoth
//		Adaptive Behavior, in press.
//
//	Usage:	deilephila	[1 | 2 | 3]
//
//	Use command line argument 1, 2, 3 to select simulation to run
//	or redefine the variable DEFAULTEXP
//
//	(The code may look as C but is actually C++)
//
//	The code has been changed from the original to work with gcc
//	random() now used instead of rand() at a critical place
//
//	The simulation could have been written more compactly by adding all probablities to the
//	state transition table - maybe next time...
//


#include <stdio.h>
#include <stdlib.h>
#include <math.h>

#define	EXP1			0
#define	EXP2			1
#define	EXP3			2

#define	DEFAULTEXP	EXP1

const float	n			=	4;
const float	beta_inc		=	0.3;
const float	beta_dec		=	0.9999;

const int		animals		=	500;


//
//	The perception table defines what the model moth
//	perceives in each state, i. e. the colour of the flower
//	in front of it and the distance to the stimulus array
//
//	The first index is the experiment to run
//	The colours are coded as (relative) receptor responses
//

float perception[3][17][4] =
{
	//	EXP1: SIMILAR COLOURS
	//	Colour table used in Simulation 1 - SAB 2002
	//	Receptor responses were calculated as described in the references

	{
		{0, 0, 0, 3},	
		{0, 0, 0, 2},
		{0, 0, 0, 2},
		{0, 0, 0, 2},
		{0, 0, 0, 2},
		{0, 0, 0, 2},
		{1.2E10, 4.1E10, 1.7E11, 1},		// yellow
		{1.6E10, 4.0E10, 1.4E11, 1},		// yellow-orange
		{2.4E9, 	6.2E10, 3.7E10, 1},		// yellow-green
		{1.2E10, 2.6E10, 9.1E10, 1},		// orange
		{2.8E9, 	3.3E10, 2.3E10, 1},		// green
		{1.2E10, 4.1E10, 1.7E11, 0},		// yellow
		{1.6E10, 4.0E10, 1.4E11, 0},		// yellow-orange
		{2.4E9, 	6.2E10, 3.7E10, 0},		// yellow-green
		{1.2E10, 2.6E10, 9.1E10, 0},		// orange
		{2.8E9, 	3.3E10, 2.3E10, 0},		// green
		{0, 0, 0, 4}
	},
	
	//	EXP2: VERY DISSIMILAR (MONOCHROMATIC) COLOURS
	//	Colour table used in Simulation 2 - SAB 2002	
	//	Receptor responses were read directly from the sensitivity curves

	{
		{0, 0, 0, 3},	
		{0, 0, 0, 2},
		{0, 0, 0, 2},
		{0, 0, 0, 2},
		{0, 0, 0, 2},
		{0, 0, 0, 2},
		
		{0.0002521, 0.9631, 0.358, 1},	// 450 nm
		{0.1671, 0.6466, 0.2664, 1},		// 400 nm
		{1, 	0.3277, 0.2031, 1},			// 350 nm
		{9.62E-10, 0.1984, 0.86161,	1},	// 500 nm
		{4.28E-18, 0.002152, 0.8394, 1},	// 550 nm
		
		{0.0002521, 0.9631, 0.358, 0},	// 450 nm
		{0.1671, 0.6466, 0.2664, 0},		// 400 nm
		{1, 	0.3277, 0.2031, 0},			// 350 nm
		{9.62E-10, 0.1984, 0.86161, 0},	// 500 nm
		{4.28E-18, 0.002152, 0.8394, 0},	// 550 nm
		
		{0, 0, 0, 4}
	},
	
	//	EXP 3: DISSIMILAR COLOURS USED IN EXPERIMENT WITH REAL MOTHS
	//	Colour table used in Simulation 2 in Adaptive Behavior article, 2003	
	//	Receptor responses were calucated for the colour used in the experiment
	//	Not in SAB article

	{
		{0, 0, 0, 3},	
		{0, 0, 0, 2},
		{0, 0, 0, 2},
		{0, 0, 0, 2},
		{0, 0, 0, 2},
		{0, 0, 0, 2},
		
		{1.2E10, 4.1E10, 1.7E11, 1},			// yellow
		{1.37E10, 3.19E10, 2.54E10, 1},		// black
		{1.63E10, 1.07E11, 5.60E10, 1},		// blue
		{1.99E10, 4.26E10, 7.10E10, 1},		// green
		{2.14E10, 3.39E10, 4.09E10, 1},		// orange
		
		{1.2E10, 4.1E10, 1.7E11, 0},			// yellow
		{1.37E10, 3.19E10, 2.54E10, 0},		// black
		{1.63E10, 1.07E11, 5.60E10, 0},		// blue
		{1.99E10, 4.26E10, 7.10E10, 0},		// green
		{2.14E10, 3.39E10, 4.09E10, 0},		// orange
		
		{0, 0, 0, 4}
	}
};



// Q for white used for von Kries adaptation

float UVmax	= 1.657E10;
float Bmax	= 2.916E11;
float Gmax	= 2.859E11;



//
//	State transition table
// 	Maps a <state, action> pair to the next state
//

int trans[17][7] =	
{
	// search, approach, left, right, chose, forage, rest

	{0, 0, 0, 0, 0, 0, 16},

	{0, 6, 1, 2, 1, 1, 16},
	{0, 7, 1, 3, 2, 2, 16},
	{0, 8, 2, 4, 3, 3, 16},
	{0, 9, 3, 5, 4, 4, 16},
	{0, 10, 4, 5, 5, 5, 16},

	{0, 6,  6, 7, 11, 6, 16},
	{0, 7,  6, 8, 12, 7, 16},
	{0, 8,  7, 9, 13, 8, 16},
	{0, 9,  8, 10, 14, 9, 16},
	{0, 10, 9, 10, 15, 10, 16},
	
	{0, 11, 6, 6, 11, 11, 16},
	{0, 12, 6, 7, 12, 12, 16},
	{0, 13, 7, 8, 13, 13, 16},
	{0, 14, 8, 9, 14, 14, 16},
	{0, 15, 9, 9, 15, 15, 16},

	{16, 16, 16, 16, 16, 16, 16}	// Welcome to the Hotel California!
};




// Random returns a random number between low and high (inclusive)

float
Random(float low, float high)
{
	return low + (float(rand())/float(RAND_MAX))*(high-low);
}										



//
//	get_perception
//
//	returns the color in a state and distance to flower with von Kries adaptation
//

void
get_perception(int e, int state, float & UV, float & Blue, float & Green, int & dist)
{
	// Get colour and apply von Kries adaptation (i. e. divide by values for white)
	
	UV 		= perception[e][state][0]/UVmax;
	Blue 		= perception[e][state][1]/Bmax;
	Green 	= perception[e][state][2]/Gmax;
	
	// Get distance to stimulus array (called A, B, C in Adaptive Behavior article; A, B, C, D before)

	dist 		= int(perception[e][state][3]);

	// Normalize colour

	float s 	= sqrt(UV*UV + Blue*Blue + Green*Green);
	
	if(s != 0)
	{
		UV /= s;
		Blue /= s;
		Green /= s;
	}
}



//
//	world
//
//	returns new state from <state, action> pair
//

int
world(int state, int action)
{
	if(action < 0)				// what is this?
		return -action;

	if(state == 0 && action == 1)	// randomize what flower the moth approaches
		return(random() % 5 + 1);

	return trans[state][action];
}



//
//	Moth Simulation Code
//

float 	UVMemory 	= 0;
float 	BMemory 		= 0;
float 	GMemory 		= 0;

int		PlaceMemory 	= 0;
float		PlaceValue	= 0.0;


//
//	stimulus_value
//
//	calculates the value/similarity of the stimulus in front and the rewarded colour
//

float
stimulus_value(float UV, float B, float G)
{
	float v = UV*UVMemory + B*BMemory + G*GMemory;
	return pow(v, n);
}



//
//	print_stimulus_values
//
//	prints some useful info
//

void
print_stimulus_values(int e)
{
	printf("\n\nStimulus Values:\n");
	for(int i=11; i<16; i++)
	{
		float UV0, B0, G0;
		int d0;
		get_perception(e, i, UV0, B0, G0, d0);
		printf("\tStimulus %d = %.2f %.2f %.2f: %f\n", i, UV0, B0, G0, stimulus_value(UV0, B0, G0));
	}
	printf("\n\n\tMemory = %.2f %.2f %.2f\n\n", UVMemory, BMemory, GMemory);

}



//
//	init_color_memory
//
//	Teaches the mot hthe colour or the rewarding stimulus
//	Assumed to be done the day before the test
//

void
init_color_memory(int e)
{
	int d;
	int state = 11;
	get_perception(e, state, UVMemory, BMemory, GMemory, d);
}



//
//	init_place_memory
//

void
init_place_memory()
{		
	PlaceValue	= 0;
	PlaceMemory 	= 0;
}



//
//	moth
//
//	the actual simulation of the moth
//	arguments: place, last action and perception 
//	

int
moth(int place, int last_action, int flower_distance, float UV, float Blue, float Green, float Reward)
{
	PlaceValue	*=	beta_dec;	// will decay until tomorrow (but we will stop the simulation long before that...)

	//
	// 	FEEADING and LEARNING
	//
	//	Update the memory/learning variables if last action was 'eat'
	//	
	
	if(last_action == 5)
	{
		if(Reward > 0)	// Reward
		{
			UVMemory	=	UV;
			BMemory		=	Blue;
			GMemory		=	Green;
			
			PlaceValue	+=	beta_inc;
			PlaceValue	=	(PlaceValue > 1 ? 1 : PlaceValue);
			PlaceMemory 	=	place-5;	// Changed after SAB article

			return 0;	// Leave
		}
				
		else
		{
			if(Random(0,1) < 0.5)	// Leave 50%
				return 0; 
			if(Random(0,1) < 0.5)	// Go left 25%
				return 2; 
			return 3; 				// Go right 25%
		}
	}
	
	//
	// 	OTHER BEHAVIOR
	//
	//	Select what to do depnding on the distance to the flowers etc
	//
	//	distance 0		=>		eat
	//	distance 1		=>		choose according to similarity or go left or right
	//	distance 2		=>		approach
	//	otherwise		=>		move to place, stay or approach
	//
	
	switch(flower_distance)
	{
		case 0:	
			return 5; // Eat

		case 1:
			if(stimulus_value(UV, Blue, Green) > Random(0, 1))	// Choose flower in front?
				return 4; 
			else if(Random(0,1) > 0.5)
				return 2; 
			else
				return 3; 
		
		case 2:
			return 1;
		
		default:
			if(Random(0, 1) < PlaceValue) // Use Place or Color Strategy?
				return -PlaceMemory;
			else
				return random() % 2;
	}
}



//	main
//
//	Main simulation code mainly collects statistics and control the simulated experiment
//

int
main(int argc, char * argv[])
{	
	int	choice[5][5] = {{0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}};	
	int	experiment = (argc > 1 ? atoi(argv[1])-1 :  DEFAULTEXP);

	if(experiment < 0 || experiment >2)
	{
		printf("Usage: deilephila [experiment] (where experiment = 1, 2 or 3)\n");
		exit(1);
	}
	
	for(int animal = 0; animal < animals; animal++)
	{
		// Reset simulation for new experiment
		
		float UV = 0;
		float B = 0;
		float G = 0;
		
		int state = 0;
		int action = 0;
		int distance = 0;
		
		int touch_count = 0;		// How many times has the moth touched a flower? Time for a reward?
		int rew_count = 0;		// How many times has themoth been rewarded?

		// Reset moth

		init_color_memory(experiment);
		init_place_memory();

//		print_stimulus_values(experiment);

		// Run one experiment until the moth has been rewarded 3 times
		// Only the first two were reported in the articles

		while(rew_count < 4)
		{
			// Reward trial
			
			if(touch_count > 3) 
			{
				get_perception(experiment, 11, UV, B, G, distance);
				action = moth(15, 5, 0, UV, B, G, 1.0);
				
				rew_count++;
				touch_count = 0;
				state = 0;
			}
		
			// Test trial

			else	
			{
				get_perception(experiment, state, UV, B, G, distance);	
				action = moth(state, action, distance, UV, B, G, 0);

				if(action == 0 && state != 0)					// The moth leaves the array; new trial
					touch_count++;
				
				if(action == 5 && 11<= state && state <=15)		// Collect statistics when eating
					choice[rew_count][state-11]++;
			
				state = world(state, action);
			}
		}
	}
	
	// PRINT RESULTS
	
	printf("\nRESULTS SIMULATION %d\n", experiment+1);
	
	printf("\nChoices (SUMS)\n");
	for(int p=0; p<4; p++)
		printf("%d: %5d %5d %5d %5d %5d\n", p, choice[p][0], choice[p][1], choice[p][2], choice[p][3], choice[p][4]);

	printf("\nChoices (%%)\n");
	float sum;
	for(int p=0; p<4; p++)
		if(sum = choice[p][0] + choice[p][1] + choice[p][2] + choice[p][3] + choice[p][4])
			printf("%d: %5.1f %5.1f %5.1f %5.1f %5.1f\n", p, 100*choice[p][0]/sum, 100*choice[p][1]/sum, 100*choice[p][2]/sum, 100*choice[p][3]/sum, 100*choice[p][4]/sum);
				
	printf("\nN = %d, n = %.2f, beta_inc = %.2f, beta_dec = %.6f\n", animals, n, beta_inc, beta_dec);
}