/*******************************************************************************
* McStas, neutron ray-tracing package
*         Copyright 1997-2003, All rights reserved
*         Risoe National Laboratory, Roskilde, Denmark
*         Institut Laue Langevin, Grenoble, France
*
* Component: SANSPDBFast
*
* %I
* Written by: Martin Cramer Pedersen (mcpe@nbi.dk) and Søren Kynde (kynde@nbi.dk)
* Date: October 17, 2012
* Origin: KU-Science
*
* A sample describing a thin solution of proteins using linear interpolation
* to increase computational speed. This components must be compiled with the
* -lgsl and -lgslcblas flags (and possibly linked to the appropriate libraries).
*
* %D
* This components expands the formfactor amplitude of the protein on spherical
* harmonics and computes the scattering profile using these. The expansion is
* done on amino-acid level and does not take hydration layer into account.
* The component must have a valid .pdb-file as an argument.
*
* %P
* RhoSolvent: [AA]               Scattering length density of the buffer - default is 100% D2O.
* Concentration: [mM]            Concentration of sample.
* AbsorptionCrosssection: [1/m]  Absorption cross section of the sample.
* xwidth: [m]                    Dimension of component in the x-direction.
* yheight: [m]                   Dimension of component in the y-direction.
* zdepth: [m]                    Dimension of component in the z-direction.
* SampleToDetectorDistance: [m]  Distance from sample to detector (for focusing the scattered neutrons).
* DetectorRadius: [m]            Radius of the detector (for focusing the scattered neutrons).
* qMin: [AA^-1]                  Lowest q-value, for which a point is generated in the scattering profile
* qMax: [AA^-1]                  Highest q-value, for which a point is generated in the scattering profile
* NumberOfQBins: []		Number of points generated in initalscattering profile.
* PDBFilepath: []		Path to the file describing the high resolution structure of the protein.
*
* %E
*******************************************************************************/

DEFINE COMPONENT SANSPDBFast

DEFINITION PARAMETERS ()

SETTING PARAMETERS (int NumberOfQBins = 200,
	RhoSolvent = 6.4e-14, Concentration = 0.01, AbsorptionCrosssection = 0.0,
	xwidth, yheight, zdepth,
	SampleToDetectorDistance, DetectorRadius,
	qMin = 0.001, qMax = 0.5,
	string PDBFilepath = "PDBfile.pdb")

OUTPUT PARAMETERS ()
DEPENDENCY " @GSLFLAGS@ "
NOACC

SHARE
%{
	// GSL libraries
	#include <gsl/gsl_sf_legendre.h>
	#include <gsl/gsl_sf_bessel.h>
	#include <complex.h>

        #define OrderOfHarmonics 21
	// Structs describing the protein
	struct Bead
	{
		double x;
		double y;
		double z;

		double xv;
		double yv;
		double zv;

		double xa;
		double ya;
		double za;

		char *NameOfResidue;

		double Volume;
		double ScatteringLength;
	};
	typedef struct Bead BeadStruct;

	struct Protein
	{
		BeadStruct *Beads;
		int NumberOfResidues;
	};
	typedef struct Protein ProteinStruct;

	// Simple mathematical functions
	int Sign(double x) {
		int Sign;

		if (x > 0) {
			Sign = 1;
		} else if (x < 0) {
			Sign = -1;
		} else {
			Sign = 0;
		}

		return Sign;
	}

	double complex Polar(double R, double Concentration) {
		double complex Polar;

		Polar = R * (cos(Concentration) + _Complex_I * sin(Concentration));

		return Polar;
	}

	// Function used to determine the number of residues in the .pdb-file
	int CountResidues(char PDBFilepath[256])
	{
		// Declarations
		double Dummy1;
		double Dummy2;
		double Dummy3;
		char Line[256];
		char DummyChar;
		char Atom;
		int NumberOfResidues = 0;
		int ResidueID;
		int PreviousResidueID = 0;
		FILE *PDBFile;

		// I/O
		PDBFile = fopen(PDBFilepath, "r");

		if (PDBFile == NULL) {
			printf("Cannot open %s... \n", PDBFilepath);
		} else {
			printf(".pdb-file located... \n");
		}

		while (fgets(Line, sizeof(Line), PDBFile) != NULL) {
		    ResidueID = 0;

		    if (sscanf(Line, "ATOM%*18c%d%*4c%lf%lf%lf", &ResidueID, &Dummy1, &Dummy2, &Dummy3) == 4) {

		        if (ResidueID != PreviousResidueID && ResidueID != 0) {
		            ++NumberOfResidues;
				}

		        PreviousResidueID = ResidueID;
		    }
		}

		fclose(PDBFile);
		printf("Calculating scattering from %d residues...\n", NumberOfResidues);

        return NumberOfResidues;
	}

	// Initialization of the arrays
	double complex ** ComplexMatrix(int A, int B)
	{
		int i;
		int j;
		double complex ** Matrix;

		Matrix = (double complex **) malloc(A * sizeof(double complex *));

		for (i = 0; i < A; ++i) {
		    Matrix[i] = (double complex *) malloc(B * sizeof(double complex));

		    for (j = 0; j < B; ++j) {
		        Matrix[i][j] = 0.0;
		    }
		}

		return Matrix;
	}

	void InitializeProtein(ProteinStruct *Protein, int NumberOfResiduesInFile)
	{
		Protein->NumberOfResidues = NumberOfResiduesInFile;
		Protein->Beads = calloc((size_t) NumberOfResiduesInFile, sizeof(BeadStruct));
	}

	void InitializeArray(int Size, double** Array)
	{
		// Declarations
		int i;
		double * DummyArray;

		// Initialization
		DummyArray = (double *) calloc(Size, sizeof(double));

		for (i = 0; i < Size; ++i) {
			DummyArray[i] = 0.0;
		}

        *Array = DummyArray;
	}

	// Function used to read .pdb-file
	void ReadAminoPDB(char PDBFilename[256], ProteinStruct *Protein)
	{
		// Declarations and input
		int NumberOfResidues = Protein->NumberOfResidues;
		BeadStruct *Residue = Protein->Beads;
		FILE *PDBFile;

		int i = 0;
		int PreviousResidueID = 0;
		int ResidueID = 0;

		double Weight = 0.0;
		double W = 0.0;

		double Aweight = 0.0;
		double A = 0.0;

		double x;
		double y;
		double z;

		double X = 0.0;
		double Y = 0.0;
		double Z = 0.0;

		double XA = 0.0;
		double YA = 0.0;
		double ZA = 0.0;

		char Atom;

		char Line[256];
		char Buffer[256];
		char DummyChar;

		// Atomic weighing factors
		const double WH = 5.15;
		const double WC = 16.44;
		const double WN = 2.49;
		const double WO = 9.13;
		const double WS = 19.86;
		const double WP = 5.73;

		// Scattering lengths
		const double AH =  - 3.741e-15;
		const double AD =    6.674e-15;
		const double AC =    6.648e-15;
		const double AN =    9.360e-15;
		const double AO =    5.805e-15;
		const double AP =    5.130e-15;
		const double AS =    2.847e-15;

		// Program
		if ((PDBFile = fopen(PDBFilename, "r")) == 0) {
		    printf("Cannot open file: %s. \n", PDBFilename);
		}

		Residue[i].NameOfResidue = (char *) calloc(3, sizeof(char));

		while (fgets(Buffer, sizeof(Buffer), PDBFile) != NULL) {
		    Atom = 0;
			ResidueID = 0;

			sscanf(Buffer,"ATOM%*9c%c%*8c%d%*4c%lf%lf%lf%*23c%c", &DummyChar, &ResidueID, &x, &y, &z, &Atom);

		    if (ResidueID != PreviousResidueID && ResidueID != 0) {

		        if (PreviousResidueID != 0) {

					// Assign center of scattering
				    Residue[i].xv = X / Weight;
					Residue[i].yv = Y / Weight;
					Residue[i].zv = Z / Weight;

					// Assign center of mass
				    Residue[i].x = XA / Aweight;
					Residue[i].y = YA / Aweight;
					Residue[i].z = ZA / Aweight;

					// Other residue attributes
				    Residue[i].Volume = Weight;
				    Residue[i].ScatteringLength = Aweight;

				    X = 0.0;
					Y = 0.0;
					Z = 0.0;
					Weight = 0.0;

					XA = 0.0;
					YA = 0.0;
					ZA = 0.0;
					Aweight = 0.0;

				    ++i;

				    if (i < NumberOfResidues) {
				        Residue[i].NameOfResidue = (char *) calloc(3, sizeof(char));
				    }

				}

		        PreviousResidueID = ResidueID;
		    }

			// Finish the final amino acid
			if (i == NumberOfResidues - 1) {
				Residue[i].xv = X / Weight;
				Residue[i].yv = Y / Weight;
				Residue[i].zv = Z / Weight;

				// Assign center of mass
				Residue[i].x = XA / Aweight;
				Residue[i].y = YA / Aweight;
				Residue[i].z = ZA / Aweight;

				// Other residue attributes
				Residue[i].Volume = Weight;
				Residue[i].ScatteringLength = Aweight;
			}

		    sscanf(Buffer, "ATOM%*9cCA%*2c%*s%*10c%lf%lf%lf", &Residue[i].xa, &Residue[i].ya, &Residue[i].za);
		    sscanf(Buffer, "ATOM%*13c%s", Residue[i].NameOfResidue);

		    switch(Atom) {
		        case 'C':
		            A = AC;
		            W = WC;
		        break;

		        case 'N':
		            A = AN;
		            W = WN;
		        break;

		        case 'O':
		            A = AO;
		            W = WO;
		        break;

		        case 'S':
		            A = AS;
		            W = WS;
		        break;

		        case 'H':
		            A = AH;
		            W = WH;
		        break;

		        case 'P':
		            A = AP;
		            W = WP;
		        break;

		        default:
		            A = 0.0;
		            W = 0.0;
		    }

		    Weight += W;
		    Aweight += A;

		    X += W * x;
		    Y += W * y;
		    Z += W * z;

		    XA += A * x;
		    YA += A * y;
		    ZA += A * z;
		}

		fclose(PDBFile);
	}

	// Function used to expand protein structure on spherical harmonics
	void ExpandStructure(double complex ** Matrix, ProteinStruct *Protein, int ResidueID, double q, double RhoSolvent)
	{
		// Dummy integers
		int i;
		int j;

		// Arrays used for storing Legendre and Bessel function values
		double Legendre[OrderOfHarmonics + 1];
		double Bessel[OrderOfHarmonics + 1];

		// Residue information
		const double Volume = Protein->Beads[ResidueID].Volume;
		const double DeltaRhoProtein = Protein->Beads[ResidueID].ScatteringLength - Volume * RhoSolvent;

		const double x = (Protein->Beads[ResidueID].x * Protein->Beads[ResidueID].ScatteringLength -
						  RhoSolvent * Volume * Protein->Beads[ResidueID].xv) / DeltaRhoProtein;

		const double y = (Protein->Beads[ResidueID].y * Protein->Beads[ResidueID].ScatteringLength -
						  RhoSolvent * Volume * Protein->Beads[ResidueID].yv) / DeltaRhoProtein;

		const double z = (Protein->Beads[ResidueID].z * Protein->Beads[ResidueID].ScatteringLength -
						  RhoSolvent * Volume * Protein->Beads[ResidueID].zv) / DeltaRhoProtein;

		// Convert bead position to spherical coordinates
		const double Radius = sqrt(pow(x, 2) + pow(y, 2) + pow(z, 2));
		const double Theta  = acos(z / Radius);
		const double Concentration    = acos(x / (Radius * sin(Theta))) * Sign(y);

		// Expand protein structure on harmonics
		gsl_sf_bessel_jl_array(OrderOfHarmonics, q * Radius, Bessel);

		for (i = 0; i <= OrderOfHarmonics; ++i) {
		    gsl_sf_legendre_array(OrderOfHarmonics, i, cos(Theta), &Legendre[i]);

		    for(j = i; j <= OrderOfHarmonics; ++j) {
		        Matrix[j][i] += sqrt(4.0 * PI) * cpow(_Complex_I, j) * DeltaRhoProtein * Bessel[j] * Legendre[j] * Polar(1.0, -i * Concentration);
		    }
		}
	}

	// Function used to generate intensity for a given q-value
	double ComputeIntensity(double complex ** Matrix, int Size)
	{
		// Declarations
		int i;
		int j;
		double Intensity = 0.0;

		// Computation
		for (i = 0; i <= Size; ++i) {

		    for (j = 0; j <= i; ++j) {
		        Intensity += ((j > 0) + 1.0) * pow(cabs(Matrix[i][j]), 2);
		    }
		}

		return Intensity;
	}

	// Function used to reinitialize a matrix
	void ResetMatrix(double complex ** Matrix, int Size)
	{
		int i;
		int j;

		for (i = 0; i <= Size; ++i) {

		    for(j = 0; j <= Size; ++j) {
		        Matrix[i][j] = 0.0;
		    }
		}
	}

%}

DECLARE 
%{
	// Declarations
	double Absorption;
	double NumberDensity;

	// Arrays for storing q and I(q)
	double * qArray;
	double * IArray;
%}

INITIALIZE
%{
	// Rescale concentration into number of aggregates per m^3 times 10^-4
	NumberDensity = Concentration * 6.02214129e19;


	// Standard sample handling
	if (!xwidth || !yheight || !zdepth) {
		printf("%s: Sample has no volume - check parameters...! \n", NAME_CURRENT_COMP);
	}

	Absorption = AbsorptionCrosssection;

	// Declarations
	int i;
	int j;
	double qDummy;
	const double qStep = (qMax - qMin) / (1.0 * NumberOfQBins);

	// Initialize matrix
	double complex ** Matrix = ComplexMatrix(OrderOfHarmonics + 1, OrderOfHarmonics + 1);

	// Initialize protein
	int NumberOfResidues;
	ProteinStruct Protein;

	printf("Initializing protein structure...\n");
    NumberOfResidues = CountResidues(PDBFilepath);
	InitializeProtein(&Protein, NumberOfResidues);

	printf("Initializing arrays...\n");
	InitializeArray(NumberOfQBins, &qArray);
	InitializeArray(NumberOfQBins, &IArray);

	printf("Creating protein structure...\n");
	ReadAminoPDB(PDBFilepath, &Protein);

	// Computing scattering profile
	printf("Computing scattering from protein...\n");
	for (i = 0; i < NumberOfQBins; ++i) {
		ResetMatrix(Matrix, OrderOfHarmonics);

		qDummy = qMin + (i + 0.5) * qStep;

		for (j = 0; j < NumberOfResidues; ++j) {
			ExpandStructure(Matrix, &Protein, j, qDummy, RhoSolvent);
		}

		qArray[i] = qDummy;
	 	IArray[i] = ComputeIntensity(Matrix, OrderOfHarmonics);
	}

	printf("Initializations complete...\n");
%}

TRACE
%{
	// Declarations
	double t0;
	double t1;
	double l_full;
	double l;
	double l1;
	double Intensity;
	double Weight;
	double IntensityPart;
	double SolidAngle;
	double q;
	double qx;
	double qy;
	double qz;
	double v;
	double dt;
	double vx_i;
	double vy_i;
	double vz_i;
	char Intersect = 0;
	double Slope;
	double Offset;
	int i;

	// Computation
	Intersect = box_intersect(&t0, &t1, x, y, z, vx, vy, vz, xwidth, yheight, zdepth);

	if (Intersect) {

		if (t0 < 0.0) {
			fprintf(stderr, "Neutron already inside sample %s - absorbing...\n", NAME_CURRENT_COMP);
			ABSORB;
    	}

		// Compute properties of neutron
		v = sqrt(pow(vx, 2) + pow(vy, 2) + pow(vz, 2));
		l_full = v * (t1 - t0);
		dt = rand01() * (t1 - t0) + t0;
		PROP_DT(dt);
	    l = v * (dt - t0);

		// Store properties of incoming neutron
		vx_i = vx;
		vy_i = vy;
		vz_i = vz;

		// Generate new direction of neutron
		randvec_target_circle(&vx, &vy, &vz, &SolidAngle, 0, 0, SampleToDetectorDistance, DetectorRadius);

		NORM(vx, vy, vz);

		vx *= v;
		vy *= v;
		vz *= v;

		// Compute q
		qx = V2K * (vx_i - vx);
		qy = V2K * (vy_i - vy);
		qz = V2K * (vz_i - vz);

		q = sqrt(pow(qx, 2) + pow(qy, 2) + pow(qz, 2));

		// Discard neutron, if q is out of range
		if ((q < qArray[0]) || (q > qArray[NumberOfQBins - 1])) {
		    ABSORB;
		}

		// Find the first value of q in the curve larger than that of the neutron
		i = 1;

		while (q > qArray[i]) {
			++i;
		}

		// Do a linear interpolation
		l1 = v * t1;

		Slope = (IArray[i] - IArray[i - 1]) / (qArray[i] - qArray[i - 1]);
		Offset = IArray[i] - Slope * qArray[i];

		Intensity = (Slope * q + Offset);

		p *= l_full * SolidAngle / (4.0 * PI) * NumberDensity * Intensity * exp(- Absorption * (l + l1) / v);

		SCATTER;
	}
%}

MCDISPLAY
%{

	box(0, 0, 0, xwidth, yheight, zdepth);
%}

END
