#pragma warning(disable:4996)

#include <stdio.h>
#include <string.h>
#include <vector>
#include <random>
#include <cmath>

#define DT  1.0e-3    /* time step */
#define T_END  2.0    /* end of simulation */

#define N  100  /* total number of neurons */
#define A    8  /* the number of clusters */

#define STAT_FLAG false /* stationary(true) or non-stationary(false) */

#define SEED 1 

int ClusterSize[A] = { 2 * (N / 10), (N / 10) , (N / 10) , (N / 10) , 2 * (N / 10) , (N / 10) , (N / 10) , (N / 10) };
std::mt19937_64 RandHost(SEED);

/*====<< common parameter of neuron >>===========*/
#define I_EXT  0.3      /* external current [a.u] */
#define EXCIT_CHEM_W  0.08 /* weights of excitatory chemical synapse [a.u] */
#define INHIB_CHEM_W -0.08 /* weights of inhibitory chemical synapse [a.u] */
#define GAP_W 0.3       /* weights of gap junction [a.u.] */
#define IN_CHEM_R  0.5 /* synaptic connection probability among neurons in same ensemble */
#define OUT_CHEM_R 0.5 /* synaptic connection probability among neurons in different ensemble */
#define GAP_R 0.8      /* gap connection probability among neurons */ 

/*====<< common parameter of noise >>===========*/
#define SIGMA_NOISE  0.2 /* standard deviation of Gaussian noise to the neural activity [a.u] */

/*====<< common parameter of synaptic connection >>===========*/
#define V_half   0.3   /* half maximum point of the logistic curve (chemical synapse) [a.u] */
#define V_width  0.1   /* half width at half maximum point of the logistic curve (chemical synapse) [a.u] */

/*====<< common parameter of fluorescence intensity >>===========*/
#define Fmax  5.0    /* max fluorescence intensity */
#define K_D   0.3    /* dissociation constant (half-maximal response constant in the Hill equation) [a.u] */
#define Hc    3.0    /* Hill coefficient (2-5 for GCaMPs) */
#define F_tau 0.1    /* time delay constant of fluorescence intensity */

class Variable;
class Parameter;
double affinity(double arg_x);
double I_ext(int arg_i, double arg_t, Parameter arg_param);
Variable Euler_Maruyama(double arg_t, Variable arg_neuron, Parameter arg_param);

/*=================================*/
/* Variables of neuron             */
/*=================================*/
class Variable {
public:
	std::vector<double> V;  /* neural activity */
	std::vector<double> F;  /* fluorescence intensity */

	void Init() {
		V.resize(N);
		F.resize(N);
		/* itial value of the variable */
		std::uniform_real_distribution<> RealDist(0.0, 0.1);
		for (int i = 0; i < N; i++) {
			V[i] = RealDist(RandHost);
			F[i] = affinity(V[i]);
		}

	}
	void Free() {
		V.clear();
		V.shrink_to_fit();
		F.clear();
		F.shrink_to_fit();
	}
};

/*================================================*/
/* Parameters of neuron and synaptic connections  */
/*================================================*/
class Parameter {
public:
	std::vector<int> ClusterID;  /* Cluster ID of neuron */
	std::vector<double> VAct;    /* action potential (firing state of neural activity) */
	std::vector<double> VRest;   /* resting membrane potential */
	std::vector<double> VTh;     /* firing threshold of neuron */
	std::vector<double> VInit;   /* initial value of neural activity */
	std::vector<double> Tau;     /* membrane time constant of neuron */
	std::vector<double> R;       /* membrane resistance of neuron */
	std::vector<std::vector<double>> WChem; /* synaptic weight of chemical synapse (positive for excitatory, negative for inhibitory) */
	std::vector<std::vector<double>> WGap;  /* synaptic weight of gap junction */

	void Init() {
		/* cluster ID of neuron */
		ClusterID.resize(N);
		int l = 0; 
		for (int k = 0; k < A; k++) {
			for (int j = 0; j < ClusterSize[k]; j++) {
				ClusterID[l] = k + 1;
				l++;
			}
		}
		/* parameters of neuron */
		VAct.resize(N);
		VRest.resize(N);
		VTh.resize(N);
		VInit.resize(N);
		Tau.resize(N);
		R.resize(N);
		for (int i = 0; i < N; i++) {
			VAct[i] = 1.0;
			VTh[i] = 0.2;
			VRest[i] = 0.0;
			VInit[i] = 0.0;
			R[i] = 1.0;
			if (ClusterID[i] <= 4) {
				Tau[i] = 5.0e-2; /* fast neuron in C1-C4 [sec] */
			}
			else {
				Tau[i] = 1.0e-1; /* fast neuron in C5-C8 [sec] */
			}
		}
		/* synaptic weights of chemical synapse and gap junction */
		std::uniform_real_distribution<> RealDist(0.0, 1.0);
		WChem.resize(N);
		WGap.resize(N);
		for (int i = 0; i < N; i++) {
			WChem[i].resize(N);
			WGap[i].resize(N);
		}
		for (int i = 0; i < N; i++) {
			for (int j = 0; j < N; j++) {
				WChem[i][j] = 0.0;
				WGap[i][j] = 0.0;
				if (ClusterID[i] == ClusterID[j]) {
					if (RealDist(RandHost) < IN_CHEM_R) { /* excitatory chemical synapse [a.u.] */
						WChem[i][j] = EXCIT_CHEM_W;
					}
					if (RealDist(RandHost) < GAP_R) { /* gap junction [a.u.] */
						WGap[i][j] = GAP_W;
						WGap[j][i] = WGap[i][j]; /* bidirectional connection */
					}
					else { /* gap junction [a.u.] */
						WGap[i][j] = 0.0;
						WGap[j][i] = WGap[i][j]; /* bidirectional connection */
					}
				}
				else if ((ClusterID[i] == 2) && (ClusterID[j] == 1)) { /* C1 --> C2 */
					if (RealDist(RandHost) < OUT_CHEM_R) {
						WChem[i][j] = EXCIT_CHEM_W; /* excitatory chemical synapse [a.u.] */
					}
				}
				else if ((ClusterID[i] == 3) && (ClusterID[j] == 1)) { /* C1 --> C3 */
					if (RealDist(RandHost) < OUT_CHEM_R) {
						WChem[i][j] = INHIB_CHEM_W; /* inhibitory chemical synapse [a.u.] */
					}
				}
				else if ((ClusterID[i] == 5) && (ClusterID[j] == 6)) { /* C6 --> C5 */
					if (RealDist(RandHost) < OUT_CHEM_R) {
						WChem[i][j] = EXCIT_CHEM_W; /* excitatory chemical synapse [a.u.] */
					}
				}
				else if ((ClusterID[i] == 5) && (ClusterID[j] == 7)) { /* C7 --> C5 */
					if (RealDist(RandHost) < OUT_CHEM_R) {
						WChem[i][j] = INHIB_CHEM_W; /* inhibitory chemical synapse [a.u.] */
					}
				}
			}
		}
	}
	void Free() {
		ClusterID.clear();
		ClusterID.shrink_to_fit();
		VAct.clear();
		VAct.shrink_to_fit();
		VRest.clear();
		VRest.shrink_to_fit();
		VTh.clear();
		VTh.shrink_to_fit();
		VInit.clear();
		VInit.shrink_to_fit();
		Tau.clear();
		Tau.shrink_to_fit();
		R.clear();
		R.shrink_to_fit();
		for (int i = 0; i < N; i++) {
			WChem[i].clear();
			WChem[i].shrink_to_fit();
			WGap[i].clear();
			WGap[i].shrink_to_fit();
		}
		WChem.clear();
		WChem.shrink_to_fit();
		WGap.clear();
		WGap.shrink_to_fit();
	}
};

/*==================================================================*/
/* Affinity function (the Hill equation) for fluorescence intensity */
/* Hill equation: f(x) = fmax * x^h / ( x^h + (K_D)^h ) for x > 0   */
/*==================================================================*/
double affinity(double arg_x)
{
	double r;

	if (arg_x > 0.0) {
		r = Fmax * std::pow(arg_x, Hc) / (std::pow(arg_x, Hc) + std::pow(K_D, Hc));
	}
	else {
		r = 0.0;
	}

	return r;
}

/*==================================================================*/
/* Time-dependent external current                                  */
/*==================================================================*/
double I_ext(int arg_i, double arg_t, Parameter arg_param)
{
	/* stationary */
	if (STAT_FLAG) {
		return I_EXT;
	}
	/* non-stationary */
	else {
		double r = 0.0;
		if ((arg_param.ClusterID[arg_i] == 1)
			|| (arg_param.ClusterID[arg_i] == 5)) {
			if ((((T_END / 4.0) * 0) < arg_t)
				&& (arg_t < ((T_END / 4.0) * 1))) {
				r = I_EXT;
			}
			else if ((((T_END / 4.0) * 2) < arg_t)
				&& (arg_t < ((T_END / 4.0) * 4))) {
				r = I_EXT;
			}
		}
		else if ((arg_param.ClusterID[arg_i] == 2)
			|| (arg_param.ClusterID[arg_i] == 6)) {
			if ((((T_END / 4.0) * 1) < arg_t)
				&& (arg_t < ((T_END / 4.0) * 2))) {
				r = I_EXT;
			}
			else if ((((T_END / 4.0) * 3) < arg_t)
				&& (arg_t < ((T_END / 4.0) * 4))) {
				r = I_EXT;
			}
		}
		else if ((arg_param.ClusterID[arg_i] == 3)
			|| (arg_param.ClusterID[arg_i] == 7)) {
			if ((((T_END / 4.0) * 2) < arg_t)
				&& (arg_t < ((T_END / 4.0) * 1))) {
				r = I_EXT;
			}
		}
		else if ((arg_param.ClusterID[arg_i] == 4)
			|| (arg_param.ClusterID[arg_i] == 8)) {
			if ((((T_END / 4.0) * 0) < arg_t)
				&& (arg_t < ((T_END / 4.0) * 1))) {
				r = I_EXT;
			}
			else if ((((T_END / 4.0) * 2) < arg_t)
				&& (arg_t < ((T_END / 4.0) * 3))) {
				r = I_EXT;
			}
		}

		return r;
	}
}

/*= ==========================================================================*/
/* Stochastic differential equations are solved by the Euler-Maruyama method */
/*===========================================================================*/
Variable Euler_Maruyama(double arg_t, Variable arg_neuron, Parameter arg_param)
{
	double I_chem, I_gap, rnd_sigma;
	Variable r = arg_neuron;

	rnd_sigma = SIGMA_NOISE * std::sqrt(DT);  /* noise intensity in the Euler-Maruyama method */
	std::normal_distribution<> NormDist(0.0, rnd_sigma);
	for (int i = 0; i < N; i++) {
		/* noisy neural activity: Leaky integrate-and-fire model with noise */
		if (arg_neuron.V[i] == arg_param.VAct[i]) {  /* (1) if previous state is the firing state */
			r.V[i] = arg_param.VInit[i];  /* reset to the initial value */
		}
		else if (arg_neuron.V[i] >= arg_param.VTh[i]) {  /* (2) if neural activity exceeds the threshold */
			r.V[i] = arg_param.VAct[i];  /* neuron generates action potential */
		}
		else {   /* (3) stochastic differential equation */
			I_chem = 0.0;
			I_gap = 0.0;
			for (int j = 0; j < N; j++) {
				I_chem += arg_param.WChem[i][j] / (1.0 + std::exp(-(arg_neuron.V[j] - V_half) / V_width));
				I_gap += arg_param.WGap[i][j] * (arg_neuron.V[j] - arg_neuron.V[i]);
			}

			r.V[i] = arg_neuron.V[i]
			       + DT * (-(arg_neuron.V[i] - arg_param.VRest[i]) + arg_param.R[i] * (I_chem + I_gap + I_ext(i, arg_t, arg_param)) + NormDist(RandHost)) / arg_param.Tau[i];
			if (r.V[i] > arg_param.VTh[i]) {
				r.V[i] = arg_param.VTh[i];
			}
		}

		/* fluorescence intensity */
		r.F[i] = arg_neuron.F[i] + DT * (affinity(arg_neuron.V[i]) - arg_neuron.F[i]) / F_tau;
	}

	return r;
}


int main()
{
	Variable Neuron;
	Parameter Param;
	Neuron.Init();
	Param.Init();

	FILE* fp1, *fp2, *fp3, *fp4;

	/* output the cluster ID of neuron */
	fp1 = fopen("Result_ClusterID.txt", "w");
	fprintf(fp1, "#neuron_ID  cluster_ID\n");
	for (int i = 0; i < N; i++) {
		fprintf(fp1, "%-12d %d\n", i + 1, Param.ClusterID[i]);
	}
	fclose(fp1);

	/* generate stochastic neural activity data by the Euler-Maruyama method */
	fp1 = fopen("Result_V.csv", "w");
	fp2 = fopen("Result_F.csv", "w");
	for (double t = 0.0; t < T_END; t += DT) {
		/* output the variables */
		for (int i = 0; i < N; i++) {
			fprintf(fp1, "%lf,", Neuron.V[i]);
			fprintf(fp2, "%lf,", Neuron.F[i]);
		}
		fprintf(fp1, "\n");
		fprintf(fp2, "\n");
		/* variables at next time step (Euler-Maruyama method) */
		Neuron = Euler_Maruyama(t, Neuron, Param);
	}
	fclose(fp1);
	fclose(fp2);

	fp3 = fopen("Result_WChem.csv", "w");
	fp4 = fopen("Result_WGap.csv", "w");
	for (int i = 0; i < N; i++) {
		for (int j = 0; j < N; j++) {
			fprintf(fp3, "%lf,", Param.WChem[i][j]);
			fprintf(fp4, "%lf,", Param.WGap[i][j]);
		}
		fprintf(fp3, "\n");
		fprintf(fp4, "\n");
	}
	fclose(fp3);
	fclose(fp4);

	Neuron.Free();
	Param.Free();

	return 0;
}