MODULE PZSpectrum EXPORTS Main;

(* Writes the Fourier power spectrum OF a curve. *)

IMPORT  ParseParams, Process, FileRd, Wr, FileWr, Math;
IMPORT  OSError, Thread, Stdio, Fmt, FPut;
IMPORT  PZCurve, PZLR3Chain, PZLRChain, PZFourierChain;

FROM Stdio IMPORT stderr;
 
<* FATAL Wr.Failure, Thread.Alerted, OSError.E *>

TYPE
  LONG = LONGREAL;
  BOOL = BOOLEAN;
  
  Options = RECORD
      inFile: TEXT;         (* Input float chain file (with extension). *)
      outName: TEXT;        (* Output float/curv chain file name (without extension). *)
      lambdaMin: LONG;      (* Min scale length for spectrum computation. *)
      uniformize: BOOL;     (* TRUE to uiformize velocity before sampling *)
    END;
    
PROCEDURE Main() =
  BEGIN
    WITH
      o = GetOptions(),
      rp = PZLR3Chain.Read(FileRd.Open(o.inFile)),
      p = rp.c^,
      m = NUMBER(p),
      c = PZCurve.New(m)
    DO 
      c.setSamples(p);
      Wr.PutText(stderr, "length = " & Fmt.LongReal(c.length) & "\n");
      IF o.uniformize THEN 
        WITH diff = UniformizeVelocity(c) DO
          Wr.PutText(stderr, "UniformizeVelocity diff = " & Fmt.LongReal(diff) & "\n")
        END;
        WITH diff = UniformizeVelocity(c) DO
          Wr.PutText(stderr, "UniformizeVelocity diff = " & Fmt.LongReal(diff) & "\n")
        END
      END;
      WritePowerSpectrum(o, c);
    END
  END Main;
  
PROCEDURE GetOptions(): Options =
  VAR o: Options;
  BEGIN
    WITH
      pp = NEW(ParseParams.T).init(stderr)
    DO
      TRY
        pp.getKeyword("-inFile");
        o.inFile := pp.getNext();
        
        pp.getKeyword("-outName");
        o.outName := pp.getNext();
        
        pp.getKeyword("-lambdaMin");
        o.lambdaMin := pp.getNextLongReal();

        o.uniformize := pp.keywordPresent("-uniformize");
      
        pp.finish();      
        
      EXCEPT                                                            
      | ParseParams.Error =>                                              
          Wr.PutText(stderr, "Usage: PZSpectrum \\\n");
          Wr.PutText(stderr, "  -inFile NAME \\\n");
          Wr.PutText(stderr, "  -outName NAME  \\\n");
          Wr.PutText(stderr, "  -lambdaMin NUM  \\\n");
          Wr.PutText(stderr, "  [ -uniformize ] \n");
        Process.Exit(1);
      END;
    END;
    RETURN o
    
  END GetOptions;


PROCEDURE WritePowerSpectrum(
    READONLY o: Options;
    READONLY c: PZCurve.T
  ) =
  BEGIN
    WITH
      pwr = ComputePowerSpectrum(c, o.lambdaMin)^,
      length = c.length
    DO
      WriteFreqPower(o.outName, pwr, length);
      WriteBandPower(o.outName, pwr, length, o.lambdaMin, 8, 2.0d0)
    END
  END WritePowerSpectrum;


PROCEDURE UniformizeVelocity(c: PZCurve.T): LONG =
  (*
    Given a curve "c", this procedure adjusts the nodal times of "c"
    so as to uniformize its velocity. *)
  VAR diff, TNew: LONG;
  BEGIN
    WITH
      tOld = c.t^,
      TOld = c.tPeriod,
      tNew = NEW(REF ARRAY OF LONG, c.m)^
    DO
      (* Set the time "c.t[j]" to the length from sample "[0]" to sample "[j]": *)
      VAR s: LONG := 0.0d0;
      BEGIN
        FOR j := 0 TO c.m-1 DO tNew[j] := s; s := s + c.s[j] END;
        TNew := s
      END;
      (* Compute variation in spacing: *)
      diff := 0.0d0;
      FOR i := 1 TO c.m-1 DO
        WITH 
          an = tNew[i-1], bn = tNew[i],
          ao = tOld[i-1], bo = tOld[i]
        DO 
          diff := MAX(diff, ABS((bn-an)/TNew-(bo-ao)/TOld));
        END
      END;
      WITH 
        an = tNew[0] + TNew, bn = tNew[c.m-1],
        ao = tOld[0] + TOld, bo = tOld[c.m-1]
      DO 
        diff := MAX(diff, ABS((bn-an)/TNew-(bo-ao)/TOld));
      END;
      c.setTimes(tNew, TNew);
      RETURN diff
    END
  END UniformizeVelocity;


PROCEDURE ComputePowerSpectrum(
    READONLY c: PZCurve.T; 
    lambdaMin: LONG;
  ): REF ARRAY OF LONG =
  (*
    Returns a vector "pwr[f]" with the power (square of amplitude)
    for each frequency "f", from 0 to "fMax = c.m DIV 2".
    Elements "pwr[0] and "pwr[fMax]" reflect only one term of the
    the series, all other elements include two conjugate terms. *)
  VAR p: LONG;
  BEGIN
    WITH
      m = c.m,
      estLength = 1.5d0 * c.length,
      nMin = MAX(2*m, CEILING(8.0d0*estLength/lambdaMin)),
      n = PowerOfTwo(nMin),
      fMax = n DIV 2,
      t = NEW(REF PZLRChain.T, n)^,
      u = NEW(REF PZLR3Chain.T, n)^,
      v = NEW(REF PZLR3Chain.T, n)^,
      r = NEW(REF PZLRChain.T, n)^,
      a = NEW(REF PZFourierChain.T, 3, n)^,
      rpwr = NEW(REF ARRAY OF LONG, fMax+1), pwr = rpwr^
    DO
      FOR i := 0 TO n-1 DO t[i] := c.tPeriod * FLOAT(i,LONG)/FLOAT(n,LONG) END;
      c.sample(t, u, v, r);
      PZLR3Chain.FourierTransform(u, a);
      
      WITH ax = a[0], ay = a[1], az = a[2] DO
        FOR f := 0 TO fMax DO
          WITH
            tx = ax[f],
            ty = ay[f],
            tz = ax[f]
          DO
            p := tx*tx + ty*ty + tz*tz
          END;
          IF f > 0 AND f < fMax THEN
            WITH
              tx = ax[n-f],
              ty = ay[n-f],
              tz = az[n-f]
            DO
              p := p + tx*tx + ty*ty + tz*tz
            END
          END;
          pwr[f] := p
        END
      END;
      RETURN rpwr
    END
  END ComputePowerSpectrum;
    

PROCEDURE PowerOfTwo(n: CARDINAL): CARDINAL = 
  (* Smallest power of two not less than "n" *)
  VAR m: CARDINAL := 1;
  BEGIN
    WHILE m < n DO m := m+m END;
    RETURN m
  END PowerOfTwo;


PROCEDURE WriteFreqPower(
    name: TEXT; 
    READONLY pwr: ARRAY OF LONG; 
    length: LONG;
  ) =
  BEGIN  
    WITH 
      fMax = LAST(pwr),
      fileName = name & "-fpw.plt",
      wr = FileWr.Open(fileName)
    DO 
      Wr.PutText(stderr, "writing file " & fileName & " ...\n");
      Wr.PutText(wr, "# freq   lambda     rms  \n");
      Wr.PutText(wr, "# ----  ------- -------  \n");
      FOR f := 1 TO fMax DO
        WITH
          lambda = length/FLOAT(f, LONG)
        DO
          FPut.Int(wr, f, 6); 
          Wr.PutText(wr, "  ");
          FPut.LongReal(wr, lambda, 4, Fmt.Style.Fix); 
          Wr.PutText(wr, "  ");
          FPut.LongReal(wr, Math.sqrt(pwr[f]), 4, Fmt.Style.Fix); 
          Wr.PutText(wr, "\n");
        END;
      END;
      Wr.Flush(wr)
    END    
  END WriteFreqPower;


PROCEDURE WriteBandPower(
    name: TEXT; 
    READONLY pwr: ARRAY OF LONG;
    length: LONG;
    lambdaMin: LONG;
    nScales: CARDINAL;
    scaleRatio: LONG;
  ) =
  VAR  
    lambda, lambdaLo, lambdaHi: LONG;
    sum: LONG;
    nb: CARDINAL;
  BEGIN  
    WITH 
      fMax = LAST(pwr),
      fileName = name & "-bpw.plt",
      wr = FileWr.Open(fileName)
    DO 
      Wr.PutText(stderr, "writing file " & fileName & " ...\n");
      Wr.PutText(wr, "# lambdaLo lambdaHi    rms  \n");
      lambdaLo := 0.0d0;
      lambdaHi := lambdaMin;
      sum := 0.0d0;
      nb := 0;
      FOR f := fMax TO 0 BY -1 DO
        IF f = 0 THEN 
          lambda := LAST(LONG) 
        ELSE
          lambda := length/FLOAT(f, LONG)
        END;
        WHILE lambda > lambdaHi AND lambdaLo < length DO 
          FPut.LongReal(wr, lambdaLo, 4, Fmt.Style.Fix); 
          Wr.PutText(wr, "  ");
          FPut.LongReal(wr, lambdaHi, 4, Fmt.Style.Fix); 
          Wr.PutText(wr, "  ");
          FPut.LongReal(wr, Math.sqrt(sum), 4, Fmt.Style.Fix); 
          Wr.PutText(wr, "\n");
          INC(nb);
          lambdaLo := lambdaHi; 
          IF nb > nScales THEN 
            lambdaHi := length 
          ELSE 
            lambdaHi := scaleRatio*lambdaLo
          END;
          sum := 0.0d0;
        END;
        sum := sum + pwr[f]
      END;
      Wr.Flush(wr)
    END    
  END WriteBandPower;

BEGIN
  Main()
END PZSpectrum.