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

{-# LANGUAGE InstanceSigs #-}
-- This option was enabled, because of
{-# OPTIONS_GHC -Wno-orphans #-}

module Lorentz.FixedArith
  ( -- * Lorentz instructions
    Lorentz.FixedArith.div
  , castNFixedToFixed
  , castFixedToNFixed

    -- * Lorentz casts
  , LorentzRounding (..)
  , LorentzFixedCast (..)

  ) where

import GHC.TypeLits qualified as Lit
import Prelude hiding (and, compare, drop, natVal, some, swap)
import Prelude qualified as P

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

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

-- | Class that enables support of rounding operations for Lorentz non-integer values
-- Note: Round is implemented using "banker's rounding" strategy, rounding half-way values
-- towards nearest even value
class LorentzRounding a b where
  round_ :: a : s :-> b : s
  ceil_ :: a : s :-> b : s
  floor_ :: a : s :-> b : s

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|]

-- | Since Michelson doesn't support divide operation, we will use our own to represent
-- divison of Fixed and Rational values
data Div

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


-- | Operation that represents division of two values with a given result
div
  :: forall r n m s. ArithOpHs Div n m r
  => n : m : s :-> r : s
div = evalArithOpHs @Div


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

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
     )