/** This program solves systems of equations. */

/** TODO:  Add isSolved method

    TODO:  Factor out solution of a single equation into its own method.

    THINK ABOUT:  Format solutions in MathML?

    THINK ABOUT:  Format solutions as text-formatted equations?

    THINK ABOUT: What if symbols are simple aliases of other symbols like
    a === b ?  We could solve much faster if we folded aliases together and
    just solved for one of them.

    THINK ABOUT:  How should we represent multiple solutions of an equation
    or system of equations?  For example, when solving a === b^2, there are
    multiple solutions for b.
*/
use allTransforms.frink

class System
{
   /* An array of equations, parsed to expressions.  These use the
      === operator to separate the sides of the equations. */
   var equations = new array

   /** A set of symbols that we're *not* going to solve for, for example,
       constants, user-supplied knowns, etc. */
   var ignoreSymbols = new set

   /** An array of LocalSolution objects */
   var localSolutions = new array

   /** All symbols except for those in ignoreSet. */
   var activeSymbols = new set

   /** This is a dictionary mapping from symbol names to a set of equation
       numbers in which that symbol occurs. */
   var equationsContaining = new dict

   /** This is an array containing a set of equation numbers directly connected
       to this equation number. */
   var connectedEquations = new array

   /** This is an array containing a set of equation numbers recursively
       connected to this equation number. */
   var recursivelyConnectedEquations = new array

   /** An array of unsolved equations. */
   var unsolvedEquations = new array

   /** A list of found solutions.  This will change. */
   var foundSolutions = new array

   /** A bitmasked set of flags for debugging */
   var debugFlags = 0

   // The following are debug-level flags that can be set. */
   /** Warn when an equation cannot be solved. */
   class var DEBUG_UNSOLVED = 0b1

   /** Warn when a solution is found. */
   class var DEBUG_SOLUTION_FOUND = 0b10

   new[equationList, ignoreSet=[]] :=
   {
      setEquations[equationList]
      
      addIgnoreSymbols[ignoreSet]

      makeLocalSolutions[]

      initializeConnectedEquations[]
   }

   /** Private array of LocalSolution objects. */
   makeLocalSolutions[] :=
   {
      num = 0
      for eq = equations
      {
         loc = new LocalSolution[eq, ignoreSymbols]
         localSolutions.push[loc]
         syms = loc.getSymbols[]
         activeSymbols = union[activeSymbols, syms]
         for s = syms
            equationsContaining.addToSet[s, num]
         num = num + 1
      }
   }

   /** Private method to set up connected equations. */
   initializeConnectedEquations[] :=
   {
      len = length[equations]
      for i=0 to len-1
      {
         // Set of connected equation numbers
         s = new set

         // Get local symbols of this equation
         theseVars = (localSolutions@i).getSymbols[]
         OTHER:
         for j = 0 to len-1
         {
            if i==j
               next OTHER

            otherVars = (localSolutions@j).getSymbols[]
            if setsIntersect[theseVars, otherVars]
               s.put[j]
         }

         connectedEquations@i = s
      }

      // Now recursively set up recursively-connected equations
      for i=0 to len-1
      {
         connected = new set
         visited = new set
         initializeRecursivelyConnected[i, connected, visited]

         // Remove this equation so it's not cyclic
         connected.remove[i]
         recursivelyConnectedEquations@i = connected
      }
   }

   /** Private recursive function to set up recursively-connected equations.

       This demonstrates how to use connectedEquations to recursively
       traverse equations without visiting them twice.
   */
   initializeRecursivelyConnected[i, connected, visited] :=
   {
      connEqs = connectedEquations@i
      connected.putAll[connEqs]
      visited.put[i]
      for eq = connEqs
         if ! visited.contains[eq]
            initializeRecursivelyConnected[eq, connected, visited]
   }

   /** Given an equation number i, returns a set of directly-connected equation
       numbers. */
   directlyConnectedEquations[i] :=
   {
      return connectedEquations@i
   }

   /** Given an equation number i, returns a set of recursively-connected
       equation numbers. */
   recursivelyConnectedEquations[i] :=
   {
      return recursivelyConnectedEquations@i
   }
   
   /** Returns a set of equation numbers containing the specified symbol.
       The symbol can be specified as a symbol or a string.

       THINK ABOUT:  This currently returns undef for a symbol that does not
       exist, but should that change to an empty set? */
   equationsContaining[symbol] :=
   {
      symbol = makeSymbol[symbol]
      return equationsContaining@symbol
   }

   /** Returns a set of the symbols in the numbered equation. */
   getSymbols[eqNum] :=
   {
      return localSolutions@eqNum.getSymbols[]
   }

   /** Sets the debug flags.  These should be a bitwise AND of events we want
       to log. */
   setDebugFlags[flags] :=
   {
      debugFlags = flags
   }

   /** Solve for a single symbol. */
   solveFor[symbol] :=
   {
      symbol = makeSymbol[symbol]

      // Select the equations that contain this symbol and solve for those
      // locally.
      eqnums = equationsContaining[symbol]

      // Trigger single-equation solution solving
      for num = eqnums
      {
         sols = localSolutions@num.solveFor[symbol]
         for sol = sols
            addSolution[sol, new set[num]]
      }

      // TODO:  Now solve for combinations of equations.
      solveOrder = sort[toArray[eqnums.subsets[false, true]], byLength[]]
      //      println["solveOrder is $solveOrder"]
   }

   solveForRecursive[symbol, currEqs, i, usedEqs] :=
   {
      connEqs = connectedEquations@i
      for eq1 = connEqs
      {
         newEqs = usedEqs.copyAndPut[eq1]
         if usedEqs.contains[newEqs]
         {
            println["Already used $newEqs"]
         } else
         {
//          TODO:  Finish this
//            for leq = (localSolutions@eq1).solveFor[symbol]
//               cockpoo
         }
      }
   }

   /** Function to solve for all active symbols. */
   solveAll[] :=
   {
      for sym = activeSymbols
         solveFor[sym]
   }

   /** Private method to add the symbols in the set (or array) to the set of
       ignoreSymbols. */
   addIgnoreSymbols[ignoreSet] :=
   {
      for sym = ignoreSet
         ignoreSymbols.put[makeSymbol[sym]]
   }

   /** Adds an equation to the list of unsolved equations. */
   addUnsolvedEquation[eq] :=
   {
      // TODO:  Log a common error here?
      unsolvedEquations.push[eq]
   }

   /** Turns a string or a symbol into a symbol. */
   class makeSymbol[s] :=
   {
      if type[s] == "Symbol"
         return s
      else
         return constructExpression["Symbol", [s]]
   }

   /** Private method to add equations.  The equations can be a string, a
       multi-line string, an array of strings, an pre-parsed "Solve" equation
       with a === operator, or an array of "Solve" equations.

       TODO: Lines separated with semicolons.  What else?
   */
   setEquations[equationList] :=
   {
      if isArray[equationList]
      {
         for elem = equationList
            setEquations[elem]
         
         return
      }

      if isString[equationList]
      {
         // Split on newlines or semicolon or other vertical separators
         for line = split[%r/\r\n|\r|\n|\u2028|\u2029|\u000B|\u000C|\u0085|;/, equationList]
            addStringEquation[line]

         return
      }

      // Equation with a === operator?
      if type[equationList] == "Solve"
      {
         equations.push[equationList]
         return
      }

      // Turn an assignment statement (with a single = ) into ===
      if type[equationList] == "Assignment"
      {
         equations.push[substituteExpression[equationList, "_a=_b", "_a===_b"]]
         return
      }

      println["Solver: invalid equation type: $equationList"]
   }

   /** Private method to add a single equation as a string. */
   addStringEquation[str] :=
   {
      str = trim[str]
      
      if length[str] != 0 and ! startsWith[str, "//"]  // Exclude comments
      {
         // Replace a single equals sign with triple-equals
         str =~ %s/(^|[^<>=])(=)($|[^=])/$1===$3/
         expr = parseToExpression[str]
         if type[expr] == "Solve"
            equations.push[expr]
         else
            println["Not a proper equation for solving.  Requires === operator:  $str"]
      }
   }

   /** This is a public method to determine if a solved equation is truly
       solved.  It expects an equation of the form "a === solution".

       THINK ABOUT:  Should we pass in the variable that we expect it to be
       solved for and check that that's the case?
   */
   class isSolved[eq] :=
   {
      return type[eq] == "Solve" and getChildCount[eq] == 2
   }

   /** Prints a human-readable representation of the state of this system. */
   dump[] :=
   {
      len = length[equations]
      println["Equations:"]
      for i = 0 to len-1
         println["   $i:\t" + inputForm[equations@i]]

      if ! ignoreSymbols.isEmpty[]
         println["\nIgnore symbols:\n   $ignoreSymbols"]

      println[]
      println["Active symbols:\n   " + activeSymbols]

      println[]
      println["Equations containing:"]
      for [sym, eqNo] = equationsContaining
         println["   " + inputForm[sym] +"\t" + sortSet[eqNo]]

      println[]
      println["Directly-connected equations:"]
      for i = 0 to len-1
         println["   $i:\t" + sortSet[directlyConnectedEquations[i]]]
      
      println[]
      println["Recursively-connected equations:"]
      for i = 0 to len-1
         println["   $i:\t" + sortSet[recursivelyConnectedEquations[i]]]

      println[]
      println["Local solutions:"]
      for i = 0 to len-1
      {
         ls = (localSolutions@i).dump[]
         if ls != ""
            print["   $i:\n$ls"]
      }

      println[]
      println["Found solutions:"]
      for i = 0 to length[foundSolutions]-1
         println["   " + inputForm[foundSolutions@i@0] + "   " + foundSolutions@i@1]
   }

   /** Private method for sorting a set. */
   class sortSet[set] :=
   {
      return lexicalSort[toArray[set]]
   }

   /** Adds a new solution to the list of solutions.  This will change. */
   addSolution[eq, eqsUsed] :=
   {
      foundSolutions.push[[eq, eqsUsed]]
      
      if bitAnd[debugFlags, DEBUG_SOLUTION_FOUND] != 0
         println["Solution found:  $eq using equations $eqsUsed"]
   }
}

/** LocalSolution is a class that holds a single original equation, its
    symbols, and the solutions for each symbol.  It caches results so they are
    only calculated once.  The primary methods are the solveFor methods which
    solve the equation if necessary or return it from the cache if already
    solved.
*/
class LocalSolution
{
   /** The original equation. */
   var eq

   /** A set of symbols in this equation minus the ignored symbols. */
   var symbols

   /** A dictionary of solutions where the key is the variable as a symbol
       and the value is an array of solved equations of the form
       var === solution
   */
   var solutions = new dict

   /** A dictionary of solutions where the key is the variable as a symbol
       and the value is an array of the right-hand-side of solved equations.
   */
   var solutionsRight = new dict

   /** Construct a new LocalSolution around the specified equation.  The
       equation should be pre-normalized by the System class. */
   new[eq, ignoreSet] :=
   {
      this.eq = eq
      before = getSymbolsAsSymbols[eq]
      symbols = setDifference[before, ignoreSet]
   }

   /** Solve the equation for all of the symbols not in the ignoreSet. */
   solveAll[] :=
   {
      for s = symbols
         solveFor[s]
   }

   /** Solve for a single variable and returns an array of solutions in the
       form [s === solution1] .  This caches solutions so you can use it to
       efficiently return the solution. */
   solveFor[s] :=
   {
      s = System.makeSymbol[s]

      // Return the solution from the cache if it exists.
      sol = solutions@s
      if sol != undef
         return sol

      // This actually solves the equation.
      // TODO:  make this a separate function?
      solveEq = constructExpression["FunctionCall", ["solve", eq, s]]
      solved = transformExpression[solveEq]

      // Force even single solutions to an array
      solved = flatten[toArray[solved]]
      
      // Verify that equations were solved and extract their right-hand-sides
      rights = new array
      for eq1 = solved
      {
         // Check to see that each equation is actually solved.
         if (System.isSolved[eq1])
         {
            solutions.addToList[s, eq1]  // Add solution to list
            // Extract right-hand side
            solutionsRight.addToList[s, getChild[eq1, 1]]
         } else
         {
            // THINK ABOUT: How to get debug settings from System?
            // Have that passed in as an argument?
//            if (bitAnd[debugFlags, DEBUG_UNSOLVED] != 0)
            // TODO:  Push unsolved equation into System
            // system.addUnsolvedEquation[eq]
            println["Solver: Unsolved equation: " + eq]
         }
      }

      return solved
   }

   /** Solve for a single variable and returns an array of right-hand-side of
       solutions.  This caches solutions so you can use it to
       efficiently return the solution. */
   solveForRight[s] :=
   {
      s = System.makeSymbol[s]
      solveFor[s]
      return solutionsRight@s
   }

   /** Returns a set of the symbols included in this equation. */
   getSymbols[] :=
   {
      return symbols
   }

   /** Prints the current state of this LocalSolution.  It does not force any
       solutions so solutions may be empty unless solveFor has been called. */
   dump[] :=
   {
      str = ""
      for [sym, sols] = solutionsRight
      {
         str = str + "      " + inputForm[sym] + " =\n"
         for sol = sols
            str = str + "          " + inputForm[sol] + "\n"
      }

      return str
   }
}