MODULE TestEnergy EXPORTS Main;

(*
  Tests an energy function.  Takes as input two ".top" files with same
  topology but different coordinate vectors, and evaluates a given
  energy function for all configurations that interpolate between the
  two, at "nSteps" equal intervals.  
  
  The output is a ".plot" file that can be examined with "gnuplot".
  Optionally, the program writes the ".top" file with minimum energy.
*)

IMPORT Triang, MixedEnergy, ParseEnergyParams;
IMPORT LR3, Fmt, FileWr, OSError, ParseParams, Process, Wr, Thread;
FROM Triang IMPORT Coords;
FROM Stdio IMPORT stderr;

TYPE
  BOOLS = ARRAY OF BOOLEAN;

TYPE
  Options = RECORD
      aFile, bFile: TEXT;        (* Configuration files (minus ".top") *)
      outFile: TEXT;             (* Plot file name (minus ".plot") *)
      eFunction: MixedEnergy.T;  (* Energy function *)
      nSteps: CARDINAL;          (* Number of steps to take *)
      showMin: BOOLEAN;          (* TRUE to write interpolated ".x3d" files *)
      all: BOOLEAN;              (* TRUE to write even non-existing triangles *)
    END;
    
PROCEDURE DoIt() =
  <* FATAL Wr.Failure, Thread.Alerted, OSError.E *>
  VAR eMin: LONGREAL := LAST(LONGREAL);
      sMin: LONGREAL;
  BEGIN
    WITH 
      o = GetOptions(),
      atc = Triang.Read(o.aFile), 
      top = atc.top,
      ac = atc.c^,
      bc = Triang.Read(o.bFile).c^,
      NV = top.NV,
      c = NEW(REF ARRAY OF LR3.T, top.NV)^,
      cDs = NEW(REF ARRAY OF LR3.T, top.NV)^,
      eDc = NEW(REF ARRAY OF LR3.T, top.NV)^,
      vAll = NEW(REF BOOLS, top.NV)^,
      vSome = NEW(REF BOOLS, top.NV)^,
      pwr = FileWr.Open(o.outFile & ".plot")
    DO
      <* ASSERT NUMBER(ac) = NV *>
      <* ASSERT NUMBER(bc) = NV *>
      WriteHeader(pwr);
      o.eFunction.defTop(top);
      Triang.GetVariableVertices(top, vAll);
      FOR i := 0 TO NV-1 DO 
        cDs[i] := LR3.Sub(bc[i], ac[i]);
        vSome[i] := LR3.Norm(cDs[i]) # 0.0d0
      END;

      FOR i := 0 TO o.nSteps DO
        WITH 
          s = FLOAT(i, LONGREAL)/FLOAT(o.nSteps, LONGREAL)
        DO
          Wr.PutText(pwr, " "); 
          Wr.PutText(pwr, Fmt.Pad(Fmt.LongReal(s, Fmt.Style.Fix, 4), 6)); 
          (* Define coordinates *)
          FOR i := 0 TO top.NV-1 DO 
            c[i] := LR3.Mix(1.0d0 - s, ac[i], s, bc[i])
          END;
          
          (* Evaluate with all or only some vertices variable: *)
          WITH
            eAll  = EvalTest(pwr, o.eFunction, c, cDs, vAll,  eDc),
            eSome = EvalTest(pwr, o.eFunction, c, cDs, vSome, eDc)
          DO
            (* Compare: *)
            Wr.PutText(pwr, Fmt.Pad(Fmt.LongReal(eAll-eSome, Fmt.Style.Fix, 4), 12));
            IF eAll < eMin  THEN 
              eMin := eAll;
              sMin := s
            END;
          END;
          Wr.PutText(pwr, "\n");
          Wr.Flush(pwr);
        END
      END;
      Wr.Close(pwr);
      IF o.showMin THEN 
        WITH name = o.outFile & "-min" DO
          FOR i := 0 TO top.NV-1 DO 
            c[i] := LR3.Mix(1.0d0 - sMin, ac[i], sMin, bc[i])
          END;
          Triang.Write(name, top, c)          
        END
      END;
    END
  END DoIt;
  
PROCEDURE EvalTest(
    pwr: Wr.T;
    eFunction: MixedEnergy.T;      (* Energy function *)
    VAR c: Coords;                 (* Coordinates *)
    READONLY cDs: Coords;          (* Coordinate changes *)
    READONLY v: ARRAY OF BOOLEAN;  (* Which vertices are variable *)
    VAR eDc: Coords;               (* Work areas *)
  ): LONGREAL =
  <* FATAL Wr.Failure, Thread.Alerted *>
  VAR e, ep: LONGREAL;
      eDs: LONGREAL;
  BEGIN
    WITH
      NV = NUMBER(c)
    DO
      <* ASSERT NUMBER(v) = NV *>
      eFunction.defVar(v);
      eFunction.eval(c, e, TRUE, eDc);
      Wr.PutText(pwr, Fmt.Pad(Fmt.LongReal(e, Fmt.Style.Fix, 4), 12)); 
      
      (* Compute directional derivative from gradient: *)
      eDs := 0.0d0;
      FOR i := 0 TO NV-1 DO eDs := eDs + LR3.Dot(eDc[i], cDs[i]) END;
      Wr.PutText(pwr, Fmt.Pad(Fmt.LongReal(eDs, Fmt.Style.Fix, 4), 12));
      
      (* Check if setting "grad=FALSE" has any effect on "e": *)
      eFunction.eval(c, ep, FALSE, eDc);
      <* ASSERT ep = e *>
      
      (* Compute directional derivative numerically: *)
      WITH ds = 0.0001d0 DO
        FOR i := 0 TO NV-1 DO c[i] := LR3.Mix(1.0d0, c[i], ds, cDs[i]) END;
        eFunction.eval(c, ep, FALSE, eDc);
        eDs := (ep - e)/ds
      END;
      Wr.PutText(pwr, Fmt.Pad(Fmt.LongReal(eDs, Fmt.Style.Fix, 4), 12));
    END;
    RETURN e
  END EvalTest;

PROCEDURE WriteHeader(pwr: Wr.T) =
  <* FATAL Wr.Failure, Thread.Alerted *>
  BEGIN
    Wr.PutText(pwr, "#");
    Wr.PutText(pwr, Fmt.Pad("s", 6));
    Wr.PutText(pwr, Fmt.Pad("etot", 12));
    Wr.PutText(pwr, Fmt.Pad("etotDs", 12));
    Wr.PutText(pwr, Fmt.Pad("etotDc*cDs", 12));
    Wr.PutText(pwr, Fmt.Pad("evar", 12));
    Wr.PutText(pwr, Fmt.Pad("evarDs", 12));
    Wr.PutText(pwr, Fmt.Pad("evarDc*cDs", 12));
    Wr.PutText(pwr, Fmt.Pad("etot - evar", 12));
    Wr.PutText(pwr, "\n");

    Wr.PutText(pwr, "#");
    Wr.PutText(pwr, Fmt.Pad(" ",  6, '-', Fmt.Align.Right));
    Wr.PutText(pwr, Fmt.Pad(" ", 12, '-', Fmt.Align.Right));
    Wr.PutText(pwr, Fmt.Pad(" ", 12, '-', Fmt.Align.Right));
    Wr.PutText(pwr, Fmt.Pad(" ", 12, '-', Fmt.Align.Right));
    Wr.PutText(pwr, Fmt.Pad(" ", 12, '-', Fmt.Align.Right));
    Wr.PutText(pwr, Fmt.Pad(" ", 12, '-', Fmt.Align.Right));
    Wr.PutText(pwr, Fmt.Pad(" ", 12, '-', Fmt.Align.Right));
    Wr.PutText(pwr, Fmt.Pad(" ", 12, '-', Fmt.Align.Right));
    Wr.PutText(pwr, "\n");
  END WriteHeader;

PROCEDURE GetOptions (): Options =
  <* FATAL Thread.Alerted, Wr.Failure *>
  VAR o: Options;
  BEGIN
    WITH 
      pp = NEW(ParseParams.T).init(stderr)
    DO
      TRY
        pp.getKeyword("-aFile");                               
        o.aFile := pp.getNext();                         

        pp.getKeyword("-bFile");                               
        o.bFile := pp.getNext();                         

        pp.getKeyword("-outFile");                               
        o.outFile := pp.getNext();                         

        IF pp.keywordPresent("-nSteps") THEN
          o.nSteps := pp.getNextInt(1, 1000);
        ELSE
          o.nSteps := 100;
        END;

        o.eFunction := ParseEnergyParams.Parse(pp);
        IF o.eFunction = NIL THEN 
          pp.error("\"-energy\" not specified")
        END;

        o.showMin := pp.keywordPresent("-showMin");

        IF o.showMin THEN
          o.all := pp.keywordPresent("-all");
        ELSE
          o.all := FALSE
        END;

        pp.finish();                                       
      EXCEPT                                                            
      | ParseParams.Error =>                                              
          Wr.PutText(stderr, "Usage: TestElectEnergy\\\n");
          Wr.PutText(stderr, "  -aFile <name> -bFile <name> -outFile <name>\\\n");
          Wr.PutText(stderr, "  [ -nSteps <num> ]\\\n");
          Wr.PutText(stderr, "  [ -showMin [ -all ]]\\\n");
          Wr.PutText(stderr, ParseEnergyParams.Help);
          Wr.PutText(stderr, "\n");
          Process.Exit (1);                                              
      END
    END;
    RETURN o
  END GetOptions;

BEGIN 
  DoIt();
END TestEnergy.