/*******************************************************************************
* 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

SHARE
%{
  %include "read_table-lib"; // for Open_File

  #include <gsl/gsl_sf_legendre.h>
  #include <gsl/gsl_sf_bessel.h>
  %include "mccode-complex-lib"
  #ifndef SAXSPDB
  #define SAXSPDB
  #define SAXSPDBOrderOfHarmonics 21

  void**
  matrix2d_new (size_t esize, size_t idim, size_t jdim) {
    size_t const ptrs_size = sizeof (void*) * idim;
    size_t const row_size = esize * jdim;
    void** const rows = malloc (ptrs_size);
    for (size_t i = 0; i < idim; ++i)
      rows[i] = malloc (row_size);
    return rows;
  }
  // 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 (cdouble** Matrix, int N, int M) {
    int i, j;
    for (i = 0; i < N; ++i) {
      for (j = 0; j < M; ++j) {
        cdouble z = Matrix[i][j];
        fprintf (stderr, "(%.12e,%.12e)%s", creal (z), cimag (z), (j < M - 1) ? " " : "\n");
      }
    }
  }

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

    Polar = cplx (R * cos (Concentration), R * 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);
    cdouble** Matrix = (cdouble**)matrix2d_new (sizeof (cdouble), 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] = cplx (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] = cplx (0, 0);
          else
            Matrix[j][i] = cadd (Matrix[j][i],
                                 rmul (sqrt (4.0 * PI) * DeltaRhoProtein * Bessel[j] * Legendre[j], cmul (cpow (cplx (0, 1), cplx (j, 0)), 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