-- 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
  , Fixed (..)
  , NFixed (..)

  -- * Support types and functions
  , LorentzFixedBaseKind
  , DecBase
  , BinBase
  , resolution_
  , toFixed
  , fromFixed
  , LorentzFixedBase

  -- * Internals
  , getBase

  ) where

import Data.Fixed (Fixed(..), HasResolution(..))
import Data.Ratio ((%))
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) #-}

data LorentzFixedBaseKindTag

-- | Open kind for fixed value bases.
type LorentzFixedBaseKind = LorentzFixedBaseKindTag -> Type

-- | Represents binary base of the Lorentz fixed-point values
data BinBase :: Lit.Nat -> LorentzFixedBaseKind

-- | Represents decimal base of the Lorentz fixed-point values
data DecBase :: Lit.Nat -> LorentzFixedBaseKind

type LorentzFixedBase :: (Lit.Nat -> LorentzFixedBaseKind) -> Constraint
class Typeable a => LorentzFixedBase a where
  getBase :: Num b => b

instance LorentzFixedBase DecBase where
  getBase = 10
instance LorentzFixedBase BinBase where
  getBase = 2

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

instance KnownNat p => HasResolution (BinBase p) where
  resolution _ = getBase @BinBase ^ (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 (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 (HasResolution a) => Real (NFixed a) where
  toRational (MkNFixed x) = fromIntegral x % resolution (Proxy @a)

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' Instances
------------------------------------------------------------------------------

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 ~ Fixed (b1 (a Lit.+ b)), b1 ~ b2) => ArithOpHs Mul (Fixed (b1 a)) (Fixed (b2 b)) 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 ~ NFixed (b1 (a Lit.+ b)), b1 ~ b2) => ArithOpHs Mul (NFixed (b1 a)) (NFixed (b2 b)) 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 ~ Fixed (b1 (a Lit.+ b)), b1 ~ b2) => ArithOpHs Mul (Fixed (b1 a)) (NFixed (b2 b)) r
instance (r ~ Fixed (b1 (a Lit.+ b)), b1 ~ b2) => ArithOpHs Mul (NFixed (b1 a)) (Fixed (b2 b)) r

instance (r ~ (NFixed (BinBase a))) => ArithOpHs Lsl (NFixed (BinBase a)) Natural r
instance (r ~ (NFixed (BinBase a))) => 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, b1 ~ b2, LorentzFixedBase b1)
  => LorentzRounding (Fixed (b1 a)) (Fixed (b2 b)) where
  round_ = roundingHelper Round
  ceil_ = roundingHelper Ceil
  floor_ = roundingHelper Floor

instance (KnownNat a, KnownNat b, b1 ~ b2, LorentzFixedBase b1)
  => LorentzRounding (NFixed (b1 a)) (NFixed (b2 b)) where
  round_ = roundingHelper Round
  ceil_ = roundingHelper Ceil
  floor_ = roundingHelper Floor

-- | Convert to the corresponding integral type by banker's rounding.
fromFixed
  :: forall a f base t s.
    ( a ~ f (base t)
    , ToT (f (base 0)) ~ ToT (Unwrappabled a)
    , LorentzRounding a (f (base 0))
    )
  => a : s :-> Unwrappabled a : s
fromFixed = round_ @_ @(f (base 0)) # forcedCoerce_

-- | Convert from the corresponding integral type.
toFixed
  :: forall a f base t s.
    ( a ~ f (base t), LorentzFixedBase base
    , Unwrappable a
    , KnownNat t
    , ArithOpHs Mul Natural (Unwrappabled a) (Unwrappabled a)
    )
  => Unwrappabled a : s :-> a : s
toFixed = rebase @base @t # forcedCoerce_ @(Unwrappabled a) @(f (base t))

type DivConstraint a b t r f b1 b2 =
  ( KnownNat b
  , b1 ~ b2
  , DivConstraint1 a t r f b1
  )

type DivConstraint1 a t r f base =
  ( KnownNat a
  , KnownNat t
  , LorentzFixedBase base
  , r ~ Maybe (f (base t))
  )

instance DivConstraint a b t r Fixed b1 b2 => ArithOpHs Div (Fixed (b1 a)) (Fixed (b2 b)) r where
  evalArithOpHs = fixedDivHelper

instance DivConstraint1 a t r Fixed b1 => ArithOpHs Div Integer (Fixed (b1 a)) r where
  evalArithOpHs = toFixed @_ @_ @_ @0 # fixedDivHelper

instance DivConstraint1 a t r Fixed b1 => ArithOpHs Div Natural (Fixed (b1 a)) r where
  evalArithOpHs = int # toFixed @_ @_ @_ @0 # fixedDivHelper

-- NFixed

instance DivConstraint a b t r NFixed b1 b2 => ArithOpHs Div (NFixed (b1 a)) (NFixed (b2 b)) r where
  evalArithOpHs = fixedDivHelper

instance DivConstraint1 a t r NFixed b1 => ArithOpHs Div Natural (NFixed (b1 a)) r where
  evalArithOpHs = toFixed @_ @_ @_ @0 # fixedDivHelper

instance DivConstraint1 a t r Fixed b1 => ArithOpHs Div Integer (NFixed (b1 a)) r where
  evalArithOpHs = toFixed @_ @_ @_ @0 # dip castNFixedToFixed # fixedDivHelper

type DivIntegralConstraint r b =
  ( KnownNat r
  , LorentzFixedBase b
  )

-- NB: Instances are parametrized by the result type by design since we also
-- have rational arithmetic.

instance DivIntegralConstraint r b => ArithOpHs Div Integer Integer (Maybe (Fixed (b r))) where
  evalArithOpHs = toFixed @(Fixed (b r)) # dip (toFixed @(Fixed (b 0))) # fixedDivHelper
instance DivIntegralConstraint r b => ArithOpHs Div Natural Natural (Maybe (Fixed (b r))) where
  evalArithOpHs = int # toFixed @(Fixed (b r)) # dip (int # toFixed @(Fixed (b 0))) # fixedDivHelper
instance DivIntegralConstraint r b => ArithOpHs Div Integer Natural (Maybe (Fixed (b r))) where
  evalArithOpHs = toFixed @(Fixed (b r)) # dip (int # toFixed @(Fixed (b 0))) # fixedDivHelper
instance DivIntegralConstraint r b => ArithOpHs Div Natural Integer (Maybe (Fixed (b r))) where
  evalArithOpHs = int # toFixed @(Fixed (b r)) # dip (toFixed @(Fixed (b 0))) # fixedDivHelper

instance DivIntegralConstraint r b => ArithOpHs Div Natural Natural (Maybe (NFixed (b r))) where
  evalArithOpHs = toFixed @(NFixed (b r)) # dip (toFixed @(NFixed (b r))) # fixedDivHelper

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 (base a)), KnownNat a, LorentzFixedBase base)
  => ArithOpHs EDiv (Fixed (base a)) Integer r where
  evalArithOpHs = edivHelper

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

instance (r ~ Maybe (Integer, NFixed (base a)), KnownNat a, LorentzFixedBase base)
  => ArithOpHs EDiv (NFixed (base a)) Integer r where
  evalArithOpHs = edivHelper

instance (r ~ Maybe (Natural, NFixed (base a)), KnownNat a, LorentzFixedBase base)
  => ArithOpHs EDiv (NFixed (base a)) Natural r where
  evalArithOpHs = edivHelper

----------------------------------------------------------------------------
-- 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 base f.
    ( KnownNat a, KnownNat b
    , ForbidTicket (ToT (Unwrappabled r1))
    , MichelsonCoercible r1 r2
    , SingI (ToT (Unwrappabled r1))
    , Unwrappable r2
    , Unwrappable r1
    , ArithOpHs Add Natural (Unwrappabled r2) (Unwrappabled r2)
    , ArithOpHs Add (Unwrappabled r2) (Unwrappabled r2) (Unwrappabled r2)
    , ArithOpHs And (Unwrappabled r2) Natural Natural
    , ArithOpHs EDiv (Unwrappabled r1) Natural (Maybe (Unwrappabled r2, Natural))
    , ArithOpHs Mul Natural r1 r1
    , LorentzFixedBase base
    , r1 ~ f (base a)
    , r2 ~ f (base b)
    , NiceConstant (Unwrappabled r2)
    , Num (Unwrappabled r2)
    )
  => RoundingPattern -> (r1 : s :-> r2 : s)
roundingHelper rp =
  let halfBase :: Natural = base `P.div` 2
      powDifference :: Integer = Lit.natVal (Proxy @b) - Lit.natVal (Proxy @a)
      newPow = 2 * halfNewPow
      halfNewPow :: Natural = halfBase * (base ^ (Prelude.abs powDifference - 1))
      base = getBase @base
   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 :: Unwrappabled r2) # add) nop
                Floor -> car
              # unsafeCoerceWrap

fixedDivHelper
  :: forall t1 t2 t3 base any s f x y r.
     ( x ~ f (base t1), y ~ f (base t2), r ~ f (base t3)
     , LorentzFixedBase base
     , Each '[Unwrappable] '[x, y, r]
     , Each '[KnownNat] '[t1, t2, t3]
     , ArithOpHs
          EDiv
          (Unwrappabled x)
          (Unwrappabled y)
          (Maybe (Unwrappabled r, any))
     , IsoValue r, Typeable f
     , ArithOpHs Mul Natural x x
     , ArithOpHs Mul Natural y y
     )
  => x : y : s :-> Maybe r : s
fixedDivHelper =
  adjust #
  dip coerceUnwrap #
  coerceUnwrap #
  ediv #
  Lorentz.Instr.map @(Maybe (Unwrappabled r, any)) (car # unsafeCoerceWrap)
  where
    powDifference :: Integer
    powDifference = Lit.natVal (Proxy @t2) + Lit.natVal (Proxy @t3) - Lit.natVal (Proxy @t1)
    -- see Note [fixedDivHelper] below for an explanation.
    multiplier :: Natural
    multiplier = getBase @base ^ P.abs powDifference
    adjust = case P.compare powDifference 0 of
      P.EQ -> nop
      P.GT -> push multiplier # mul
      P.LT -> dip $ push multiplier # mul

{-
Note [fixedDivHelper]
~~~~~~~~~~~~~~~~~~~~~

This code may look a little bit mysterious without an explanation, so here goes:

We want to compute r = x / y, where x = a β⁻ⁿ¹, y = b β⁻ⁿ², and r has to be
represented as some r = c β⁻ⁿ³. Ultimately we need to compute c. Thus:

c β⁻ⁿ³ = (a β⁻ⁿ¹) / (b β⁻ⁿ²)
c β⁻ⁿ³ = (a / b) β⁻ⁿ¹⁺ⁿ²
c = (a / b) β⁻ⁿ¹⁺ⁿ²⁺ⁿ³

Let ε = | -n₁ + n₂ + n₃ |
Let μ = βᵋ

(-n₁ + n₂ + n₃) can be either positive, negative, or zero. If it's zero, μ = 1,
and no futher action is necessary.
If it's positive, we need to multiply a by μ before division.
If it's negative, technically we could divide a by μ, but division is more
costly than multiplication, so instead we can multiply b by μ to obtain the same
result.

In the code above, β = getBase @base, μ = multiplier, and
(-n₁ + n₂ + n₃) = powDifference.

-- @lierdakil
-}

edivHelper
  :: forall a base x y r1 r2 s f.
     ( KnownNat a
     , ArithOpHs Mul Natural y y
     , ArithOpHs EDiv (Unwrappabled x) y (Maybe (r1, Unwrappabled r2))
     , Unwrappable x, Unwrappable r2
     , LorentzFixedBase base
     , x ~ f (base a)
     )
  => (x : y : s) :-> (Maybe (r1, r2) : s)
edivHelper =
  dip (rebase @base @a) #
  coerceUnwrap #
  ediv #
  forcedCoerce_ @(Maybe (r1, Unwrappabled r2)) @(Maybe (r1, r2))

rebase
  :: forall base (exp :: Lit.Nat) b s.
    (KnownNat exp, ArithOpHs Mul Natural b b, LorentzFixedBase base)
  => b : s :-> b : s
rebase = case getBase @base ^ Lit.natVal (Proxy @exp) :: Natural of
  1 -> nop
  pow -> push pow # mul