package nxm.sys.prim;

import nxm.sys.inc.Units;
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;

/**
  Perform FFT based spectral functions.

  @author Jeff Schoen
*/
public class fft extends Primitive {

   /** windowing function name list */
  public static final String windowList = DataOp.windowList;

  /** PARTIAL mode (process or discard partial input data frames) - NONE, ALL or START */
  public static final String partialList= "None,All,Start";
  /** PARTIAL mode: discard any partial input data insufficient to form a whole frame */
  public static final int P_NONE        = 1;
  /** PARTIAL mode: zero-pad any partial input data to produce an output frame */
  public static final int P_ALL         = 2;
  /** PARTIAL mode: zero-pad partial input data only at the start of processing;
   *  discard partial data if we have already output at least one frame */
  public static final int P_START       = 3;

  /** Packing mode (how the nyquist bin is handled) for real data (P_UNPACK ,P_DC or P_ZERO) */
  public static final String packList = "Unpack,DC,Zero";
  /** Packing mode that unpacks the nyquist into its own bin which may cause sub-optimal transfer lengths */
  public static final int P_UNPACK    = 1;
  /** Packing mode that packs the real part of the nyquist bin into the always zero imaginary component of the P_DC bin */
  public static final int P_DC        = 2;
  /** Packing mode that drops the nyquist bin */
  public static final int P_ZERO      = 3;

  /** DataFile for inputs/outputs/cross output file */
  private DataFile hi1, ho1, hi2, ho2, hox;

  /** Data buffers for inputs/outputs/computations */
  private Data dbi1,dbi2,dbf1,dbf2,dbfx,dbo,dbw,dba1,dba2,dbax,dac,db;

  private int nfft;      // transform size
  private int nin;       // transfer length
  private int nrot;      // number rotate for complex
  private int nout;      // number out
  private int navg;      // num transforms to average
  private int nexp;      // num exponential averages
  private int iavg,iexp; // averaging counters
  private int stride;    // stride (num elements to skip)
  private int pack;      // packing mode
  private int partial;   // padding option for incomplete input

  private double factor; // windowing scale factor
  private double over;   // fraction of {@link #nfft} to overlap blocks
  private double size;   // size of input files (smaller of IN1 or IN2)
  private double rnavg;  // 1/navg or 1/{@link #navg}
  private double rnexp;  // 1/nexp/navg or 1/{@link #nexp}*{@link #navg}

  private boolean mag;   // output magnitude squared (MAG)
  private boolean psd;   // output power spectral density (PSD)
  private boolean log;   // output MAG or PSD in 20log format
  private boolean fdb;   // use fast dB log computation
  private boolean cxm;   // input data is in complex mode
  private boolean absc;  // Align input files by abscissa
  private boolean out1d; // output a type 1000 file
  private boolean in2d;  // input 1 is a type 2000 file
  private boolean dp;    // do double precision buffers & output
  private boolean ac;    // Take out DC bias before window and FFT
  private boolean cmag;  // output mag/log on cross as well (NXM-3899)
  private String  wname; // Name of windowing function
  private Fft     fft;   // FFT library

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

  @Override
  public int open () {

    hi1 = MA.getDataFile ("IN1","1000,2000","S#,C#",0);
    hi1.open(DataFile.WARNEMPTY);
    in2d = (hi1.typeClass==2);

    hi2 = MA.getDataFile ("IN2","1000,2000","S#,C#",0);
    hi2.open(DataFile.WARNEMPTY|DataFile.OPTIONAL);

    nfft  = MA.getL("NFFT");
    if (nfft == -1) nfft = (int)hi1.getSize(); //Matches X-Midas for -1 to get file size (2GB int limit)
    nin   = MA.getL("/TL",nfft);
    psd   = MA.getState("/PSD");
    log   = MA.getState("/LOG");
    fdb   = MA.getState("/DBF");
    mag   = MA.getState("/MAG",fdb||log||psd);
    cmag  = MA.getState("/CMAG");
    absc  = MA.getState("/ABSC",true);
    out1d = MA.getState("/1D");
    ac    = MA.getState("/AC");
    dp    = MA.getState("/DP");

    partial = MA.getSelectionIndex("/PARTIAL",partialList,P_START);

    if (state!=RESTART) {                             // retain realtime changes
      navg = MA.getL("NAVG");
      over = MA.getD("OVER");
      nexp = MA.getL("/NEXP",1);
      wname = MA.getS("WIN");
    }

    double delx = hi1.getXDelta();
    double delf = 1.0/(nfft*delx);
    double dely;
    double stry; // input's Start (X Start for 1D, Y Start for 2D)
    cxm = (hi1.spa==2) || (hi2.isOpen && hi2.spa==2);

    if (in2d) {
      nin  = hi1.getSubSize();
      dely = hi1.getYDelta();
      stry = hi1.getYStart();
      hi1.setFS(0);
      hi2.setFS(0);
      setOverlap(over);                               // initialize stride
    } else {
      setOverlap(over);                               // initialize stride
      dely = delx*stride;
      stry = hi1.getXStart();
      if (hi2.isOpen) hi1.align(hi2,absc,false);      // align both input files
    }

    size = hi1.size;                                  // estimate size
    if (hi2.isOpen) size = Math.min(hi2.size,size);
    if (navg<0) navg = (int)((size-nfft)/stride)+1;   // average the entire file
    setNAvg (navg);
    setNExp (nexp);

    int flags = Fft.FORWARD;
    byte dtype = (dp)? DOUBLE:FLOAT;
    String cform = (dp)? "CD":"CF";
    String sform = (dp)? "SD":"SF";
    pack = out1d? P_UNPACK : P_ZERO; // unpack only if single output frame (i.e. /1D)
    pack = MA.getSelectionIndex("/PACK",packList,pack); // since 2.7.1 for BUG 1853

    if (cxm) {
      nout = nfft;
      nrot = nfft/2-1;
      flags |= Fft.COMPLEX;
      factor = 1.0/nfft;
      pack = 0;                      // no packing for complex data
    } else {
      nout = nfft/2;
      if (nfft != nout*2) M.error("Invalid FFT size parameter: "+nfft);
      nrot = 0;
      flags |= Fft.REAL;
      if (pack==P_UNPACK) { flags |= Fft.UNPACKED; nout += 1; }
      else                { flags |= Fft.PACKED;              } // ZERO and DC are same (i.e. PACKED) except that ZERO drops the nyquist bin
      factor = 1.0/(2*nfft);
    }
    // check for scaling
    factor = MA.getD("/SCALE", factor);
    fft = new Fft(nfft,flags);
    if (psd) factor /= Math.sqrt(delf);
    if (!MA.getState("/ROTATE",cxm)) nrot=0;

    String oform =   (mag)? sform : cform;  // output file format
    double osize = (out1d)? 1.0   : Math.max(1.0,Math.floor((size-nfft)/stride/navg)); // output file size
    ho1 = MA.getDataFile("OUT1","2000",oform,0);
    ho1.propagate(hi1,"TIMECODE|KEYWORDS"); // propagate timecode and keywords from IN1 to OUT1 (Bug 2101)
    ho1.setSubSize (nout);
    ho1.setXStart(-nrot*delf);
    ho1.setXDelta(delf);
    ho1.setXUnits(Units.FREQUENCY);
    ho1.setYStart(stry + (nfft+(navg-1)*stride)/2 * delx);
    ho1.setYDelta(dely*navg);
    if (!out1d) {
      ho1.setYUnits(Units.TIME);  // if not /1D then set YU=time (since 3.1.0 for BUG 2495)
    }
    ho1.setSize(osize);
    ho1.open(DataFile.WARNEMPTY|DataFile.OPTIONAL);

    setWindow (wname);

    dac  = hi1.getDataBuffer(1,dtype);
    dbi1 = hi1.getDataBuffer(nin,dtype);
    dbf1 = new Data (cform,nout);
    dbo  = ho1.getDataBuffer(1,dtype);
    dba1 = ho1.getDataBuffer(1,dtype);

    ho2 = MA.getDataFile ("OUT2",ho1,0);    // since input files are aligned (have same data rate), use ho1's headers
    ho2.propagate(hi2,"TIMECODE|KEYWORDS"); // propagate timecode and keywords from input file 2 to output file 2
    ho2.open(DataFile.WARNEMPTY|DataFile.OPTIONAL);

    hox = MA.getDataFile ("CROSS",ho1,0);   // use headers & keywords from output 1 for cross output file
    if (cmag) hox.setFormat(oform);
    else hox.setFormat(cform);
    hox.open(DataFile.WARNEMPTY|DataFile.OPTIONAL);

    if (!hi2.isOpen && hox.isOpen) {
      M.error("CROSS parameter is invalid when only one input file supplied");
    }

    if (ho2.isOpen || hox.isOpen) {
      dbi2 = hi2.getDataBuffer(nin,dtype);
      dba2 = ho2.getDataBuffer(1,dtype);
      dbf2 = new Data (cform,nout);
      dbfx = new Data (cform,nout);
      dbax = new Data (cmag?oform:cform,nout); // Jeff NXM-3899
    }
    iavg=iexp=1;
    return NORMAL;
  }

  @Override
  public int process () {
    int navgSaved = navg; // save off navg, so it can be restored after an EOF scenario (3.9.0)
    int stat = NORMAL;
    int n1 = (hi1.isOpen)? hi1.read(dbi1) : 0;
    int n2 = (hi2.isOpen)? hi2.read(dbi2) : n1;
    if (n1<0 || n2<0) {                               // EOF
      stat = FINISH;
      if (iavg==1) return stat;                       // already a clean finish
      navg = iavg;                                    // prep for last partial average
      n1 = n2 = Math.max(0,Math.min(n1,n2));          // finish off last frame
    }
    rnavg = 1.0/navg;
    rnexp = 1.0/navg/iexp;
    if (n1 == nin || partial != P_NONE) {
      if (hi1.isOpen) processInput (n1,ho1,dbi1,dbf1,dba1);
      if (hi2.isOpen) processInput (n2,ho2,dbi2,dbf2,dba2);
      if (hox.isOpen) processCross (hox,dbf1,dbf2,dbax);
    }
    if (++iavg > navg) {          // have we done enough for the next (averaged) output frame?
      iavg  = 1;                  // reset counter
      if (out1d) stat = FINISH;   // our work here is done
      if (partial == P_START) partial = P_NONE; // don't want any more partials
      if (nexp<=1) iexp = 1; else if (++iexp > nexp) iexp = nexp;
    }
    navg = navgSaved; // restore original navg - since 3.9.0 for NXM-3210
    return stat;
  }

  private void processInput (int n, DataFile ho, Data dbi, Data dbf, Data dba) {

    dbf.setSPA(dbi.spa);                              // see dbf as input

    if (ac) {                                         // remove DC bias
      DataOp.sum(dbi,dac,n);
      DataOp.scale(dac,dac,1,(-1.0/n));
      DataOp.add(dbi,dac,dbi,n);
    }

    DataOp.mul(dbw,dbi,dbf);                          // apply window

    if (n<nfft) DataOp.zero(dbf,n,nfft-n);            // zero fill

    if (cxm&&dbi.spa==1) DataOp.cmplx(dbf,dbf,nfft,DataOp.SC);  // force complex FFT

    fft.work (dbf);                                   // perform FFTs
    dbf.setSPA(2); db=dbf;                            // see dbf as output
    if (pack==P_ZERO) dbf.setF(1,0.0F);

    if (mag) { DataOp.mag2(db,dbo); db=dbo; }         // compute PSD

    if (iexp>1) {
      DataOp.scale(dba,dba,nout,1.0-rnexp);           // scale by 1-1/(navg*nexp)
      DataOp.scale(db ,db ,nout,rnexp);               // scale by 1/(navg*nexp)
      DataOp.add(db,dba,dba);
      if (iavg<navg) return;
      db=dba;
    }
    else if (navg>1) {                                // average
      if (iavg==1) DataOp.copy(db,dba);
      else DataOp.add(db,dba,dba);
      if (iavg<navg) return;
      DataOp.scale(dba,dba,nout,rnavg);               // scale by 1/navg
      db=dba;
    }
    if (nexp>1) { DataOp.copy(dba,dbo); db=dbo; }

    if (db==dbf && hox.isOpen) { DataOp.copy(dbf,dbo); db=dbo; } // preserve buffer for cross

    if (fdb) DataOp.dbf(db,nout,0);                   // compute fast Db
    else if (log) DataOp.dblogb(db,db);               // compute log

    if (nrot!=0) fft.rotate (db);                     // rotate complex fft bins
    db.setHeader( dbi.getHeader() );                  // propagate time line
    ho.write (db);                                    // write output
  }

  private void processCross (DataFile ho, Data dbf1, Data dbf2, Data dba) {

    db=dbfx;
    DataOp.mul(dbf1,dbf2,db,nout,DataOp.GCC);         // compute cross spectra

    if (cmag) { DataOp.mag2(db,dbo); db=dbo; }        // if /CMAG compute PSD of cross spectrum output as well (NXM-3899)

    if (nexp>1) {
      DataOp.scale(dba,dba,nout,1.0-rnexp);           // scale by 1-1/(navg*nexp)
      DataOp.scale(db ,db ,nout,rnexp);               // scale by 1/(navg*nexp)
      DataOp.add(db,dba,dba);
      if (iavg<navg) return;
      db=dba;
    }
    else if (navg>1) {                                // average
      if (iavg==1) DataOp.copy(db,dba);
      else DataOp.add(db,dba,dba);
      if (iavg<navg) return;
      DataOp.scale(dba,dba,nout,rnavg);               // scale by 1/navg
      db=dba;
    }

    // IF /CMAG apply /LOG, and /DBF to cross spectrum output as well (NXM-3899)
    if (!cmag);
    else if (fdb) DataOp.dbf(db,nout,0);              // compute fast Db
    else if (log) DataOp.dblogb(db,db);               // compute log

    if (nrot!=0) fft.rotate (db);                     // rotate complex fft bins
    ho.write(db);                                     // write output
  }

  @Override
  public int close () {
    fft.free();
    if (out1d && ho1.isOpen) ho1.setFS(0);
    if (out1d && ho2.isOpen) ho2.setFS(0);
    if (out1d && hox.isOpen) hox.setFS(0);
    hi1.close();
    ho1.close();
    hi2.close();
    ho2.close();
    hox.close();
    return NORMAL;
  }

  @Override
  public int processReady () {
    if (state!=PROCESS) return 0;
    int n = hi1.processReady(nin);
    if (hi2.isOpen) n = Math.min(n, hi2.processReady(nin) );
    if (ho1.isOpen) n = Math.min(n, ho1.processReady(1) );
    if (ho2.isOpen) n = Math.min(n, ho2.processReady(1) );
    if (hox.isOpen) n = Math.min(n, hox.processReady(1) );
    return n;
  }

  /** get the number of input frames to block average for each output frame
   *  @return the number of input frames to block average for each output frame
   */
  public int getNAvg () { return navg; }
  /** get the number of output frames to exponentially average for each output frame
   *  @return the number of output frames to exponentially average for each output frame
   */
  public int getNExp () { return nexp; }
  /** get the FFT size
   *  @return fft size
   */
  public int getNFft () { return nfft; }
  /** get the fractional overlap, zero for no overlap, negative to skip
   *  @return the fractional overlap
   */
  public double getOverlap () { return over; }
  /** get the windowing function name
   *  @return the windowing function name
   */
  public String getWindow () { return wname; }

  /** set the number of input frames to block average for each output frame
   *  @param value number of input frames to block average for each output frame
   */
  public void setNAvg (int value) {
    navg = Math.max(1,value);
  }

  /** set the number of output frames to exponentially average for each output frame
   *  @param value number of output frames to exponentially average for each output frame
   */
  public void setNExp (int value) {
    if (value<0) { iexp=1; value=-value; }
    nexp = Math.max(1,value);
  }

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

  /** set the fractional overlap, zero for no overlap, negative to skip
   *  @param value fractional overlap
   */
  public void setOverlap (double value) {
    if (in2d && stride%nin!=0) M.warning("Type 2000 inputs don't have integral overlap="+over);
    over = value;        // note stride of 0 can be useful
    stride = Math.max(0,(int)((1.0-over)*nin));
    hi1.setXferLength(nin);
    hi1.setConsLength(stride);
    hi2.setXferLength(nin);
    hi2.setConsLength(stride);
  }

  /** set the windowing function
   *  @param value name of windowing function
   */
  public void setWindow (String value) {
    int nwind = Math.min(nin,nfft);
    wname = value;
    dbw = DataOp.getWindow(wname,nwind,factor,M);
  }

}
