//-----------------------------------------------------------------------------
//
// Copyright (C) Microsoft Corporation.  All Rights Reserved.
//
//-----------------------------------------------------------------------------
using System;
using System.Text;
using System.IO;
using System.Collections;
using System.Collections.Generic;
using System.Diagnostics.Contracts;
using System.Linq;
using Microsoft.Basetypes;
using Microsoft.Boogie.VCExprAST;

// Method to turn VCExprs into strings that can be fed into SMT
// solvers. This is currently quite similar to the
// SimplifyLikeLineariser (but the code is independent)

namespace Microsoft.Boogie.SMTLib
{

  // Options for the linearisation
  public class LineariserOptions
  {
    public static readonly LineariserOptions Default = new LineariserOptions();
  }


  ////////////////////////////////////////////////////////////////////////////////////////

  // Lineariser for expressions. The result (bool) is currently not used for anything
  public class SMTLibExprLineariser : IVCExprVisitor<bool, LineariserOptions/*!*/>
  {

    public static string ToString(VCExpr e, UniqueNamer namer, SMTLibProverOptions opts)
    {
      Contract.Requires(e != null);
      Contract.Requires(namer != null);
      Contract.Ensures(Contract.Result<string>() != null);

      StringWriter sw = new StringWriter();
      SMTLibExprLineariser lin = new SMTLibExprLineariser(sw, namer, opts);
      Contract.Assert(lin != null);
      lin.Linearise(e, LineariserOptions.Default);
      return cce.NonNull(sw.ToString());
    }

    ////////////////////////////////////////////////////////////////////////////////////////

    private readonly TextWriter wr;
    [ContractInvariantMethod]
    void ObjectInvariant()
    {
      Contract.Invariant(wr != null);
      Contract.Invariant(Namer != null);
    }

    private SMTLibOpLineariser OpLinObject = null;
    private IVCExprOpVisitor<bool, LineariserOptions>/*!>!*/ OpLineariser
    {
      get
      {
        Contract.Ensures(Contract.Result<IVCExprOpVisitor<bool, LineariserOptions>>() != null);

        if (OpLinObject == null)
          OpLinObject = new SMTLibOpLineariser(this, wr);
        return OpLinObject;
      }
    }

    internal readonly UniqueNamer Namer;
    internal readonly SMTLibProverOptions ProverOptions;

    public SMTLibExprLineariser(TextWriter wr, UniqueNamer namer, SMTLibProverOptions opts)
    {
      Contract.Requires(wr != null); Contract.Requires(namer != null);
      this.wr = wr;
      this.Namer = namer;
      this.ProverOptions = opts;
    }

    public void Linearise(VCExpr expr, LineariserOptions options)
    {
      Contract.Requires(expr != null);
      Contract.Requires(options != null);
      expr.Accept<bool, LineariserOptions>(this, options);
    }

    /////////////////////////////////////////////////////////////////////////////////////

    private static void TypeToStringHelper(Type t, StringBuilder sb)
    {
      Contract.Requires(t != null);

      TypeSynonymAnnotation syn = t as TypeSynonymAnnotation;
      if (syn != null) {
        TypeToStringHelper(syn.ExpandedType, sb);
      } else {
        if (t.IsMap && CommandLineOptions.Clo.UseArrayTheory) {
          MapType m = t.AsMap;
          Contract.Assert(m.MapArity == 1);
          sb.Append("(Array ");
          TypeToStringHelper(m.Arguments[0], sb);
          sb.Append(" ");
          TypeToStringHelper(m.Result, sb);
          sb.Append(")");
        } else if (t.IsMap) {
          MapType m = t.AsMap;
          sb.Append('[');
          for (int i = 0; i < m.MapArity; ++i) {
            if (i != 0)
              sb.Append(',');
            TypeToStringHelper(m.Arguments[i], sb);
          }
          sb.Append(']');
          TypeToStringHelper(m.Result, sb);
        } else if (t.IsBool || t.IsInt || t.IsBv) {
          sb.Append(TypeToString(t));
        } else {
          System.IO.StringWriter buffer = new System.IO.StringWriter();
          using (TokenTextWriter stream = new TokenTextWriter("<buffer>", buffer, false)) {
            t.Emit(stream);
          }
          sb.Append(buffer.ToString());
        }
      }

    }

    public static string TypeToString(Type t)
    {
      Contract.Requires(t != null);
      Contract.Ensures(Contract.Result<string>() != null);

      if (t.IsBool)
        return "Bool";
      else if (t.IsInt)
        return "Int";
      else if (t.IsBv) {
        return "(_ BitVec " + t.BvBits + ")";
      } else {
        StringBuilder sb = new StringBuilder();        
        TypeToStringHelper(t, sb);
        var s = sb.ToString();
        if (s[0] == '(')
          return s;
        else
          return SMTLibNamer.QuoteId("T@" + s);
      }
    }

    public static string ExtractBuiltin(Function f)
    {
      Contract.Requires(f != null);
      string retVal = null;
      retVal = f.FindStringAttribute("bvbuiltin");
      
      // It used to be "sign_extend 12" in Simplify, and is "sign_extend[12]" with SMT
      if (retVal != null && retVal.StartsWith("sign_extend "))
        return retVal.Replace(" ", "[") + "]";

      if (retVal == null) {
        retVal = f.FindStringAttribute("builtin");
      }
      return retVal;
    }

    /////////////////////////////////////////////////////////////////////////////////////

    public bool Visit(VCExprLiteral node, LineariserOptions options)
    {
      if (node == VCExpressionGenerator.True)
        wr.Write("true");
      else if (node == VCExpressionGenerator.False)
        wr.Write("false");
      else if (node is VCExprIntLit) {
        BigNum lit = ((VCExprIntLit)node).Val;
        if (lit.IsNegative)
          // In SMT2 "-42" is an identifier (SMT2, Sect. 3.2 "Symbols")
          wr.Write("(- 0 {0})", lit.Abs);
        else
          wr.Write(lit);
      } else {
        Contract.Assert(false);
        throw new cce.UnreachableException();
      }

      return true;
    }

    /////////////////////////////////////////////////////////////////////////////////////

    public bool Visit(VCExprNAry node, LineariserOptions options)
    {
      //Contract.Requires(node != null);
      //Contract.Requires(options != null);

      VCExprOp op = node.Op;
      Contract.Assert(op != null);

      if (op.Equals(VCExpressionGenerator.AndOp) ||
           op.Equals(VCExpressionGenerator.OrOp)) {
        // handle these operators without recursion

        wr.Write("({0}",
          op.Equals(VCExpressionGenerator.AndOp) ? "and" : "or");
        foreach (var ch in node.UniformArguments) {
          wr.Write("\n");
          Linearise(ch, options);
        }
        wr.Write(")");

        return true;
      }

      return node.Accept<bool, LineariserOptions/*!*/>(OpLineariser, options);
    }

    /////////////////////////////////////////////////////////////////////////////////////

    public bool Visit(VCExprVar node, LineariserOptions options)
    {
      wr.Write(Namer.GetQuotedName(node, node.Name));
      return true;
    }

    /////////////////////////////////////////////////////////////////////////////////////

    public bool Visit(VCExprQuantifier node, LineariserOptions options)
    {
      Contract.Assert(node.TypeParameters.Count == 0);

      Namer.PushScope(); try {

        string kind = node.Quan == Quantifier.ALL ? "forall" : "exists";
        wr.Write("({0} (", kind);

        for (int i = 0; i < node.BoundVars.Count; i++) {
          VCExprVar var = node.BoundVars[i];
          Contract.Assert(var != null);
          string printedName = Namer.GetQuotedLocalName(var, var.Name);
          Contract.Assert(printedName != null);
          wr.Write("({0} {1}) ", printedName, TypeToString(var.Type));
        }

        wr.Write(") ");

        VCQuantifierInfos infos = node.Infos;
        var weight = QKeyValue.FindIntAttribute(infos.attributes, "weight", 1);
        if (!ProverOptions.UseWeights)
          weight = 1;
        var hasAttrs = node.Triggers.Count > 0 || infos.qid != null || weight != 1 || infos.uniqueId != -1;

        if (hasAttrs)
          wr.Write("(! ");

        Linearise(node.Body, options);

        if (hasAttrs) {
          wr.Write("\n");
          if (infos.qid != null)
            wr.Write(" :named {0}\n", SMTLibNamer.QuoteId(infos.qid));
          if (weight != 1)
            wr.Write(" :weight {0}\n", weight);
          if (infos.uniqueId != -1)
            wr.Write(" :skolemid {0}\n", infos.uniqueId);
          WriteTriggers(node.Triggers, options);
          
          wr.Write(")");
        }
        
        wr.Write(")");

        return true;

      } finally {
        Namer.PopScope();
      }
    }

    private void WriteTriggers(List<VCTrigger/*!>!*/> triggers, LineariserOptions options)
    {
      Contract.Requires(options != null);
      Contract.Requires(triggers != null);
      // first, count how many neg/pos triggers there are
      int negTriggers = 0;
      int posTriggers = 0;
      foreach (VCTrigger vcTrig in triggers) {
        Contract.Assert(vcTrig != null);
        if (vcTrig.Pos) {
          posTriggers++;
        } else {
          negTriggers++;
        }
      }

      if (posTriggers > 0) {
        foreach (VCTrigger vcTrig in triggers) {
          Contract.Assert(vcTrig != null);
          if (vcTrig.Pos) {
            wr.Write(" :pattern (");
            foreach (VCExpr e in vcTrig.Exprs) {
              Contract.Assert(e != null);
              wr.Write(" ");
              var subPat = e;
              var nary = e as VCExprNAry;
              if (nary != null && (nary.Op == VCExpressionGenerator.NeqOp || nary.Op == VCExpressionGenerator.EqOp)) {
                if (nary[0] is VCExprLiteral)
                  subPat = nary[1];
                else if (nary[1] is VCExprLiteral)
                  subPat = nary[0];
              }
              Linearise(subPat, options);
            }
            wr.Write(")\n");
          }
        }
      } else if (negTriggers > 0) {
        // if also positive triggers are given, the SMT solver (at least Z3)
        // will ignore the negative patterns and output a warning. Therefore
        // we never specify both negative and positive triggers
        foreach (VCTrigger vcTrig in triggers) {
          Contract.Assert(vcTrig != null);
          if (!vcTrig.Pos) {
            wr.Write(" :no-pattern (");
            Contract.Assert(vcTrig.Exprs.Count == 1);
            Linearise(vcTrig.Exprs[0], options);
            wr.Write(")\n");
          }
        }
      }
    }

    /////////////////////////////////////////////////////////////////////////////////////

    public bool Visit(VCExprLet node, LineariserOptions options)
    {
      Namer.PushScope();
      try {
        
        foreach (VCExprLetBinding b in node) {
          wr.Write("(let (");
          Contract.Assert(b != null);
          wr.Write("({0} ", Namer.GetQuotedName(b.V, b.V.Name));
          Linearise(b.E, options);
          wr.Write("))\n");
        }        
        Linearise(node.Body, options);
        foreach (VCExprLetBinding b in node)
          wr.Write(")");
        return true;
      } finally {
        Namer.PopScope();
      }
    }

    /////////////////////////////////////////////////////////////////////////////////////

    // Lineariser for operator terms. The result (bool) is currently not used for anything
    internal class SMTLibOpLineariser : IVCExprOpVisitor<bool, LineariserOptions/*!*/>
    {
      private readonly SMTLibExprLineariser ExprLineariser;
      private readonly TextWriter wr;
      [ContractInvariantMethod]
      void ObjectInvariant()
      {
        Contract.Invariant(wr != null);
        Contract.Invariant(ExprLineariser != null);
      }


      public SMTLibOpLineariser(SMTLibExprLineariser ExprLineariser, TextWriter wr)
      {
        Contract.Requires(ExprLineariser != null);
        Contract.Requires(wr != null);
        this.ExprLineariser = ExprLineariser;
        this.wr = wr;
      }

      ///////////////////////////////////////////////////////////////////////////////////
      private void WriteApplication(string opName, IEnumerable<VCExpr>/*!>!*/ args, LineariserOptions options)
      {
        Contract.Requires(cce.NonNullElements(args));
        Contract.Requires(options != null);
        Contract.Assert(opName != null);

        bool hasArgs = false;
        foreach (VCExpr e in args) {
          Contract.Assert(e != null);
          if (!hasArgs)
            wr.Write("({0}", opName);
          wr.Write(" ");
          ExprLineariser.Linearise(e, options);
          hasArgs = true;
        }

        if (hasArgs)
          wr.Write(")");
        else
          wr.Write("{0}", opName);
      }

      ///////////////////////////////////////////////////////////////////////////////////

      public bool VisitNotOp(VCExprNAry node, LineariserOptions options)
      {
        WriteApplication("not", node, options);      // arguments can be both terms and formulas
        return true;
      }

      private bool PrintEq(VCExprNAry node, LineariserOptions options)
      {
        Contract.Requires(node != null);
        Contract.Requires(options != null);

        // not sure if this is needed
        if (node[0].Type.IsBool) {
          Contract.Assert(node[1].Type.IsBool);
          // use equivalence
          WriteApplication("iff", node, options);
        } else {
          WriteApplication("=", node, options);
        }

        return true;
      }

      public bool VisitEqOp(VCExprNAry node, LineariserOptions options)
      {
        return PrintEq(node, options);
      }

      public bool VisitNeqOp(VCExprNAry node, LineariserOptions options)
      {
        //Contract.Requires(node != null);
        //Contract.Requires(options != null);
        wr.Write("(not ");
        PrintEq(node, options);
        wr.Write(")");
        return true;
      }

      public bool VisitAndOp(VCExprNAry node, LineariserOptions options)
      {
        WriteApplication("and", node, options);
        return true;
      }

      public bool VisitOrOp(VCExprNAry node, LineariserOptions options)
      {
        WriteApplication("or", node, options);
        return true;
      }

      public bool VisitImpliesOp(VCExprNAry node, LineariserOptions options)
      {
        WriteApplication("=>", node, options);
        return true;
      }

      public bool VisitIfThenElseOp(VCExprNAry node, LineariserOptions options)
      {
        WriteApplication("ite", node, options);
        return true;
      }

      public bool VisitCustomOp(VCExprNAry node, LineariserOptions options)
      {
        VCExprCustomOp op = (VCExprCustomOp)node.Op;
        WriteApplication(op.Name, node, options);
        return true;
      }

      public bool VisitDistinctOp(VCExprNAry node, LineariserOptions options)
      {
        //Contract.Requires(node != null);
        //Contract.Requires(options != null);

        if (node.Length < 2) {
          ExprLineariser.Linearise(VCExpressionGenerator.True, options);
        } else {
          var groupings = node.GroupBy(e => e.Type).Where(g => g.Count() > 1).ToArray();

          if (groupings.Length > 1)
            wr.Write("(and ");

          foreach (var g in groupings) {
            wr.Write("(distinct");
            foreach (VCExpr e in g) {
              Contract.Assert(e != null);
              wr.Write(" ");
              ExprLineariser.Linearise(e, options);
            }
            wr.Write(")");
          }

          if (groupings.Length > 1)
            wr.Write(")");
          wr.Write("\n");
        }

        return true;
      }

      public bool VisitLabelOp(VCExprNAry node, LineariserOptions options)
      {
        var op = (VCExprLabelOp)node.Op;
        if (ExprLineariser.ProverOptions.UseLabels) {
          // Z3 extension
          wr.Write("({0} {1} ", op.pos ? "lblpos" : "lblneg", SMTLibNamer.QuoteId(op.label));
        }

        if (!op.pos)
          wr.Write("(or {0} ", SMTLibNamer.QuoteId(SMTLibNamer.BlockedLabel(op.label)));
            
        ExprLineariser.Linearise(node[0], options);

        if (!op.pos) wr.Write(")");

        if (ExprLineariser.ProverOptions.UseLabels)
          wr.Write(")");

        return true;
      }

      public bool VisitSelectOp(VCExprNAry node, LineariserOptions options)
      {        
        var name = SimplifyLikeExprLineariser.SelectOpName(node);
        name = ExprLineariser.Namer.GetQuotedName(name, name);
        if (CommandLineOptions.Clo.UseArrayTheory)
          name = "select";
        WriteApplication(name, node, options);
        return true;
      }

      public bool VisitStoreOp(VCExprNAry node, LineariserOptions options)
      {
        var name = SimplifyLikeExprLineariser.StoreOpName(node);
        name = ExprLineariser.Namer.GetQuotedName(name, name);
        if (CommandLineOptions.Clo.UseArrayTheory)
          name = "store";
        WriteApplication(name, node, options);
        return true;
      }

      static char[] hex = { '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'a', 'b', 'c', 'd', 'e', 'f' };
      public bool VisitBvOp(VCExprNAry node, LineariserOptions options)
      {
        var lit = (VCExprIntLit)node[0];
        var bytes = lit.Val.ToByteArray();
        if (node.Type.BvBits % 8 == 0) {
          wr.Write("#x");
          for (var pos = node.Type.BvBits / 8 - 1; pos >= 0; pos--) {
            var k = pos < bytes.Length ? bytes[pos] : 0;
            wr.Write(hex[k >> 4]);
            wr.Write(hex[k & 0xf]);
          }
        } else {
          wr.Write("#b");
          for (var pos = node.Type.BvBits - 1; pos >= 0; pos--) {
            var i = pos >> 3;
            var k = i < bytes.Length ? bytes[i] : 0;
            wr.Write((k & (1 << (pos & 7))) == 0 ? '0' : '1');
          }
        }
        return true;
      }

      public bool VisitBvExtractOp(VCExprNAry node, LineariserOptions options)
      {
        var op = (VCExprBvExtractOp)node.Op;
        wr.Write("((_ extract {0} {1}) ", op.End - 1, op.Start);
        ExprLineariser.Linearise(node[0], options);
        wr.Write(")");
        return true;
      }

      public bool VisitBvConcatOp(VCExprNAry node, LineariserOptions options)
      {
        WriteApplication("concat", node, options);
        return true;
      }

      public bool VisitAddOp(VCExprNAry node, LineariserOptions options)
      {
        WriteApplication("+", node, options);
        return true;
      }

      public bool VisitSubOp(VCExprNAry node, LineariserOptions options)
      {
        WriteApplication("-", node, options);
        return true;
      }

      public bool VisitMulOp(VCExprNAry node, LineariserOptions options)
      {
        WriteApplication("*", node, options);
        return true;
      }

      public bool VisitDivOp(VCExprNAry node, LineariserOptions options)
      {
        WriteApplication("int_div", node, options);
        return true;
      }

      public bool VisitModOp(VCExprNAry node, LineariserOptions options)
      {
        WriteApplication("int_mod", node, options);
        return true;
      }

      public bool VisitLtOp(VCExprNAry node, LineariserOptions options)
      {
        WriteApplication("<", node, options);
        return true;
      }

      public bool VisitLeOp(VCExprNAry node, LineariserOptions options)
      {
        WriteApplication("<=", node, options);
        return true;
      }

      public bool VisitGtOp(VCExprNAry node, LineariserOptions options)
      {
        WriteApplication(">", node, options);
        return true;
      }

      public bool VisitGeOp(VCExprNAry node, LineariserOptions options)
      {
        WriteApplication(">=", node, options);
        return true;
      }

      public bool VisitSubtypeOp(VCExprNAry node, LineariserOptions options)
      {
        WriteApplication("UOrdering2", node, options);
        return true;
      }

      public bool VisitSubtype3Op(VCExprNAry node, LineariserOptions options)
      {
        WriteApplication("UOrdering3", node, options);
        return true;
      }

      public bool VisitBoogieFunctionOp(VCExprNAry node, LineariserOptions options)
      {
        VCExprBoogieFunctionOp op = (VCExprBoogieFunctionOp)node.Op;
        Contract.Assert(op != null);
        string printedName;

        var builtin = ExtractBuiltin(op.Func);
        if (builtin != null)
          printedName = builtin;
        else
          printedName = ExprLineariser.Namer.GetQuotedName(op.Func, op.Func.Name);
        Contract.Assert(printedName != null);

        WriteApplication(printedName, node, options);

        return true;
      }

    }
  }

}