package nxm.dsp.prim;

import java.util.Arrays;

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

/**
    <pre>
  ###############
  #             #
  #  CORRELATE  #
  #             #
  ###############

  NAME:  CORRELATE

  PURPOSE:
     To correlate the data in file &lt;in&gt; with that of file
     &lt;impulse&gt; to effect a linear time-invariant filter.

  DESCRIPTION:
     Get command line parameters
     Set-up necessary MIDAS files
     Read buffer of data
     Perform correlation
     Write buffer of data
     Post headers, close files and return to MIDAS

  CALLING SEQUENCE:

     CORRELATE &lt;in&gt; &lt;impulse&gt; &lt;out&gt;


  INPUT FILES:
     MIDAS input file containing data
     MIDAS input file containing impulse response

  OUTPUT FILES:
     MIDAS file of results

  </pre>

  For explanation of running sums, (inMeans and inStdDevFactors)
  see http://scribblethink.org/Work/nvisionInterface/nip.html#eq3:xform

  @author Scott Fernandez
 */
public class correlate extends Primitive {
  /** Indicates whether the native correlate library has been loaded properly. */
  private static boolean loaded = Shell.loadLibrary("dsp", "prim", "correlate");
  /** Transfer length.  Specify on command line with the '/TL=' switch */
  private int tl = 4096;
  /** Specify on command line with /NATIVE */
  private boolean useNative;
  /** Specify on command line with /NORMALIZE */
  private boolean normalize;
  /** Specify on command line with /QUIET */
  private boolean quiet;

  private DataFile inputFile, in1File;
  private DataFile outputFile;
  private DataFile impulseFile, in2File;
  private Data inputData;
  private Data impulseData;
  private Data outputData;

  private double outXStart;
  private float[] inNormFactors;
  private float[] input;
  private float[] impulse;
  private float[] output;
  private float max;
  private int maxIndex;
  private int impulseTransfer;
  private int numInput;
  private int numImpulse;
  private int unprocessed;
  private int totalProcessed;
  private int shouldProcess;
  private boolean done, noop, complexIn, complexImp;

  private static native int correlateNative (
      Object inbuf, Object impbuf, Object outbuf, int bufsize, int ngot);
  private static native int normalizedCorrelateNative (
      Object inbuf, Object impbuf, Object outbuf, int bufsize, int ngot,
      Object in_norm_factors_buf);

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

  public int open () {

    max = maxIndex = totalProcessed = 0;

    /** Decide if using native */
    useNative = MA.getState("/NATIVE", true) && Native.useMath && loaded;
    /** See if doing normalized correlation */
    normalize = MA.getState("/NORMALIZE", false);
    /** See if outputting maximum */
    quiet = MA.getState("/QUIET", false);

    /* Get and open files */
    inputFile = MA.getDataFile ("IN", "1000", "S#, C#",0);
    impulseFile = MA.getDataFile ("IMPULSE", "1000", "S#, C#", 0);
    inputFile.open(DataFile.WARNEMPTY);
    impulseFile.open(DataFile.WARNEMPTY);

    complexIn  = (inputFile.spa == 2);
    complexImp = (impulseFile.spa == 2);

    /* Get sizes of files */
    int inputSize = (int)inputFile.getSize();
    int impulseSize = (int)impulseFile.getSize();
    if (inputSize <= 0 || impulseSize <= 0) {
      M.error("An input file is empty.  Check IN and IMPULSE.");
    }
    /* Switch files internally if impulse is larger than input */
    if ( inputSize >= impulseSize ) {
      in1File = inputFile;
      in2File = impulseFile;
    } else {
      in1File = impulseFile;
      in2File = inputFile;
      int tempSize = inputSize;
      inputSize = impulseSize;
      impulseSize = tempSize;
    }

    if (in1File.getXDelta() != in2File.getXDelta()) {
      M.warning("Deltas are not the same");
    }

    /* Format output file */
    outputFile = MA.getDataFile ("OUT",in1File,0);
    outXStart = in1File.getXStart() - in2File.getXStart() - (in2File.getSize() - 1);
    outputFile.setXStart( outXStart );
    outputFile.setSize( inputSize + impulseSize - 1.0 );
    outputFile.open();

    /* Set up transfer length */
    tl = MA.getL("/TL", tl);
    if (tl <= 0) {
      M.error("Invalid transfer length.");
      return FINISH;
    }
    tl = Math.min(inputSize, tl);
    impulseTransfer = impulseSize;

    /* Set up buffers */
    outputData = outputFile.getDataBuffer(tl, FLOAT);
    impulseData = in2File.getDataBuffer(impulseTransfer, FLOAT);
    inputData = in1File.getDataBuffer(tl, FLOAT);

    /* Set up impulse array */
    setupImpulseArray();

    /* Create and initialize input and, if normalizing, norm array */
    input = new float[tl+impulseTransfer];
    if (normalize) { inNormFactors = new float[tl]; }
    unprocessed = numImpulse-1;
    for(int i = 0; i < unprocessed; i++)
      input[i] = 0;

    return NORMAL;
  }

  /* Oversees the work */
  public int process () {
    /* Setup the input array and figure out how many to process */
    int shouldProcess = setupInput();
    /* If no input received, noop */
    if (noop) {
      noop = false;
      return NOOP;
    }
    /* Do work in transfer length chunks if not noop */
    while (shouldProcess > tl) {
      int processed = work(tl);
      shouldProcess -= processed;
    }
    while (shouldProcess > 0) {
      int processed = work(shouldProcess);
      shouldProcess -= processed;
    }

    /* If done, finish */
    return done ? FINISH : NORMAL;
  }

  /* Close files */
  public int close () {
    if (!quiet) { M.info("max was " + max + ", at abcissa " + maxIndex + " in the larger file"); }
    inputFile.close();
    impulseFile.close();
    outputFile.close();
    return (NORMAL);
  }

  /*
   =============================
   |      SETUP FUNCTIONS      |
   =============================
   */
  /** Sets up the input
   * @return The number of elements that should be processed
   */
  private int setupInput () {
    int shouldProcess = 0;
    /* Read input file */
    int read = Math.min(input.length - unprocessed, tl);
    numInput = in1File.read(inputData, read);
    /* Stall if no data */
    if (numInput == 0 && read > 0) {
      noop = true;
      return shouldProcess;
    }
    /* EOF means we're done */
    done = (numInput < 0);
    if (!done) {
      /* If not done, store new data in input array */
      float[] temp = inputData.castF(true);
      for(int i = 0, j = unprocessed; i < numInput; i++, j++)
        input[j] = temp[i];
      inputData.uncast(temp, false);
      unprocessed += numInput;
      shouldProcess = Math.max(unprocessed - numImpulse + 1, 0);
    } else {
      /* When finishing, set to process all of the elements that are left
           and clear out the end of the input array */
      shouldProcess = unprocessed;
      int length = input.length;
      for(int i = unprocessed; i < length; i++) {
        input[i] = 0;
      }
    }
    return shouldProcess;
  }

  /** Sets up the impulse array... as the name would suggest */
  private void setupImpulseArray() {
    /* Read in impulse file. If reached EOF, we're done. */
    numImpulse = in2File.read(impulseData);

    /* Cast impulse buffer to temporary float array, copy into a new array, then uncast */
    float[] temp = impulseData.castF(true);
    impulse = Arrays.copyOf(temp, temp.length);
    /* If normalizing, normalize the impulse array */
    if (normalize) { normalize(impulse); }
    impulseData.uncast(temp, false);
  }

  /*
  ============================
  |      WORK FUNCTIONS      |
  ============================
  */
  /** Does work for the process function
      @param toProcess The number of elements to calculate the correlation for
      @return The number of elements processed
   */
  private int work (int toProcess) {
    int i, j, processed;
    /* Return if no elements to process */
    if (toProcess == 0)
      return 0;

    /* Get output array */
    output = outputData.castF(false);

    /* Correlate and reduce unprocessed by number of elements processed */
    if (normalize) {
      /* Normalize input */
      getNormalizationFactors(toProcess);
      if (useNative) {
        processed = normalizedCorrelateNative(input, impulse, output,
            toProcess, numImpulse, inNormFactors);
      } else {
        processed = normalizedCorrelation(toProcess);
      }
    } else {
      if (useNative) {
        processed = correlateNative(input, impulse, output, toProcess, numImpulse);
      } else {
        processed = correlation(toProcess);
      }
    }
    unprocessed -= processed;
    totalProcessed += processed;

    /* Update max if not quiet */
    if (!quiet) { updateMax(processed); }
    /* Uncast to buffer and write to file */
    outputData.uncast(output, true);
    outputFile.write(outputData, processed);

    /* Store data for next time */
    for(i = 0, j = processed; i < unprocessed; i++, j++) {
      input[i] = input[j];
    }
    /* If finishing up, zero out the end */
    if (done) {
      int end = unprocessed+processed;
      for(i = unprocessed; i < end; i++) {
        input[i] = 0;
      }
    }
    return processed;
  }

  /** Calculate normalized cross-correlation
      @param toProcess The number of elements to calculate the correlation for
      @return The number of elements processed without any verification.
   */
  private int normalizedCorrelation (int toProcess) {
    double sum, sum1, sum2, sum3, sum4;
    int unrolledEnd = toProcess-4;
    int i, j, ij;
    /* Unrolled outer looping--do 4 values at once */
    for (i = 0; i < unrolledEnd; i+=4) {
      sum1 = sum2 = sum3 = sum4 = 0.0f;
      for (ij = i, j = 0; j < numImpulse; j++, ij++) {
        float impVal = impulse[j];
        sum1 += input[ij]   * impVal;
        sum2 += input[ij+1] * impVal;
        sum3 += input[ij+2] * impVal;
        sum4 += input[ij+3] * impVal;
      }

      output[i]   = (float)(sum1/inNormFactors[i]);
      output[i+1] = (float)(sum2/inNormFactors[i+1]);
      output[i+2] = (float)(sum3/inNormFactors[i+2]);
      output[i+3] = (float)(sum4/inNormFactors[i+3]);
    }
    /* Finish up the last few, if any */
    for (; i < toProcess; i++) {
      sum = 0.0f;
      for(ij = i, j = 0; j < numImpulse; j++, ij++) {
        sum += input[ij] * impulse[j];
      }
      output[i] =  (float)sum/inNormFactors[i];
    }
    return toProcess;
  }

  /** Calculate cross-correlation
      @param toProcess The number of elements to calculate the correlation for
      @return The number of elements processed without any verification.
   */
  private int correlation (int toProcess) {
    double sum, sum1, sum2, sum3, sum4;
    int unroll = 4;
    int unrolledEnd = toProcess-4;
    int i, j, ij;
    if (complexIn && complexImp) {
      for (i = 0; i < unrolledEnd; i+=unroll) {
        sum1 = sum2 = sum3 = sum4 = 0.0f;
        for (ij = 2*i, j = 0; j < 2*numImpulse; j+=2) {
          float impVal1 = impulse[j];
          float impVal2 = impulse[j+1];
          float in1 = input[ij];
          float in2 = input[ij+1];
          sum1 += in1 * impVal1 - in2 * impVal2;
          sum2 += in1 * impVal2 + in2 * impVal1;

          ij += 2;
          in1 = input[ij];
          in2 = input[ij+1];

          sum3 += in1 * impVal1 - in2 * impVal2;
          sum4 += in1 * impVal2 + in2 * impVal1;
        }
        output[i]   = (float)sum1;
        output[i+1] = (float)sum2;
        output[i+2] = (float)sum3;
        output[i+3] = (float)sum4;
      }
      /* Finish up the last few, if any */
      for (; i < toProcess; i+=2) {
        sum1 = sum2 = 0.0f;
        for(ij = 2*i, j = 0; j < 2*numImpulse; j+=2, ij+=2) {
          float impVal1 = impulse[j];
          float impVal2 = impulse[j+1];
          float in1 = input[ij];
          float in2 = input[ij+1];
          sum1 += in1 * impVal1 - in2 * impVal2;
          sum2 += in1 * impVal2 + in2 * impVal1;
        }
        output[i]   = (float)sum1;
        output[i+1] = (float)sum2;
      }
    } else if (complexIn) {

    } else if (complexImp) {

    }
    /* Unrolled outer looping--do 4 values at once */
    for (i = 0; i < unrolledEnd; i+=unroll) {
      sum1 = sum2 = sum3 = sum4 = 0.0f;
      for (ij = i, j = 0; j < numImpulse; j++, ij++) {
        float impVal = impulse[j];
        sum1 += input[ij]   * impVal;
        sum2 += input[ij+1] * impVal;
        sum3 += input[ij+2] * impVal;
        sum4 += input[ij+3] * impVal;
      }
      output[i]   = (float)sum1;
      output[i+1] = (float)sum2;
      output[i+2] = (float)sum3;
      output[i+3] = (float)sum4;
    }
    /* Finish up the last few, if any */
    for (; i < toProcess; i++) {
      sum = 0.0f;
      for(ij = i, j = 0; j < numImpulse; j++, ij++) {
        sum += input[ij] * impulse[j];
      }
      output[i] = (float)sum;
    }
    return toProcess;
  }

  /*
   ================================
   |      AUXILARY FUNCTIONS      |
   ================================
   */
  /** Updates the maximum value of the array
      @param processed The number of elements to look through for the new max
   */
  private void updateMax (int processed) {
    for(int i = 0; i < processed; i++) {
      float temp = output[i];
      if (temp > max) {
        max = temp;
        maxIndex = totalProcessed - processed + i + (int)outXStart;
      }
    }
  }

  /** Reduces calculation time for normalization
      @param toProcess The number of elements to get the normalization factor for
   */
  private void getNormalizationFactors (int toProcess) {
    /* these are double for precision in calculations, especially the sums */
    double temp, temp2, sum, sumsquared;
    sum = sumsquared = 0;
    for(int i = 0; i < numImpulse; i++) {
      temp = input[i];
      sum += temp;
      sumsquared += temp*temp;
    }

    for(int i = 0, j = numImpulse; i < toProcess; i++, j++) {
      /* Get mean and standard deviation */
      temp = sum/numImpulse;
      temp2 = Math.sqrt(sumsquared - sum*temp);
      inNormFactors[i] = (temp2 == 0) ? 1 : (float)temp2;

      /* Calculate next frame's sum and sumsquared */
      temp = input[i];
      temp2 = input[j];
      /* Subtract out last elements, add new elements */
      sum += temp2 - temp;
      sumsquared += temp2*temp2 - temp*temp;
    }
  }

  /** Copies the given array into a new array of the same size, then returns the new array.
      @param array input array to copy
      @return array containing a copy of input array
      @deprecated Since 3.5.1, use {@link Arrays#copyOf(float[], int)}
   */
  @Deprecated
  public float[] copyArray (float[] array) {
    return Arrays.copyOf(array, array.length);
  }

  /** Normalizes the given array.
      @param array The array to normalize
   */
  public void normalize (float[] array) {
    int length = array.length;
    /* Calculate mean */
    double mean = 0;
    for (int i = 0; i < length; i++) {
      mean += array[i];
    }
    mean /= length;

    /* Store the elements shifted by the mean in a new array.
         Meanwhile, calculate the standard deviation factor, the standard
         deviation multiplied by the square root of the number of elements */
    double stdDevFactorFactor = 0;
    for (int i = 0; i < length; i++) {
      double centered = array[i] - mean;
      array[i] = (float)centered;
      stdDevFactorFactor += centered*centered;
    }
    stdDevFactorFactor = Math.sqrt(stdDevFactorFactor);

    /* Normalize impulse array and return it */
    if (stdDevFactorFactor != 0) { // If the standard dev is 0, the array is all 0s
      for (int i = 0; i < length; i++) {
        array[i] /= stdDevFactorFactor;
      }
    }
  }

  /** For debugging */
  static String array2String (float[] arr) {
    String str = ""+arr[0];
    for(int i = 1; i < arr.length; i++) {
      str += ", "+arr[i];
    }
    return str;
  }
}
