content: /*************************************************************************
  *                                                                        *
  *  This file is part of the 20n/act project.                             *
  *  20n/act enables DNA prediction for synthetic biology/bioengineering.  *
  *  Copyright (C) 2017 20n Labs, Inc.                                     *
  *                                                                        *
  *  Please direct all queries to act@20n.com.                             *
  *                                                                        *
  *  This program is free software: you can redistribute it and/or modify  *
  *  it under the terms of the GNU General Public License as published by  *
  *  the Free Software Foundation, either version 3 of the License, or     *
  *  (at your option) any later version.                                   *
  *                                                                        *
  *  This program 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 this program.  If not, see <http://www.gnu.org/licenses/>. *
  *                                                                        *
  *************************************************************************/
  
  package com.act.lcms.v2.fullindex;
  
  import com.act.lcms.LCMSNetCDFParser;
  import com.act.lcms.LCMSSpectrum;
  import com.act.lcms.MS1;
  import com.act.utils.CLIUtil;
  import com.act.utils.rocksdb.DBUtil;
  import com.act.utils.rocksdb.RocksDBAndHandles;
  import org.apache.commons.cli.CommandLine;
  import org.apache.commons.cli.Option;
  import org.apache.commons.lang3.StringUtils;
  import org.apache.commons.lang3.tuple.Pair;
  import org.apache.logging.log4j.LogManager;
  import org.apache.logging.log4j.Logger;
  import org.joda.time.DateTime;
  import org.rocksdb.RocksDB;
  import org.rocksdb.RocksDBException;
  import org.rocksdb.WriteOptions;
  
  import javax.xml.parsers.ParserConfigurationException;
  import javax.xml.stream.XMLStreamException;
  import java.io.File;
  import java.io.IOException;
  import java.nio.ByteBuffer;
  import java.nio.charset.Charset;
  import java.nio.charset.StandardCharsets;
  import java.util.ArrayList;
  import java.util.Collections;
  import java.util.HashSet;
  import java.util.Iterator;
  import java.util.LinkedList;
  import java.util.List;
  import java.util.Set;
  
  public class Builder {
    private static final Logger LOGGER = LogManager.getFormatterLogger(Builder.class);
    private static final Charset UTF8 = StandardCharsets.UTF_8;
    /* TIMEPOINTS_KEY is a fixed key into a separate column family in the index that just holds a list of time points.
     * Within that column family, there is only one entry:
     *   "timepoints" -> serialized array of time point doubles
     * and we use this key to write/read those time points.  Since time points are shared across all traces, we can
     * maintain this one copy in the index and reconstruct the XZ pairs as we read trace intensity arrays. */
    // All of these are intentionally package private.
    static final byte[] TIMEPOINTS_KEY = "timepoints".getBytes(UTF8);
  
    static final Double WINDOW_WIDTH_FROM_CENTER = MS1.MS1_MZ_TOLERANCE_DEFAULT;
    /* This step size should make it impossible for us to miss any readings in the index due to FP error.
     * We could make the index more compact by spacing these windows out a bit, but I'll leave that as a TODO. */
    static final Double MZ_STEP_SIZE = MS1.MS1_MZ_TOLERANCE_DEFAULT / 2.0;
    static final Double MIN_MZ = 50.0;
    static final Double MAX_MZ = 950.0;
  
    public static final String OPTION_INDEX_PATH = "x";
    public static final String OPTION_SCAN_FILE = "i";
  
    public static final String HELP_MESSAGE = StringUtils.join(new String[]{
        "This class extracts and indexes readings from an LCMS scan files, ",
        "and writes them to an on-disk index for later processing."
    }, "");
  
    public static final List<Option.Builder> OPTION_BUILDERS = new ArrayList<Option.Builder>() {{
      add(Option.builder(OPTION_INDEX_PATH)
          .argName("index path")
          .desc("A path to the directory where the on-disk index will be stored; must not already exist")
          .hasArg().required()
          .longOpt("index")
      );
      add(Option.builder(OPTION_SCAN_FILE)
          .argName("scan file")
          .desc("A path to the LCMS NetCDF scan file to read")
          .hasArg().required()
          .longOpt("input")
      );
      add(Option.builder("h")
          .argName("help")
          .desc("Prints this help message")
          .longOpt("help")
      );
    }};
  
    public static class Factory {
      public static Builder makeBuilder(File indexDir) throws RocksDBException{
        RocksDB.loadLibrary();
        LOGGER.info("Creating index at %s", indexDir.getAbsolutePath());
        RocksDBAndHandles<ColumnFamilies> dbAndHandles = DBUtil.createNewRocksDB(indexDir, ColumnFamilies.values());
        return new Builder(dbAndHandles);
      }
    }
  
    private RocksDBAndHandles<ColumnFamilies> dbAndHandles;
  
    Builder(RocksDBAndHandles<ColumnFamilies> dbAndHandles) {
      this.dbAndHandles = dbAndHandles;
    }
  
    public static void main(String[] args) throws Exception {
      CLIUtil cliUtil = new CLIUtil(Builder.class, HELP_MESSAGE, OPTION_BUILDERS);
      CommandLine cl = cliUtil.parseCommandLine(args);
  
      File inputFile = new File(cl.getOptionValue(OPTION_SCAN_FILE));
      if (!inputFile.exists()) {
        cliUtil.failWithMessage("Cannot find input scan file at %s", inputFile.getAbsolutePath());
      }
  
      File indexDir = new File(cl.getOptionValue(OPTION_INDEX_PATH));
      if (indexDir.exists()) {
        cliUtil.failWithMessage("Index file at %s already exists--remove and retry", indexDir.getAbsolutePath());
      }
  
      Builder indexBuilder = Factory.makeBuilder(indexDir);
      try {
        indexBuilder.processScan(indexBuilder.makeTargetMasses(), inputFile);
      } finally {
        if (indexBuilder != null) {
          indexBuilder.close();
        }
      }
  
      LOGGER.info("Done");
    }
  
    public void close() throws RocksDBException {
      dbAndHandles.close();
    }
  
    private List<Double> makeTargetMasses() {
      List<Double> targets = new ArrayList<>();
      for (double m = MIN_MZ - MZ_STEP_SIZE; m <= MAX_MZ + MZ_STEP_SIZE; m += MZ_STEP_SIZE) {
        targets.add(m);
      }
      return targets;
    }
  
    public void processScan(List<Double> targetMZs, File scanFile)
        throws RocksDBException, ParserConfigurationException, XMLStreamException, IOException {
      DateTime start = DateTime.now();
      LOGGER.info("Accessing scan file at %s", scanFile.getAbsolutePath());
      LCMSNetCDFParser parser = new LCMSNetCDFParser();
      Iterator<LCMSSpectrum> spectrumIterator = parser.getIterator(scanFile.getAbsolutePath());
  
      WriteOptions writeOptions = new WriteOptions();
      /* The write-ahead log and disk synchronization features are useful when we need our writes to be durable (i.e. to
       * survive a crash).  However, our index construction is effectively a one-shot deal: if it doesn't succeed, we'll
       * just start from scratch.  Since we don't care about durability while we're constructing the index, the WAL and
       * sync features eat a lot of disk space and I/O bandwidth, which slows us down.  So long as we cleanly close the
       * index once it's built, nobody has to know that we disabled these features. */
      writeOptions.setDisableWAL(true);
      writeOptions.setSync(false);
      dbAndHandles.setWriteOptions(writeOptions);
  
      // TODO: split targetMZs into batches of ~100k and extract incrementally to allow huge input sets.
  
      LOGGER.info("Extracting traces");
      List<MZWindow> windows = targetsToWindows(targetMZs);
      extractTriples(spectrumIterator, windows);
  
      LOGGER.info("Writing search targets to on-disk index");
      writeWindowsToDB(windows);
  
      DateTime end = DateTime.now();
      LOGGER.info("Index construction completed in %dms", end.getMillis() - start.getMillis());
    }
  
    private List<MZWindow> targetsToWindows(List<Double> targetMZs) {
      List<MZWindow> windows = new ArrayList<MZWindow>() {{
        int i = 0;
        for (Double targetMZ : targetMZs) {
          add(new MZWindow(i, targetMZ));
          i++;
        }
      }};
  
      /* We *must* ensure the windows are sorted in m/z order for the sweep line to work.  However, we don't know anything
       * about the input targetMZs list, which may be immutable or may be in some order the client wants to preserve.
       * Rather than mess with that array, we'll sort the windows in our internal array and leave be the client's targets.
       */
      Collections.sort(windows, (a, b) -> a.getTargetMZ().compareTo(b.getTargetMZ()));
  
      return windows;
    }
  
    protected void extractTriples(
        Iterator<LCMSSpectrum> iter,
        List<MZWindow> windows)
        throws RocksDBException, IOException {
      /* Warning: this method makes heavy use of ByteBuffers to perform memory efficient collection of values and
       * conversion of those values into byte arrays that RocksDB can consume.  If you haven't already, go read this
       * tutorial on ByteBuffers: http://mindprod.com/jgloss/bytebuffer.html
       *
       * ByteBuffers are quite low-level structures, and they use some terms you need to watch out for:
       *   capacity: The total number of bytes in the array backing the buffer.  Don't write more than this.
       *   position: The next index in the buffer to read or write a byte.  Moves with each read or write op.
       *   limit:    A mark of where the final byte in the buffer was written.  Don't read past this.
       *             The remaining() call is affected by the limit.
       *   mark:     Ignore this for now, we don't use it.  (We'll always, always read buffers from 0.)
       *
       * And here are some methods that we'll use often:
       *   clear:     Set position = 0, limit = 0.  Pretend the buffer is empty, and is ready for more writes.
       *   flip:      Set limit = position, then position = 0.  This remembers how many bytes were written to the buffer
       *              (as the current position), and then puts the position at the beginning.
       *              Always call this after the write before a read.
       *   rewind:    Set position = 0.  Buffer is ready for reading, but unless the limit was set we might now know how
       *              many bytes there are to read.  Always call flip() before rewind().  Can rewind many times to re-read
       *              the buffer repeatedly.
       *   remaining: How many bytes do we have left to read?  Requires an accurate limit value to avoid garbage bytes.
       *   reset:     Don't use this.  It uses the mark, which we don't need currently.
       *
       * Write/read patterns look like:
       *   buffer.clear(); // Clear out anything already in the buffer.
       *   buffer.put(thing1).put(thing2)... // write a bunch of stuff
       *   buffer.flip(); // Prep for reading.  Call *once*!
       *
       *   while (buffer.hasRemaining()) { buffer.get(); } // Read a bunch of stuff.
       *   buffer.rewind(); // Ready for reading again!
       *   while (buffer.hasRemaining()) { buffer.get(); } // Etc.
       *   buffer.reset(); // Forget what was written previously, buffer is ready for reuse.
       *
       * We use byte buffers because they're fast, efficient, and offer incredibly convenient means of serializing a
       * stream of primitive types to their minimal binary representations.  The same operations on objects + object
       * streams require significantly more CPU cycles, consume more memory, and tend to be brittle (i.e. if a class
       * definition changes slightly, serialization may break).  Since the data we're dealing with is pretty simple, we
       * opt for the low-level approach.
       */
  
      /* Because we'll eventually use the window indices to map a mz range to a list of triples that fall within that
       * range, verify that all of the indices are unique.  If they're not, we'll end up overwriting the data in and
       * corrupting the structure of the index. */
      ensureUniqueMZWindowIndices(windows);
  
      // For every mz window, allocate a buffer to hold the indices of the triples that fall in that window.
      ByteBuffer[] mzWindowTripleBuffers = new ByteBuffer[windows.size()];
      for (int i = 0; i < mzWindowTripleBuffers.length; i++) {
        /* Note: the mapping between these buffers and their respective mzWindows is purely positional.  Specifically,
         * mzWindows.get(i).getIndex() != i, but mzWindowTripleBuffers[i] belongs to mzWindows.get(i).  We'll map windows
         * indices to the contents of mzWindowTripleBuffers at the very end of this function. */
        mzWindowTripleBuffers[i] = ByteBuffer.allocate(Long.BYTES * 4096); // Start with 4096 longs = 8 pages per window.
      }
  
      // Every TMzI gets an index which we'll use later when we're querying by m/z and time.
      long counter = -1; // We increment at the top of the loop.
      // Note: we could also write to an mmapped file and just track pointers, but then we might lose out on compression.
  
      // We allocate all the buffers strictly here, as we know how many bytes a long and a triple will take.  Then reuse!
      ByteBuffer counterBuffer = ByteBuffer.allocate(Long.BYTES);
      ByteBuffer valBuffer = ByteBuffer.allocate(TMzI.BYTES);
      List<Float> timepoints = new ArrayList<>(2000); // We can be sloppy here, as the count is small.
  
      /* We use a sweep-line approach to scanning through the m/z windows so that we can aggregate all intensities in
       * one pass over the current LCMSSpectrum (this saves us one inner loop in our extraction process).  The m/z
       * values in the LCMSSpectrum become our "critical" or "interesting points" over which we sweep our m/z ranges.
       * The next window in m/z order is guaranteed to be the next one we want to consider since we address the points
       * in m/z order as well.  As soon as we've passed out of the range of one of our windows, we discard it.  It is
       * valid for a window to be added to and discarded from the working queue in one application of the work loop. */
      LinkedList<MZWindow> tbdQueueTemplate = new LinkedList<>(windows); // We can reuse this template to init the sweep.
  
      int spectrumCounter = 0;
      while (iter.hasNext()) {
        LCMSSpectrum spectrum = iter.next();
        float time = spectrum.getTimeVal().floatValue();
  
        // This will record all the m/z + intensity readings that correspond to this timepoint.  Exactly sized too!
        ByteBuffer triplesForThisTime = ByteBuffer.allocate(Long.BYTES * spectrum.getIntensities().size());
  
        // Batch up all the triple writes to reduce the number of times we hit the disk in this loop.
        // Note: huge success!
        RocksDBAndHandles.RocksDBWriteBatch<ColumnFamilies> writeBatch = dbAndHandles.makeWriteBatch();
  
        // Initialize the sweep line lists.  Windows go follow: tbd -> working -> done (nowhere).
        LinkedList<MZWindow> workingQueue = new LinkedList<>();
        LinkedList<MZWindow> tbdQueue = (LinkedList<MZWindow>) tbdQueueTemplate.clone(); // clone is in the docs, so okay!
        for (Pair<Double, Double> mzIntensity : spectrum.getIntensities()) {
          // Very important: increment the counter for every triple.  Otherwise we'll overwrite triples = Very Bad (tm).
          counter++;
  
          // Brevity = soul of wit!
          Double mz = mzIntensity.getLeft();
          Double intensity = mzIntensity.getRight();
  
          // Reset the buffers so we end up re-using the few bytes we've allocated.
          counterBuffer.clear(); // Empty (virtually).
          counterBuffer.putLong(counter);
          counterBuffer.flip(); // Prep for reading.
  
          valBuffer.clear(); // Empty (virtually).
          TMzI.writeToByteBuffer(valBuffer, time, mz, intensity.floatValue());
          valBuffer.flip(); // Prep for reading.
  
          // First, shift any applicable ranges onto the working queue based on their minimum mz.
          while (!tbdQueue.isEmpty() && tbdQueue.peekFirst().getMin() <= mz) {
            workingQueue.add(tbdQueue.pop());
          }
  
          // Next, remove any ranges we've passed.
          while (!workingQueue.isEmpty() && workingQueue.peekFirst().getMax() < mz) {
            workingQueue.pop(); // TODO: add() this to a recovery queue which can then become the tbdQueue.  Edge cases!
          }
          /* In the old indexed trace extractor world, we could bail here if there were no target m/z's in our window set
           * that matched with the m/z of our current mzIntensity.  However, since we're now also recording the links
           * between timepoints and their (t, m/z, i) triples, we need to keep on keepin' on regardless of whether we have
           * any m/z windows in the working set right now. */
  
          // The working queue should now hold only ranges that include this m/z value.  Sweep line swept!
  
          /* Now add this intensity to the buffers of all the windows in the working queue.  Note that since we're only
           * storing the *index* of the triple, these buffers are going to consume less space than they would if we
           * stored everything together. */
          for (MZWindow window : workingQueue) {
            // TODO: count the number of times we add intensities to each window's accumulator for MS1-style warnings.
            counterBuffer.rewind(); // Already flipped.
            mzWindowTripleBuffers[window.getIndex()] = // Must assign when calling appendOrRealloc.
                Utils.appendOrRealloc(mzWindowTripleBuffers[window.getIndex()], counterBuffer);
          }
  
          // We flipped after reading, so we should be good to rewind (to be safe) and write here.
          counterBuffer.rewind();
          valBuffer.rewind();
          writeBatch.put(ColumnFamilies.ID_TO_TRIPLE, Utils.toCompactArray(counterBuffer), Utils.toCompactArray(valBuffer));
  
          // Rewind again for another read.
          counterBuffer.rewind();
          triplesForThisTime.put(counterBuffer);
        }
  
        writeBatch.write();
  
        assert(triplesForThisTime.position() == triplesForThisTime.capacity());
  
        ByteBuffer timeBuffer = ByteBuffer.allocate(Float.BYTES).putFloat(time);
        timeBuffer.flip(); // Prep both bufers for reading so they can be written to the DB.
        triplesForThisTime.flip();
        dbAndHandles.put(ColumnFamilies.TIMEPOINT_TO_TRIPLES,
            Utils.toCompactArray(timeBuffer), Utils.toCompactArray(triplesForThisTime));
  
        timepoints.add(time);
  
        spectrumCounter++;
        if (spectrumCounter % 1000 == 0) {
          LOGGER.info("Extracted %d time spectra", spectrumCounter);
        }
      }
      LOGGER.info("Extracted %d total time spectra", spectrumCounter);
  
      // Now write all the mzWindow to triple indexes.
      RocksDBAndHandles.RocksDBWriteBatch<ColumnFamilies> writeBatch = dbAndHandles.makeWriteBatch();
      ByteBuffer idBuffer = ByteBuffer.allocate(Integer.BYTES);
      for (int i = 0; i < mzWindowTripleBuffers.length; i++) {
        idBuffer.clear();
        idBuffer.putInt(windows.get(i).getIndex());
        idBuffer.flip();
  
        ByteBuffer triplesBuffer = mzWindowTripleBuffers[i];
        triplesBuffer.flip(); // Prep for read.
  
        writeBatch.put(ColumnFamilies.WINDOW_ID_TO_TRIPLES, Utils.toCompactArray(idBuffer), Utils.toCompactArray(triplesBuffer));
      }
      writeBatch.write();
  
      dbAndHandles.put(ColumnFamilies.TIMEPOINTS, TIMEPOINTS_KEY, Utils.floatListToByteArray(timepoints));
      dbAndHandles.flush(true);
    }
  
    private void ensureUniqueMZWindowIndices(List<MZWindow> windows) {
      Set<Integer> ids = new HashSet<>(windows.size());
      for (MZWindow window : windows) {
        if (ids.contains(window.getIndex())) {
          String msg = String.format("Assumption violation: found duplicate mzWindow index when all should be unique: %d",
              window.getIndex());
          LOGGER.error(msg);
          // A hard invariant has been violated, so crash the program.
          throw new RuntimeException(msg);
        }
        ids.add(window.getIndex());
      }
    }
  
    /**
     * Writes all MZWindows to the DB as serialized objects.  Windows are keyed by target MZ (as a Double) for easy lookup
     * in the case that we have a target m/z that exactly matches a window.
     *
     * @param windows The windows to write.
     * @throws RocksDBException
     * @throws IOException
     */
    protected void writeWindowsToDB(List<MZWindow> windows) throws RocksDBException, IOException {
      for (MZWindow window : windows) {
        byte[] keyBytes = Utils.serializeObject(window.getTargetMZ());
        byte[] valBytes = Utils.serializeObject(window);
  
        dbAndHandles.put(ColumnFamilies.TARGET_TO_WINDOW, keyBytes, valBytes);
      }
  
      dbAndHandles.flush(true);
      LOGGER.info("Done writing window data to index");
    }
  }
  