// This program solves systems of equations.
// See "systemSolverTest.frink" or
//  https://frinklang.org/fsp/solve2.fsp
// for examples of using it.

use allTransforms.frink

showApproximations[false]

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

   // A dictionary mapping symbol names (as strings)
   // to an array of the equation numbers
   // in which they occur.  (Or should it be the equation objects?)
   var variableOccurrences

   // A list of symbols in each equation.
   var symbols

   // A set of symbols that we're not going to solve for, because
   // they're intended to be units.  This is a set of strings.
   var ignoreSymbolsStrings = new set

   // A set of symbols that we're not going to solve for, because
   // they're intended to be units.  This is a set of symbols
   var ignoreSymbols = new set
   
   // A dictionary containing solved equations.
   // The key is the variable name (as a string)
   // The value is an array of equations that comprise solutions for that
   // variable.
   var solutions

   // Create a new System with the given equations, passed in as a list.
   new[equationList, ignoreSet=[], debug=false] :=
   {
      if debug
         println["orig ignoreSet is " + ignoreSet]
      for sym = ignoreSet
      {
         if debug
            println["Adding $sym"]
         addIgnoreSymbol[sym]
      }

      if debug
         println["ignoreSymbolsStrings is " + ignoreSymbolsStrings]
      
      equationTemp = deepCopy[equationList]
      equationTemp = solveSimultaneous[equationTemp, debug]
      //      equationTemp = solveSimultaneous[equationTemp, debug]

      equations = new array
      variableOccurrences = new dict
      symbols = new array

      for e = equationTemp
         addEquation[e]

      solutions = new dict
   }

   // Adds a symbol to the ignore list.  This adds it both to the
   // ignoreSymbols and ignoreSymbolsStrings sets.
   addIgnoreSymbol[sym is string] :=
   {
      ignoreSymbolsStrings.put[sym]
      ignoreSymbols.put[constructExpression["Symbol", sym]]
   }

   // Adds a new equation to the system.  Normally we would want all of the
   // equations to be known at construction time, but this may allow us to
   // add more equations to the system later, to provide a "refinement" of
   // the system we're using.  We generally don't want people to add equations
   // after we've partially solved the system, or if we do, we'll have to
   // clear any cache (like the "solutions" dictionary) and re-solve.
   addEquation[e] :=
   {
      i = length[equations]
      equations.push[e]
      
      syms = getSymbols[e]
      symbols@i = syms
      for sym = syms
      {
         if variableOccurrences@sym == undef
            variableOccurrences@sym = new array

         (variableOccurrences@sym).push[i]
      }
   }

   // Find the solutions for the symbol x (specified as a string.)
   // This function sets up for a recursive call.
   solveFor[x, debug=false] :=
   {
      prev = solutions@x
      if prev
         return prev
      
      results = new array       // Array of expressions containing results.
      usedEquations = new set   // Set of numbers of used equations.
      usedSymbols = new set     // Set of names of used symbols.
      results = solveFor[x, results, usedEquations, usedSymbols, debug]
      
      solutions@x = flatten[results]     // Put into dictionary
      return results
   }

   // The private recursive call.
   solveFor[x, results, usedEquations, usedSymbols, debug=false] :=
   {
      usedSymbols.put[x]
      // println["Used symbol $x"]
      for ei = variableOccurrences@x  // For each equation that x is used in...
      {
         if usedEquations.contains[ei]
            next                // Skip equation if it's already been used.
            
         usedEquations.put[ei]
         // println["Used equation $ei"]
         eq = equations@ei

         results.push[solveSingle[eq, x]]
         // println["The equation is $solveEq"]
         // println[result]
      }

      // println["Results before recursion:\n" + results]

      // We count the number of results currently on the list and only
      // substitute into the existing results, as we'll be adding to the
      // results list.
      s = length[results]
      for n = 0 to s-1
      {
         res = results@n
         symbols = array[getSymbols[res]]

         for sym = symbols
            if ! usedSymbols.contains[sym]
            {
               results2 = new array
//               println["Solving for $sym"]
               results2 = solveFor[sym, results2, deepCopy[usedEquations], deepCopy[usedSymbols], debug]
//               println["Solved for $sym, got back " + results2]
               symSymbol = constructExpression["Symbol", sym]

               // TODO:  Only substitute below if the result is a fully solved
               // === expression with the appropriate form.
               for eq2 = flatten[results2]
               {
                  if structureEquals[_a === _b, eq2]
                  {
                     results.push[transformExpression[substituteExpression[res, symSymbol, child[eq2, 1]]]]
//                   println["Used $eq2"]
                  } else
                       if debug
                           println["Discarded unsolved $eq2"]
               }
            }
      }

      return eliminateOverconstrained[results, debug]
   }

   // Find the solutions for the symbol x (specified as a string.)
   // values is an array of [varname, value] pairs.
   solveForValues[x, values, allValues=false, debug=false] :=
   {
      origSolutions = solveFor[x, debug]
      newSolutions = new array
//      if debug
//         println["solveForValues: origSolutions are:" + origSolutions]
      
      for sol = origSolutions
      {
         newSol = sol
         for [symName, value] = values
         {
            if debug
            {
               println["\nReplacing $symName with $value"]
               println["old is:" + newSol]
            }
            
            newSol = replaceVar[newSol, symName, value]
            
            if debug
               println["new is:" + newSol]
         }

         newSolutions.push[transformExpression[newSol]]
      }

      if (allValues == false)
         newSolutions = eliminateOverconstrained[newSolutions, debug]

      solvedVals = new array
      allVals = new array
      for sol = newSolutions
      {
         right = getChild[sol,1]
         allVals.push[right]
         if length[setDifference[getSymbols[right], ignoreSymbolsStrings]] == 0 // Solved?
            solvedVals.push[right]
      }

      if (allValues==false and length[solvedVals] > 0)  // Have fully-solved?
         return solvedVals
      else
         return allVals
   }

   // Find the solutions for the symbol x (specified as a string.)
   // values is an array of [varname, value] pairs.
   // This just sets up arguments for the recursive call.
   solveForValuesOld[x, values, allValues=false, debug=false] :=
   {
      // Create a new set of values we've already substituted for.
      substituted = new set
      // Dictionary<set, equations>
      cachedSolutions = new dict
      return solveForValuesRecursive[x, values, substituted, cachedSolutions, allValues, debug]
   }

   // Puts the specified values into the cache.
   class putCache[cachedSolutions, varName, substituted, vals] :=
   {
      if ! cachedSolutions.containsKey[varName]
         cachedSolutions@varName = new dict

      cachedSolutions@varName@substituted = vals
   }

   // This is the "private" method that does recursive solution of
   // solutions for the symbol x.
   solveForValuesRecursive[x, values, substituted, cachedSolutions, allValues, debug=false] :=
   {
      if cachedSolutions.containsKey[x]
         if (cachedSolutions@x).containsKey[substituted]
            return cachedSolutions@x@substituted

      if debug
         println["Solving for $x, substituted is $substituted"]

      sols = solveFor[x, debug]
      valset = new set
      for [varname, val] = values
      {
         if varname == x        // User passed in solution?
         {
            putCache[cachedSolutions, x, substituted, [val]]
            return [val]          // TODO:  Warn as overspecified?
         }
//         sols = replaceVar[sols, varname, val]  // THINK ABOUT:  Do this later?
         valset.put[varname]
      }

      if debug
         println["Initial solutions for $x are $sols"]
      if (allValues == false)
         sols = eliminateOverconstrained[flatten[sols]]
      sort[sols, {|a,b| length[getSymbols[a]] <=> length[getSymbols[b]]}]

      // Ignore the variable we're solving for and things in the ignore list.
      ignoreVals = new set
      ignoreVals.put[x]
      ignoreVals = union[ignoreVals, ignoreSymbolsStrings]
      res = new array
      for sol = sols
         if length[setDifference[getSymbols[sol], ignoreVals]] == 0  // Solved?
            for s = array[sol]
               res.push[child[s,1]]  // Eval here?

      if length[res] > 0        // Got fully-solved solutions?
      {
         putCache[cachedSolutions, x, substituted, res]
         if debug
            println["Solutions for $x (fully solved) are $res"]
         return res
      } else    // Not fully solved, return symbolic solutions 
      {
          xSym = constructExpression["Symbol", x]
          allIgnores = deepCopy[ignoreSymbolsStrings]
          allIgnores = union[allIgnores, substituted]
          allIgnores = union[allIgnores, valset]
          allIgnores.put[x]

          // Get all symbols in all solutions, minus the things we're
          // ignoring.
          allSyms = setDifference[getSymbols[sols], allIgnores]

          // solsDict is a dictionary <solveFor, [equations]>
          solsDict = new dict
          for u = allSyms
          {
             newsub = deepCopy[substituted]
             newsub.put[x]
             solsDict@u = solveForValuesRecursive[u, values, newsub, cachedSolutions, allValues, debug]
          }

//          println["solsDict is $solsDict"]

          // Now iterate through each solution and substitute in each 
          for s = sols
             if structureEquals[_a === _b, s]
             {
                right = child[s,1]
                unknowns = setDifference[getSymbols[right], allIgnores]
                newsub = deepCopy[substituted]
                newsub.put[x]

                for u = unknowns
                {
                   res2 = solveForValuesRecursive[u, values, newsub, cachedSolutions, allValues, debug]
                   uSym = constructExpression["Symbol", u]
                   for r2 = flatten[res2]
                   {
                      if debug
                         println["Substituting in $s, from $u, to $r2"]
                      s = substituteExpression[s, uSym, r2]
                   }
                }
                
                s = transformExpression[s]
                
                g = getSymbols[child[s, 1]]
                if ! g.contains[x]  // Solution contains x?
                {
                   res.push[s]
                   if debug
                      println["Adding solution $s"]
                } else
                   if debug
                      println["Solution $s rejected because it contains $x"]
             } else
                if debug
                   println["Solution $s rejected because unsolved."]
             

          for [varname, val] = values
             res = replaceVar[res, varname, val]  // THINK ABOUT:  Do this later?

          if (allValues == false)
             res = eliminateOverconstrained[res, debug]
          size = length[res]
          for i=0 to size-1
          {
             child = child[res@i, 1]  // Return only right-hand-sides
             if isSubset[getSymbols[child], ignoreSymbolsStrings]
                res@i = child  // Eval here?
             else
                res@i = child
          }
          if debug
             println["Solutions for $x are $res"]

          putCache[cachedSolutions, x, substituted, res]
          return res  //, false, debug]
      }
   }
   
   // Solve for all unknowns in the system.
   // Setting fullReplace=true makes this work harder at performing full
   // substitutions of all variables in all solutions, but it may be
   // prohibitively slow for a system with a large number of solutions and
   // variables.
   solveAll[fullReplace=false, debug=false] :=
   {
      results = new array
      for [sym, appear] = variableOccurrences
         if ! ignoreSymbolsStrings.contains[sym]
            results.push[solveFor[sym, debug]]

      if debug
         println["Before fullReplace:\n" + join["\n", flatten[results]]]

         // TODO:  Comment back in fullReplace.
         //println["Results are" + results]

      if fullReplace
         return fullReplace[flatten[results], true, debug]
      else
         return flatten[results]
   }

   
   // Once all equations have been solved, this function will redistribute all
   // solutions through all equations.
   fullReplace[equations, allValues=false, debug=false] :=
   {
      // This is a dictionary with variable names as keys and a list of 
      // equations that are the solutions with those variables.  It will
      // replace the solutions member when we're complete.
      newSolutions = new dict
      solArray = new array

      for eq = equations
      {
         if debug
            println["Equation is: $eq"]
         left = getChild[eq, 0]
         right = getChild[eq, 1]
         leftStr = toString[left]
//         println["leftStr is $leftStr, type of leftStr is " + type[leftStr]]
         allSyms = getSymbols[right]
//         println["Symbols on right are " + allSyms]
//         println["ignoreSymbols are " + ignoreSymbols]
//         println["ignoreSymbolsStrings are " + ignoreSymbolsStrings]
         replaceSyms = setDifference[allSyms, ignoreSymbolsStrings]

         symList = toArray[replaceSyms]
         if debug
            println["symList is " + symList]
         stateList = new array
         solList = new array

         for sym = symList
         {
            stateList.push[new range[false, true]]
            
            // Make an array of right-hand-sides of solutions for that symbol
            solList2 = new array
            for sol = solutions@sym
            {
               solRight = child[sol, 1]
//               println["$sym, " + getSymbols[solRight]]
               if ! (getSymbols[solRight].contains[leftStr])
                  solList2.push[solRight]
//               else
//                  println["Rejected $sol because it contains $leftStr"]
            }
            solList.push[solList2]
         }

         numStates = length[symList]

         newEq = eq
         if debug
            println["Before replace: $newEq"]

         multifor states = stateList
         {
            tempSolArray = [eq]
            if debug
            {
               println["states is $states"]
               println["solList is $solList"]
            }

            newSolArray = new array
            for i = 0 to numStates-1
            {
               replaced = false
               if states@i == true
               {
                  for sol = solList@i
                  {
                     replaced = true
                     existingSols = length[tempSolArray]
                     
                     for k = 0 to existingSols-1
                     {
                        curreq = tempSolArray@k
                        curreq = substituteExpression[curreq, constructExpression["Symbol", symList@i], sol]
                        newSolArray.push[transformExpression[curreq]]
                     }
                  }
               }

               if replaced
                  tempSolArray = newSolArray
            }

            tempSolArray = eliminateOverconstrained[tempSolArray]
            if debug
               println["after replace:  $tempSolArray\n"]
            // TODO:  Change to set to avoid dups?
            solArray.pushAll[tempSolArray]
         }
      }
      
      println["Pre-Reduced is $solArray"]
      solArray = eliminateOverconstrained[solArray]
      
      if debug
         println["Reduced is $solArray"]

      for sol = solArray
      {
         key = toString[child[sol,0]]
         if newSolutions.containsKey[key]
            (newSolutions@key).push[sol]
         else
            newSolutions@key = [sol]
      }

      solutions = newSolutions
      
      return solArray
   }

   // Once all equations have been solved, this function will redistribute all
   // solutions through all equations.
   fullReplaceOld[equations, debug=false] :=
   {
      newSolutions = new dict
      res = deepCopy[equations]
      for eq = equations
      {
         if debug
            println["Working on equation $eq"]

         eqLeft = child[eq,0]
         eqRight = child[eq,1]
         lsymbols = getSymbols[eqLeft]
         lsymbolStr = array[lsymbols]@0  // Kludge to get symbol on left
         rsymbols = getSymbols[eqRight]
         syms = getSymbols[eqRight]  // Symbols on right-hand-side of equation
         syms = setDifference[syms, ignoreSymbolsStrings]
         symArray = array[syms]
         len = length[symArray]
         if debug
            println[" symArray is $symArray"]

         if len > 0
         {
            // Now go through and replace all permutations of variables.
            // This essentially counts through all the 2^n binary permutations
            states = new array 
            for i=0 to len-2
               states@i = false

            states@(len-1) = true

            i = len-1
            r = eq
            while i>=0
            {
               if (debug)
                  println["states is $states, $symArray, j is $j, symArray@j is " + symArray@j ]
               for repval = flatten[solveFor[symArray@j]]
               {
                  replaceVal = child[repval,1]
                  for j=0 to len-1
                     if states@j
                        if ! (getSymbols[replaceVal].contains[lsymbolStr])
                        {
                           if debug
                              println["    Replacing " + symArray@j + " with " + replaceVal + " from $repval in " + r]
                           r = transformExpression[replaceVar[r, symArray@j, replaceVal]]
                           if debug
                              println["     Solution is $r"]
                        }
               }

               // Eliminate x === x
               if ! structureEquals[eqLeft, child[r,1]]
                  res.push[r]

               if structureEquals[r, eq]
                  println["Expression $r unchanged."]
               else
                  println["Expression $eq changed to $r"]
  
               // Advance to next binary state
               flipped = false
               i = len-1
               while i>=0 and !flipped
               {
                  // Enter next state
                  if states@i == false
                  {
                     states@i = true
                     flipped = true
                  } else
                  {  // Carry
                     states@i = false
                     i = i - 1
                  }
               }
            }
         }
      }

      res = eliminateOverconstrained[res]
      println["\nAfter fullReplace:\n" + join["\n", res] + "\n"]

      return res
   }

   // Attempts to solve simultaneous equations in the system.
   solveSimultaneous[equationList, debug=true] :=
   {
      // Sort so that shortest equations are first.  This simplifies
      // substitutions and makes resulting equations shorter.
      // We can't do the latter because the sort function doesn't have
      // access to ignoreSymbols!  We have to create a closure to make this
      // happen?
      //      sort[equationList, {|a,b| length[setDifference[getSymbols[a],
      //      ignoreSymbols]] <=> length[setDifference[getSymbols[b],
      //      ignoreSymbols]]}]
      sort[equationList, {|a,b| length[getSymbols[a]] <=> length[getSymbols[b]]}]
      
      s = length[equationList]
      for i = 0 to s-2
      {
         ei = equationList@i
         ui = getSymbols[ei]
         aui = array[ui]
         lui = length[aui]
         
         JLOOP:
         for j = i+1 to s-1
         {
            ej = equationList@j
            uj = getSymbols[ej]

            commonSet = intersection[ui, uj]
            if length[commonSet] >= 1 or lui == 1
            {
               commonMinusIgnore = setDifference[commonSet, ignoreSymbols]
               if length[commonMinusIgnore] > 0
                  solveFor = sort[array[commonMinusIgnore]]@0
               else
                  solveFor = array[ui]@0
               
               eqSolved = solveSingle[ei, solveFor]

               if (debug)
                  println["Solving $ei for $solveFor, result is $eqSolved"]

               for part = flatten[eqSolved]
               {
                  // See if part is fully solved.
                  if structureEquals[_a === _b, part]
                  {
                     // Right-hand-side of solved equation
                     eqSolution = child[part,1]  

                     xSymbol = constructExpression["Symbol", solveFor]

                     // Substitute result into other equations.
                     for k=0 to s-1
                        if k != i     // Don't substitute into this equation.
                        {
                           orig = equationList@k
                           subst = substituteExpression[equationList@k, xSymbol, eqSolution]
                           if (orig != subst)
                           {
                              subst = transformExpression[subst]
                              // Only replace the equation if we actually simplified it.
                              if length[getSymbols[subst]] <= length[getSymbols[orig]]
                              {
                                 if debug
                                    println["Useful substitution: $part TO $subst"]
                                 equationList@k = subst
                              } else
                              if debug
                                 println["Unuseful substitution: $part TO $subst"]
                              }
                           }
                           break JLOOP
                  } else
                      println["Unsolved: $part"]
               }
            }
         }
      }
      return equationList
   }

   // This function eliminates overconstrained equations.  For example, a
   // system containing the solutions a===1/2 c r  and  a===c d^-1 r^2  is
   // overconstrained because a value can always be obtained with the first
   // equation.  The second is not necessary, and could lead to
   // inconsistent results.   This method ignores any symbols listed in the
   // ignoreSymbols list, (these are probably units,) eliminating them from
   // the equations.
   eliminateOverconstrained[eqArray, debug=false] :=
   {
      size = length[eqArray]
      unknowns = new array
      lefts = new array
      for i = 0 to size-1
      {
         lefts@i = child[eqArray@i, 0]
         unknowns@i = setDifference[getSymbols[child[eqArray@i,1]], ignoreSymbols]
      }

      res = new array

      // Check for duplicates.
      for i=0 to size-1
      {
         overconstrained = false
         j = 0
         do
         {
            if i != j and structureEquals[lefts@i, lefts@j]
            {
               // If the unknowns in j are a proper subset of the unknowns in
               // i, then it's probably a simpler form of the equation.  The
               // counterexample is with single solutions or degenerate cases
               // like 0 where we want to see all solutions.  For now, if the
               // right-hand-side is zero, we don't consider it overconstrained
               overconstrained = isProperSubset[unknowns@j, unknowns@i] && child[eqArray@j, 1] != 0 || (i<j and structureEquals[eqArray@i, eqArray@j])
               if (overconstrained)
                  if debug
                     println[eqArray@j + " is a proper subset or match of " + eqArray@i]
            }
            j=j+1
         } while (j < size) and ! overconstrained

         if (! overconstrained)
            res.push[eqArray@i]  // If we got here, no j is a proper subset of i.
      }

      return res
   }
}

// Solve a single equation for the specified symbol.
solveSingle[eq, symbol] :=
{
   xSymbol = constructExpression["Symbol", symbol]
   // We could use this in symbolic mode, otherwise it warns.
   // solveEq = solve[eq, xSymbol]
   solveEq = constructExpression["FunctionCall", ["solve", eq, xSymbol]]
   
   return transformExpression[solveEq]
}

// Replace a variable in the specified equation with the specified value.
// You should probably wrap the "eq" and "value" in a noEval[] block if
// you're not passing them in from a variable.
replaceVar[eq, varName, value] :=
{
   sym = constructExpression["Symbol", varName]
   res = substituteExpression[eq, sym, value]
   return res
}

/** This filters a list, returning only positive results. */
positive[x] :=
{
   result = new array
   for sol = x
   {
      rs = realSignum[sol]
      if rs != undef && rs > 0
         result.push[sol]
   }
   return result      
}

"systemSolver2.frink included OK"