/*******************************************************************************
* McXtrace, X-ray tracing package
*           Copyright, All rights reserved
*           Risoe National Laboratory, Roskilde, Denmark
*           Institut Laue Langevin, Grenoble, France
*
* Component: SAXSPDBFast
*
* %I
* Written by: Martin Cramer Pedersen (mcpe@nbi.dk) and Søren Kynde (kynde@nbi.dk)
* Date: May 2, 2012
* Origin: KU-Science
* Release: McXtrace 1.0
*
* 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.
*
* This is fast implementation of the SAXSPDB sample component.
*
* %P
* RhoSolvent: [AA]    Scattering length density of the buffer.
* 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 x-rays).
* DetectorRadius: [m] Radius of the detector (for focusing the scattered x-rays).
* 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 inital scattering profile.
* PDBFilepath: []     Path to the file describing the high resolution structure of the protein.
*
* %E
*******************************************************************************/

DEFINE COMPONENT SAXSPDBFast



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


DEPENDENCY " @GSLFLAGS@ "
NOACC
/*X-ray PARAMETERS (x, y, z, kx, ky, kz, phi, t, Ex, Ey, Ez, p)*/

SHARE
%{
  %include "read_table-lib"; // for Open_File
  
	#include <gsl/gsl_sf_legendre.h>
	#include <gsl/gsl_sf_bessel.h>
	#include <complex.h>
#ifndef SAXSPDB
#define SAXSPDB
	#define SAXSPDBOrderOfHarmonics 21

	// 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;
	}
	
	void complex_print_matrix(double complex **Matrix, int N, int M)
	{
	  int i,j;
	  for (i = 0; i < N; ++i)
    {
      for (j = 0; j < M; ++j)
        {
          double complex z = Matrix[i][j];
          fprintf(stderr,
                  "(%.12e,%.12e)%s",
                  creal(z),
                  cimag(z),
                  (j < M - 1) ? " " : "\n");
        }
    }
	}

	double complex Polar(double R, double Concentration) {
		double complex Polar;
		
		Polar = R * (cos(Concentration) + _Complex_I * sin(Concentration));

		return Polar;
	} 

	// Protein structs
	struct Bead
	{
		double x;
		double y;
		double z;

		double xv;
		double yv;
		double zv;

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

	struct Protein
	{
		BeadStruct *Beads;
		int NumberOfResidues;
	};
	typedef struct Protein ProteinStruct;
	
	// functions for the INITIALIZE ----------------------------------------------

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

		// I/O
		PDBFile = Open_File(PDBFilepath, "r",NULL);
		if (PDBFile == NULL) {
			exit(fprintf(stderr, "SAXSPDBFast: %s: ERROR: Cannot open %s... \n", __FILE__, PDBFilepath));
		}

		while (fgets(Line, sizeof(Line), PDBFile) != NULL) {
		    ResidueID = 0;
        if (strncmp(Line, "ATOM", 4)) continue;
		    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);
    return NumberOfResidues;
	} // CountResidues

	// Function used to read .pdb-file
	int ReadAminoPDB(char *PDBFilename, 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 Buffer[65535];
		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 =  1 * 2.82e-13;
		const double AD =  1 * 2.82e-13;
		const double AC =  6 * 2.82e-13;
		const double AN =  7 * 2.82e-13;
		const double AO =  8 * 2.82e-13;
		const double AP = 15 * 2.82e-13;
		const double AS = 16 * 2.82e-13;

		// Program
		if (NumberOfResidues <= 0 || (PDBFile = Open_File(PDBFilename, "r",NULL)) == 0) {
		    exit(printf("ERROR: Cannot open file: %s. \n", PDBFilename));
		}
		
		while (fgets(Buffer, sizeof(Buffer), PDBFile) != NULL) {
		    // a typical line is:
		    // ATOM   8726  N   VAL B 576      76.450  47.214  58.026  1.00111.85           N
		    Atom = 0; 
        ResidueID = 0;
        if (strncmp(Buffer, "ATOM", 4)) continue;
        if (!sscanf(Buffer,"ATOM%*9c%c%*8c%d%*4c%lf%lf%lf%*23c%c", &DummyChar, &ResidueID, &x, &y, &z, &Atom)) {
          fprintf(stderr, "SAXSPDBFast: %s: ReadAminoPDB: [%i] invalid PDB line %s\n", __FILE__, i, Buffer);
          continue;
        }

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

            if (PreviousResidueID != 0 && Aweight && Weight) {

                // 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;
                Residue[i].Atom             = Atom;

                X  = Y  = Z  = Weight  = 0.0;
                XA = YA = ZA = Aweight = 0.0;

                ++i;

            }

		        PreviousResidueID = ResidueID;
		    }

        // Finish the final amino acid
        if (i == NumberOfResidues - 1 && Aweight && Weight) {
            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;
            Residue[i].Atom             = 'X';
        }

		    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);
		
		return(NumberOfResidues);
	} // ReadAminoPDB\

	
#endif /*SAXSPDB*/
%}

DECLARE
%{
	double Absorption;
	double NumberDensity;

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

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

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

  // count the number of residues
	Absorption               = AbsorptionCrosssection;
  Protein.NumberOfResidues = CountResidues(PDBFilepath);
  Protein.Beads            = calloc(Protein.NumberOfResidues,sizeof(BeadStruct));
  if (Protein.Beads == NULL) 
    exit(fprintf(stderr, "SAXSPDB: %s: ERROR: memory allocation\n",  NAME_CURRENT_COMP));

	qArray = create_darr1d(NumberOfQBins);
  IArray = create_darr1d(NumberOfQBins);

  // initialize the protein from the PDB
	ReadAminoPDB(PDBFilepath, &Protein);
	MPI_MASTER(
	printf("SAXSPDBFast: %s: Initializing scattering from %s with %d residues on %d Q-values\n", 
	  NAME_CURRENT_COMP, PDBFilepath, Protein.NumberOfResidues, NumberOfQBins);
  );

	// Computing scattering profile I(q)
	for (qbin = 0; qbin < NumberOfQBins; ++qbin) {
	    int    i,j,ResidueID;
	    double qStep = (qMax - qMin) / (1.0 * NumberOfQBins);
	    double q     =  qMin + qStep * (qbin + 0.5);
	    double complex Matrix[SAXSPDBOrderOfHarmonics+1][SAXSPDBOrderOfHarmonics+1];
	    // init Matrix = 0
	    // ResetMatrix(Matrix, SAXSPDBOrderOfHarmonics);
	    for (i = 0; i <= SAXSPDBOrderOfHarmonics; ++i)
		    for (j = 0; j <= SAXSPDBOrderOfHarmonics; ++j)
		        Matrix[i][j] = 0.0;
      

	    for (ResidueID = 0; ResidueID < Protein.NumberOfResidues; ++ResidueID) {
	      // ExpandStructure(Matrix, &Protein, ResidueID, qDummy, RhoSolvent);		
	      double Legendre[SAXSPDBOrderOfHarmonics + 1];
		    double Bessel[SAXSPDBOrderOfHarmonics + 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 C      = acos(x / (Radius * sin(Theta))) * Sign(y);

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

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

		        for(j = 0; j <= SAXSPDBOrderOfHarmonics; ++j) {
		            if (j < i) Matrix[j][i] = 0;
		            else
		            Matrix[j][i] += sqrt(4.0 * PI) * cpow(_Complex_I, j) * DeltaRhoProtein * Bessel[j] * Legendre[j] * Polar(1.0, -i * C);
		        }
		    }
		    
	    } // for ResidueID

	    qArray[qbin] = q;
	    // IArray[qbin] = ComputeIntensity(Matrix, SAXSPDBOrderOfHarmonics);
	    IArray[qbin] = 0;
	    for (i = 0; i <= SAXSPDBOrderOfHarmonics; ++i) {
		    for (j = 0; j <= i; ++j) {
		        IArray[qbin] += ((j > 0) + 1.0) * pow(cabs(Matrix[i][j]), 2);
		    }
		  }
		  // printf("I(q=%g) = %g\n", q, IArray[qbin]);
	} // for qbin
	MPI_MASTER(
  	printf("SAXSPDBFast: %s: %s initialization I(q) done\n", NAME_CURRENT_COMP, PDBFilepath);
  );

%}

TRACE
%{
	// Declarations	
	double l0; 
	double l1; 
	double l_full;
	double l;
	double l_1;
	double q;
	double Intensity;
	double Weight;
	double IntensityPart;
	double SolidAngle;
	double qx; 
	double qy; 
	double qz;
	double k;
	double dl;
	double kx_i;
	double ky_i;
	double kz_i;
	char Intersect = 0;
	double Slope;	
	double Offset;
	int i;

	// Computation
  Intersect = box_intersect(&l0, &l1, x, y, z, kx, ky, kz, xwidth, yheight, zdepth);
    
	if (Intersect) {

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

		// Compute properties of photon
		k = sqrt(pow(kx, 2) + pow(ky, 2) + pow(kz, 2));
		l_full = l1 - l0;
		dl = rand01() * (l1 - l0) + l0; 
		PROP_DL(dl);                  
		l = dl - l0;

		// Store properties of incoming photon
		kx_i = kx;
		ky_i = ky;
		kz_i = kz;

		// Generate new direction of photon
		randvec_target_circle(&kx, &ky, &kz, &SolidAngle, 0, 0, SampleToDetectorDistance, DetectorRadius);

		NORM(kx, ky, kz);

		kx *= k;
		ky *= k;
		kz *= k;

		// Compute q
		qx = kx_i - kx;
		qy = ky_i - ky;
		qz = kz_i - kz;

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

		// Discard photon, 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 photon
		i = 1;

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

		// Do a linear interpolation
		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));

		SCATTER;
  }
  %}

MCDISPLAY
%{
  box(0, 0, 0, xwidth, yheight, zdepth,0, 0, 1, 0);
%}

END