package nxm.dsp.prim;

import nxm.sys.lib.Data;
import nxm.sys.lib.DataFile;
import nxm.sys.lib.DataOp;
import nxm.sys.lib.Primitive;
import nxm.sys.libm.Fft;
import nxm.sys.libm.MMath;

/** Use an FFT convolution algorithm to do cross-correlation quickly.
    @author Scott Fernandez / Neon Ngo
    @since NeXtMidas 3.1.0
 */
public class fastcorrelate extends Primitive {

  /** input (long) file (IN=) */
  private DataFile hi;
  /** input (short) filter file (FIR=) */
  private DataFile hf;
  /** output file (OUT=) */
  private DataFile ho;
  /** data buffer for work */
  private Data dbw;
  /** data buffer for filter */
  private Data dbf;
  /** data buffer for mux */
  private Data dbm;
  /** FFT size */
  private int nfft;
  /** (short) filter input (FIR=) file size */
  private int nfir;
  /** number of bins (i.e. size) for data buffers */
  private int nbins;
  /** flags for FFT library */
  private int flags;
  /** (-bias * 4) offset into {@link #dbw} for copying to {@link #dbm} when {@link #mux} is enabled */
  private int bias;
  /** number of elements to keep in each process loop ({@link #nfft} - ({@link #over}) */
  private int keep;
  /** overlap amount ({@link #nfir} - 1) */
  private int over;
  /** number of zero elements to prepend to input data */
  private int pad1;
  /** number of zero elements to append to input data */
  private int pad2;
  /** MUX in & out for hilbert */
  private boolean mux;
  /** pad transient with zeros */
  private boolean pad;
  /** complex mode (input file or input impulse is complex) */
  private boolean cxm;
  /** FFT library for doing forward FFT */
  private Fft fft;
  /** FFT library for doing inverse FFT */
  private Fft ifft;
  /** scale factor for FIR filter */
  private double scale;

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

  public int open () {

    hi = MA.getDataFile ("IN","1000","S#,C#,-SD,-CD",0);
    hi.open();

    hf = MA.getDataFile ("FIR","1000","S#,C#,-SD,-CD",0);
    hf.open();

    cxm  = (hi.spa==2) || (hf.spa==2);        // complex mode
    mux  = MA.getState("/MUX");               // mux in & out for hilbert
    pad  = MA.getState("/PAD", false) || mux; // pad transient with zeros
    if (mux && cxm) M.error("Cannot /MUX with complex input file");

    // determine size of input files
    double size1 = hi.getSize();
    double size2 = hf.getSize();
    nfir = (int) size2;

    // Get fft size from command line switch (default = 32K)
    nfft = Math.max(32768, MMath.power2GE(2*nfir));
    if (MA.getState("/FFT")) {
      nfft = MA.getL("/FFT");
      M.deprecate("FASTCORRELATE: Use of /FFT is deprecated. Use /NFFT instead.");
    }
    nfft = MA.getL("/NFFT", nfft);

    // Perform error checking and defaulting
    double size3 = (size1 - size2) + 1; // size of output
    if (size3 <= 0) {
      M.error("input length ["+size1+"] is not longer than short (FIR) input ["+size2+"]");
    }

    double hiXDelta = hi.getXDelta();
    double hfXDelta = hf.getXDelta();
    if (hiXDelta != hfXDelta) { // Are abscissa deltas the same?
      M.warning("Abscissa deletas of inputs are unequal: " + hiXDelta + " vs "+hfXDelta);
    }

    over = nfir - 1;
    keep = nfft - over;
    bias = (int)Math.round(hf.getXStart()/hfXDelta);

    if (pad) { pad1 = pad2 = over; }
    else     { pad1 = pad2 = 0; }

    if (cxm) {
      flags = Fft.COMPLEX;
      scale = 1.0/nfft;
      nbins = nfft;
    } else {
      flags = Fft.REAL|Fft.UNPACKED;
      scale = 1.0/(4*nfft);
      nbins = nfft/2 + 1;
    }
    fft  = new Fft(nfft,flags|Fft.FORWARD);
    ifft = new Fft(nfft,flags|Fft.INVERSE);

    ho = MA.getDataFile("OUT", hi, 0);
    ho.setSize( size1-over+pad1+pad2 );

    double outXStart = (pad)? hi.getXStart() - hf.getXStart() + (pad1*hi.getXDelta())
                            : hi.getXStart() - hf.getXStart(); // <-- matches XM's fastcorl
    ho.setXStart(outXStart);

    if (cxm||mux) ho.setFormatMode(COMPLEX);
    String outFormat = (char)ho.getFormatMode() + "F";
    outFormat = Data.promoteFormat(hi.getFormat(), hf.getFormat(), outFormat);
    ho.setFormat(outFormat);
    ho.open();

    hi.setXferLength(nfft);
    hi.setConsLength(keep);

    // promote to at least complex mode for FFT calculation on data buffers
    String bufFormat = Data.promoteFormat(outFormat, "CF", null);
    dbw = new Data(bufFormat, nbins);         // data buffer for read hi & work
    dbf = new Data(bufFormat, nbins);         // data buffer for filter
    if (mux) dbm = new Data(bufFormat, nfft); // data buffer for mux

    // if (verbose) M.info("DEBUG: outFormat="+outFormat +" bufFormat="+bufFormat+ " nfft="+nfft+" nbins="+nbins+" keep="+keep+" over="+over+" bias="+bias+" scale="+scale+" outXStart="+outXStart+" pad1="+pad1+" pad2="+pad2+" size3="+size3);
    return NORMAL;
  }

  public int process () {

    if (hf.avail() > 0) {
      loadFilter();                      // new filter available
    }

    dbw.setSPA(hi.spa);
    if (pad1 > 0) {                      // handle front pad I/O
      DataOp.zero(dbw, 0, pad1);         // insert pad1 samples of zero
      dbw.boff = (pad1 * dbw.bpe);       // adjust byte offset for read buffer
      hi.setXferLength(nfft - pad1);     // short the transfer
      hi.setConsLength(keep - pad1);     // short the consume
    }

    int n = hi.read(dbw);                // read input data
    if (n == 0) return NOOP;
    if (n < 0)  return FINISH;

    if (pad1 > 0) {                      // restore std I/O parameters
      dbw.boff = 0;
      hi.setXferLength(nfft);
      hi.setConsLength(keep);
      n += pad1;
      pad1 = 0;
    }

    if (n < nfft) {                                  // handle back end I/O
      if (n <= over && pad2 <= 0) { return FINISH; } // no good points
      DataOp.zero(dbw, n, nfft - n);                 // zero pad to read buffer to xfer (FFT) size
      n = Math.min(nfft, n + pad2);
    }

    if (cxm && hi.spa == 1)
      DataOp.cmplx(dbw, dbw, nfft, DataOp.SC);
    if (mux)
      DataOp.copy(dbw, -bias * 4, 4, dbm, 0, 8, 4, n - over);

    fft.work(dbw); // perform forward FFT
    dbw.setSPA(2);

    // complex multiply with complex conjugate
    DataOp.mul(dbw, dbf, dbw, nbins, DataOp.CGC);

    ifft.work(dbw); // perform inverse FFT

    if (mux) {
      dbw.setSPA(1);           // set for I/O libs
      DataOp.copy(dbw, 0, 4, dbm, 4, 8, 4, n - over);
      ho.write(dbm, n - over); // write just the good points
    }
    else {
      dbw.setSPA(ho.spa);      // set for I/O libs
      ho.write(dbw, n - over); // write just the good points
    }

    return NORMAL;
  }

  public int close () {
    fft.free();
    ifft.free();
    hi.close();
    hf.close();
    ho.close();
    return NORMAL;
  }

  private void loadFilter () {
    float fscale = (float) scale;
    dbf.setSPA(hf.spa);
    int read = hf.read(dbf, nfir); // get the FIR taps
    over = read - 1;
    keep = nfft - over;
    if (pad) {
      pad1 = pad2 = over;
    }
    fscale *= MA.getF("/SCALE", 1);

    Data dbs = new Data(fscale); // scale factor
    DataOp.mul(dbf, dbs, dbf, read);
    if (cxm && hf.spa == 1)
      DataOp.cmplx(dbf, dbf, nfft, DataOp.SC);
    fft.work(dbf); // perform forward FFT of impulse response filter
    dbf.setSPA(2);
  }

  /** get the FFT size
      @return FFT size
   */
  public int getNFft () {
    return nfft;
  }

  /** set the FFT size.
      @param value new FFT size to use
   */
  public void setNFft (int value) {
    if (value == nfft) return;
    MA.put("/NFFT", "" + value);
    setState(RESTART);
  }

}
