/****************************************************************************
*
* McStas, neutron ray-tracing package
*         Copyright 1997-2003, All rights reserved
*         Risoe National Laboratory, Roskilde, Denmark
*         Institut Laue Langevin, Grenoble, France
*
* Component: Pol_bender
*
* %I
* Written by: Peter Christiansen
* Date: August 2006
* Origin: RISOE
*
* Polarising bender.
*
* %D
* Based on Guide_curved written by Ross Stewart.
* Models a rectangular curved guide tube with entrance centered on the Z axis.
* The entrance lies in the X-Y plane.  Draws a true depiction
* of the guide with multiple slits (but without spacers), and trajectories.
* It relies on similar physics as the Monochromator_pol.
* The reflec function and parameters are passed to this component to
* give a bigger freedom.
* The up direction is hardcoded to be along the y-axis (0, 1, 0)
*
* The guide is asummed to have half a spacer on each side:
*    slit1  slit2  slit3
*  |+     ++     ++     +|
*  |+     ++     ++     +|
*  <---------------------> xwidth
*  <------> xwidth/nslit (nslit=3)
*         <> d
*
* The reflection functions and parameters defaults as follows:
* Bot defaults to Top, Left defaults to Top, Right defaults to left.
* Down defaults to down and up defaults to up for all functions and
* Top(Up and Down) defaults to StdReflecFunc and {0.99,0.0219,6.07,2.0,0.003}
* which stands for {R0, Qc, alpha, m, W}.
*
* Example:
* Pol_bender(xwidth = 0.08, yheight = 0.08, length = 1.0, radius= 10.0,
* 	     nslit=5, d=0.0, endFlat=0, drawOption=2,
* 	     rTopUpPar={0.99, 0.0219, 6.07, 3.0, 0.003},
* 	     rTopDownPar={0.99, 0.0219, 6.07, 1.0, 0.003})
*
* See also the example instruments Test_Pol_Bender and
* Test_Pol_Bender_Vs_Guide_Curved (under tests).
*
* %BUGS
* This component has been against tested Guide_curved and found to
* give the same intensities. Gravity option has not been tested.
*
* GRAVITY: YES (when gravity is along y-axis)
*
* %P
* INPUT PARAMETERS:
*
* xwidth: [m]        Width at the guide entry 
* yheight: [m]       Height at the guide entry 
* length: [m]        length of guide along center 
* radius: [m]        Radius of curvature of the guide (+:curve left/-:right) 
* G: [m/s^2]         Gravitational constant
* nslit: [1]         Number of slits 
* d: [m]             Width of spacers (subdividing absorbing walls) 
* endFlat: [1]       If endflat>0 then entrance and exit planes are parallel. 
* rTopUpPar: [1]     Top mirror Parameters for spin up standard reflectivity function
* rTopDownPar: [1]   Top mirror Parameters for spin down standard reflectivity function
* rBotUpPar: [1]     Bottom mirror Parameters for spin up standard reflectivity function
* rBotDownPar: [1]   Bottom mirror Parameters for spin down standard reflectivity function
* rLeftUpPar: [1]    Left mirror Parameters for spin up standard reflectivity function
* rLeftDownPar: [1]  Left mirror Parameters for spin down standard reflectivity function
* rRightUpPar: [1]   Right mirror Parameters for spin up standard reflectivity function
* rRightDownPar: [1] Right mirror Parameters for spin down standard reflectivity function
* rTopUpData: [1]     Reflectivity file for top mirror, spin up 
* rTopDownData: [1]   Reflectivity file for top mirror, spin down 
* rBotUpData: [1]     Reflectivity file for bottom mirror, spin up
* rBotDownData: [1]   Reflectivity file for bottom mirror, spin down 
* rLeftUpData: [1]    Reflectivity file for left mirror, spin up
* rLeftDownData: [1]  Reflectivity file for left mirror, spin down 
* rRightUpData: [1]   Reflectivity file for right mirror, spin up
* rRightDownData: [1] Reflectivity file for right mirror, spin down 
* drawOption: [1]    1: fine(all slits/90 points per arc), 2: normal (max 20/40), 3: rough (max 5/10) 
* debug: [1]         if debug > 0 print out some internal parameters 
*
* CALCULATED PARAMETERS:
*
* localG: [m/s/s]    Gravity vector in guide reference system 
* normalXXX: [1]     Several normal vector used for defining the geometry 
* pointXXX: [1]      Several points used for defining the geometry 
* rXXXParPtr: []     Pointers to reflection parameters used with ref. functions.
*
* %L
*
* %E
*******************************************************************************/

DEFINE COMPONENT Pol_bender

SETTING PARAMETERS (xwidth, yheight, length, radius, G=9.8, int nslit=1, d=0.0, int debug=0, int endFlat=0,
vector rTopUpPar={0.99,0.0219,6.07,2.0,0.003},
vector rTopDownPar={0.99,0.0219,6.07,2.0,0.003},
vector rBotUpPar={0.99,0.0219,6.07,2.0,0.003},
vector rBotDownPar={0.99,0.0219,6.07,2.0,0.003},
vector rLeftUpPar={0.99,0.0219,6.07,2.0,0.003},
vector rLeftDownPar={0.99,0.0219,6.07,2.0,0.003},
vector rRightUpPar={0.99,0.0219,6.07,2.0,0.003},
vector rRightDownPar={0.99,0.0219,6.07,2.0,0.003},
string rTopUpData="", string rTopDownData="",string rBotUpData="",string rBotDownData="",
string rLeftUpData="", string rLeftDownData="",string rRightUpData="",string rRightDownData="",
int drawOption=1)
/* Neutron parameters: (x,y,z,vx,vy,vz,t,sx,sy,sz,p) */

SHARE
%{
  %include "pol-lib"
  %include "ref-lib"
%}

DECLARE
%{
  Coords localG;
  Coords normTopBot;
  Coords normIn;
  Coords normOut;
  Coords pointTop;
  Coords pointBot;
  Coords pointIn;
  Coords pointOut;

  t_Table rTopUpTable;
  t_Table rTopDownTable;
  t_Table rBotUpTable;
  t_Table rBotDownTable;
  t_Table rLeftUpTable;
  t_Table rLeftDownTable;
  t_Table rRightUpTable;
  t_Table rRightDownTable;
  int useTables;
%}

INITIALIZE
%{
  double angle;

  if (strlen (rTopUpData) && strcmp (rTopUpData, "NULL")) {
    useTables = 1;
    /*if rUpTopData is set assume we'll be usning tabled data for all reflectivities*/
    if (Table_Read (&rTopUpTable, rTopUpData, 1) <= 0) {
      fprintf (stderr, "Pol_bender: %s: can not read file %s\n", NAME_CURRENT_COMP, rTopUpData);
      exit (1);
    }
    if (Table_Read (&rTopDownTable, rTopDownData, 1) <= 0) {
      fprintf (stderr, "Pol_bender: %s: can not read file %s\n", NAME_CURRENT_COMP, rTopDownData);
      exit (1);
    }
    if (Table_Read (&rBotUpTable, rBotUpData, 1) <= 0) {
      fprintf (stderr, "Pol_bender: %s: can not read file %s\n", NAME_CURRENT_COMP, rBotUpData);
      exit (1);
    }
    if (Table_Read (&rBotDownTable, rBotDownData, 1) <= 0) {
      fprintf (stderr, "Pol_bender: %s: can not read file %s\n", NAME_CURRENT_COMP, rBotDownData);
      exit (1);
    }
    if (Table_Read (&rLeftUpTable, rLeftUpData, 1) <= 0) {
      fprintf (stderr, "Pol_bender: %s: can not read file %s\n", NAME_CURRENT_COMP, rLeftUpData);
      exit (1);
    }
    if (Table_Read (&rLeftDownTable, rLeftDownData, 1) <= 0) {
      fprintf (stderr, "Pol_bender: %s: can not read file %s\n", NAME_CURRENT_COMP, rLeftDownData);
      exit (1);
    }
    if (Table_Read (&rRightUpTable, rRightUpData, 1) <= 0) {
      fprintf (stderr, "Pol_bender: %s: can not read file %s\n", NAME_CURRENT_COMP, rRightUpData);
      exit (1);
    }
    if (Table_Read (&rRightDownTable, rRightDownData, 1) <= 0) {
      fprintf (stderr, "Pol_bender: %s: can not read file %s\n", NAME_CURRENT_COMP, rRightDownData);
      exit (1);
    }
  }
  if ((xwidth <= 0) || (yheight <= 0) || (length <= 0) || (radius == 0)) {
    fprintf (stderr,
             "Pol_bender: %s: NULL or negative length scale!\n"
             "ERROR      (xwidth,yheight,length, radius). Exiting\n",
             NAME_CURRENT_COMP);
    exit (1);
  }

  if (drawOption < 1 || drawOption > 3) {
    fprintf (stderr, "Pol_bender: %s: drawOption %ld not supported. Exiting.\n", NAME_CURRENT_COMP, drawOption);
    exit (1);
  }

  if (mcgravitation) {

    localG = rot_apply (ROT_A_CURRENT_COMP, coords_set (0, -GRAVITY, 0));
    fprintf (stdout, "Pol_bender %s: Gravity is on!\n", NAME_CURRENT_COMP);
    if (localG.x != 0 || localG.z != 0)
      fprintf (stderr,
               "WARNING: Pol_Bender: %s: "
               "This component only gives correct resulta with gravitation,\n"
               "when gravity is strictly along the y-axis!\n",
               NAME_CURRENT_COMP);

  } else
    localG = coords_set (0, 0, 0);

  // To be able to handle the situation properly where a component of
  // the gravity is along the z-axis we also define entrance (in) and
  // exit (out) planes

  angle = length / radius;
  normIn = coords_set (0, 0, 1);
  if (endFlat)
    normOut = coords_set (0, 0, 1);
  else
    normOut = coords_set (sin (angle), 0, cos (angle));
  pointIn = coords_set (0, 0, 0);
  pointOut = coords_set (radius - radius * cos (angle), 0, radius* sin (angle));

  // Top and bot plane (+y dir) can be spanned by (1, 0, 0) & (0, 0, 1)
  // and the top point (0, yheight/2, 0) and bot point (0, -yheight/2, 0)
  // A normal vector is (0, 1, 0)
  normTopBot = coords_set (0, 1, 0);
  pointTop = coords_set (0, yheight / 2, 0);
  pointBot = coords_set (0, -yheight / 2, 0);
%}

TRACE
%{
  const double whalf = 0.5 * xwidth;                   /* half width of guide */
  const double hhalf = 0.5 * yheight;                  /* half height of guide */
  const double z_off = radius * sin (length / radius); /* z-comp of guide length */
  const double dThreshold = 1e-10;                     /* distance threshold */
  const double tThreshold = dThreshold / sqrt (vx * vx + vy * vy + vz * vz);
  double angle_z_vout; /* angle between z-axis and v_out */

  // Variables used in the case of multiple slits
  const double slitWidth = xwidth / nslit; // slitwidth
  const double spacerhalf = 0.5 * d;       /* half width of spacers */
  int slitHit;                             // decide which slit is hit
  double posInSlit;                        // position in slit

  double t11, t12, t21, t22, theta, alpha, endtime, phi;
  int i_bounce;
  int nerr = 0;

  // Pol variables
  double FN, FM, Rup, Rdown, weight;

  double Rleft;  /* radius of curvature of left mirror */
  double Rright; /* radius of curvature of right mirror */
  double absR = fabs (radius);
  double sign = 1;
  if (radius < 0)
    sign = -1;

  /* Propagate neutron to guide entrance. */
  PROP_Z0;
  if (!inside_rectangle (x, y, xwidth, yheight))
    ABSORB;

  if (nslit > 1) {
    // check if neutron is absorbed on a spacer
    posInSlit = fmod (x + whalf, slitWidth);
    if (posInSlit <= spacerhalf || posInSlit >= slitWidth - spacerhalf)
      ABSORB;

    // check which slat is hit
    slitHit = (int)((x + whalf) / slitWidth);

    // Modify R1 and R2 according to which slat was hit
    Rleft = absR + sign * whalf - sign * (slitHit + 1) * slitWidth + sign * spacerhalf;
    Rright = absR + sign * whalf - sign * slitHit * slitWidth - sign * spacerhalf;

    if (debug > 0)
      printf ("\nslitHit: %d/%f, Rleft: %f, Rright: %f\n", slitHit, (x + whalf) / slitWidth, Rleft, Rright);
  } else { // only 1 slit

    Rleft = absR - sign * whalf;
    Rright = absR + sign * whalf;
  }

  for (;;) {

    double tLeft, tRight, tTop, tBot, tIn, tOut, tMirror;
    double tUp, tSide, time, endtime;
    double R, Q;
    Coords vVec, xVec;
    int isPolarising;
    double vel_xz;

    isPolarising = 0;

    xVec = coords_set (x, y, z);
    vVec = coords_set (vx, vy, vz);

    solve_2nd_order (&tTop, NULL, 0.5 * coords_sp (normTopBot, localG), coords_sp (normTopBot, vVec), coords_sp (normTopBot, coords_sub (xVec, pointTop)));

    solve_2nd_order (&tBot, NULL, 0.5 * coords_sp (normTopBot, localG), coords_sp (normTopBot, vVec), coords_sp (normTopBot, coords_sub (xVec, pointBot)));

    solve_2nd_order (&tIn, NULL, 0.5 * coords_sp (normIn, localG), coords_sp (normIn, vVec), coords_sp (normIn, coords_sub (xVec, pointIn)));

    solve_2nd_order (&tOut, NULL, 0.5 * coords_sp (normOut, localG), coords_sp (normOut, vVec), coords_sp (normOut, coords_sub (xVec, pointOut)));

    /* Find itersection points with inside and outside guide walls */
    if (!cylinder_intersect (&t11, &t12, x - radius, y, z, vx, vy, vz, Rleft, 2 * yheight)) {
      /*neutron did not hit the cylinder*/
      t11 = t12 = 0;
    }
    if (!cylinder_intersect (&t21, &t22, x - radius, y, z, vx, vy, vz, Rright, 2 * yheight)) {
      /*neutron did not hit the cylinder*/
      t21 = t22 = 0;
    }

    /* Choose appropriate reflection time */
    tLeft = (t11 < tThreshold) ? t12 : t11;
    tRight = (t21 < tThreshold) ? t22 : t21;

    /* Choose appropriate reflection time */
    if (tTop > tThreshold && (tTop < tBot || tBot <= tThreshold))
      tUp = tTop;
    else
      tUp = tBot;

    if (tLeft > tThreshold && (tLeft < tRight || tRight <= tThreshold))
      tSide = tLeft;
    else
      tSide = tRight;

    if (tUp > tThreshold && (tUp < tSide || tSide <= tThreshold))
      time = tUp;
    else
      time = tSide;

    if (time <= tThreshold) {
      nerr++;
      if (nerr < 10) {
        fprintf (stdout,
                 "tTop: %e, tBot:%e, tRight: %e, tLeft: %e\n"
                 "tUp: %e, tSide: %e, time: %e\n",
                 tTop, tBot, tRight, tLeft, tUp, tSide, time);
      } else {
        fprintf (stdout, "Found 10 propagation error for this neutron, terminating!\n");
        break;
      }
    }

    /* Has neutron left the guide? */
    if (tOut > tThreshold && (tOut < tIn || tIn <= tThreshold))
      endtime = tOut;
    else
      endtime = tIn;

    if (time > endtime)
      break;

    PROP_DT (time);
    SCATTER;

    /* Find reflection surface */
    if (time == tSide) { /* Left or right side */

      if (time == tLeft)
        R = sign * Rleft;
      else
        R = sign * Rright;

      phi = atan (vx / vz);       /* angle of neutron trajectory */
      alpha = asin (z / R);       /* angle of guide wall */
      theta = fabs (phi - alpha); /* angle of reflection */
      angle_z_vout = 2.0 * alpha - phi;

      vel_xz = sqrt (vx * vx + vz * vz); /* in plane velocity */
      vz = vel_xz * cos (angle_z_vout);
      vx = vel_xz * sin (angle_z_vout);

    } else { /* Top or Bottom wall */
      theta = fabs (atan (vy / vz));
      vy = -vy;
    }

    /* Now compute reflectivity. */
    Q = 2.0 * sin (theta) * sqrt (vx * vx + vy * vy + vz * vz) * V2K;

    // calculate reflection probability
    if (time == tTop) {
      if (useTables) {
        Rup = Table_Value (rTopUpTable, Q, 1);
        Rdown = Table_Value (rTopDownTable, Q, 1);
      } else {
        StdReflecFunc (Q, rTopUpPar, &Rup);
        StdReflecFunc (Q, rTopDownPar, &Rdown);
      }
      if (debug > 0)
        fprintf (stdout, "\tTop hit:\n");
    } else if (time == tBot) {
      if (useTables) {
        Rup = Table_Value (rBotUpTable, Q, 1);
        Rdown = Table_Value (rBotDownTable, Q, 1);
      } else {
        StdReflecFunc (Q, rBotUpPar, &Rup);
        StdReflecFunc (Q, rBotDownPar, &Rdown);
      }
      if (debug > 0)
        fprintf (stdout, "\tBot hit:\n");
    } else if (time == tRight) {
      if (useTables) {
        Rup = Table_Value (rRightUpTable, Q, 1);
        Rdown = Table_Value (rRightDownTable, Q, 1);
      } else {
        StdReflecFunc (Q, rRightUpPar, &Rup);
        StdReflecFunc (Q, rRightDownPar, &Rdown);
      }
      if (debug > 0)
        fprintf (stdout, "\tRight hit:\n");
    } else if (time == tLeft) {
      if (useTables) {
        Rup = Table_Value (rLeftUpTable, Q, 1);
        Rdown = Table_Value (rLeftDownTable, Q, 1);
      } else {
        StdReflecFunc (Q, rLeftUpPar, &Rup);
        StdReflecFunc (Q, rLeftDownPar, &Rdown);
      }
      if (debug > 0)
        fprintf (stdout, "\tLeft hit:\n");
    }
    if (Rup != Rdown) {
      isPolarising = 1;
      GetMonoPolFNFM (Rup, Rdown, &FN, &FM);
      GetMonoPolRefProb (FN, FM, sy, &weight);
      /* Output of PW discussions with Hal Lee 2024/03/08
         We have now done our QM "measurement", thus
         forcing the spin to assume up/down: */
      sx = 0;
      sz = 0;
    } else
      weight = Rup;

    if (debug > 0)
      printf ("\tlambda: %.2f AA, Q: %.4f, Rup: %.4f, Rdown: %.4f,"
              " weight: %.4f\n",
              2 * PI / (sqrt (vx * vx + vy * vy + vz * vz) * V2K), Q, Rup, Rdown, weight);

    // check that refWeight is meaningfull
    if (weight <= 0)
      ABSORB;
    if (weight > 1)
      weight = 1;

    if (isPolarising) {
      SetMonoPolRefOut (FN, FM, weight, &sx, &sy, &sz);
      if (sx * sx + sy * sy + sz * sz > 1.000001)
        fprintf (stderr, "Pol_bender: %s: Warning: polarisation |s| = %g > 1\n", NAME_CURRENT_COMP,
                 sx * sx + sy * sy + sz * sz); // check that polarisation is meaningfull
    }

    p *= weight;

    if (p == 0) {
      ABSORB;
      break;
    }
  }
%}

MCDISPLAY
%{
  double x1, x2, z1, z2;
  x1 = x2 = z1 = z2 = 0;
  const int n = 90;
  double* xplot = malloc (n * sizeof (double));
  double* zplot = malloc (n * sizeof (double));
  int ns = 0;
  int j = 1;
  const double lengthOfGuide = sin (length / radius) * radius;
  const double slitWidth = xwidth / nslit;
  double R = 0; /* radius of arc */
  int nSlitsMax = nslit;
  int nMax = n;

  if (lengthOfGuide <= 0)
    exit (fprintf (stdout, "Pol_bender: %s: Negative guide length ! lengthOfGuide=%g\n", NAME_CURRENT_COMP, lengthOfGuide));

  if (drawOption == 2) {

    if (nSlitsMax > 20)
      nSlitsMax = 20;
    nMax = 40;
  } else if (drawOption == 3) {

    if (nSlitsMax > 5)
      nSlitsMax = 5;
    nMax = 10;
  }

  // draw opening
  rectangle ("xy", 0, 0, 0, xwidth, yheight);

  for (ns = 0; ns < nSlitsMax + 1; ns++) {

    // to make sure the sides are drawn properly
    if (ns == nSlitsMax && nSlitsMax < nslit)
      ns = nslit;

    // calculate x for this R
    R = radius - 0.5 * xwidth + ns * slitWidth;

    for (j = 0; j < nMax; j++) {

      if (endFlat) {

        if (ns == 0) // only calculate once
          zplot[j] = j * lengthOfGuide / (double)(nMax - 1);
      } else
        zplot[j] = R * sin (length / radius * (double)j / (double)(nMax - 1));

      if (radius > 0)
        xplot[j] = radius - sqrt (R * R - zplot[j] * zplot[j]);
      else
        xplot[j] = radius + sqrt (R * R - zplot[j] * zplot[j]);
    }

    // To be able to draw end we store some of the point values
    if (ns == 0) { // first wall

      x1 = xplot[nMax - 1];
      z1 = zplot[nMax - 1];
    } else if (ns == nslit) { // last wall

      x2 = xplot[nMax - 1];
      z2 = zplot[nMax - 1];
    }

    for (j = 0; j < nMax - 1; j++) {
      line (xplot[j], 0.5 * yheight, zplot[j], xplot[j + 1], 0.5 * yheight, zplot[j + 1]);
      line (xplot[j], -0.5 * yheight, zplot[j], xplot[j + 1], -0.5 * yheight, zplot[j + 1]);
    }
  }

  // draw end gap
  line (x1, 0.5 * yheight, z1, x2, 0.5 * yheight, z2);
  line (x1, 0.5 * yheight, z1, x1, -0.5 * yheight, z1);
  line (x2, -0.5 * yheight, z2, x2, 0.5 * yheight, z2);
  line (x1, -0.5 * yheight, z1, x2, -0.5 * yheight, z2);
  free (xplot);
  free (zplot);
  %}

END
