MODULE BITAverage EXPORTS Main;
(* Last edited on 1999-12-31 10:33:32 by hcgl *)

(* Computes the average and standard deviation of a bunch of real chains. *)

IMPORT  ParseParams, Process, Text, Rd, Wr, FileWr, FileRd;
IMPORT  OSError, Thread, TextWr;
FROM Math IMPORT sqrt, exp, log;
IMPORT PZLRChain, PZLR3Chain;
FROM Stdio IMPORT stderr;
FROM PZTypes IMPORT LONG, LONGS, NAT, BOOL;

<* FATAL Rd.Failure, Wr.Failure, Thread.Alerted, OSError.E *>

TYPE
  Options = RECORD
      outFile: TEXT;         (* Output file name (WITH extension).  *)
      nSamples: NAT;         (* Num samples in each file.  *)   
      inFile: REF FileNames; (* Input file names (WITH extension). *)
      log: BOOL;             (* TRUE computes the log-mean-exp of the chains. *)
    END;
  
  FileNames = ARRAY OF TEXT;
  
PROCEDURE Main() =
  VAR v: LONG;
  BEGIN
    WITH
      o = GetOptions(),
      nFiles = NUMBER(o.inFile^),
      nf = FLOAT(nFiles, LONG),
      sum1 = NEW(REF LONGS, o.nSamples)^,
      sum2 = NEW(REF LONGS, o.nSamples)^,
      res = NEW(REF PZLR3Chain.T, o.nSamples)^
    DO 
      FOR k := 0 TO o.nSamples-1 DO sum1[k] := 0.0d0; sum2[k] := 0.0d0 END;
      FOR i := 0 TO nFiles-1 DO
        Wr.PutText(stderr, o.inFile[i] & " ...");
        WITH
          rd = FileRd.Open(o.inFile[i]),
          rp = PZLRChain.Read(rd, headerOnly := FALSE), p = rp.c^
        DO
          Rd.Close(rd);
          <* ASSERT NUMBER(p) = o.nSamples *>
          FOR k := 0 TO o.nSamples-1 DO 
            IF o.log THEN 
              v := 0.5d0*log(rp.unit*rp.unit/4.0d0 + p[k]*p[k])
            ELSE
              v := p[k]
            END;
            sum1[k] := sum1[k] + v;
            sum2[k] := sum2[k] + v*v
          END;
        END;
        Wr.PutText(stderr, "\n");
      END;
      FOR k := 0 TO o.nSamples-1 DO
        WITH
          avg = sum1[k]/nf,
          dev = sqrt(MAX(sum2[k] - avg*avg*nf, 0.0d0)/(nf - 1.0d0)),
          lo = avg - dev,
          hi = avg + dev
        DO
          IF o.log THEN 
            res[k][0] := exp(avg); 
            res[k][1] := exp(lo); 
            res[k][2] := exp(hi)
          ELSE
            res[k][0] := avg;
            res[k][1] := lo;
            res[k][2] := hi
          END
        END
      END;
      PZLR3Chain.Write(
        FileWr.Open(o.outFile), 
        cmt := "average AND confidence interval of the following files:\n" &
               CatFiles(o.inFile^), 
        c := res, 
        unit := SelectUnit(res)
      )
    END 
  END Main;
  
PROCEDURE SelectUnit(READONLY s: PZLR3Chain.T): LONG =
  CONST MaxScaled = 256.0d0*256.0d0*256.0d0;  (* 2^24 *)
  VAR unit: LONG := 1.0d0/(MaxScaled*MaxScaled*MaxScaled);
  BEGIN
    FOR i := 0 TO LAST(s) DO
      FOR k := 0 TO 2 DO
        WITH eps = ABS(s[i][k])/MaxScaled DO
          WHILE unit < eps DO unit := 2.0d0*unit END
        END
      END
    END;
    RETURN unit
  END SelectUnit;


PROCEDURE CatFiles(READONLY file: ARRAY OF TEXT): TEXT =
  BEGIN
    WITH
      wr = NEW(TextWr.T).init()
    DO
      FOR i := 0 TO LAST(file) DO
        Wr.PutText(wr, "  "); Wr.PutText(wr, file[i]); Wr.PutText(wr, "\n"); 
      END;
      RETURN TextWr.ToText(wr)
    END
  END CatFiles;

 
PROCEDURE GetOptions(): Options =
  VAR o: Options;
  BEGIN
    WITH
      pp = NEW(ParseParams.T).init(stderr)
    DO
      TRY
        pp.getKeyword("-outFile");
        o.outFile := pp.getNext();
        
        pp.getKeyword("-nSamples");
        o.nSamples := pp.getNextInt(1);
        
        o.log := pp.keywordPresent("-log");
        
        pp.skipParsed();
        
        WITH nFiles = NUMBER(pp.arg^) - pp.next DO
          IF nFiles < 2 THEN pp.error("needs at least two files") END;
          o.inFile := NEW(REF ARRAY OF TEXT, nFiles);
          FOR i := 0 TO nFiles-1 DO
            WITH f = pp.getNext() DO
              o.inFile[i] := f;
              IF Text.Length(f) > 1 AND Text.GetChar(f,0) = '-' THEN 
                pp.error("bad file name \"" & f & "\"")
              END
            END
          END
        END;
        
        pp.finish();      
        
      EXCEPT                                                            
      | ParseParams.Error =>                                              
          Wr.PutText(stderr, "Usage: BITAverage \\\n");
          Wr.PutText(stderr, "  -outFile FILE \\\n");
          Wr.PutText(stderr, "  -nSamples NAT  \\\n");
          Wr.PutText(stderr, "  [ -log ] \\\n");
          Wr.PutText(stderr, "  FILE1 FILE2 ... FILEn\n");
          Process.Exit(1);
      END;
    END;
    RETURN o
  END GetOptions;


BEGIN
  Main()
END BITAverage.