package nxm.dsp.prim;

import nxm.sys.lib.Data;
import nxm.sys.lib.DataFile;
import nxm.sys.lib.Native;
import nxm.sys.lib.Parser;
import nxm.sys.lib.Primitive;
import nxm.sys.lib.Shell;

/** Normalizes an input file by frames and according to the mode specified.
  <pre>
  DESCRIPTION:
     Get command line parameters
     Set-up necessary MIDAS files
     Read buffer of data
     Perform normalization
     Write buffer of data
     Post headers, close files and return to MIDAS

  CALLING SEQUENCE:

     NORMALIZE &lt;in&gt; &lt;out&gt; &lt;fs&gt; &lt;mode&gt; &lt;scale&gt; &lt;shift&gt;

  INPUT FILES:
     MIDAS input file containing data

  OUTPUT FILES:
     MIDAS file of normalized data

  </pre>

  @author Scott Fernandez / Neon Ngo
  @since NeXtMidas 3.1.0
 */
public class normalize extends Primitive {
  /** This list of normalization modes.  Options are MEAN, SD, PEAK, and NONE */
  public static final String modeList = "MEAN,SD,PEAK,NONE";
  /** mean only normalization mode */
  private static final int MEAN = 1;
  /** mean + standard deviation + scale normalization mode */
  private static final int SD   = 2;
  /** peak-to-peak normalization mode */
  private static final int PEAK = 3;
  /** no normalization--just scale mode */
  private static final int NONE = 4;

  /** Indicates whether the native correlate library has been loaded properly. */
  private static boolean nativeLoaded = Shell.loadLibrary("dsp", "prim", "normalize");
  /** Specify on command line with /NATIVE */
  private boolean useNative;
  /** Frame size.  Default is 16k. */
  private int fs = 16384;

  /** Mode of normalization to perform. Options are {@link #MEAN}, {@link #SD}, {@link #PEAK}, and {@link #NONE}. */
  private int mode;
  /** Value to scale the output by (multiply with output) */
  private double scale;
  /** Value to shift the output by (add to output) */
  private double shift;

  private DataFile inputFile;
  private DataFile outputFile;
  private Data data;

  /** atoms per element from input file */
  private int ape;
  /** Elements read from input file */
  private int numRead;
  /** Real mean of elements in buffer */
  private double rmean;
  /** Imaginary mean of elements in buffer */
  private double imean;
  /** Real sum of elements in buffer */
  private double rsum;
  /** Imaginary sum of elements in buffer */
  private double isum;
  /** Sum of real squared + imaginary squared (only for complex) of elements in buffer */
  private double magsum;
  /** Standard deviation of elements in buffer */
  private double stdDev;
  /** Maximum value (magnitude for complex numbers) of elements in buffer */
  private double max;
  /** Minimum value (magnitude for complex numbers) of elements in buffer */
  private double min;
  /** Whether the format is complex */
  private boolean complex;

  private static native int nativeScaleAndShift (Object buf, double scale, double shift, int start, int end);

  /** default constructor */
  public normalize () {}

  public int open () {
    /** Specify on command line with /NATIVE */
    useNative = MA.getState("/NATIVE", true) && Native.useMath && nativeLoaded;
    /** Mode of normalization.  Options SD, MEAN, PEAK, and NONE correspond to 1, 2, 3, and 4 */
    mode = MA.getChoice("MODE", modeList, SD);
    /** Value to scale the output by (multiply with output) */
    scale = MA.getD("SCALE", 1);
    /** Value to shift the output by (add to output) */
    shift = MA.getD("SHIFT", 0);

    /* Get files */
    inputFile = MA.getDataFile("IN", "1000,2000", "S#,C#",0);
    inputFile.open(DataFile.WARNEMPTY);
    outputFile = MA.getDataFile("OUT", inputFile, 0);

    complex = (inputFile.getFormatMode() == COMPLEX);
    ape     = inputFile.getAPE();

    /* Format and open output file */
    byte outType = Data.promoteType(inputFile.dataType, FLOAT);
    outputFile.setFormatType(outType);
    outputFile.setSize(inputFile.getSize());
    outputFile.open(DataFile.WARNEMPTY|DataFile.OPTIONAL);

    // Set up frame size (normalize width/span)
    if (inputFile.getTypeCodeClass() == 2) {
      fs = inputFile.getFrameSize(); // use Type 2000 input file's frame size
      xfer = 1;
    }
    else {
      fs = MA.getL("FS", fs);
      if (fs <= 0) {
        M.error("Invalid frame size specified: "+fs);
        return FINISH;
      }
      xfer = fs = Math.max(Math.min(fs, (int)inputFile.getSize()), 1);
    }

    /* Set up buffer */
    data = outputFile.getDataBuffer(xfer);

    resetAccumulators();

    return NORMAL;
  }

  /* Oversees the work */
  public int process () {
    /* Read input file */
    numRead = inputFile.read(data);
    /* If no input data received, noop */
    if (numRead == 0) { return NOOP; }
    /* If reached EOF, finish */
    if (numRead == -1) { return FINISH; }
    /* Otherwise, reset variables */
    resetAccumulators();

    /* Get the statistics if mode isn't NONE */
    if (mode != NONE) {
      computeStatistics(data, numRead, ape, complex);
    }

    /* Do work */
    switch (mode) {
    case SD:   sdNormalize    (data, scale, shift); break;
    case MEAN: zeroMean       (data, scale, shift); break;
    case PEAK: normalizePeaks (data, scale, shift); break;
    case NONE: scaleAndShift  (data, scale, shift); break;
    }

    /* Uncast and write normalized data */
    outputFile.write(data, numRead);

    return NORMAL;
  }

  /* Close files */
  public int close () {
    inputFile.close();
    outputFile.close();
    return NORMAL;
  }

  // ******************** Begin get/set run-time properties ********************
  /** get current normalize mode.
      @return normalize mode
   */
  public String getMode () { return Parser.get(modeList, mode); }
  /** set normalize mode.
      @param modeStr new normalize mode
   */
  public void setMode (String modeStr) {
    mode = Parser.find(modeList, modeStr, mode, 0, Parser.E_ERROR);
  }

  /** get factor to scale normalized output.
      @return scale factor
   */
  public double getScale () { return scale; }
  /** set factor to scale normalized output.
      @param scaleFactor new scale factor
   */
  public void setScale (double scaleFactor) { scale = scaleFactor; }

  /** get amount to shift the normalized and scaled output.
      @return amount to shift
   */
  public double getShift () { return shift; }
  /** set amount to shift the normalized and scaled output.
      @param shiftAmount amount to shift
   */
  public void setShift (double shiftAmount) { shift = shiftAmount; }
  // ********************* End get/set run-time properties *********************

  /** Resets the accumulators. */
  protected void resetAccumulators () {
    rsum = isum = magsum = stdDev = 0;
    max  = DMIN;
    min  = DMAX;
  }

  protected void computeStatistics (Data data, int ngot, int atomsPerElement, boolean complex) {
    double[] dbuf = data.castD(true);
    if (complex) {
      for (int i=0; i < (ngot*atomsPerElement*2); ) {
        double rval = dbuf[i++];
        double ival = dbuf[i++];
        rsum += rval;
        isum += ival;
        double mag  = (rval*rval) + (ival*ival);
        magsum += mag;
        // leave min/max as squared until end for speed
        if (mag > max) { max = mag; }
        if (mag < min) { min = mag; }
      }
    } else { // scalar data
      for (int i=0; i<ngot*atomsPerElement; ) {
        double dval = dbuf[i++];
        rsum   += dval;
        magsum += dval*dval; // no imaginary component
        // leave min/max as squared until end for speed
        if (dval > max) { max = dval; }
        if (dval < min) { min = dval; }
      }
    }
    data.uncast(dbuf, false);
    finishCalculations();
  }

  /** finish the statistics calculations */
  protected void finishCalculations () {
    int nelem = numRead*ape;
    rmean  = rsum / nelem;
    imean  = isum / nelem;
    stdDev = (magsum/nelem) - (rmean*rmean) - (imean*imean);
    stdDev = Math.sqrt( Math.max(0.0,stdDev) );

    if (complex) {  // revert complex squared back to magnitude for min/max
      min = Math.sqrt(min);
      max = Math.sqrt(max);
    }
  }

  /** Normalizes the given data.
      @param data The data to normalize
      @param scale The value to multiply the output by
      @param shift The value to add to the output
   */
  private void sdNormalize (Data data, double scale, double shift) {
    /* If the standard deviation is 0, every element in the data will be zero
       after adjusting for the mean, so scaling will do nothing. */
    boolean noScale = (stdDev == 0);
    /* Otherwise, the standard deviation simply another factor to scale by */
    scale = noScale ? 1 : scale / stdDev;

    zeroMean(data, scale, shift);
  }

  /** Centers the first numElem elements, making the mean of them 0
      @param data The data to zero the mean of
      @param scale The value to multiply the output by
      @param shift The value to add to the output
   */
  private void zeroMean (Data data, double scale, double shift) {
    /* result = (original-mean)*scale + shift = original*scale + (shift - mean*scale) */
    shift -= rmean*scale;

    scaleAndShift(data, scale, shift);

    /* If complex data, add back the real mean and subtract out the
       complex mean from the complex data */
    if (complex) {
      shiftImaginary(data, rmean*scale - imean*scale);
    }
  }

  private void shiftImaginary (Data data, double shift) {
    double[] dbuf = data.castD(true);
    scalarAddToArray(dbuf, shift, 1, 2);
    data.uncast(dbuf, true);
  }

  /** Normalizes the peaks of the data, shifting and scaling so that the
      maximum is 0.5 and the minimum is -0.5 before scaling and shifting
      by the user input.
      @param data The data to normalize the peaks of
      @param scale The value to multiply the output by
      @param shift The value to add to the output
   */
  private void normalizePeaks (Data data, double scale, double shift) {
    double range = max - min;

    /* If range is 0, the data is constant, and we can't really do anything */
    if (range == 0) { scaleAndShift(data, 1, shift-max); return; }

    if (complex) {
      Data norm = data.copy();
      sdNormalize(norm, 1, 0);
      resetAccumulators();
      computeStatistics(norm, numRead, ape, true);

      range = max-min;

      scale = (range == 0) ? 1 : scale/range;
      shift -= scale*(max+min)/2;

      scaleAndShift(data, scale, shift);
    } else {
      /* Subtract out the max, setting max to 0, then add 0.5 times user scaling */
      shift += 0.5*scale - max*scale/range;
      /* Divide scale by range to set range to 1 */
      scale /= range;

      scaleAndShift(data, scale, shift);
    }
  }

  /** Scales and shifts data.
      @param data The data to scale and shift
      @param scale The value to multiply the output by
      @param shift The value to add to the output
   */
  private void scaleAndShift (Data data, double scale, double shift) {
    float[] array = data.castF(true);
    if (useNative) {
      nativeScaleAndShift(array, scale, shift, 0, array.length);
    } else if (scale == 1 || shift == 0) {
      scalarMultiplyToArray(array, scale);
      scalarAddToArray(array, shift);
    } else {
      /* Otherwise, do them both at once */
      int unroll = 4;
      int end = array.length;
      int unrolledEnd = end - (end % unroll);
      int i;
      for (i = 0; i < unrolledEnd; i+=unroll) {
        array[i]   = (float)(array[i]  *scale + shift);
        array[i+1] = (float)(array[i+1]*scale + shift);
        array[i+2] = (float)(array[i+2]*scale + shift);
        array[i+3] = (float)(array[i+3]*scale + shift);
      }
      for (; i < end; i++) {
        array[i] = (float)(array[i]*scale + shift);
      }
    }
    data.uncast(array, true);
  }

  /**
      @param array The float array to multiply the scalar into
      @param scale The value to multiply each element by
   */
  private void scalarMultiplyToArray (float[] array, double scale) {
    /* Return if nothing to do */
    if (scale == 1) { return; }

    /* Unroll to four assignments per iteration */
    int unroll = 4;
    int end = array.length;
    int unrolledEnd = end - (end % unroll);
    int i;
    for (i = 0; i < unrolledEnd; i+=unroll) {
      array[i]   *= scale;
      array[i+1] *= scale;
      array[i+2] *= scale;
      array[i+3] *= scale;
    }
    for (; i < end; i++) {
      array[i] *= scale;
    }
  }

  /**
      @param array The float array to add to
      @param shift The value to add to the array
   */
  private void scalarAddToArray (float[] array, double shift) {
    /* Return if nothing to do */
    if (shift == 0) { return; }

    /* Unroll to four assignments per iteration */
    int unroll = 4;
    int end = array.length;
    int unrolledEnd = end - (end % unroll);
    int i;
    for (i = 0; i < unrolledEnd; i+=unroll) {
      array[i]   += shift;
      array[i+1] += shift;
      array[i+2] += shift;
      array[i+3] += shift;
    }
    for (; i < end; i++) {
      array[i] += shift;
    }
  }

  /**
      @param data The double data array to add to
      @param shift The value to add
      @param start Where to start changing data
      @param stride The stride to take
   */
  private void scalarAddToArray (double[] data, double shift, int start, int stride) {
    /* Return if nothing to do */
    if (shift == 0) { return; }

    /* Unroll to four assignments per iteration */
    int unroll = 4;
    int move = unroll*stride;
    int end = data.length-start;
    int unrolledEnd = end - (end % move);
    int i;
    for (i = start; i < unrolledEnd; i+=move) {
      data[i]          += shift;
      data[i+stride]   += shift;
      data[i+2*stride] += shift;
      data[i+3*stride] += shift;
    }
    for (; i < end; i+=stride) {
      data[i] += shift;
    }
  }
}
