// ******************************************************************************
   //
   // Title:       Force Field X.
   // Description: Force Field X - Software for Molecular Biophysics.
   // Copyright:   Copyright (c) Michael J. Schnieders 2001-2025.
   //
   // This file is part of Force Field X.
   //
   // Force Field X is free software; you can redistribute it and/or modify it
  // under the terms of the GNU General Public License version 3 as published by
  // the Free Software Foundation.
  //
  // Force Field X is distributed in the hope that it will be useful, but WITHOUT
  // ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
  // FOR A PARTICULAR PURPOSE. See the GNU General Public License for more
  // details.
  //
  // You should have received a copy of the GNU General Public License along with
  // Force Field X; if not, write to the Free Software Foundation, Inc., 59 Temple
  // Place, Suite 330, Boston, MA 02111-1307 USA
  //
  // Linking this library statically or dynamically with other modules is making a
  // combined work based on this library. Thus, the terms and conditions of the
  // GNU General Public License cover the whole combination.
  //
  // As a special exception, the copyright holders of this library give you
  // permission to link this library with independent modules to produce an
  // executable, regardless of the license terms of these independent modules, and
  // to copy and distribute the resulting executable under terms of your choice,
  // provided that you also meet, for each linked independent module, the terms
  // and conditions of the license of that module. An independent module is a
  // module which is not derived from or based on this library. If you modify this
  // library, you may extend this exception to your version of the library, but
  // you are not obligated to do so. If you do not wish to do so, delete this
  // exception statement from your version.
  //
  // ******************************************************************************
  package ffx.potential.nonbonded;
  
  import edu.rit.pj.IntegerForLoop;
  import edu.rit.pj.IntegerSchedule;
  import edu.rit.pj.ParallelRegion;
  import ffx.crystal.Crystal;
  import ffx.potential.bonded.Atom;
  
  import java.nio.DoubleBuffer;
  import java.util.logging.Level;
  import java.util.logging.Logger;
  
  import static java.lang.String.format;
  import static java.util.Arrays.fill;
  import static org.apache.commons.math3.util.FastMath.floor;
  
  /**
   * This class implements a spatial decomposition based on partitioning a grid into octants.
   *
   * @author Michael J. Schnieders
   * @since 1.0
   */
  public class SpatialDensityRegion extends ParallelRegion {
  
    /** Constant <code>logger</code> */
    protected static final Logger logger = Logger.getLogger(SpatialDensityRegion.class.getName());
  
    public final int nThreads;
    protected final int nSymm;
    public int[] actualCount;
    public int nAtoms;
    /** The number of divisions along the A-axis. */
    protected int nA;
    /** The number of divisions along the B-axis. */
    protected int nB;
    /** The number of divisions along the C-Axis. */
    protected int nC;
    /**
     * Number of octant work cells with at least one atom (actualWork is less than or equal to nWork).
     */
    protected int actualWork;
  
    protected double[][][] coordinates;
    protected boolean[][] select;
    protected Crystal crystal;
    protected SpatialDensityLoop[] spatialDensityLoop;
    /** The A index of each octant (0..nA - 1) that has atoms. */
    int[] actualA;
    /** The B index of each octant (0..nB - 1) that has atoms. */
    int[] actualB;
    /** The C index of each octant (0..nC - 1) that has atoms. */
    int[] actualC;
    /** The list of atoms in each cell. [nsymm][natom] = atom index */
    int[][] cellList;
    /** The number of atoms in each cell. [nsymm][ncell] */
    int[][] cellCount;
    /** The index of the first atom in each cell. [nsymm][ncell] */
    int[][] cellStart;
    /** Basis size. */
    private int basisSize;
    /** Minimum amount of work. */
    private int minWork;
   /** The number of cells in one plane (nDivisions^2). */
   private int nAB;
   /** The number of cells (nDivisions^3). */
   private int nCells;
   /** Number of octant work cells. */
   private int nWork;
   /** A temporary array that holds the index of the cell each atom is assigned to. */
   private int[][] cellIndex;
   /** The A index of each octant (0..nA - 1) that may not have any atoms. */
   private int[] workA;
   /** The B index of each octant (0..nB - 1) that may not have any atoms. */
   private int[] workB;
   /** The C index of each octant (0..nC - 1) that may not have any atoms. */
   private int[] workC;
   /**
    * The offset of each atom from the start of the cell. The first atom atom in the cell has 0
    * offset. [nsymm][natom] = offset of the atom
    */
   private int[][] cellOffset;
 
   private double[] xf;
   private double[] yf;
   private double[] zf;
   private int gridSize;
   private double[] grid = null;
   private DoubleBuffer gridBuffer;
   private double initValue = 0.0;
   private GridInitLoop gridInitLoop;
 
   /**
    * Constructor for SpatialDensityRegion.
    *
    * @param gX a int.
    * @param gY a int.
    * @param gZ a int.
    * @param grid an array of double.
    * @param basisSize a int.
    * @param nSymm a int.
    * @param minWork a int.
    * @param threadCount a int.
    * @param crystal a {@link ffx.crystal.Crystal} object.
    * @param atoms an array of {@link ffx.potential.bonded.Atom} objects.
    * @param coordinates an array of double.
    */
   public SpatialDensityRegion(
       int gX,
       int gY,
       int gZ,
       double[] grid,
       int basisSize,
       int nSymm,
       int minWork,
       int threadCount,
       Crystal crystal,
       Atom[] atoms,
       double[][][] coordinates) {
     this(gX, gY, gZ, basisSize, nSymm, minWork, threadCount, crystal, atoms, coordinates);
     this.grid = grid;
     if (grid != null) {
       gridBuffer = DoubleBuffer.wrap(grid);
     }
   }
 
   /**
    * Chop up the 3D unit cell domain into fractional coordinate chunks to allow multiple threads to
    * put charge density onto the grid without needing the same grid point. First, we partition the
    * X-axis, then the Y-axis, and finally the Z-axis if necessary.
    *
    * @param gX a int.
    * @param gY a int.
    * @param gZ a int.
    * @param basisSize a int.
    * @param nSymm a int.
    * @param minWork a int.
    * @param threadCount a int.
    * @param crystal a {@link ffx.crystal.Crystal} object.
    * @param atoms an array of {@link ffx.potential.bonded.Atom} objects.
    * @param coordinates an array of double.
    */
   private SpatialDensityRegion(
       int gX,
       int gY,
       int gZ,
       int basisSize,
       int nSymm,
       int minWork,
       int threadCount,
       Crystal crystal,
       Atom[] atoms,
       double[][][] coordinates) {
     this.coordinates = coordinates;
     this.nSymm = nSymm;
     this.nAtoms = atoms.length;
     this.nThreads = threadCount;
     this.basisSize = basisSize;
     this.minWork = minWork;
     gridInitLoop = new GridInitLoop();
     setCrystal(crystal.getUnitCell(), gX, gY, gZ);
   }
 
   /**
    * Assign asymmetric and symmetry mate atoms to cells. This is very fast; there is little to be
    * gained from parallelization at this point.
    */
   public void assignAtomsToCells() {
     // Call the selectAtoms method of subclasses.
     selectAtoms();
 
     for (int iSymm = 0; iSymm < nSymm; iSymm++) {
       final int[] cellIndexs = cellIndex[iSymm];
       final int[] cellCounts = cellCount[iSymm];
       final int[] cellStarts = cellStart[iSymm];
       final int[] cellLists = cellList[iSymm];
       final int[] cellOffsets = cellOffset[iSymm];
       final boolean[] selected = select[iSymm];
 
       // Zero out the cell counts.
       for (int i = 0; i < nCells; i++) {
         cellCounts[i] = 0;
       }
 
       // Convert to fractional coordinates.
       final double[][] xyz = coordinates[iSymm];
       final double[] x = xyz[0];
       final double[] y = xyz[1];
       final double[] z = xyz[2];
       crystal.toFractionalCoordinates(nAtoms, x, y, z, xf, yf, zf);
 
       // Assign each atom to a cell using fractional coordinates.
       for (int i = 0; i < nAtoms; i++) {
         if (!selected[i]) {
           continue;
         }
 
         final int index = getCellIndex(i);
         cellIndexs[i] = index;
 
         // The offset of this atom from the beginning of the cell.
         cellOffsets[i] = cellCounts[index]++;
       }
 
       // Define the starting indices.
       cellStarts[0] = 0;
       for (int i = 1; i < nCells; i++) {
         final int i1 = i - 1;
         cellStarts[i] = cellStarts[i1] + cellCounts[i1];
       }
 
       // Move atom locations into a list ordered by cell.
       for (int i = 0; i < nAtoms; i++) {
         if (!selected[i]) {
           continue;
         }
         final int index = cellIndexs[i];
         cellLists[cellStarts[index]++] = i;
       }
 
       // Redefine the starting indices again.
       cellStarts[0] = 0;
       for (int i = 1; i < nCells; i++) {
         final int i1 = i - 1;
         cellStarts[i] = cellStarts[i1] + cellCounts[i1];
       }
     }
 
     // Loop over work chunks and get rid of empty chunks.
     actualWork = 0;
     for (int icell = 0; icell < nWork; icell++) {
       int ia = workA[icell];
       int ib = workB[icell];
       int ic = workC[icell];
       int ii = count(ia, ib, ic);
       // Fractional chunks along the C-axis.
       if (nC > 1) {
         ii += count(ia, ib, ic + 1);
         // Fractional chunks along the B-axis.
         if (nB > 1) {
           ii += count(ia, ib + 1, ic);
           ii += count(ia, ib + 1, ic + 1);
           // Fractional chunks along the A-axis.
           if (nA > 1) {
             ii += count(ia + 1, ib, ic);
             ii += count(ia + 1, ib, ic + 1);
             ii += count(ia + 1, ib + 1, ic);
             ii += count(ia + 1, ib + 1, ic + 1);
           }
         }
       }
 
       // If there is work in this chunk, include it.
       if (ii > 0) {
         actualA[actualWork] = ia;
         actualB[actualWork] = ib;
         actualC[actualWork] = ic;
         actualCount[actualWork++] = ii;
       }
     }
 
     if (logger.isLoggable(Level.FINEST)) {
       logger.finest(format(" Empty chunks: %d out of %d.", nWork - actualWork, nWork));
     }
   }
 
   private int getCellIndex(int i) {
     double xu = xf[i];
     double yu = yf[i];
     double zu = zf[i];
 
     // Move the atom into the range 0.0 <= x < 1.0
     while (xu >= 1.0) {
       xu -= 1.0;
     }
     while (xu < 0.0) {
       xu += 1.0;
     }
     while (yu >= 1.0) {
       yu -= 1.0;
     }
     while (yu < 0.0) {
       yu += 1.0;
     }
     while (zu >= 1.0) {
       zu -= 1.0;
     }
     while (zu < 0.0) {
       zu += 1.0;
     }
 
     // The cell indices of this atom.
     int a = (int) floor(xu * nA);
     int b = (int) floor(yu * nB);
     int c = (int) floor(zu * nC);
 
     // Check to make sure a, b and c are less than nA, nB and nC, respectively.
     if (a >= nA) {
       a = nA - 1;
     }
     if (b >= nB) {
       b = nB - 1;
     }
     if (c >= nC) {
       c = nC - 1;
     }
 
     // The cell index of this atom.
     final int index = a + b * nA + c * nAB;
     return index;
   }
 
   /**
    * Getter for the field <code>grid</code>.
    *
    * @return an array of double.
    */
   public double[] getGrid() {
     return grid;
   }
 
   /**
    * getNsymm
    *
    * @return a int.
    */
   public int getNsymm() {
     return nSymm;
   }
 
   /**
    * index
    *
    * @param ia a int.
    * @param ib a int.
    * @param ic a int.
    * @return a int.
    */
   public int index(int ia, int ib, int ic) {
     return ia + ib * nA + ic * nAB;
   }
 
   /** {@inheritDoc} */
   @Override
   public void run() {
     int ti = getThreadIndex();
     int actualWork1 = actualWork - 1;
     SpatialDensityLoop loop = spatialDensityLoop[ti];
 
     // This lets the same SpatialDensityLoops be used with different SpatialDensityRegions.
     loop.setNsymm(nSymm);
 
     try {
       execute(0, gridSize - 1, gridInitLoop);
       execute(0, actualWork1, loop.setOctant(0));
       // Fractional chunks along the C-axis.
       if (nC > 1) {
         execute(0, actualWork1, loop.setOctant(1));
         // Fractional chunks along the B-axis.
         if (nB > 1) {
           execute(0, actualWork1, loop.setOctant(2));
           execute(0, actualWork1, loop.setOctant(3));
           // Fractional chunks along the A-axis.
           if (nA > 1) {
             execute(0, actualWork1, loop.setOctant(4));
             execute(0, actualWork1, loop.setOctant(5));
             execute(0, actualWork1, loop.setOctant(6));
             execute(0, actualWork1, loop.setOctant(7));
           }
         }
       }
     } catch (Exception e) {
       String message = " Exception in SpatialDensityRegion.";
       logger.log(Level.SEVERE, message, e);
     }
   }
 
   /**
    * Select atoms that should be assigned to cells. The default is to include all atoms, which is
    * set up in the constructor. This function should be over-ridden by subclasses that want finer
    * control.
    */
   public void selectAtoms() {}
 
   /**
    * setAtoms.
    *
    * @param atoms an array of {@link ffx.potential.bonded.Atom} objects.
    */
   public void setAtoms(Atom[] atoms) {
     nAtoms = atoms.length;
     if (select == null || select.length < nSymm || select[0].length < nAtoms) {
       select = new boolean[nSymm][nAtoms];
       for (int i = 0; i < nSymm; i++) {
         fill(select[i], true);
       }
       cellList = new int[nSymm][nAtoms];
       cellIndex = new int[nSymm][nAtoms];
       cellOffset = new int[nSymm][nAtoms];
     }
     if (xf == null || xf.length < nAtoms) {
       xf = new double[nAtoms];
       yf = new double[nAtoms];
       zf = new double[nAtoms];
     }
   }
 
   /**
    * Setter for the field <code>crystal</code>.
    *
    * @param crystal a {@link ffx.crystal.Crystal} object.
    * @param gX a int.
    * @param gY a int.
    * @param gZ a int.
    */
   public final void setCrystal(Crystal crystal, int gX, int gY, int gZ) {
     // If the crystal is unchanged, return.
     if (this.crystal != null && this.crystal.equals(crystal.getUnitCell())) {
       return;
     }
 
     this.crystal = crystal.getUnitCell();
     if (xf == null || xf.length < nAtoms) {
       xf = new double[nAtoms];
       yf = new double[nAtoms];
       zf = new double[nAtoms];
     }
 
     gridSize = gX * gY * gZ * 2;
     int nX = gX / basisSize;
     int nY = gY / basisSize;
     int nZ = gZ / basisSize;
     if (nThreads > 1 && nZ > 1) {
       if (nZ % 2 != 0) {
         nZ--;
       }
       nC = nZ;
       int div = 2;
       int currentWork = nC / div / nThreads;
       // If we have enough work per thread, stop dividing the domain.
       if (currentWork >= minWork || nY < 2) {
         nA = 1;
         nB = 1;
         // Reduce the number of divisions along the Z-axis if possible
         while (currentWork >= 2 * minWork) {
           nC -= 2;
           currentWork = nC / div / nThreads;
         }
 
       } else {
         if (nY % 2 != 0) {
           nY--;
         }
         nB = nY;
         div = 4;
         currentWork = nB * nC / div / nThreads;
         // If we have 4 * threadCount * minWork chunks, stop dividing the domain.
         if (currentWork >= minWork || nX < 2) {
           nA = 1;
           while (currentWork >= 2 * minWork) {
             nB -= 2;
             currentWork = nB * nC / div / nThreads;
           }
         } else {
           if (nX % 2 != 0) {
             nX--;
           }
           nA = nX;
           div = 8;
           currentWork = nA * nB * nC / div / nThreads;
           while (currentWork >= 2 * minWork) {
             nA -= 2;
             currentWork = nA * nB * nC / div / nThreads;
           }
         }
       }
       nAB = nA * nB;
       nCells = nAB * nC;
       nWork = nCells / div;
     } else {
       nA = 1;
       nB = 1;
       nC = 1;
       nAB = 1;
       nCells = 1;
       nWork = 1;
     }
 
     logger.fine(format("   Spatial cells:                 (%3d,%3d,%3d)", nA, nB, nC));
     logger.fine(format("   Spatial work:                           %4d", nWork));
 
     if (workA == null || workA.length < nWork) {
       workA = new int[nWork];
       workB = new int[nWork];
       workC = new int[nWork];
       actualA = new int[nWork];
       actualB = new int[nWork];
       actualC = new int[nWork];
       actualCount = new int[nWork];
       int index = 0;
       for (int h = 0; h < nA; h += 2) {
         for (int k = 0; k < nB; k += 2) {
           for (int l = 0; l < nC; l += 2) {
             workA[index] = h;
             workB[index] = k;
             workC[index++] = l;
           }
         }
       }
     }
     if (select == null || select.length < nSymm || select[0].length < nAtoms) {
       select = new boolean[nSymm][nAtoms];
       for (int i = 0; i < nSymm; i++) {
         fill(select[i], true);
       }
       cellList = new int[nSymm][nAtoms];
       cellIndex = new int[nSymm][nAtoms];
       cellOffset = new int[nSymm][nAtoms];
     }
     if (cellStart == null || cellStart.length < nSymm || cellStart[0].length < nCells) {
       cellStart = new int[nSymm][nCells];
       cellCount = new int[nSymm][nCells];
     }
   }
 
   /**
    * setDensityLoop
    *
    * @param loops an array of {@link ffx.potential.nonbonded.SpatialDensityLoop} objects.
    */
   public void setDensityLoop(SpatialDensityLoop[] loops) {
     spatialDensityLoop = loops;
   }
 
   /**
    * Setter for the field <code>initValue</code>.
    *
    * @param initValue a double.
    */
   public void setInitValue(double initValue) {
     this.initValue = initValue;
   }
 
   /**
    * Setter for the field <code>gridBuffer</code>.
    *
    * @param grid a {@link java.nio.DoubleBuffer} object.
    */
   void setGridBuffer(DoubleBuffer grid) {
     gridBuffer = grid;
   }
 
   private int count(int ia, int ib, int ic) {
     int count = 0;
     for (int iSymm = 0; iSymm < nSymm; iSymm++) {
       final int index = index(ia, ib, ic);
       final int start = cellStart[iSymm][index];
       final int stop = start + cellCount[iSymm][index];
       count += (stop - start);
     }
     return count;
   }
 
   private class GridInitLoop extends IntegerForLoop {
 
     private final IntegerSchedule schedule = IntegerSchedule.fixed();
 
     @Override
     public void run(int lb, int ub) {
       if (gridBuffer != null) {
         // if (grid != null) {
         for (int i = lb; i <= ub; i++) {
           // grid[i] = initValue;
           gridBuffer.put(i, initValue);
         }
       }
     }
 
     @Override
     public IntegerSchedule schedule() {
       return schedule;
     }
   }
 }