#include <stdio.h>
#include <math.h>
#include <string.h> // gives us memcpy
#include <stdlib.h> // for rand, system
#include <time.h>
#include <unistd.h> // for sleep

//////////////////////////////////////////////////////////////////
//				1D Hartree minimizer:							//
//	We always have Dirichlet conditions ar the very boundaries. //
//	Play with the external potential by editing the function    //
//	"potential". Play with positive and negative coupling		//
//	constants for the Hartree term.								//
//////////////////////////////////////////////////////////////////
// Please feel free to use, improve and redistribute.			//
// Do not make commercial use of this template, no selling!		//
// Created by Niels Benedikter.									//
// Please keep this document head. Feel free to add your name	//
// if you make modifications.									//
//////////////////////////////////////////////////////////////////

// physical constants:
#define COUPLINGCONSTANT 0.7	//good: 0.7 (this is the Hartree term)
#define EXTERNALPOTENTIAL 0.1		//0.1 works well

// affects the printout of energys but nothing else
#define ENERGYSCALE 1000		//good: 1000

// constants affecting the 'random trials' - generally no need to play with those.
#define SIZE 100	// good: 100 or 200 (this is the number of lattice points)
#define RESOLUTION 0.00001		//good: 0.00001
#define N_ITERATIONS 500000	//good: 500000

#define PI 3.14159265
#define WIGGLEWIDTH 8			//good: 6
#define WIGGLESTRENGTH 600		//good: 600

int max(int a, int b)
{
	if (a > b)
	{
		return a;
	}
	return b; // only if a was not larger than b
}

int min(int a, int b)
{
	if (a < b)
	{
		return a;
	}
	return b; // only if a was not smaller than b
}

double maxoffield(double varphi[SIZE+4]) // finds the maximum of a field configuration varphi
{
	double maxofvarphi = 0;
	for (int x=1; x<= SIZE; x++) {
		if (fabs(varphi[x]) > maxofvarphi)
		{
			maxofvarphi = fabs(varphi[x]);
		}
	}
	return maxofvarphi;
}

double potential (int x) // defines the external potential at position x, e.g. harmonic oscillator or box - pay attention to the required lattice size
{
// no extra potential
	return 0;

//  harmonic oscillator: lives on SIZE = 100
//	return ((double)x-50)*((double)x-50);

// 	double well: lives on SIZE = 200
//	return min((x-50)*(x-50),(x-150)*(x-150));

// box with Dirichlet conditions: lives on SIZE = 100
//	if (x > 20 && x < 80) {return 0;}
//	else {return 100000;}

//  box with Dirichlet conditions and small middle wall: lives on SIZE = 200
//	if (x > 20 && x < 80) {return 0;}
//	if (x > 120 && x < 180) {return 0;}
//	if (x >= 80 && x <= 120) {return 50;}
//	return 100000;

// mexican hat:
//	return ((double)x-60)*((double)x-60)*((double)x-140)*((double)x-140)/(81000);

//  asymmetric box with Dirichlet conditions: lives on SIZE = 200
//	if (x > 20 && x < 80) {return 20;}
//	if (x > 120 && x < 180) {return 0;}
//	if (x >= 80 && x <= 120) {return 10;}
//	return 100000;
}

void energy (double varphi[SIZE+4]) // calculates the energy of a given field configuration and stores it in place SIZE+2
{
	varphi[SIZE+2] =0;	

	for (int x=1; x<= SIZE; x++) {	// x is the lattice site
		// kinetic term, using the discrete Laplacian
		varphi[SIZE+2] += ENERGYSCALE*varphi[x]*(-varphi[x+1]-varphi[x-1]+2*varphi[x]);
		// interaction term, quartic nonlinearity
		varphi[SIZE+2] += ENERGYSCALE*COUPLINGCONSTANT*varphi[x]*varphi[x]*varphi[x]*varphi[x];
		// external potential
		varphi[SIZE+2] += EXTERNALPOTENTIAL*potential(x)*varphi[x]*varphi[x];
	}
}

double norm (double varphi[SIZE+4]) // calculates the norm of a given field configuration and returns it
{
	double squarednorm =0;	

	for (int x=1; x<= SIZE; x++) {	// x is the lattice site
		// norm of L^2
		squarednorm += varphi[x]*varphi[x];
	}
	return sqrt(squarednorm);
}

void normalize (double varphi[SIZE+4]) // normalizes a field configuration to one
{
	double normvalue;
	normvalue = norm(varphi);

	for (int x=1; x<= SIZE; x++) {	// x is the lattice site
			// divide every point by the norm
			varphi[x] /= normvalue;
		}
}

void initialize (double varphi[SIZE+4]) // initializes a field configuration with a (not normalized) sine, i.e. non-interacting system
{
	for (int x=1; x<= SIZE; x++) {	// x is the lattice site
			// divide every point by the norm
			varphi[x] = sin(PI*x/(SIZE+1));
		}
	varphi[0] = 0;
	varphi[101] = 0;
}

void wiggle (double varphi[SIZE+4], int wigglew) // wiggles a field configuration by some randomness
{
	int xnot = rand() % SIZE;
	int xminus = max(xnot-wigglew,1);
	int xplus = min(xnot+wigglew,SIZE);
	
	for (int x=xminus; x<= xplus; x++) {
		varphi[x] += RESOLUTION*((double)(rand() % WIGGLESTRENGTH -((double)WIGGLESTRENGTH/2.)));
	}
}

int saveconfig (double varphi[SIZE+4], int iterationnumber) // save field to file
{
	char filename[150];
	sprintf(filename,"field/field%i.dat",iterationnumber);
	// open file for saving
 	FILE* filestream = fopen(filename, "wt");
    if (filestream == NULL)  {
        printf("Folder named ->field<- is required!\n");
       	return(1);
    }

	for (int x=1; x<= SIZE; x++) {	// x is the lattice site
		fprintf(filestream, "%f\n", varphi[x]);
	}
	
	fclose(filestream);
	return(0);
}

int saveinitialconfig (double varphi[SIZE+4]) // save field to file
{
	// open file for saving
 	FILE* filestream = fopen("field/initialfield.dat", "wt");
    if (filestream == NULL)  {
        printf("Folder named ->field<- is required!\n");
       	return(1);
    }

	for (int x=1; x<= SIZE; x++) {	// x is the lattice site
		fprintf(filestream, "%f\n", varphi[x]);
	}
	
	fclose(filestream);
	return(0);
}

int savepotential (double energyscale, double maximum) // save field to file
{
	// open file for saving
 	FILE* filestream = fopen("field/potential.dat", "wt");
    if (filestream == NULL)  {
        printf("Folder named ->field<- is required!\n");
       	return(1);
    }

	// we use the energy of the final state to decide where we cutoff the plotting of the external potential; we rescale the height to the solution

	for (int x=1; x<= SIZE; x++) {	// x is the lattice site
		if (fabs(EXTERNALPOTENTIAL*potential(x)) <= fabs(energyscale)) {
			fprintf(filestream, "%f\n", (EXTERNALPOTENTIAL*potential(x)/energyscale)*maximum);
		}
		else {
			fprintf(filestream, "%f\n", copysign(maximum,EXTERNALPOTENTIAL*potential(x)));
//			fprintf(filestream, "?\n");
		}
	}
	
	fclose(filestream);
	return(0);
}

void createscript (int limit)	// write a gnuplot string for our obtained data
{
 	FILE* filestream = fopen("script.gp", "wt");
    if (filestream == NULL)  {
        printf("Failed to create gnuplot script!\n");
       	return;
    }

	fprintf(filestream, "set datafile missing \"?\"\n");
	fprintf(filestream, "do for [i=1:%d] { plot sprintf('field/field%%d.dat', i), \"field/initialfield.dat\", \"field/potential.dat\" with lines; pause 0.01 }", limit);
	printf("Gnuplot script 'script.gp' successfully written.\n");
	
	fclose(filestream);
}

int main (void)
{
	// prepare a file that will hold the statistics
 	FILE* filestreamconvergence = fopen("field/convergence.dat", "wt");
    if (filestreamconvergence == NULL)  {
        printf("Folder named ->field<- is required!\n");
       	return(1);
    }

	// here the main stuff starts:

	int goodsteps = 0;
	srand(time(NULL));

	printf("Numerical Solver for 1D Hartree problem with Dirichlet boundary conditions.\n");
printf("Please wait. If you start from a very good initial guess, the first step often takes long. (E.g. if you take both coupling constants to be vanishing, the initial data is already optimal and the program cannot find any improvement for the first step.)\n");
	// field configuration:
	// point 0, SIZE+1 are boundary conditions, here: = 0
	// point SIZE+2 stores the energy
	// IN LATER OPTIMIZED VERSIONS: point SIZE+3 stores whether the field configuration is the energetically better one (0 = bad, 1 = good)
	double fields[3][SIZE+4];
	// the first field configuration is the currently found optimal, the second is used to do the variation

	initialize(fields[0]);
	initialize(fields[2]);	// just keeping the initial wave function in this field
	normalize(fields[2]);
	saveinitialconfig(fields[2]);	// saving the initial wave function to a file for drawing
	normalize(fields[0]);
	energy(fields[0]);
	double initialenergy = fields[0][SIZE+2];

	int stepbegin=0; int laststepsize = 0;
	int i = 0;

	int wigglew = WIGGLEWIDTH;

	while (laststepsize < N_ITERATIONS) {
		memcpy(fields[1],fields[0],sizeof(double)*(SIZE+4)); // copy field 0 into field 1, the one to be wiggled
		wiggle(fields[1],wigglew);
		normalize(fields[1]);
		energy(fields[1]);

		// if the new energy is better, overwrite the current optimal field config held in fields[0] by the new one
		if (fields[1][SIZE+2] < fields[0][SIZE+2])	{
			memcpy(fields[0],fields[1],sizeof(double)*(SIZE+4));
			
			if (1 == saveconfig(fields[1],goodsteps)) {		// we save only when the configuration improves
				return 1;			
			}
			laststepsize = i-stepbegin;
			printf("Improvement %d (after %d tries):\t Energy = %.10lf.\n", goodsteps, laststepsize, fields[0][SIZE+2]);	
			fprintf(filestreamconvergence,"%.10lf %.10lf\n", (double)(laststepsize % N_ITERATIONS)/(double)N_ITERATIONS*initialenergy, fields[0][SIZE+2]);	// using a modolus to cutoff the high values where the algorithm breaks off
			goodsteps++;
			stepbegin = i;
		}

	// reached convergence break-off -> reset wigglewidth and continue once more
	if ((laststepsize >= N_ITERATIONS) && (wigglew > 1)) {
		wigglew /= 2;
		printf("\tWigglewidth=%d, iteration continues.\n",wigglew);
		sleep(1);
		laststepsize = 0;
	}

	i++;
//	if (i % 10000 == 0) { printf(".");}
	}

	// savin the external potential with cutoff
	double maximum;
	maximum = maxoffield(fields[0]);
	savepotential(fields[0][SIZE+2],maximum);

	fclose(filestreamconvergence);

	createscript(goodsteps-1);

	char d;
	printf("Run gnuplot to show simulation? (y=yes,other=no)");
	scanf("%c",&d);
	if (d=='y') {
		system("gnuplot -p 'script.gp'");
		printf("Showing convergence of energy and tries needed for an improvement step.\n");
		sleep(1);
		system("gnuplot -p 'statisticsscript.gp'");
	}

	return 0;
}