// Functions for taking roots of a number using Newton's method to arbitrary
// precision using the root[] function.
//
// This function also performs integer and rational exponents to arbitrary
// precision using its pow[] function.
//
// These methods shouldn't be used if you're already working with imprecise
// floating-point numbers.
//
// Newton's method converges quadratically if you have a good guess.
// It can be horribly slow if you don't start with a good guess.
//
// This uses a few tricks:
//  * Uses floating-point math (when possible) for a good initial guess.
//  * Dynamically increases the working precision as estimates get better.


// Convenience function to do square root to current working precision.
sqrt[n] := sqrt[n, getPrecision[]]

// Convenience method for sqrt.
sqrt[n, digits, debug=false] := root[n, 2, digits, debug]

// Convenience function to do power to current working precision.
pow[n, p] := pow[n, p, getPrecision[]]

// Arbitrary-precision power
pow[n, p, digits, debug=false] :=sssssssssssssssssssssssssssssssssssss
{
   if debug
      println["Entering pow $n $p $digits"]
   
   if p == 0
      return 1
   else
   {
      origPrec = getPrecision[]
      try
      {
         setPrecision[digits+4]
         return root[n, 1/p, digits, debug]
      }
      finally
         setPrecision[origPrec]
   }
}

// Convenience function to do root to current working precision.
root[n, p] := root[n, p, getPrecision[]]

// General arbitrary-precision root finder
// n is the number, p is the root (e.g. 2 for square root, 3 for cube root)
root[n, p, digits, debug=false] :=
{
   if debug
      println["in root[$n, $p]"]
   if p == 1
      return n
   if p == 0
      return 1
   if p < 0
   {
      origPrec = getPrecision[]
      try
      {
         setPrecision[digits+4]
         return 1/root[n, -p]
      }
      finally
         setPrecision[origPrec]
   }
   
   alter = false
   if n<0
   {
      if p == 2
      {
         alter=true
         factor = i
         n = negate[n]
      } else
        if p mod 2 == 1         // Negative base, odd power
        {
           alter = true
           factor = -1
           n = negate[n]        // Negative base, even power
        } else
        {
           println["Arbitrary-precision root cannot produce complex numbers.  Arguments were root[$n, $p, $digits]"]
           return undef
        }
   }

   origPrec = getPrecision[]
   try
   {
      // Handle exact rational numbers
      if isRational[p]
      {
         prec = getPrecision[]
         setPrecision[digits+3]

         // TODO:  This needs to use arbitary-precision powers!
         // We can't apparently replace it with call to pow[] because
         // the stack never terminates.
         retval = root[n, numerator[p], digits, debug]^denominator[p]
         if alter
            retval = retval * factor
         setPrecision[prec]
         return retval
      }
      
      prec = getPrecision[]
      err = 10.^-(ceil[digits]+1)

      // Initial guess
      setPrecision[5]
      if (n<1e+308 and n>1e-323)
      {
         x = n^(1/p)               // Use floating-point if we can for guess
         err = err * x / 10.
         scale = approxLog2[x] / 3.219   // Approx. log 10
      } else
      {
         x = 2.^(approxLog2[n]/p)  // Dumb guess; could use introot function below
         err = err * x
         scale = approxLog2[x] / 3.219  // Approx. log 10
      }

      if scale == 0
         scale = 1
      
      if (debug)
      {
         println["First guess: $x"]
         println["Err is: $err"]
         println["Scale is: $scale"]
      }

      newWorkingPrecision = ceil[min[30,digits+3]]
      if newWorkingPrecision < 30
         newWorkingPrecision = 30
      workingPrecision = 15
      diff = abs[x]

      scaledErr = abs[err*scale]
      
      while (workingPrecision < digits+p) || (abs[diff] > scaledErr)
      {
         workingPrecision = newWorkingPrecision
         if debug
            println["precision is $workingPrecision"]
         setPrecision[workingPrecision]
         oldx = x
         if p==2
         {
            x = (x + n / x) / 2
            diff = oldx - x
         } else
         {
            // TODO:  This needs to use arbitary-precision powers!
            // We can't apparently replace it with call to pow[] because
            // the stack never terminates.
            errterm = (x^p - n)/ (p x^(p-1))
            x = x - errterm
            diff = errterm
         }

         if debug
         {
            println["x is $x"]
            println["diff is $diff"]
         }
         // This slightly more than doubles  working digits.
         setPrecision[10]
         if diff == 0
            goodDigits = workingPrecision * 2
         else
         {
            if debug
               println["approxLog2 is " + approxLog2[abs[diff]]]
            
            goodDigits = -approxLog2[abs[diff]]/3.219+scale
         }

         if debug
            println["Good digits: $goodDigits"]
         if (goodDigits < 30)
            goodDigits = 30
         newWorkingPrecision = min[ceil[digits+p+1], ceil[goodDigits*2.1]] 
      }

      if (debug)
         println["Final diff is $diff"]

      if alter
      {
         setPrecision[digits+p+1]
         retval = factor * x
      } else
      retval = x
      
      return retval
   }
   finally
      setPrecision[origPrec]
}

// Inverse square root.  1/sqrt[n].
inverseSquareRoot[n, digits, debug=false] :=
{
   alter = false
   if n<0
   {
      alter=true
      factor = i
      n = -n
   }

   origPrec = getPrecision[]
   try
   {
      prec = getPrecision[]
      err = 10.^-(ceil[digits]+1)

      // Initial guess
      setPrecision[5]
      if (n<1e+308 and n>1e-323)
      {
         x = 1/(n^(1/2))           // Use floating-point if we can for guess
         err = err * x / 10.
         scale = approxLog2[x] / 3.219   // Approx. log 10
      } else
      {
         x = 1/(2.^(approxLog2[n]/2))// Dumb guess; could use introot function below
         err = err * x
         scale = approxLog2[x] / 3.219  // Approx. log 10
      }

      if scale == 0
         scale = 1
      
      if (debug)
      {
         println["First guess: $x"]
         println["Err is: $err"]
         println["Scale is: $scale"]
      }

      newWorkingPrecision = ceil[min[30,digits+3]]
      if newWorkingPrecision < 30
         newWorkingPrecision = 30
      workingPrecision = 15
      diff = abs[x]

      scaledErr = abs[err*scale]
      
      while (workingPrecision < digits+2) || (diff > scaledErr)
      {
         workingPrecision = newWorkingPrecision
         if debug
            println["precision is $workingPrecision"]
         setPrecision[workingPrecision]
         oldx = x

         diff = 1 - n * x^2
         x = x + diff * x/2

         if debug
         {
            println["x is $x"]
            println["diff is $diff"]
         }
         // This slightly more than doubles  working digits.
         setPrecision[10]
         if diff == 0
            goodDigits = workingPrecision * 2
         else
            goodDigits = -approxLog2[abs[diff]]/3.219+scale
         if debug
            println["Good digits: $goodDigits"]
         if (goodDigits < 30)
            goodDigits = 30
         newWorkingPrecision = min[ceil[digits+3], ceil[goodDigits*1.8]] 
      }

      if (debug)
         println["Final diff is $diff"]

      if alter
      {
         setPrecision[digits+2+1]
         retval = factor * x
      } else
      retval = x
      
      return retval
   }
   finally
      setPrecision[origPrec]
}


// Integer square root--returns the greatest integer less than or equal to the 
// to the square root of n.
// This is Exercise 5.7 in Bressoud with my own modifications for better
// initial guess.
introot[n] :=
{
   a = 2^((bitLength[n]+1) div 2)
   //a = 2^((ceil[approxLog2[n]+1]) div 2)
   b = a - 1

   while b<a
   {
//      println["$a $b"]
      a = b
      b = (a*a + n) div (2*a)
   }

   return a
}