{-|
Module      : What4.Protocol.PolyRoot
Description : Representation for algebraic reals
Copyright   : (c) Galois Inc, 2016-2020
License     : BSD3
Maintainer  : jhendrix@galois.com

Defines a numeric data-type where each number is represented as the root of a
polynomial over a single variable.

This currently only defines operations for parsing the roots from the format
generated by Yices, and evaluating a polynomial over rational coefficients
to the rational derived from the closest double.
-}

{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE ScopedTypeVariables #-}
module What4.Protocol.PolyRoot
  ( Root
  , approximate
  , fromYicesText
  , parseYicesRoot
  ) where

import           Control.Applicative
import           Control.Lens
import qualified Data.Attoparsec.Text as Atto
import qualified Data.Map as Map
import           Data.Ratio
import           Data.Text (Text)
import qualified Data.Text as Text

import qualified Data.Vector as V
import           Prettyprinter as PP

atto_angle :: Atto.Parser a -> Atto.Parser a
atto_angle p = Atto.char '<' *> p <* Atto.char '>'

atto_paren :: Atto.Parser a -> Atto.Parser a
atto_paren p = Atto.char '(' *> p <* Atto.char ')'

-- | A polynomial with one variable.
newtype SingPoly coef = SingPoly (V.Vector coef)
  deriving (Functor, Foldable, Traversable, Show)

instance (Ord coef, Num coef, Pretty coef) => Pretty (SingPoly coef) where
  pretty (SingPoly v) =
    case V.findIndex (/= 0) v of
      Nothing -> pretty "0"
      Just j -> go (V.length v - 1)
        where ppc c | c < 0 = parens (pretty c)
                    | otherwise = pretty c

              ppi 1 = pretty "*x"
              ppi i = pretty "*x^" <> pretty i

              go 0 = ppc (v V.! 0)
              go i | seq i False = error "pretty SingPoly"
                   | i == j = ppc (v V.! i) <> ppi i
                   | v V.! i == 0 = go (i-1)
                   | otherwise = ppc (v V.! i) <> ppi i <+> pretty "+" <+> go (i-1)

fromList :: [c] -> SingPoly c
fromList = SingPoly . V.fromList

-- | Create a polyomial from a map from powers to coefficient.
fromMap :: (Eq c, Num c) => Map.Map Int c -> SingPoly c
fromMap m0 = SingPoly (V.generate (n+1) f)
  where m = Map.filter (/= 0) m0
        (n,_) = Map.findMax m
        f i   = Map.findWithDefault 0 i m

-- | Parse a positive monomial
pos_mono :: Integral c => Atto.Parser (c, Int)
pos_mono = (,) <$> Atto.decimal <*> times_x
  where times_x :: Atto.Parser Int
        times_x = (Atto.char '*' *> Atto.char 'x' *> expon) <|> pure 0

        -- Parse explicit exponent or return 1
        expon :: Atto.Parser Int
        expon = (Atto.char '^' *> Atto.decimal) <|> pure 1


-- | Parses a monomial and returns the coefficient and power
mono :: Integral c => Atto.Parser (c, Int)
mono = atto_paren (Atto.char '-' *> (over _1 negate <$> pos_mono))
     <|> pos_mono

parseYicesPoly :: Integral c => Atto.Parser (SingPoly c)
parseYicesPoly = do
     (c,p) <- mono
     go (Map.singleton p c)
  where go m = next m <|> pure (fromMap m)
        next m = seq m $ do
          _ <- Atto.char ' ' *> Atto.char '+' *> Atto.char ' '
          (c,p) <- mono
          go (Map.insertWith (+) p c m)


-- | Evaluate polynomial at a specific value.
--
-- Note that due to rounding, the result may not be exact when using
-- finite precision arithmetic.
eval :: forall c . Num c => SingPoly c -> c -> c
eval (SingPoly v) c = f 0 1 0
  where -- f takes an index, the current power, and the current sum.
        f :: Int -> c -> c -> c
        f i p s
          | seq p $ seq s $ False = error "internal error: Poly.eval"
          | i < V.length v = f (i+1) (p * c) (s + p * (v V.! i))
          | otherwise = s

data Root c = Root { rootPoly :: !(SingPoly c)
                   , rootLbound :: !c
                   , rootUbound :: !c
                   }
  deriving (Show)

-- | Construct a root from a rational constant
rootFromRational :: Num c => c -> Root c
rootFromRational r = Root { rootPoly = fromList [ negate r, 1 ]
                          , rootLbound = r
                          , rootUbound = r
                          }

instance (Ord c, Num c, Pretty c) => Pretty (Root c) where
  pretty (Root p l u) = langle <> pretty p <> comma <+> bounds <> rangle
    where bounds = parens (pretty l <> comma <+> pretty u)

-- | This either returns the root exactly, or it computes the closest double
-- precision approximation of the root.
--
-- Underneath the hood, this uses rational arithmetic to guarantee precision,
-- so this operation is relatively slow.  However, it is guaranteed to provide
-- an exact answer.
--
-- If performance is a concern, there are faster algorithms for computing this.
approximate :: Root Rational -> Rational
approximate r
    | l0 == u0       = l0
    | init_lval == 0 = l0
    | init_uval == 0 = u0
    | init_lval < 0 && init_uval > 0 = bisect (fromRational l0) (fromRational u0)
    | init_lval > 0 && init_uval < 0 = bisect (fromRational u0) (fromRational l0)
    | otherwise = error "Closest root given bad root."
  where p_rat = rootPoly r
        l0 = rootLbound r
        u0 = rootUbound r

        init_lval = eval p_rat l0
        init_uval = eval p_rat u0

        -- bisect takes a value that evaluates to a negative value under the 'p',
        -- and a value that evalautes to a positive value, and runs until it
        -- converges.
        bisect :: Double -> Double -> Rational
        bisect l u   -- Stop if mid point is at bound.
                   | m == l || m == u = toRational $
                      -- Pick whichever bound is cl oser to root.
                      if l_val <= u_val then l else u
                   | m_val == 0 = toRational m -- Stop if mid point is exact root.
                   | m_val <  0 = bisect m u -- Use mid point as new lower bound
                   | otherwise  = bisect l m -- Use mid point as new upper bound.
          where m = (l + u) / 2
                m_val = eval p_rat (toRational m)
                l_val = abs (eval p_rat (toRational l))
                u_val = abs (eval p_rat (toRational u))


atto_pair :: (a -> b -> r) -> Atto.Parser a -> Atto.Parser b -> Atto.Parser r
atto_pair f x y = f <$> x <*> (Atto.char ',' *> Atto.char ' ' *> y)

atto_sdecimal :: Integral c => Atto.Parser c
atto_sdecimal = Atto.char '-' *> (negate <$> Atto.decimal)
              <|> Atto.decimal

atto_rational :: Integral c => Atto.Parser (Ratio c)
atto_rational = (%) <$> atto_sdecimal <*> denom
  where denom = (Atto.char '/' *> Atto.decimal) <|> pure 1

parseYicesRoot :: Atto.Parser (Root Rational)
parseYicesRoot = atto_angle (atto_pair mkRoot (fmap fromInteger <$> parseYicesPoly) parseBounds)
             <|> (rootFromRational <$> atto_rational)
  where mkRoot :: SingPoly c -> (c, c) -> Root c
        mkRoot = uncurry . Root
        parseBounds :: Atto.Parser (Rational, Rational)
        parseBounds = atto_paren (atto_pair (,) atto_rational atto_rational)

-- | Convert text to a root
fromYicesText :: Text -> Maybe (Root Rational)
fromYicesText t = resolve (Atto.parse parseYicesRoot t)
  where resolve (Atto.Fail _rem _ _msg) = Nothing
        resolve (Atto.Partial f) =
          resolve (f Text.empty)
        resolve (Atto.Done i r)
          | Text.null i = Just $! r
          | otherwise = Nothing