%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%% mode 1 == Slab reactor with no energy dependence
%% mode 2 == Infinite homogeneous system with energy dependence

function simple_monte_carlo_blank(mode)

% --- Initialization and pre-processing

nBatches   =    20;
nInactive  =    10;
nParticles =  1000;

maxSteps = 1000;

nu = 3;

Emin = 1e1; % eV
Emax = 1e7; % eV

% --- Initialize keffVector

keffVector = zeros(nBatches,1);

% --- Set number of x-bins for fission profile

xBins = 100;

% --- Set number of E-bins for flux profile

eBins = 200;

% --- Create bin boundaries

fissionBoundaries = linspace(0, 100, xBins+1); % cm
energyBoundaries  = logspace(log(Emin), log(Emax), eBins+1); % eV

% --- Initialize tally matrices

fissionProfile = zeros(nBatches, xBins);
fluxProfile = zeros(nBatches, eBins);
densityProfile = zeros(nBatches, eBins);

% --- Fission sites for first batch (random between zero and 100 cm)

prevFissionSites = zeros(nParticles,1)+50.0;

% --- Simulate nBatch batches of neutrons

for iBatch = 1:nBatches

  disp(['Batch: ' num2str(iBatch)]);
  fflush(stdout);

  % --- Reset batch wise neutron production and loss

  batchProd = 0;
  batchLoss = 0;

  % --- Initialize vector to store fission sites
  %     max 1 fission per neutron
  %     Initialize everything to x = 50

  newFissionSites = zeros(nParticles,1) + 50;
  nFissionSites = 0;

  % --- Simulate nParticles particles in current batch

  for N = 1:nParticles

    % --- Get new particle from source distribution

    [x, u, E] = NewNeutron(prevFissionSites);

    % --- In the infinite case just set the direction
    %     to be perependicular to x-axis

    if (mode == 2)
      u = 0.0;
    end

    % --- Random walk until neutron track is terminated

    for step=1:maxSteps

      % --- Get material region

      regionIndex = GetMaterialRegion(x);

      % --- Get local cross sections

      if (mode == 1)
        SigmaVector = MacroXS(regionIndex, E);
      else
        SigmaVector = MacroXSInf(regionIndex, E);
      end

      % --- Sample path length to next interaction

      dS = SamplePathLength(SigmaVector);

      % --- Check surface crossing and modify dS if needed

      if (mode == 1)
        dS1 = StopAtSurface(x, u, dS);
      else
        dS1 = dS;
      end

      % --- Move neutron to new position

      x = MoveNeutron(x, u, dS1);

      %% --- Tally path-length estimators

      %% --- Get energy bin index

      [~, index] = max(energyBoundaries > E);

      % --- Add to flux profile (path-length estimator)

      fluxProfile(iBatch, index - 1)  = fluxProfile(iBatch, index - 1) + ...;

      % --- Add to neutron density profile (path-length estimator)

      densityProfile(iBatch, index - 1)  = densityProfile(iBatch, index - 1) + ...;

      % --- Get material region

      regionIndex = GetMaterialRegion(x);

      % --- Handle leakage

      if (regionIndex == 0)
        if (mode == 1)

        % --- Score loss (analog estimator)

        batchLoss = batchLoss + ...;

        % --- End random walk

        break;

        else

          error("No leakage in infinite system allowed");

        end
      end

      % --- Check surface crossing

      if (dS1 < dS)
         continue;
      end

      % --- Get local cross sections

      if (mode == 1)
        SigmaVector = MacroXS(regionIndex, E);
      else
        SigmaVector = MacroXSInf(regionIndex, E);
      end

      % --- SampleReaction

      reactionIndex = SampleReaction(SigmaVector);

      % --- Model reaction (and score analog reaction estimators)

      if (reactionIndex == 1) % Capture

        % --- Score loss (analog estimator)

        batchLoss = batchLoss + ...;

        % --- End random walk

        break;

      elseif (reactionIndex == 2) % Scattering

        % --- Update neutron energy and angle
        if (mode == 1)
          [u, E] = ScatteringSlab(u, E);
        else
          [u, E] = ScatteringInf(u, E);
        end

        if (E < energyBoundaries(1))
          warning('Neutron scattered below energy region, killing')
          break;
        end

      elseif (reactionIndex == 3) % Fission

        % --- Score production (analog estimator)

        batchProd = batchProd + ...;

        % --- Score loss (analog estimator)

        batchLoss = batchLoss + ...;

        % --- Score analog estimate for fission distribution
        % --- Get spatial index of current position

        [~, index] = max(fissionBoundaries > x);

        % --- Add to fission profile (analog estimator)

        fissionProfile(iBatch, index - 1)  = fissionProfile(iBatch, index - 1) + ...;

        % --- Add to new fission sites (will serve as a source for next batch)

        nFissionSites = nFissionSites + 1;
        newFissionSites(nFissionSites) = x;

        % --- End random walk

        break;
      else
        error('Unknown reaction index');
      end % End reaction handling

    end % End random walk loop

    % --- Check that our neutron did not do maxSteps

    if (step == maxSteps)
      error(['Possibly infinite tracking loop for neutron at ' num2str(x) ' ' num2str(u) ' ' num2str(E)]);
    end

  end % End loop over neutrons

  % --- All neutron histories have been run for this batch
  %     Collect results from this batch

  % --- Calculate batch-wise mean value of result estimators
  %     and store to memory

  % --- keff = production/loss

  keffVector(iBatch) = ... / ...;

  % --- Fission distribution (this is already calculated)

  % --- Energy spectrum (this is already calculated)

  % --- Check that we got some fission sites for next batch

  if (nFissionSites == 0)
    if (mode == 1)
      error('No fissions on previous batch, cannot continue');
    else
       nFissionSites = 1;
    end
  end

  % --- Remove unused fission sites from fission site vector

  newFissionSites(nFissionSites+1:end) = [];

  % --- New fission sites will serve as the source for next Batch

  prevFissionSites = newFissionSites;

end % End Batches

% --- Calculate mean values

keffMean  = mean(keffVector(nInactive+1:end));
keffStd   = std(keffVector(nInactive+1:end));
keffSerr  = keffStd/sqrt(nBatches-nInactive);

if (mode == 1)

  % --- Print out value of k-eff

  disp(['k-eff was ' num2str(keffMean) ' +- ' num2str(keffSerr)]);

  % --- Remove profiles from inactive cycles

  fissionProfile(1:nInactive,:) = [];

  % --- Calculate mean of fission profile

  fissionProfileMean = mean(fissionProfile,1);

  % --- Normalize fission profile

  fissionProfile = fissionProfile/max(fissionProfileMean);

  fissionProfileMean = fissionProfileMean/max(fissionProfileMean);

  % --- Calculate standard deviation and error

  fissionProfileStd  = std(fissionProfile,0,1);
  fissionProfileSerr = fissionProfileStd/sqrt(nBatches-nInactive);

  figure()

  hold on;

  xlim([0.0 100.0])
  ylim([0.0 1.1])

  val = 0.0;
  sig = 0;

  for i=1:xBins

    % --- Get bin value

    valPrev = val;
    val = fissionProfileMean(i);

    % --- Get uncertainty

    sigPrev = sig;
    sig = fissionProfileSerr(i);

    % --- Get bin limits

    x0 = fissionBoundaries(i);
    x1 = fissionBoundaries(i+1);

    % --- Horizontal line for bin value

    h = line([x0 x1],[val val]);
    set(h,'LineWidth',2);

    % --- Horizontal line for minus two sigma

    h = line([x0 x1],[val-2*sig val-2*sig]);
    set(h,'LineWidth',1);

    % --- Horizontal line for plus two sigma

    h = line([x0 x1],[val+2*sig val+2*sig]);
    set(h,'LineWidth',1);

    % --- Vertical lines to show step profile

    if (i > 1)
      h = line([x0 x0],[valPrev val]);
    end

  end

  % --- Vertical line for material boundaries

  h = line([25 25],[0 1.1]);
  set(h,'LineWidth',1,'LineStyle','--','Color','r');

  h = line([75 75],[0 1.1]);
  set(h,'LineWidth',1,'LineStyle','--','Color','r');

  xlabel("X-coordinate (cm)")
  ylabel("Fission rate (a.u.)")

  hold off
  grid on
else

  % --- Normalize flux profile with bin width

  binWidths = energyBoundaries(2:end)-energyBoundaries(1:end-1);

  fluxProfile = fluxProfile./repmat(binWidths,nBatches, 1);
  densityProfile = densityProfile./repmat(binWidths,nBatches, 1);

  % --- Calculate mean of flux profile

  fluxProfileMean = mean(fluxProfile,1);
  densityProfileMean = mean(densityProfile,1);

  % --- Normalize flux profile

  fluxProfile = fluxProfile/max(fluxProfileMean);
  densityProfile = densityProfile/max(densityProfileMean);

  fluxProfileMean = fluxProfileMean/max(fluxProfileMean);
  densityProfileMean = densityProfileMean/max(densityProfileMean);

  % --- Calculate standard deviation and error

  fluxProfileStd  = std(fluxProfile,0,1);
  fluxProfileSerr = fluxProfileStd/sqrt(nBatches-nInactive);
  densityProfileStd  = std(densityProfile,0,1);
  densityProfileSerr = densityProfileStd/sqrt(nBatches-nInactive);

  figure()

  hold on;

  val = 0.0;
  sig = 0;

  for i=1:xBins

    % --- Get bin value

    valPrev = val;
    val = fluxProfileMean(i);

    % --- Get uncertainty

    sigPrev = sig;
    sig = fluxProfileSerr(i);

    % --- Get bin limits

    x0 = energyBoundaries(i);
    x1 = energyBoundaries(i+1);

    % --- Horizontal line for bin value

    h = line([x0 x1],[val val]);
    set(h,'LineWidth',2);

    % --- Horizontal line for minus two sigma

    h = line([x0 x1],[val-2*sig val-2*sig]);
    set(h,'LineWidth',1);

    % --- Horizontal line for plus two sigma

    h = line([x0 x1],[val+2*sig val+2*sig]);
    set(h,'LineWidth',1);

    % --- Vertical lines to show step profile

    if (i > 1)
      h = line([x0 x0],[valPrev val]);
    end

  end

  set(gca,'xscale','log')

  xlim([Emin Emax])
  ylim([0.0 1.1])

  xlabel("Energy (eV)")
  ylabel("Neutron flux / energy bin width (a.u.)")

  hold off
  grid on

  figure()

  hold on;

  val = 0.0;
  sig = 0;

  for i=1:xBins

    % --- Get bin value

    valPrev = val;
    val = densityProfileMean(i);

    % --- Get uncertainty

    sigPrev = sig;
    sig = densityProfileSerr(i);

    % --- Get bin limits

    x0 = energyBoundaries(i);
    x1 = energyBoundaries(i+1);

    % --- Horizontal line for bin value

    h = line([x0 x1],[val val]);
    set(h,'LineWidth',2);

    % --- Horizontal line for minus two sigma

    h = line([x0 x1],[val-2*sig val-2*sig]);
    set(h,'LineWidth',1);

    % --- Horizontal line for plus two sigma

    h = line([x0 x1],[val+2*sig val+2*sig]);
    set(h,'LineWidth',1);

    % --- Vertical lines to show step profile

    if (i > 1)
      h = line([x0 x0],[valPrev val]);
    end

  end

  set(gca,'xscale','log')

  xlim([Emin Emax])
  ylim([0.0 1.1])

  xlabel("Energy (eV)")
  ylabel("Neutron density / energy bin width (a.u.)")

  hold off
  grid on

end
% --- Post processing

end % End main function

%%%%%%%%%%%%%%%%%%%%
% Helper functions %
%%%%%%%%%%%%%%%%%%%%

% --- New neutron from fission source

function [x, u, E] = NewNeutron(prevFissionSites)

  % --- Position of neutron based on previous batches fission sites.

  nSites = size(prevFissionSites,1);

  % --- Sample site

  iSite = randi(nSites);

  % --- Get x-coordinate of the site

  x = prevFissionSites(iSite);

  % --- Direction cosine between -1 and 1 (isotropic emission)

  u = ...;

  % --- Energy from Watt spectrum

  a = 9.8e5; % eV

  b = 2.532e-6; % 1/ev

  % --- Probability density distribution is always smaller than one
  %     We can use rejection sampling to sample neutron energy

  maxResample = 100;

  for i=1:maxResample

     % --- Sample energy between 0 and 10 MeV

     E = ...;

     % --- Get acceptance probability (Watt pdf at current energy)

     P = ...;

     % --- Sample acceptance

     if (rand(1,1) < P)

       % --- Accepted

       break;
     end
  end

  % disp(['Took ' num2str(i) ' samples, E = ' num2str(E)]);
  % fflush(stdout);

  if (i == maxResample)
    error('Could not sample neutron energy')
  end

end % End function NewNeutron



% --- Gets material region based on coordinate (slab case)

function regionIndex = GetMaterialRegion(x)

  % --- Check region

  if ((x < 0) || (x > 100))

    % --- Out of slab

    regionIndex = 0;

  elseif ((x > 25) && (x < 75))

    % --- Center region

    regionIndex = 1;

  else

    % --- Edge region

    regionIndex = 2;

  end

end % End function regionIndex



% --- Cross section vector for the current location and energy
%     Slab geometry

function SigmaVector = MacroXS(regionIndex, E)

  % --- Check region

  if (regionIndex == 0)
    error('Tried to get sigma vector even though x is out of slab')
  elseif (regionIndex > 2)
    error(['Unknown region index: ' num2str(regionIndex)]);
  end

  % --- Matrix of the cross sections (row is reaction, column is region)

  SigmaMatrix = [0.06, 0.06
                 0.05, 0.06
                 0.06, 0.04];

  % --- Compose sigma vector

  SigmaVector = SigmaMatrix(:,regionIndex);

end % End function MacroXS



% --- Cross section vector for the current location and energy
%     Infinite geometry

function SigmaVector = MacroXSInf(regionIndex, E)

  % --- Check region

  if (regionIndex == 0)
    error('Tried to get sigma vector even though x is out of slab')
  elseif (regionIndex > 1)
    error(['Unknown region index: ' num2str(regionIndex)]);
  end

  %% --- Atomic density

  N = 1.0;

  %% --- Neutron mass

  m = 1.674927471e-27;

  %% --- Conversion factor from joules to eV

  JtoeV = 6.24150913e18;

  %% --- Kinetic energy of a 2200 m/s neutron in eV

  E0 = 0.5*m*2200^2*JtoeV;

  %% --- Capture cross section at E0

  x0 = 300e-24*N;

  %% --- Capture cross section at energy E (1/v cross section)

  SigmaC = ...;

  % --- Scattering cross section

  Gamman = 3.76e2;
  Gammat = 3.7634e2;
  E0 = 2.81e3;
  sigmap = 3.68e-24;
  sigma0 = 377e-24;
  y = 2*(E - E0)/Gammat;

  % --- Calculate SigmaS (Eq. 4.1)

  SigmaS = ...;

  % --- Fission cross section is zero

  SigmaF = ...;

  % --- Compose sigma vector (column indicates region, row indicates reaction)

  SigmaVector = [SigmaC; SigmaS; SigmaF];
end % End function SigmaVector



% --- Sample path length to next interaction

function dS = SamplePathLength(SigmaVector)

  % --- Calculate total macroscopic cross section

  SigmaTot = ...;

  % --- Sample path length

  dS = ...;

end % End function SamplePathLength



% --- Function that modifies sampled path length to stop
%     neutrons at material boundaries (surface tracking)

function dSout = StopAtSurface(x, u, dSin)

  % --- We can now check all surfaces since there are
  %     only a couple of them

  % --- Vector of surface x-coordinates

  surfaces = [0.0 25.0 50.0 75.0]';

  % --- Initialize minimum distance to a surface to something larger
  % --- than dSin (the sampled path length)

  mindS = ...;

  % --- Loop over all surfaces

  for i=1:size(surfaces,1)

    % --- Get x-coordinate of current surface

    xSurf = surfaces(i);

    % --- Distance to surface at x = xSurf cm
    %     Solve equation x + u*distance = xSurf

    if (u == 0)

      % --- Going perpendicular to x-axis

      distance = ...;

    else

      % --- Distance will be finite

      distance = ...;

    end

    % --- Check if this was the closest surface
    %     Negative distance means that this surface is
    %     behind the neutron and not ahead

    if ((mindS > distance) && (distance > 0))
      mindS = distance;
    end

  end % End loop over surfaces

  % --- Return smaller of the initial dS and the distance
  %     to the closest surface + 1 nanometre

  dSout = min(dSin, mindS + 1e-7);

end % End function StopAtSurface


% --- Move neutron to new position

function x = MoveNeutron(xin, u, dS)

  % --- Move neutron

  x = ...;

end % End function MoveNeutron



% --- Sample reaction mode based on local cross sections

function reactionIndex = SampleReaction(SigmaVector)

  % --- Sample one of the reaction modes
  %     We don't really have to know the reaction modes
  %     here, just the cross sections

  %% Calculate total cross section

  SigmaTot = ...;

  % --- Sample proportion of total cross section

  randSigma = rand(1,1);

  % --- Loop over reaction modes to check, which reaction index we
  %     sampled

  for reactionIndex = 1:size(SigmaVector,1)

    % --- Subtract current reactions proportion from sampled proportion

    randSigma = ...;

    % --- If randSigma went negative, we sampled this reaction mode

    if (randSigma < 0)
      break;
    end

  end % End loop over reactions

  % --- Check that we really managed to sample some reaction mode

  if (randSigma > 0)
    error('Reaction sampling failed');
  end

end % End function reactionIndex



% --- Model scattering (slab)

function [u, E] = ScatteringSlab(u0, E0)

  % --- Sample new direction cosine between -1 and 1 (lab frame)
  %     NB: In the general case we should sample scattering cosine
  %     in the CM-frame and the new direction would be the sum of
  %     the current direction angle and the scattering angle

  u = ...;

  % --- Do not modify incoming energy

  E = E0;

end % End function scattering

% --- Model scattering (infinite medium)

function [u, E] = ScatteringInf(u0, E0)

  % --- Sample scattering cosine between -1 and 1 (CM frame)

  uC = ...;

  % --- Mass number is now 23

  A = 23;

  % --- Calculate new energy

  E = ...;

  % --- Do not modify incoming direction now
  %     Typically we would need to modify the neutron direction (lab frame)
  %     based on the scattering angle sampled in the CM frame, which
  %     requires two coordinate transformations.

  u = u0;

end % End function scattering