MODULE OptShape EXPORTS Main;

(*
  Adjusts the vertex coordinates of a triangulation
  so as to minimize some energy function.
  
  Created 1994 by Rober Marcone Rosi and J.Stolfi.
*)

IMPORT Wr, FileWr, Text, TextWr, OSError, Fmt, Process, Thread, ParseParams;
IMPORT LRN, Math, CPUTime;
IMPORT ParseEnergyParams, ParseMinimizerParams;
IMPORT Energy, MixedEnergy, Triang; 
IMPORT Minimizer;
FROM Triang IMPORT Coords, Topology;
FROM Energy IMPORT Gradient;
FROM Stdio IMPORT stderr;

IMPORT Debug;
 
TYPE
  Options = RECORD
      inFile: TEXT;             (* Initial guess file name (minus ".top") *)
      outFile: TEXT;            (* Output file name prefix *)
      eFunction: MixedEnergy.T; (* Energy function to minimize *)
      minimizer: Minimizer.T;   (* Minimization method to use *)
      vertexPasses: CARDINAL;   (* by-vertex min8n passes; 0 for simultaneous *)
      maxEvals: CARDINAL;       (* Max energy evaluations per min *)
      showEvery: CARDINAL;      (* When to write ".top" file; *)
    END;

TYPE
  LONGS = ARRAY OF LONGREAL;
  BOOLS = ARRAY OF BOOLEAN;
  CARDS = ARRAY OF CARDINAL;

TYPE
  EvalRec = RECORD           (* A configuration and its energy data: *)
      c: REF Coords;         (* Vertex coordinates *)
      e: LONGREAL;           (* Total energy *)
      termValue: REF LONGS;  (* Unweighted energy terms *)
      eDc: REF Gradient      (* Gradient of "e" relto "c" *)
    END;
    
PROCEDURE DoIt() =
  <* FATAL OSError.E, Thread.Alerted, Wr.Failure *>
  BEGIN
    WITH 
      o = GetOptions(), 
      tc = Triang.Read(o.inFile)
    DO
      WITH gnuWr = FileWr.Open(o.outFile & ".gnuplot") DO
        WriteGNUPlotCommands(gnuWr, o.outFile, o.eFunction.term^);
        Wr.Close(gnuWr)
      END;
      (* Write triangles for animation: *)
      Triang.WriteRLuisSuMa(o.outFile, tc.top, all := FALSE, comments := tc.comments);
      WITH 
        plotWr = FileWr.Open(o.outFile & ".plot"),
        stWr = FileWr.Open(o.outFile & ".st")
      DO
        WritePlotComments(plotWr, o.eFunction, o.minimizer);
        WritePlotHeader(plotWr, o.eFunction.term^);
        Minimize(
          tc, o.minimizer, o.eFunction,
          vertexPasses := o.vertexPasses,
          maxEvals := o.maxEvals, 
          showEvery := o.showEvery,
          outFile := o.outFile,
          plotWr := plotWr,
          stWr := stWr
        );
        Wr.Close(plotWr);
        Wr.Close(stWr)
      END
    END
  END DoIt;
  
PROCEDURE Minimize(
    READONLY tc: Triang.Configuration;  (* In: initial guess, Out: minimum *)
    m: Minimizer.T;         (* Minimization method *)
    e: MixedEnergy.T;       (* Energy function *)
    vertexPasses: CARDINAL; (* by-vertex min8n passes; 0 for simultaneous *)
    maxEvals: CARDINAL;     (* Evaluation budget per minimization *)
    showEvery: CARDINAL;    (* When to write ".top" file *)
    outFile: TEXT;          (* Output file name prefix *)
    plotWr: Wr.T;           (* Energy plots *)
    stWr: Wr.T;             (* Configurations for animation *)
  ) =
  
  VAR
    totEvals: CARDINAL := 0;            (* Count calls to "e.eval". *)
    cpuTime: LONGREAL := 0.0d0;         (* Accum minimization CPU time, seconds. *)
    cpuWhen: LONGREAL := CPUTime.Now(); (* When minimization was (re)started. *)
  BEGIN
    WITH
      top = tc.top,
      NV = top.NV,
      NT = NUMBER(e.term^),
      
      rMin = NewEvalRec(NV, NT)^,             (* Best configuration found so far *) 
      
      variable = VariableVertices(top)^,      (* Oficially variable vertices *)
      NVV = CountTrue(variable),              (* Number of variable vertices *)
      nSteps = MAX(1, vertexPasses * NV),     (* Number of minimizations *)
      selected = NEW(REF BOOLS, NV)^,         (* Vertices being minimized *)

      (* The following are work areas for "DoMinimizeStep": *) 
      
      rWork = NewEvalRec(NV, NT)^,            (* Probe configuration *)
      cIndex = NEW(REF CARDS, 3*NV)^,         (* Maps minimizer vars to coords *)
      xm = NEW(REF Minimizer.Point, 3*NV)^,   (* Argument vector for minimizer *)
      fm = NEW(REF Minimizer.Value)^,         (* Function value for minimizer *)
      gm = NEW(REF Minimizer.Gradient, 3*NV)^,(* Gradient for minimizer *) 

      (* Used by "ReportEval: *)
      minComments = 
        "energy function: " & e.name() & "\n" &
        "minimization algorithm: " & m.name() & "\n" &
        "vertex-by-vertex passes = " & Fmt.Int(vertexPasses) & "\n" &
        "max evals = " & Fmt.Int(totEvals) & "\n" &
        "--------------------------------------------------------\n" &
        tc.comments
    DO

      (* Tell the energy function what topolgy we are evaluationg: *)
      e.defTop(tc.top);
      
      (* Tell the minimizer how many variables are we going to minimize over; *)
      IF vertexPasses = 0 THEN 
        EVAL m.setSize(3*NVV)
      ELSE
        EVAL m.setSize(3)
      END;
      
      (* Initialize the minimum configuration: *)
      rMin.c^ := tc.c^;

      FOR step := 0 TO nSteps-1 DO

        (* Do a full evaluation to get the total energy right. *)
        (* This should be necessary only once at the beginning, *)
        (*   but let's do it once every pass for good measure. *)
        IF step MOD NV = 0 THEN
          VAR eOld: LONGREAL := rMin.e;
          BEGIN
            e.defVar(variable);
            INC(totEvals);
            e.eval(rMin.c^, rMin.e, FALSE, rMin.eDc^);
            rMin.termValue^ := e.termValue^;
            IF step # 0 AND rMin.e - eOld > 1.0d-6 * MAX(ABS(rMin.e), ABS(eOld)) THEN
              Debug.Line("pass = " & Fmt.Int(step DIV NV));
              Debug.LongReal("energy discrepancy = ", LRN.T{rMin.e - eOld});
              Debug.Line("")
            END;

            IF step = 0 THEN
              (* Write initial configuration as ".st": *)
              WriteCoords(
                stWr, e, rMin, 
                cpuTime := cpuTime, totEvals := totEvals, step := step, 
                comments := minComments
              );
              (* Write initial configuration as ".top": *)
              WITH
                name = outFile & "-initial"
              DO
                WriteConfiguration(
                  name, top, e, rMin, 
                  cpuTime := cpuTime, totEvals := totEvals, step := step, 
                  comments := minComments
                )
              END;
            END            
          END
        END;

        (* Decide which vertices to optimize: *)
        IF vertexPasses = 0 THEN
          (* Simultaneous minimimization: *)
          selected := variable;
        ELSE
          (* Single-vertex minimization: *)
          FOR v := 0 TO NV-1 DO selected[v] := FALSE END;
          WITH v = step MOD NV DO selected[v] := variable[v] END;
        END;

        (* Now optimize the selected vertices. *)
        VAR NC: CARDINAL := 0;

        PROCEDURE ReportEval(READONLY rWork, rMin: EvalRec) =
          BEGIN
            (* Stop the clock: *)
            cpuTime := cpuTime + (CPUTime.Now() - cpuWhen);
            INC (totEvals);
            PlotEnergy(plotWr, cpuTime, totEvals, step, rWork, e.weight^);
            IF NiceNumberCoords(totEvals, maxEvals, showEvery) THEN
              WriteCoords(
                stWr, e, rMin, 
                cpuTime := cpuTime, totEvals := totEvals, step := step, 
                comments := minComments
              )
            END;
            (* Restart the clock: *)
            cpuWhen := CPUTime.Now();
          END ReportEval;

        BEGIN
          (* Start the clock: *)
          cpuWhen := CPUTime.Now();
          FOR v := 0 TO NV-1 DO 
            IF selected[v] THEN
              cIndex[NC] := 3*v+0; INC(NC);
              cIndex[NC] := 3*v+1; INC(NC);
              cIndex[NC] := 3*v+2; INC(NC);
            END
          END;
          IF NC # 0 THEN
            e.defVar(selected);
            DoMinimizeStep(
              m, e, rMin,
              cIndex := SUBARRAY(cIndex, 0, NC),
              maxEvals := maxEvals,
              reportEval := ReportEval,
              rWork := rWork,
              xm := SUBARRAY(xm, 0, NC),
              fm := fm,
              gm := SUBARRAY(gm, 0, NC)
            )
          END;
        END
      END;
      
      (* Stop the clock: *)
      cpuTime := cpuTime + (CPUTime.Now() - cpuWhen);
      
      (* Write final configuration as ".st": *)
      WriteCoords(
        stWr, e, rMin, 
        cpuTime := cpuTime, totEvals := totEvals, step := nSteps-1, 
        comments := minComments
      );
      
      (* Write final configuration as ".top" too: *)
      WITH
        name = outFile & "-final"
      DO
        WriteConfiguration(
          name, top, e, rMin, 
          cpuTime := cpuTime, totEvals := totEvals, step := nSteps-1, 
          comments := minComments
        )
      END;

      (* Return minimum: *)
      tc.c^ := rMin.c^
    END
  END Minimize;
  
PROCEDURE VariableVertices(READONLY top: Topology): REF BOOLS =
  BEGIN
    WITH 
      r = NEW(REF BOOLS, top.NV)
    DO
      Triang.GetVariableVertices(top, r^);
      RETURN r
    END
  END VariableVertices;
  
PROCEDURE NewEvalRec(NV, NT: CARDINAL): REF EvalRec =
  BEGIN
    WITH 
      r = NEW(REF EvalRec)
    DO
      r^ := EvalRec{
        c := NEW(REF Coords, NV),
        e := 0.0d0,
        termValue := NEW(REF LONGS, NT),
        eDc := NEW(REF Gradient, NV)
      };
      RETURN r
    END
  END NewEvalRec;
  
TYPE
  ReportEvalProc = PROCEDURE (READONLY rWork, rMin: EvalRec);

PROCEDURE DoMinimizeStep(
    m: Minimizer.T;             (* Minimization method *)
    e: MixedEnergy.T;           (* Energy function *)
    VAR rMin: EvalRec;          (* In: initial guess, Out: minimum *)
    READONLY cIndex: CARDS;     (* Maps minimizer args to vertex coords *)
    maxEvals: CARDINAL;         (* Evaluation budget *)
    reportEval: ReportEvalProc; (* Called after every energy evaluation *)
    (* Work areas: *)
    VAR rWork: EvalRec;         (* Probe configuration *)
    VAR xm: Minimizer.Point;    (* Argument vector for minimizer *)
    VAR fm: Minimizer.Value;    (* Function value for minimizer *)
    VAR gm: Minimizer.Gradient; (* Gradient vector for minimizer *)
  ) =
  (*
    Performs an energy minimization for the vertices defined by the
    "cIndex" vector.
    Assumes that "m.setSize", "e.setTop", "e.setVar"
    have already been called.
  *)
  BEGIN
    WITH
      NT = NUMBER(e.term^),
      NC = NUMBER(cIndex),
      grad = m.needsGradient(),         (* TRUE to compute gradients *)
      eOffset = NEW(REF LONGREAL)^,     (* "eval(variable)-eval(selected)" *)
      termOffset = NEW(REF LONGS, NT)^, (* Unweigthed terms of "eOffset" *)
      
      cMin = rMin.c^, 
      eMin = rMin.e,
      eDcMin = rMin.eDc^,
      termMin = rMin.termValue^,
      
      cWork = rWork.c^, 
      eWork = rWork.e,
      eDcWork = rWork.eDc^,
      termWork = rWork.termValue^
      
    DO

      PROCEDURE Initialize(
          VAR x: Minimizer.Point; 
          VAR f: Minimizer.Value; 
          VAR g: Minimizer.Gradient
        ) =
        BEGIN
          cWork := cMin;
          e.eval(cWork, eWork, grad, eDcWork);
          termWork := e.termValue^;

          (* Save offsets between full and partial evaluations: *)
          eOffset := eMin - eWork;
          eWork := eMin;
          FOR t := 0 TO NT-1 DO
            termOffset[t] := termMin[t] - termWork[t];
          END;
          termWork := termMin;
          
          reportEval(rWork, rMin);

          (* Set function and gradient: *)
          f := eWork + eOffset;
          FOR i := 0 TO LAST(x) DO 
            WITH vk = cIndex[i], v = vk DIV 3, k = vk MOD 3 DO
              x[i] := cWork[v][k];
              g[i] := eDcWork[v][k]
            END
          END;
        END Initialize;

      PROCEDURE Eval(
          READONLY x: Minimizer.Point; 
          VAR f: Minimizer.Value; 
          VAR g: Minimizer.Gradient
        ) =
        BEGIN
          (* Stuff the argument into the "cWork" vector: *)
          FOR i := 0 TO LAST(x) DO 
            WITH vk = cIndex[i], v = vk DIV 3, k = vk MOD 3 DO
              cWork[v][k] := x[i];
            END
          END;

          (* Evaluate the partial energy for "cWork": *)
          e.eval(cWork, eWork, grad, eDcWork);
          termWork := e.termValue^;

          (* Correct to full energy: *)
          eWork := eWork + eOffset;
          FOR t := 0 TO NT-1 DO
            termWork[t] := termWork[t] + termOffset[t]
          END;

          (* See if we got something good: *)
          IF eWork < eMin THEN
            eMin := eWork;
            cMin := cWork;
            eDcMin := eDcWork;
            termMin := termWork;
          END;

          reportEval(rWork, rMin);

          (* Return energy and gradient to optimizer: *)
          f := eWork;
          IF grad THEN
            FOR i := 0 TO LAST(g) DO 
              WITH vk = cIndex[i], v = vk DIV 3, k = vk MOD 3 DO
                g[i] := eDcWork[v][k]
              END
            END
          END;
        END Eval;

      (* Now do a partial one to get us started: *)
      BEGIN
        Initialize(xm, fm, gm);
        m.minimize(
          eval := Eval,
          x := xm,
          f := fm,
          g := gm,
          dist := Math.sqrt(FLOAT(NC, LONGREAL)),
          tol := 0.01d0,
          flat := 0.0d0,
          maxEvals := maxEvals
        );
        (* Sanity check: *)
        IF fm # eMin THEN
          Debug.LongReal("** bug: fm, eMin = ", LRN.T{fm, eMin})
        END;
      END
    END
  END DoMinimizeStep;
          
PROCEDURE CountTrue(READONLY b: BOOLS): CARDINAL = 
  VAR n: CARDINAL := 0;
  BEGIN
    FOR i := 0 TO LAST(b) DO IF b[i] THEN INC(n) END END;
    RETURN n
  END CountTrue;

PROCEDURE NiceNumberCoords(
    totEvals: CARDINAL;
    maxEvals: CARDINAL;
    showEvery: CARDINAL;
  ): BOOLEAN =
  BEGIN
    RETURN (totEvals MOD showEvery = 0)
      OR totEvals <= 10
      OR totEvals = 10 OR totEvals = 20 OR totEvals = 50
      OR (totEvals MOD 100 = 0 AND totEvals <= 5000) 
      OR (totEvals MOD 500 = 0)
      OR (totEvals = maxEvals)
  END NiceNumberCoords;
  
PROCEDURE WriteCoords(
    wr: Wr.T;
    e: MixedEnergy.T;
    READONLY r: EvalRec;
    cpuTime: CPUTime.T;
    totEvals: CARDINAL;
    step: CARDINAL;
    comments: TEXT;
  ) =
  (* 
    Writes only coords values to file "wr" in ".st" format
  *)    
  BEGIN
    Triang.WriteRLuisSt(
      wr, r.c^, r.eDc^,
      comments := 
        IterationComments(cpuTime, totEvals, step, r, e) & "\n" & 
        comments
    )
  END WriteCoords;

PROCEDURE WriteConfiguration(
    name: TEXT;
    READONLY top: Topology;
    e: MixedEnergy.T;
    READONLY r: EvalRec;
    cpuTime: CPUTime.T;
    totEvals: CARDINAL;
    step: CARDINAL;
    comments: TEXT;
  ) =
  BEGIN
    Triang.Write(
      name, top, r.c^,
      comments := 
        IterationComments(cpuTime, totEvals, step, r, e) & "\n" &
        comments
    )
  END WriteConfiguration;
  
PROCEDURE IterationComments(
    cpuTime: LONGREAL;
    READONLY totEvals: CARDINAL;
    READONLY step: CARDINAL;
    READONLY r: EvalRec;
    READONLY e: MixedEnergy.T;
  ): TEXT =
  BEGIN
    WITH
      wr = NEW(TextWr.T).init()
    DO
      WritePlotHeader(wr, e.term^);
      PlotEnergy(wr, cpuTime, totEvals, step, r, e.weight^);
      RETURN TextWr.ToText(wr)
    END
  END IterationComments;
  
PROCEDURE PlotEnergy(
    wr: Wr.T;
    READONLY cpuTime: LONGREAL;
    READONLY totEvals: CARDINAL;
    READONLY step: CARDINAL;
    READONLY r: EvalRec;
    READONLY weight: ARRAY OF REAL;
  ) =
  <* FATAL Thread.Alerted, Wr.Failure *>
  BEGIN
    Wr.PutText(wr, " ");
    Wr.PutText(wr, Fmt.Pad(Fmt.LongReal(cpuTime, Fmt.Style.Fix, prec := 2), 8));
    Wr.PutText(wr, " ");
    Wr.PutText(wr, Fmt.Pad(Fmt.Int(totEvals), 8));
    Wr.PutText(wr, " ");
    Wr.PutText(wr, Fmt.Pad(Fmt.LongReal(r.e, Fmt.Style.Fix, 5), 10));
    Wr.PutText(wr, " ");
    Wr.PutText(wr, Fmt.Pad(Fmt.Int(step), 8));
    FOR i := 0 TO LAST(r.termValue^) DO 
      Wr.PutText(wr, " ");
      WITH
        t = r.termValue[i] * FLOAT(weight[i], LONGREAL)
      DO
        Wr.PutText(wr, Fmt.Pad(Fmt.LongReal(t, Fmt.Style.Fix, 5), 10))
      END;
    END;
    Wr.PutText(wr, "\n");
    Wr.Flush(wr);
  END PlotEnergy;

PROCEDURE WritePlotComments(wr: Wr.T; e: MixedEnergy.T; m: Minimizer.T) =
  <* FATAL Thread.Alerted, Wr.Failure *>
  BEGIN
    Wr.PutText(wr, "#");
    Wr.PutText(wr, "energy function: " & e.name());
    Wr.PutText(wr, "\n");
              
    Wr.PutText(wr, "#");
    Wr.PutText(wr, "minimizer: " & m.name());
    Wr.PutText(wr, "\n");

    Wr.Flush(wr);
  END WritePlotComments;

PROCEDURE WritePlotHeader(wr: Wr.T; READONLY term: ARRAY OF Energy.T) =
  <* FATAL Thread.Alerted, Wr.Failure *>
  BEGIN
    Wr.PutText(wr, "#");
    Wr.PutText(wr, Fmt.Pad("cpuTime", 8));
    Wr.PutText(wr, " ");
    Wr.PutText(wr, Fmt.Pad("evals", 8));
    Wr.PutText(wr, " ");
    Wr.PutText(wr, Fmt.Pad("energy", 10));
    Wr.PutText(wr, " ");
    Wr.PutText(wr, Fmt.Pad("step", 8));
    FOR i := 0 TO LAST(term) DO 
      Wr.PutText(wr, " ");
      Wr.PutText(wr, Fmt.Pad(EnergyTag(term[i].name()), 10));
    END;
    Wr.PutText(wr, "\n");
    Wr.Flush(wr);
  END WritePlotHeader;
  
PROCEDURE WriteGNUPlotCommands(
    wr: Wr.T; 
    outFile: TEXT;
    READONLY term: ARRAY OF Energy.T;
  ) =
  <* FATAL Thread.Alerted, Wr.Failure *>
  BEGIN
    Wr.PutText(wr, "set terminal X11\n");
    Wr.PutText(wr, "set title \"" & outFile & "\"\n");
    Wr.PutText(wr, 
      "plot \"" & outFile & ".plot\" using 1:3 title \"etot\" with lines, \\\n"
    );
    FOR i := 0 TO LAST(term) DO 
      WITH col = i + 5 DO
        Wr.PutText(wr, 
          "     \"" & outFile & ".plot\" using 1:" & 
          Fmt.Int(col) & " title \"" &
          EnergyTag(term[i].name()) & 
          "\" with linespoints"
        );
        IF i < LAST(term) THEN Wr.PutText(wr, ", \\") END;
        Wr.PutText(wr, "\n");
      END
    END;
    Wr.PutText(wr, "pause 300\n");
    Wr.PutText(wr, "quit\n");
    Wr.Flush(wr);
  END WriteGNUPlotCommands;
  
PROCEDURE EnergyTag(name: TEXT): TEXT =
  BEGIN
    WITH
      n = Text.FindChar(name, '(')
    DO
      IF n = -1 THEN 
        RETURN Text.Sub(name, 0, 8)
      ELSE
        RETURN Text.Sub(name, 0, MIN(n, 8))
      END
    END
  END EnergyTag;

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

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

        IF pp.keywordPresent("-vertexPasses") THEN
          o.vertexPasses := pp.getNextInt(0, 9999);
        ELSE
          o.vertexPasses := 0
        END;
        
        IF pp.keywordPresent("-maxEvals") THEN
          o.maxEvals := pp.getNextInt(0, 99999);
        ELSE
          IF o.vertexPasses = 0 THEN
            o.maxEvals := 1000
          ELSE
            pp.error("\"-maxEvals\" per vertex must be given")
          END
        END;

        IF pp.keywordPresent("-showEvery") THEN
          o.showEvery := pp.getNextInt(0, 1000000)
        ELSIF pp.keywordPresent("-showAll") THEN
          o.showEvery := 1
        ELSE
          o.showEvery := LAST(CARDINAL)
        END;

        o.eFunction := ParseEnergyParams.Parse(pp);
        IF o.eFunction = NIL THEN 
          pp.error("no energy specified")
        END;
        
        o.minimizer := ParseMinimizerParams.Parse(pp);
        IF o.minimizer = NIL THEN
          pp.error("no minimizer specified")
        END;

        pp.finish();                                       
      EXCEPT                                                            
      | ParseParams.Error =>                                              
          Wr.PutText(stderr, "Usage: OptShape \\\n");
          Wr.PutText(stderr, "  -inFile <name>   -outFile <name> \\\n");
          Wr.PutText(stderr, "  [ -vertexPasses <num> ] \\\n");
          Wr.PutText(stderr, "  [ -maxEvals <num> ] \\\n");
          Wr.PutText(stderr, "  [ -showEvery <num> | -showAll ] \\\n");
          Wr.PutText(stderr, ParseMinimizerParams.Help);
          Wr.PutText(stderr, " \\\n");
          Wr.PutText(stderr, ParseEnergyParams.Help);
          Wr.PutText(stderr, "\n");
          Process.Exit (1);                                              
      END;
    END;
    RETURN o
  END GetOptions;

BEGIN
  DoIt()
END OptShape.