MODULE EnergySlice EXPORTS Main;

(*
  Generates a two-dimensional plot of the energy function,
  sliced along the plane determined by three given configurations.
  
  Created Jan 1996 by Rober Marcone Rosi and J.Stolfi.
*)

IMPORT Rd, Wr, FileWr, Text, TextRd, OSError, Fmt, Process, Thread, ParseParams;
IMPORT ParseEnergyParams;
IMPORT Energy, MixedEnergy, Triang; 
FROM Triang IMPORT Coords, Topology;
FROM Energy IMPORT Gradient;
FROM Stdio IMPORT stderr;

(* MPORT Debug; *)
 
TYPE
  Options = RECORD
      oFile: TEXT;          (* Configuration at (0,0) (minus ".top") *)
      xFile: TEXT;          (* Configuration at (1,0) (minus ".top") *)
      xMin, xMax: LONGREAL; (* X range to plot *)
      yFile: TEXT;          (* Configuration at (0,1) (minus ".top") *)
      yMin, yMax: LONGREAL; (* Y range to plot *)
      outFile: TEXT;        (* Output file name prefix *)
      erg: MixedEnergy.T;   (* Energy function to plot *)
      nSteps: CARDINAL;     (* Number of steps per axis *)
    END;

TYPE
  BOOLS = ARRAY OF BOOLEAN;

PROCEDURE DoIt() =
  BEGIN
    WITH 
      o = GetOptions(), 
      otc = Triang.Read(o.oFile),
      xtc = Triang.Read(o.xFile),
      ytc = Triang.Read(o.yFile)
    DO
      <* ASSERT Triang.TriviallyIsomorphic(otc.top, xtc.top) *>
      <* ASSERT Triang.TriviallyIsomorphic(otc.top, ytc.top) *>
      WITH 
        top = otc.top,
        comments = CombineComments(otc.comments, xtc.comments, ytc.comments)
      DO
        GeneratePlot(
          top, 
          oc := otc.c^, 
          xc := xtc.c^, xMin := o.xMin, xMax := o.xMax,
          yc := ytc.c^, yMin := o.yMin, yMax := o.yMax,
          nSteps := o.nSteps,
          erg := o.erg,
          outFile := o.outFile,
          comments := comments
        );
      END
    END
  END DoIt;
  
PROCEDURE GeneratePlot(
    READONLY top: Topology;   (* Common topology *)
    READONLY oc: Coords;      (* The "(0,0)" configuration *)
    READONLY xc: Coords;      (* The "(1,0)" configuration *)
    xMin, xMax: LONGREAL;     (* X range to plot *)
    READONLY yc: Coords;      (* The "(0,1)" configuration *)
    yMin, yMax: LONGREAL;     (* Y range to plot *)
    nSteps: CARDINAL;         (* Steps per axis *)
    erg: MixedEnergy.T;       (* Energy function *)
    outFile: TEXT;            (* Output filename prefix *)
    comments: TEXT;           (* Comments for ".top" files *)
  ) =
  <* FATAL OSError.E, Thread.Alerted, Wr.Failure *>
  BEGIN
    WITH gnuWr = FileWr.Open(outFile & ".gnuplot") DO
      WriteGNUPlotCommands(gnuWr, outFile, erg, xMin, xMax, yMin, yMax);
      Wr.Close(gnuWr)
    END;
    WITH
      NV = top.NV,
      c = NEW(REF Coords, NV)^,
      variable = NEW(REF BOOLS, NV)^,
      e = NEW(REF LONGREAL)^,
      eDc = NEW(REF Gradient, NV)^,
      fSteps = FLOAT(nSteps, LONGREAL),
      plotWr = FileWr.Open(outFile & ".plot")
    DO
      (* Write header: *)
      WritePlotComments(
        plotWr, 
        erg, xMin, xMax, yMin, yMax,
        comments
      );
      WritePlotHeader(plotWr, erg);

      (* Determine the vertices that change: *)
      FOR i := 0 TO LAST(c) DO
        variable[i] := 
          (oc[i][0] # xc[i][0]) OR (oc[i][0] # yc[i][0]) OR
          (oc[i][1] # xc[i][1]) OR (oc[i][1] # yc[i][1]) OR
          (oc[i][2] # xc[i][2]) OR (oc[i][2] # yc[i][2])
      END;
      (* Tell the energy function what topolgy we are evaluating: *)
      erg.defTop(top);
      erg.defVar(variable);
      FOR xi := 0 TO nSteps DO
        BeginIsoline(plotWr);
        FOR yi := 0 TO nSteps DO
          WITH
            sx = FLOAT(xi, LONGREAL)/fSteps,
            x = sx * xMax + (1.0d0 - sx)*xMin,
            sy = FLOAT(yi, LONGREAL)/fSteps,
            y = sy * yMax + (1.0d0 - sy)*yMin,
            h = 1.0d0 - x - y
          DO
            FOR i := 0 TO LAST(c) DO
              c[i][0] := h * oc[i][0] + x * xc[i][0] + y * yc[i][0];
              c[i][1] := h * oc[i][1] + x * xc[i][1] + y * yc[i][1];
              c[i][2] := h * oc[i][2] + x * xc[i][2] + y * yc[i][2]
            END;
            IF xi MOD nSteps = 0 AND yi MOD nSteps = 0 THEN
              WriteConfiguration(
                outFile, 
                top, c,
                x := x, y := y,
                comments := comments
              )
            END;
            erg.eval(c, e, FALSE, eDc);
            PlotEnergy(plotWr, x, y, e, erg);
          END
        END;
        EndIsoline(plotWr)
      END;
      Wr.Close(plotWr)
    END
  END GeneratePlot;
  
PROCEDURE PlotEnergy(
    wr: Wr.T;
    x, y: LONGREAL;
    READONLY e: LONGREAL;
    READONLY erg: MixedEnergy.T
  ) =
  <* FATAL Thread.Alerted, Wr.Failure *>
  BEGIN
    Wr.PutText(wr, " ");
    Wr.PutText(wr, Fmt.Pad(Fmt.LongReal(x, Fmt.Style.Fix, prec := 4), 8));
    Wr.PutText(wr, " ");
    Wr.PutText(wr, Fmt.Pad(Fmt.LongReal(y, Fmt.Style.Fix, prec := 4), 8));
    Wr.PutText(wr, " ");
    Wr.PutText(wr, Fmt.Pad(Fmt.LongReal(e, Fmt.Style.Fix, 5), 10));
    Wr.PutText(wr, " ");
    FOR i := 0 TO LAST(erg.termValue^) DO 
      Wr.PutText(wr, " ");
      WITH
        t = erg.termValue[i] * FLOAT(erg.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 BeginIsoline(<*UNUSED*> wr: Wr.T) =
  BEGIN
    (* Ok *)  
  END BeginIsoline;

PROCEDURE EndIsoline(wr: Wr.T) =
  <* FATAL Thread.Alerted, Wr.Failure *>
  BEGIN
    Wr.PutText(wr, "\n");    
  END EndIsoline;

PROCEDURE WritePlotComments(
    wr: Wr.T; 
    erg: MixedEnergy.T;
    xMin, xMax: LONGREAL;     (* X range to plot *)
    yMin, yMax: LONGREAL;     (* Y range to plot *)
    comments: TEXT;           (* Other comments *)
  ) =
  <* FATAL Thread.Alerted, Wr.Failure *>
  BEGIN
    
    Wr.PutText(wr, "# energy function: " & erg.name() & "\n");

    Wr.PutText(wr, "# x range = [ ");
    Wr.PutText(wr, Fmt.Pad(Fmt.LongReal(xMin, Fmt.Style.Fix, prec := 4), 8));
    Wr.PutText(wr, " __ ");
    Wr.PutText(wr, Fmt.Pad(Fmt.LongReal(xMax, Fmt.Style.Fix, prec := 4), 8));
    Wr.PutText(wr, " ]\n");

    Wr.PutText(wr, "# y range = [ ");
    Wr.PutText(wr, Fmt.Pad(Fmt.LongReal(yMin, Fmt.Style.Fix, prec := 4), 8));
    Wr.PutText(wr, " __ ");
    Wr.PutText(wr, Fmt.Pad(Fmt.LongReal(yMax, Fmt.Style.Fix, prec := 4), 8));
    Wr.PutText(wr, " ]\n");

    WriteComments(wr, comments, "# ");

    Wr.Flush(wr);
  END WritePlotComments;
  
PROCEDURE CombineComments(ocm, xcm, ycm: TEXT): TEXT =
  BEGIN
    RETURN
      "================================================================\n" &
      "configuration (0,0): \n\n" &
      ocm &
      "================================================================\n" &
      "configuration (1,0): \n\n" &
      xcm &
      "================================================================\n" &
      "configuration (0,1): \n\n" &
      ycm
  END CombineComments;

PROCEDURE WritePlotHeader(wr: Wr.T; READONLY erg: MixedEnergy.T) =
  <* FATAL Thread.Alerted, Wr.Failure *>
  BEGIN
    Wr.PutText(wr, "#");
    Wr.PutText(wr, Fmt.Pad("x", 8));
    Wr.PutText(wr, " ");
    Wr.PutText(wr, Fmt.Pad("y", 8));
    Wr.PutText(wr, " ");
    Wr.PutText(wr, Fmt.Pad("energy", 10));
    FOR i := 0 TO LAST(erg.term^) DO 
      Wr.PutText(wr, " ");
      Wr.PutText(wr, Fmt.Pad(EnergyTag(erg.term[i].name()), 10));
    END;
    Wr.PutText(wr, "\n");
    Wr.Flush(wr);
  END WritePlotHeader;
  
PROCEDURE WriteGNUPlotCommands(
    wr: Wr.T; 
    outFile: TEXT; 
    erg: MixedEnergy.T;
    xMin, xMax: LONGREAL;     (* X range to plot *)
    yMin, yMax: LONGREAL;     (* Y range to plot *)
  ) =
  <* FATAL Thread.Alerted, Wr.Failure *>
  BEGIN
    WITH
      NT = NUMBER(erg.term^)
    DO
      Wr.PutText(wr, "set parametric\n");
      Wr.PutText(wr, "set terminal X11\n");
      Wr.PutText(wr, "set title \"" & outFile & "\"\n");
      Wr.PutText(wr, "set view 60,315,1,1\n");
      Wr.PutText(wr, "set hidden3d\n");
      Wr.PutText(wr, "set autoscale xy\n");
      FOR i := 0 TO NT DO
        Wr.PutText(wr, "splot [");
        Wr.PutText(wr, Fmt.Pad(Fmt.LongReal(xMin, Fmt.Style.Fix, prec := 4), 8));
        Wr.PutText(wr, ":");
        Wr.PutText(wr, Fmt.Pad(Fmt.LongReal(xMax, Fmt.Style.Fix, prec := 4), 8));
        Wr.PutText(wr, "] [");
        Wr.PutText(wr, Fmt.Pad(Fmt.LongReal(yMin, Fmt.Style.Fix, prec := 4), 8));
        Wr.PutText(wr, ":");
        Wr.PutText(wr, Fmt.Pad(Fmt.LongReal(yMax, Fmt.Style.Fix, prec := 4), 8));
        Wr.PutText(wr, "] \\\n");
        Wr.PutText(wr, "  \"" & outFile & ".plot\" \\\n");
        IF i < NT THEN
          Wr.PutText(wr, "  using 1:2:" & Fmt.Int(4 + i) & " \\\n");
          WITH title = EnergyTag(erg.term[i].name()) DO
            Wr.PutText(wr, "  title \"" & title & "\" with lines \n")
          END;
          Wr.PutText(wr, "pause 60\n");
        ELSE
          Wr.PutText(wr, "  using 1:2:3 \\\n");
          WITH title = "tot energy" DO
            Wr.PutText(wr, "  title \"" & title & "\" with lines \n")
          END;
          Wr.PutText(wr, "pause 300\n");
        END;
      END;
      Wr.PutText(wr, "quit\n");
      Wr.Flush(wr);
    END;
  END WriteGNUPlotCommands;
  
PROCEDURE WriteConfiguration(
    outFile: TEXT;
    READONLY top: Topology;
    READONLY c: Coords;
    x, y: LONGREAL;
    comments: TEXT;
  ) =
  BEGIN
    Triang.Write(
      outFile & "-" & PositionTag(x) & "-" & PositionTag(y), 
      top, c,
      comments := 
        PositionComment(x, y) & "\n" &
        comments
    )
  END WriteConfiguration;
  
PROCEDURE WriteComments(wr: Wr.T; comments: TEXT; prefix: TEXT) =
  (* 
    Writes the given "comments" text to "wr", with a "prefix"
    in front of every line.  Supplies a final '\n' if the text is 
    non-empty but does not end with newline. *)
  
  VAR rd: Rd.T := TextRd.New(comments);
  
  PROCEDURE CopyLine() RAISES {Rd.EndOfFile} =
  (*
    Copy one line from "rd" to "wr", prefixed by CommentPrefix. 
    Supplies a final '\n' if the next line exists but does NOT
    end with newline. Raises Rd.EndOfFile if there are no more
    lines in "rd". *)
    
    <* FATAL Rd.Failure, Wr.Failure, Thread.Alerted *>
    VAR c: CHAR;
    BEGIN
      c := Rd.GetChar(rd); (* If EOF here, propagate to caller *)
      Wr.PutText(wr, prefix);
      Wr.PutChar(wr, c);
      WHILE c # '\n' DO
        TRY c := Rd.GetChar(rd) EXCEPT Rd.EndOfFile => c := '\n' END;
        Wr.PutChar(wr, c)
      END
    END CopyLine;

  BEGIN
    TRY LOOP CopyLine() END EXCEPT Rd.EndOfFile => (* Ok *) END;
  END WriteComments;

PROCEDURE PositionTag(x: LONGREAL): TEXT =
  VAR sgn: TEXT;
  BEGIN
    IF x < 0.0d0 THEN sgn := "m"
    ELSIF x > 0.0d0 THEN sgn := "p"
    ELSE sgn := "0"
    END;
    RETURN sgn & Fmt.Pad(Fmt.Int(ROUND(ABS(x)*10000.0d0)), 5, '0')
  END PositionTag;
  
PROCEDURE PositionComment(x, y: LONGREAL): TEXT =
  BEGIN
    RETURN
      "position: x = " &
      Fmt.LongReal(x, Fmt.Style.Fix, prec := 4) &
      " y = " &
      Fmt.LongReal(y, Fmt.Style.Fix, prec := 4)
  END PositionComment;
  
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("-oFile");                               
        o.oFile := pp.getNext();                         

        pp.getKeyword("-xFile");                               
        o.xFile := pp.getNext();   
        
        IF pp.testNext("-range") THEN
          o.xMin := pp.getNextLongReal(-9.9999d0, +9.9999d0);
          o.xMax := pp.getNextLongReal(o.xMin + 0.0001d0, +9.9999d0);
        ELSE
          o.xMin := 00.0000d0;
          o.xMax := +1.0000d0;
        END;

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

        IF pp.testNext("-range") THEN
          o.yMin := pp.getNextLongReal(-9.9999d0, +9.9999d0);
          o.yMax := pp.getNextLongReal(o.yMin + 0.0001d0, +9.9999d0);
        ELSE
          o.yMin := 00.0000d0;
          o.yMax := +1.0000d0;
        END;

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

        IF pp.keywordPresent("-nSteps") THEN
          o.nSteps := pp.getNextInt(1, 9999);
        ELSE
          o.nSteps := 50
        END;
        
        o.erg := ParseEnergyParams.Parse(pp);
        IF o.erg = NIL THEN 
          pp.error("no energy specified")
        END;
        
        pp.finish();                                       
      EXCEPT                                                            
      | ParseParams.Error =>                                              
          Wr.PutText(stderr, "Usage: EnergySlice \\\n");
          Wr.PutText(stderr, "  -oFile <name> \\\n");
          Wr.PutText(stderr, "  -xFile <name> [-range <xmin> <xmax> ] \\\n");
          Wr.PutText(stderr, "  -yfile <name> [-range <ymin> <ymax> ] \\\n");
          Wr.PutText(stderr, "  -outFile <name> \\\n");
          Wr.PutText(stderr, "  [ -nSteps <num> ] \\\n");
          Wr.PutText(stderr, ParseEnergyParams.Help);
          Wr.PutText(stderr, "\n");
          Process.Exit (1);                                              
      END;
    END;
    RETURN o
  END GetOptions;

BEGIN
  DoIt()
END EnergySlice.