-- SPDX-FileCopyrightText: 2021 Oxhead Alpha
--
-- SPDX-License-Identifier: LicenseRef-MIT-OA

{-# OPTIONS_GHC -Wno-orphans #-}

module Lorentz.CustomArith.FixedArith
  ( -- * Lorentz instructions
    castNFixedToFixed
  , castFixedToNFixed
  , unsafeCastFixedToNFixed
    -- * Typeclasses
  , LorentzFixedCast (..)
  , Fixed (..)
  , NFixed (..)

  -- * Support types and functions
  , DecBase (..)
  , BinBase (..)
  , resolution_

  ) where

import Data.Fixed (Fixed(..), HasResolution(..))
import GHC.Num qualified (fromInteger)
import GHC.TypeLits qualified as Lit
import Prelude hiding (and, compare, drop, natVal, some, swap)
import Prelude qualified as P
import Text.Show qualified

import Lorentz.Arith
import Lorentz.Base
import Lorentz.Coercions
import Lorentz.Constraints.Scopes
import Lorentz.CustomArith.Common
import Lorentz.Errors
import Lorentz.Instr
import Lorentz.Macro
import Lorentz.Value
import Morley.Michelson.Typed

import Unsafe qualified

{-# ANN module ("HLint: ignore Use 'natVal' from Universum" :: Text) #-}

-- | Datatypes, representing base of the fixed-point values
data DecBase p where
  DecBase :: KnownNat p => DecBase p
data BinBase p where
  BinBase :: KnownNat p => BinBase p

instance KnownNat p => HasResolution (DecBase p) where
  resolution _ = 10 ^ (Lit.natVal (Proxy @p))

instance KnownNat p => HasResolution (BinBase p) where
  resolution _ = 2 ^ (Lit.natVal (Proxy @p))

-- | Special function to get resolution without argument
resolution_ :: forall a. HasResolution a => Natural
resolution_ =
  let r = resolution (Proxy @a)
    in if r <= 0
       then error "Lorentz Rationals support only positive resolutions"
       else Unsafe.fromIntegral @Integer @Natural r


-- | Like @Fixed@ but with a @Natural@ value inside constructor
newtype NFixed p = MkNFixed Natural deriving stock (Eq, Ord)

convertNFixedToFixed :: NFixed a -> Fixed a
convertNFixedToFixed (MkNFixed a) = MkFixed (Unsafe.fromIntegral @Natural @Integer a)

instance (HasResolution a) => Show (NFixed a) where
  show = show . convertNFixedToFixed

-- Note: This instances are copies of those in Data.Fixed for Fixed datatype
instance (HasResolution a) => Num (NFixed a) where
  (MkNFixed a) + (MkNFixed b) = MkNFixed (a + b)
  (MkNFixed a) - (MkNFixed b) = MkNFixed (a - b)
  fa@(MkNFixed a) * (MkNFixed b) = MkNFixed (P.div (a * b) (fromInteger (resolution fa)))
  negate (MkNFixed a) = MkNFixed (negate a)
  abs = id
  signum (MkNFixed a) = MkNFixed (signum a)
  fromInteger i = withResolution (\res -> MkNFixed ((fromInteger i) * res))

instance (HasResolution a) => Fractional (NFixed a) where
  fa@(MkNFixed a) / (MkNFixed b) = MkNFixed (P.div (a * (fromInteger (resolution fa))) b)
  recip fa@(MkNFixed a) = MkNFixed (P.div (res * res) a) where
      res = fromInteger $ resolution fa
  fromRational r = withResolution (\res -> MkNFixed (floor (r * (toRational res))))

instance IsoValue (NFixed p) where
  type ToT (NFixed p) = 'TNat
  toVal (MkNFixed x) = VNat x
  fromVal (VNat x) = MkNFixed x

instance Unwrappable (NFixed a) where
  type Unwrappabled (NFixed a) = Natural

-- Helpers copied from Data.Fixed, because they are not exported from there
withResolution :: forall a f. (HasResolution a) => (Natural -> f a) -> f a
withResolution foo = foo . fromInteger . resolution $ Proxy @a

------------------------------------------------------------------------------
-- Arithmetic operations' Insances
------------------------------------------------------------------------------

instance (r ~ (Fixed p)) => ArithOpHs Add (Fixed p) (Fixed p) r
instance (r ~ (Fixed p)) => ArithOpHs Add (Fixed p) Integer r
instance (r ~ (Fixed p)) => ArithOpHs Add (Fixed p) Natural r
instance (r ~ (Fixed p)) => ArithOpHs Add Integer (Fixed p) r
instance (r ~ (Fixed p)) => ArithOpHs Add Natural (Fixed p) r

instance (r ~ (NFixed p)) => ArithOpHs Add (NFixed p) (NFixed p) r
instance (r ~ (Fixed p)) => ArithOpHs Add (NFixed p) Integer r
instance (r ~ (NFixed p)) => ArithOpHs Add (NFixed p) Natural r
instance (r ~ (Fixed p)) => ArithOpHs Add Integer (NFixed p) r
instance (r ~ (NFixed p)) => ArithOpHs Add Natural (NFixed p) r

instance (r ~ (Fixed p)) => ArithOpHs Add (Fixed p) (NFixed p) r
instance (r ~ (Fixed p)) => ArithOpHs Add (NFixed p) (Fixed p) r

instance (r ~ (Fixed p)) => ArithOpHs Sub (Fixed p) (Fixed p) r
instance (r ~ (Fixed p)) => ArithOpHs Sub (Fixed p) Integer r
instance (r ~ (Fixed p)) => ArithOpHs Sub (Fixed p) Natural r
instance (r ~ (Fixed p)) => ArithOpHs Sub Integer (Fixed p) r
instance (r ~ (Fixed p)) => ArithOpHs Sub Natural (Fixed p) r

instance (r ~ (Fixed p)) => ArithOpHs Sub (NFixed p) (NFixed p) r
instance (r ~ (Fixed p)) => ArithOpHs Sub (NFixed p) Integer r
instance (r ~ (Fixed p)) => ArithOpHs Sub (NFixed p) Natural r
instance (r ~ (Fixed p)) => ArithOpHs Sub Integer (NFixed p) r
instance (r ~ (Fixed p)) => ArithOpHs Sub Natural (NFixed p) r

instance (r ~ (Fixed p)) => ArithOpHs Sub (Fixed p) (NFixed p) r
instance (r ~ (Fixed p)) => ArithOpHs Sub (NFixed p) (Fixed p) r


instance (r ~ (a Lit.+ b)) => ArithOpHs Mul (Fixed (DecBase a)) (Fixed (DecBase b)) (Fixed (DecBase r))
instance (r ~ (a Lit.+ b)) => ArithOpHs Mul (Fixed (BinBase a)) (Fixed (BinBase b)) (Fixed (BinBase r))
instance (r ~ (Fixed p)) => ArithOpHs Mul (Fixed p) Integer r
instance (r ~ (Fixed p)) => ArithOpHs Mul (Fixed p) Natural r
instance (r ~ (Fixed p)) => ArithOpHs Mul Integer (Fixed p) r
instance (r ~ (Fixed p)) => ArithOpHs Mul Natural (Fixed p) r

instance (r ~ (a Lit.+ b)) => ArithOpHs Mul (NFixed (DecBase a)) (NFixed (DecBase b)) (NFixed (DecBase r))
instance (r ~ (a Lit.+ b)) => ArithOpHs Mul (NFixed (BinBase a)) (NFixed (BinBase b)) (NFixed (BinBase r))
instance (r ~ (Fixed p)) => ArithOpHs Mul (NFixed p) Integer r
instance (r ~ (NFixed p)) => ArithOpHs Mul (NFixed p) Natural r
instance (r ~ (Fixed p)) => ArithOpHs Mul Integer (NFixed p) r
instance (r ~ (NFixed p)) => ArithOpHs Mul Natural (NFixed p) r

instance (r ~ (a Lit.+ b)) => ArithOpHs Mul (NFixed (DecBase a)) (Fixed (DecBase b)) (Fixed (DecBase r))
instance (r ~ (a Lit.+ b)) => ArithOpHs Mul (Fixed (DecBase a)) (NFixed (DecBase b)) (Fixed (DecBase r))
instance (r ~ (a Lit.+ b)) => ArithOpHs Mul (NFixed (BinBase a)) (Fixed (BinBase b)) (Fixed (BinBase r))
instance (r ~ (a Lit.+ b)) => ArithOpHs Mul (Fixed (BinBase a)) (NFixed (BinBase b)) (Fixed (BinBase r))

instance (r ~ (NFixed (BinBase b))) => ArithOpHs Lsl (NFixed (BinBase a)) Natural r
instance (r ~ (NFixed (BinBase b))) => ArithOpHs Lsr (NFixed (BinBase a)) Natural r

instance UnaryArithOpHs Neg (Fixed p) where
  type UnaryArithResHs Neg (Fixed p) = (Fixed p)
instance UnaryArithOpHs Neg (NFixed p) where
  type UnaryArithResHs Neg (NFixed p) = (Fixed p)

instance ToIntegerArithOpHs (NFixed a)

-- | Round is implemented using "banker's rounding" strategy, rounding half-way values
-- towards nearest even value
instance (KnownNat a, KnownNat b) => LorentzRounding (Fixed (DecBase a)) (Fixed (DecBase b)) where
  round_ = roundingHelper @a @b Round 10
  ceil_ = roundingHelper @a @b Ceil 10
  floor_ = roundingHelper @a @b Floor 10

instance (KnownNat a, KnownNat b) => LorentzRounding (Fixed (BinBase a)) (Fixed (BinBase b)) where
  round_ = roundingHelper @a @b Round 2
  ceil_ = roundingHelper @a @b Ceil 2
  floor_ = roundingHelper @a @b Floor 2

instance (LorentzRounding (Fixed (DecBase a)) (Fixed (DecBase b)))
  => LorentzRounding (NFixed (DecBase a)) (NFixed (DecBase b)) where
  round_ = castNFixedToFixed # round_ # unsafeCastFixedToNFixed
  ceil_ = castNFixedToFixed # ceil_ # unsafeCastFixedToNFixed
  floor_ = castNFixedToFixed # floor_ # unsafeCastFixedToNFixed

instance (KnownNat a, KnownNat b)
  => LorentzRounding (NFixed (BinBase a)) (NFixed (BinBase b)) where
  round_ = castNFixedToFixed # round_ # unsafeCastFixedToNFixed
  ceil_ = castNFixedToFixed # ceil_ # unsafeCastFixedToNFixed
  floor_ = castNFixedToFixed # floor_ # unsafeCastFixedToNFixed

-- | Class that allows casting @Fixed@ values to Integer in vice versa
class LorentzFixedCast a where
  fromFixed :: a : s :-> Integer : s
  toFixed :: Integer : s :-> a : s

instance (KnownNat a) => LorentzFixedCast (Fixed (DecBase a)) where
  fromFixed :: (Fixed (DecBase a)) : s :-> Integer : s
  fromFixed = (round_ @(Fixed (DecBase a)) @(Fixed (DecBase 0))) # forcedCoerce_
  toFixed = let pow :: Integer = 10 ^ Lit.natVal (Proxy @a) in push pow # mul # forcedCoerce_ @Integer @(Fixed (DecBase a))

instance (KnownNat a) => LorentzFixedCast (Fixed (BinBase a)) where
  fromFixed :: (Fixed (BinBase a)) : s :-> Integer : s
  fromFixed = (round_ @(Fixed (BinBase a)) @(Fixed (BinBase 0))) # forcedCoerce_
  toFixed = let pow :: Integer = 2 ^ Lit.natVal (Proxy @a) in push pow # mul # forcedCoerce_ @Integer @(Fixed (BinBase a))

instance (LorentzFixedCast (Fixed a)) => LorentzFixedCast (NFixed a) where
  fromFixed = castNFixedToFixed # fromFixed
  toFixed = toFixed # castFixedToNFixed # assertSome [mt|casting negative value to NFixed|]

instance ( KnownNat a
         , KnownNat b
         , KnownNat r
         ) => ArithOpHs Div (Fixed (DecBase a)) (Fixed (DecBase b)) (Maybe (Fixed (DecBase r))) where
  evalArithOpHs = fixedDivHelper @a @b @r 10

instance ( KnownNat a
         , KnownNat b
         , KnownNat r
         ) => ArithOpHs Div (Fixed (BinBase a)) (Fixed (BinBase b)) (Maybe (Fixed (BinBase r))) where
  evalArithOpHs = fixedDivHelper @a @b @r 2

instance ( KnownNat a
         , KnownNat b
         , KnownNat r
         ) => ArithOpHs Div (NFixed (DecBase a)) (NFixed (DecBase b)) (Maybe (NFixed (DecBase r))) where
  evalArithOpHs = fixedDivHelper @a @b @r 10

instance ( KnownNat a
         , KnownNat b
         , KnownNat r
         ) => ArithOpHs Div (NFixed (BinBase a)) (NFixed (BinBase b)) (Maybe (NFixed (BinBase r))) where
  evalArithOpHs = fixedDivHelper @a @b @r 2


castNFixedToFixed :: NFixed p : s :-> Fixed p : s
castNFixedToFixed = int # forcedCoerce_

castFixedToNFixed :: Fixed p : s :-> Maybe (NFixed p) : s
castFixedToNFixed = coerceUnwrap # isNat # forcedCoerce_

unsafeCastFixedToNFixed ::  Fixed p : s :-> NFixed p : s
unsafeCastFixedToNFixed = coerceUnwrap # Lorentz.Instr.abs # forcedCoerce_

instance (r ~ Maybe (Integer, NFixed (DecBase a)), KnownNat a) => ArithOpHs EDiv (Fixed (DecBase a)) Integer r where
  evalArithOpHs = edivHelper @a @Integer 10

instance (r ~ Maybe (Integer, NFixed (DecBase a)), KnownNat a) => ArithOpHs EDiv (Fixed (DecBase a)) Natural r where
  evalArithOpHs = edivHelper @a @Integer 10

instance (r ~ Maybe (Integer, NFixed (BinBase a)), KnownNat a) => ArithOpHs EDiv (Fixed (BinBase a)) Integer r where
  evalArithOpHs = edivHelper @a @Natural 2

instance (r ~ Maybe (Integer, NFixed (BinBase a)), KnownNat a) => ArithOpHs EDiv (Fixed (BinBase a)) Natural r where
  evalArithOpHs = edivHelper @a @Natural 2

instance (r ~ Maybe (Integer, NFixed (DecBase a)), KnownNat a) => ArithOpHs EDiv (NFixed (DecBase a)) Integer r where
  evalArithOpHs = edivHelper @a @Integer 10

instance (r ~ Maybe (Natural, NFixed (DecBase a)), KnownNat a) => ArithOpHs EDiv (NFixed (DecBase a)) Natural r where
  evalArithOpHs = edivHelper @a @Natural 10

instance (r ~ Maybe (Integer, NFixed (BinBase a)), KnownNat a) => ArithOpHs EDiv (NFixed (BinBase a)) Integer r where
  evalArithOpHs = edivHelper @a @Natural 2

instance (r ~ Maybe (Natural, NFixed (BinBase a)), KnownNat a) => ArithOpHs EDiv (NFixed (BinBase a)) Natural r where
  evalArithOpHs = edivHelper @a @Natural 2

----------------------------------------------------------------------------
-- Helpers
----------------------------------------------------------------------------

data RoundingPattern = Round | Ceil | Floor

-- Note: Round is implemented using "banker's rounding" strategy, rounding half-way values
-- towards the nearest even value.
roundingHelper
  :: forall a b r1 r2 s.
    ( KnownNat a, KnownNat b
    , FailOnTicketFound (ContainsTicket (ToT (Unwrappabled r1)))
    , MichelsonCoercible r1 r2
    , Typeable (Unwrappabled r2)
    , IsoValue (Unwrappabled r2)
    , SingI (ToT (Unwrappabled r1))
    , Unwrappable r2
    , Unwrappable r1
    , ArithOpHs Add Integer (Unwrappabled r2) (Unwrappabled r2)
    , ArithOpHs Add Natural (Unwrappabled r2) (Unwrappabled r2)
    , ArithOpHs And (Unwrappabled r2) Natural Natural
    , ArithOpHs EDiv (Unwrappabled r1) Integer (Maybe (Unwrappabled r2, Natural))
    , ArithOpHs Mul Integer r1 r1
    )
  => RoundingPattern -> Natural -> (r1 : s :-> r2 : s)
roundingHelper rp base =
  let halfBase :: Natural = base `P.div` 2
      powDifference :: Integer = (Lit.natVal (Proxy @b)) - (Lit.natVal (Proxy @a))
      newPow :: Integer = fromIntegral $ 2 * halfNewPow
      halfNewPow :: Natural = halfBase * (base ^ (Prelude.abs powDifference - 1))
   in case () of
        _ | powDifference == 0 -> (forcedCoerce_ :: (r1 : s :-> r2 : s))
          | powDifference > 0 ->
              push newPow # mul # (forcedCoerce_ :: (r1 : s :-> r2 : s))
          | otherwise ->
              push newPow #
              swap #
              coerceUnwrap # ediv #
              assertSome (Impossible @"Division by zero impossible here") #
              case rp of
                Round ->
                  unpair #
                  swap #
                  push halfNewPow #
                  compare #
                  dup #
                  ifGe0 drop (
                    -- rem >= halfNewPow
                    dip (push (1 :: Natural)) #
                    -- if rem == halfNewPow, check if quot is odd
                    ifEq0 (dupN @2 # and) nop #
                    -- if quot is odd or rem > halfNewPow, add 1 to quot.
                    add
                   )
                Ceil ->
                  unpair #
                  swap #
                  ifNeq0 (push (1 :: Integer) # add) nop
                Floor -> car
              # unsafeCoerceWrap

fixedDivHelper
  :: forall t1 t2 t3 base b1 b2 a r s.
     ( KnownNat t1, KnownNat t2, KnownNat t3
     , Num base, NiceConstant base
     , Unwrappable r, Unwrappable b1, Unwrappable a
     , Unwrappabled a ~ base
     , ArithOpHs Mul base r r
     , ArithOpHs EDiv (Unwrappabled r) (Unwrappabled b1) (Maybe (base, b2))
     , IsoValue a, Typeable a
     )
  => base -> (r : b1 : s) :-> (Maybe a : s)
fixedDivHelper base =
  let powA = Lit.natVal (Proxy @t1)
      powB = Lit.natVal (Proxy @t2)
      powR = base ^ (Lit.natVal (Proxy @t3) - 1)
   in
     push powB #
     push powA #
     ifGt (push powR # mul) nop #
     dip coerceUnwrap #
     coerceUnwrap #
     ediv @_ @_ @(Maybe (base, b2)) #
     ifNone (none) (car # unsafeCoerceWrap # some)

edivHelper
  :: forall a base x y r1 r2 m s.
     ( KnownNat a, Num base, NiceConstant base, KnownValue r1, KnownValue r2, Dupable m
     , UnaryArithOpHs Eq' m, ArithOpHs Mul base y m
     , ArithOpHs EDiv (Unwrappabled x) m (Maybe (r1, Unwrappabled r2))
     , Unwrappable x, Unwrappable r2, UnaryArithResHs Eq' m ~ Bool
     )
  => base -> (x : y : s) :-> (Maybe (r1, r2) : s)
edivHelper base =
  let powA = base ^ (Lit.natVal (Proxy @a))
   in
     swap  #
     push powA #
     mul #
     dup #
     ifEq0 @m (dropN @2 # none)
     (
       swap #
       coerceUnwrap #
       ediv #
       assertSome (Impossible @"Division by zero impossible here") #
       unpair #
       dip unsafeCoerceWrap #
       pair #
       some
     )