-- SPDX-FileCopyrightText: 2021 Tocqueville Group
--
-- SPDX-License-Identifier: LicenseRef-MIT-TQ

{-# LANGUAGE RebindableSyntax #-}
{-# LANGUAGE NoApplicativeDo #-}
{-# OPTIONS_GHC -Wno-unused-do-bind #-}
{-# OPTIONS_GHC -Wno-orphans #-}

module Lorentz.CustomArith.RationalArith
  ( Rational
  , NRational

  -- * Common functions
  , oppositeRational
  , gcdEuclid
  , euclidExtendedNormalization
  , reduce
  , tripleMul

    -- * Constructor functions for Rational
  , mkRational
  , (%!)

   -- * Constructor functions for NRational
  , mkNRational
  , (%%!)

    -- * Rational specific typeclasses
  , LorentzRational

    -- * Constructors
  , mkRational_
  , numerator
  , denominator
  , constructRational
  , deconstructRational
  , uncheckedPairToRational
  , unsafePairToRational
  , pairToRational
  ) where

import Prelude (Either(..), Eq, Integer, Maybe, Natural, Show, Text, show, unsafe, ($), (++), (==))
import Text.Show qualified

import Lorentz.Annotation
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.Rebinded
import Morley.Michelson.Text
import Morley.Michelson.Typed
import Morley.Util.Named
import Morley.Util.Type

newtype Rational =
  UnsafeRational (Integer, Natural)
  deriving newtype (Eq)

infixl 7 %!
(%!) :: Integer -> Natural -> Rational
(%!) x y = unsafe $ mkRational (x, y)

mkRational :: (Integer, Natural) -> Either Text Rational
mkRational (x, y) = case y of
  0 -> Left $ "Attempt to construct Rational with zero denominator."
  _ -> Right $ UnsafeRational (x, y)

instance Show Rational where
  show (UnsafeRational (a, b)) = show a ++ " / " ++ show b

instance HasAnnotation Rational where
  getAnnotation _ = starNotes

instance IsoValue Rational where
  type ToT Rational = 'TPair 'TInt 'TNat
  toVal (UnsafeRational (a, b)) = VPair (VInt a, VNat b)
  fromVal (VPair (VInt a, VNat b)) = UnsafeRational (a, b)

instance Unwrappable Rational where
  type Unwrappabled Rational = (Integer, Natural)

newtype NRational =
  UnsafeNRational (Natural, Natural)
  deriving newtype (Eq)

mkNRational :: (Natural, Natural) -> Either Text NRational
mkNRational (x, y) = case y of
  0 -> Left "Attempt to construct NRational with zero denominator."
  _ -> Right $ UnsafeNRational (x, y)

infixl 7 %%!
(%%!) :: Natural -> Natural -> NRational
(%%!) x y = unsafe $ mkNRational (x, y)

instance Show NRational where
  show (UnsafeNRational (a, b)) = show a ++ " / " ++ show b

instance HasAnnotation NRational where
  getAnnotation _ = starNotes

instance IsoValue NRational where
  type ToT NRational = 'TPair 'TNat 'TNat
  toVal (UnsafeNRational (a, b)) = VPair (VNat a, VNat b)
  fromVal (VPair (VNat a, VNat b)) = UnsafeNRational (a, b)

instance Unwrappable NRational where
  type Unwrappabled NRational = (Natural, Natural)

class Unwrappable r => LorentzRational r where
  numerator :: Unwrappabled r ~ (a, Natural) => r : s :-> a : s
  numerator = coerceUnwrap # car

  denominator :: Unwrappabled r ~ (a, Natural) => r : s :-> Natural : s
  denominator = coerceUnwrap # cdr

  deconstructRational :: Unwrappabled r ~ (a, Natural) => r : s :-> a : Natural : s
  deconstructRational = coerceUnwrap # unpair

  constructRational :: Unwrappabled r ~ (a, Natural) => a : Natural : s :-> r : s
  constructRational = pair # uncheckedPairToRational

  unsafePairToRational :: Unwrappabled r ~ (a, Natural) => (a, Natural) : s :-> r : s
  unsafePairToRational = unpair # dip (dup # assertNeq0 [mt|Attempt to construct rational with zero denominator|]) # pair # unsafeCoerceWrap

  uncheckedPairToRational :: Unwrappabled r ~ (a, Natural) => (a, Natural) : s :-> r : s
  uncheckedPairToRational = unsafeCoerceWrap

  pairToRational :: (Unwrappabled r ~ (a, Natural), KnownValue r) => (a, Natural) : s :-> Maybe r : s
  pairToRational = do
    unpair; swap
    dup
    ifEq0
      do
        dropN @2
        none
      do
        swap
        pair
        unsafeCoerceWrap
        some

  oppositeRational :: r : s :-> r : s

  reduce :: (KnownList s, KnownList (ToTs s)) => r : s :-> r : s

instance LorentzRational Rational where
  oppositeRational = do
    deconstructRational
    dup
    ifLt0 (dip neg) (dip int)
    abs
    swap
    constructRational
  reduce = reduceRationalHelper

instance LorentzRational NRational where
  oppositeRational = do
    deconstructRational
    swap
    constructRational
  reduce = reduceNRationalHelper

instance (r ~ Rational) => ArithOpHs Add Rational Integer r where
  evalArithOpHs = swap # add
instance (r ~ Rational) => ArithOpHs Add Integer Rational r where
  evalArithOpHs = addSubHelper add
instance (r ~ Rational) => ArithOpHs Add Rational Natural r where
  evalArithOpHs = swap # add
instance (r ~ Rational) => ArithOpHs Add Natural Rational r where
  evalArithOpHs = addSubHelper add
instance (r ~ Rational) => ArithOpHs Add Rational Rational r where
  evalArithOpHs = addSubRationalHelper add

instance (r ~ Rational) => ArithOpHs Sub Rational Integer r where
  evalArithOpHs = swap # sub # negRational
instance (r ~ Rational) => ArithOpHs Sub Integer Rational r where
  evalArithOpHs = addSubHelper sub
instance (r ~ Rational) => ArithOpHs Sub Rational Natural r where
  evalArithOpHs = swap # sub # negRational
instance (r ~ Rational) => ArithOpHs Sub Natural Rational r where
  evalArithOpHs = addSubHelper sub
instance (r ~ Rational) => ArithOpHs Sub Rational Rational r where
  evalArithOpHs = addSubRationalHelper sub

instance (r ~ Rational) => ArithOpHs Mul Rational Integer r where
  evalArithOpHs = swap # mul
instance (r ~ Rational) => ArithOpHs Mul Integer Rational r where
  evalArithOpHs = dip deconstructRational # mul # constructRational
instance (r ~ Rational) => ArithOpHs Mul Rational Natural r where
  evalArithOpHs = swap # mul
instance (r ~ Rational) => ArithOpHs Mul Natural Rational r where
  evalArithOpHs = dip deconstructRational # mul # constructRational

instance (r ~ Rational) => ArithOpHs Mul Rational Rational r where
  evalArithOpHs = mulRationalHelper @Rational @Rational @Rational

instance (r ~ Rational) => ArithOpHs Div Integer Rational r where
  evalArithOpHs = dip oppositeRational # mul
instance (r ~ Rational) => ArithOpHs Div Natural Rational r where
  evalArithOpHs = dip oppositeRational # mul
instance (r ~ Rational) => ArithOpHs Div Rational Integer r where
  evalArithOpHs = oppositeRational # mul # oppositeRational
instance (r ~ Rational) => ArithOpHs Div Rational Natural r where
  evalArithOpHs = oppositeRational # mul
instance (r ~ Rational) => ArithOpHs Div Rational Rational r where
  evalArithOpHs = dip oppositeRational # mul

instance (r ~ Rational) => ArithOpHs Add NRational Integer r where
  evalArithOpHs = swap # add
instance (r ~ Rational) => ArithOpHs Add Integer NRational r where
  evalArithOpHs = addSubHelper add
instance (r ~ NRational) => ArithOpHs Add NRational Natural r where
  evalArithOpHs = swap # add
instance (r ~ NRational) => ArithOpHs Add Natural NRational r where
  evalArithOpHs = addSubHelper add
instance (r ~ NRational) => ArithOpHs Add NRational NRational r where
  evalArithOpHs = addSubRationalHelper add

instance (r ~ Rational) => ArithOpHs Sub NRational Integer r where
  evalArithOpHs = swap # sub # negRational
instance (r ~ Rational) => ArithOpHs Sub Integer NRational r where
  evalArithOpHs = addSubHelper sub
instance (r ~ Rational) => ArithOpHs Sub NRational Natural r where
  evalArithOpHs = swap # sub # negRational
instance (r ~ Rational) => ArithOpHs Sub Natural NRational r where
  evalArithOpHs = addSubHelper sub
instance (r ~ Rational) => ArithOpHs Sub NRational NRational r where
  evalArithOpHs = addSubRationalHelper sub

instance (r ~ Rational) => ArithOpHs Mul NRational Integer r where
  evalArithOpHs = swap # mul
instance (r ~ Rational) => ArithOpHs Mul Integer NRational r where
  evalArithOpHs = dip deconstructRational # mul # constructRational
instance (r ~ NRational) => ArithOpHs Mul NRational Natural r where
  evalArithOpHs = swap # mul
instance (r ~ NRational) => ArithOpHs Mul Natural NRational r where
  evalArithOpHs = dip deconstructRational # mul # constructRational

instance (r ~ NRational) => ArithOpHs Mul NRational NRational r where
  evalArithOpHs = mulRationalHelper @NRational @NRational @NRational

instance (r ~ Rational) => ArithOpHs Div Integer NRational r where
  evalArithOpHs = dip oppositeRational # mul
instance (r ~ NRational) => ArithOpHs Div Natural NRational r where
  evalArithOpHs = dip oppositeRational # mul
instance (r ~ Rational) => ArithOpHs Div NRational Integer r where
  evalArithOpHs = oppositeRational # mul # oppositeRational
instance (r ~ NRational) => ArithOpHs Div NRational Natural r where
  evalArithOpHs = oppositeRational # mul
instance (r ~ NRational) => ArithOpHs Div NRational NRational r where
  evalArithOpHs = dip oppositeRational # mul

instance (r ~ Rational) => ArithOpHs Add Rational NRational r where
  evalArithOpHs = addSubRationalHelper add
instance (r ~ Rational) => ArithOpHs Add NRational Rational r where
  evalArithOpHs = addSubRationalHelper add
instance (r ~ Rational) => ArithOpHs Sub Rational NRational r where
  evalArithOpHs = addSubRationalHelper sub
instance (r ~ Rational) => ArithOpHs Sub NRational Rational r where
  evalArithOpHs = addSubRationalHelper sub
instance (r ~ Rational) => ArithOpHs Mul Rational NRational r where
  evalArithOpHs = mulRationalHelper @Rational @NRational @Rational
instance (r ~ Rational) => ArithOpHs Mul NRational Rational r where
  evalArithOpHs = mulRationalHelper @NRational @Rational @Rational
instance (r ~ Rational) => ArithOpHs Div Rational NRational r where
  evalArithOpHs = dip oppositeRational # mul
instance (r ~ Rational) => ArithOpHs Div NRational Rational r where
  evalArithOpHs = dip oppositeRational # mul


instance {-# OVERLAPPABLE #-} ( LorentzRational r
         , Unwrappabled r ~ (a, Natural)
         , ArithOpHs EDiv a Natural (Maybe (a, Natural))
         , ArithOpHs Add a Natural a
         , ArithOpHs Add Natural a a
         ) => LorentzRounding r a where
  round_ = do
    deconstructRational; dip dup
    ediv @a @Natural @(Maybe (a, Natural)); assertSome (Impossible @"Division by zero impossible here")
    unpair
    dip do
      push (2 :: Natural)
      mul
      ifGt (push (1 :: Natural)) (push (0 :: Natural))
    add
  ceil_ = do
    deconstructRational
    ediv @a @Natural @(Maybe (a, Natural)); assertSome (Impossible @"Division by zero impossible here")
    unpair
    swap
    ifEq0 nop $ do
      push (1 :: Natural)
      add
  floor_ = do
    deconstructRational
    ediv @a @Natural @(Maybe (a, Natural)); assertSome (Impossible @"Division by zero impossible here")
    car

-- | Special multiplication helper, in case you want to multiply three 'Rational values'
-- given values @(a / b) * (c / d) * (e / f)@ performs @(a * c * e) / (b * d * f)@.
tripleMul :: forall r a s. (ArithOpHs Mul a a a, Unwrappabled r ~ (a, Natural), LorentzRational r) => r : r : r : s :-> r : s
tripleMul = do
  dip do
    dip deconstructRational
    deconstructRational
  deconstructRational
  dip (swap)
  dipN @3 (swap)
  dipN @2 (swap)
  mul @a @a @a; mul @a @a @a
  dip (do mul; mul)
  constructRational

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

-- | For a given 'Natural' values, calculates their gcd, using Euclid algorithm.
gcdEuclid :: Natural : Natural : s :-> Natural : s
gcdEuclid = do
  dupTop2
  ifLt swap nop
  dupN @2
  neq0
  loop do
    dip dup
    ediv
    assertSome (Impossible @"Division by zero impossible here")
    cdr
    dup
    neq0
    dip swap
  dip drop

-- | Reduce 'Rational' value to the point, when numerator's and denominator's gcd is 1.
-- This operation should be commonly used after several arithmetic operations, because
-- numerator and denominator can become quite big. That in order, can lead to serious
-- gas consumption.
reduceRationalHelper :: (KnownList s, KnownList (ToTs s)) => Rational : s :-> Rational : s
reduceRationalHelper = do
  deconstructRational
  dup; ifLt0 (push (-1 :: Integer)) (push (1 :: Integer)); dug @2; abs;
  reduceHelper; dig @2; mul; pair; uncheckedPairToRational

reduceNRationalHelper :: NRational : s :-> NRational : s
reduceNRationalHelper = do
  deconstructRational
  reduceHelper;
  constructRational

reduceHelper :: Natural : Natural : s :-> Natural : Natural : s
reduceHelper = do
  dupTop2; gcdEuclid; dup; push (1 :: Natural)
  ifEq drop $ do
    dup; dug @3; swap;
    ediv; assertSome (Impossible @"Division by zero impossible here"); car; dug @2
    ediv; assertSome (Impossible @"Division by zero impossible here"); car
    swap

-- | Reduce 'Rational' value, using extended Euclid algorithm.
-- Consumes slightly more gas, than @reduce@, but contract with it is cheaper in terms of origination.
euclidExtendedNormalization :: forall s. (KnownList s, KnownList (ToTs s)) => Rational : s :-> Rational : s
euclidExtendedNormalization = do
  deconstructRational; dip int
  dipN @2 do
    push (1 :: Integer)
    push (0 :: Integer)
    push (0 :: Integer)
    push (1 :: Integer)
  swap; dup; neq0; dip swap;
  loop do
    dupTop2
    ediv
    assertSome (Impossible @"Division by zero impossible here")
    car
    dup
    dup
    dupN @9; mul; dig @7; sub; dug @2
    dupN @7; mul; dig @5; sub; swap
    dupN @5; mul; dig @3; sub
    dig @3; dig @4; dug @2; dipN @4 (swap)
    swap; dup; neq0; dip swap
  dropN @3; dip drop
  dup; ifLt0 (dip neg) (nop); abs; swap
  neg; pair; uncheckedPairToRational

addSubHelper
  :: forall a1 a2 a3 a4 a5 a6 b r s1 s2.
     (  b ~ Natural
     , NiceConstant a1
     , ArithOpHs Mul a2 a1 r
     , Unwrappable a3
     , Unwrappabled a3 ~ (a5, a1)
     , Unwrappabled a4 ~ (a6, Natural)
     , LorentzRational a4
     )
  => ((r : a5 : a1 : s1) :-> (a6 : b : s2))
  -> (a2 : a3 : s1) :-> (a4 : s2)
addSubHelper oper = do
  dip do
    coerceUnwrap
    unpair
    dupN @2
  mul
  oper
  pair
  uncheckedPairToRational

addSubRationalHelper
  :: forall a1 a2 a3 a4 a5 a6 a7 a8 r1 r2 r3 n m s1 s2.
    ( NiceConstant a1
    , NiceConstant a2
    , LorentzRational a7
    , r3 ~ Natural
    , ArithOpHs Mul a1 a3 r1, ArithOpHs Mul a2 a4 r2
    , ArithOpHs Mul n m r3, Unwrappable a5, Unwrappable a6
    , Unwrappabled a7 ~ (a8, r3)
    , Unwrappabled a5 ~ (a3, a2), Unwrappabled a6 ~ (a4, a1)
    )
  => ((r2 : r1 : a1 : a2 : s1) :-> (a8 : n : m : s2))
  -> (a6 : a5 : s1) :-> (a7 : s2)
addSubRationalHelper oper = do
  dip do
    coerceUnwrap
    unpair
  coerceUnwrap; unpair
  dip do
    dup
    dug @2
    mul
  dupN @4
  mul
  oper
  dip mul
  pair
  uncheckedPairToRational

mulRationalHelper
  :: forall r1 r2 r3 a b c s.
     ( LorentzRational r1
     , LorentzRational r2
     , LorentzRational r3
     , Unwrappabled r1 ~ (a, Natural)
     , Unwrappabled r2 ~ (b, Natural)
     , Unwrappabled r3 ~ (c, Natural)
     , ArithOpHs Mul a b c
     )
  => r1 : r2 : s :-> r3 : s
mulRationalHelper = do
  dip deconstructRational
  deconstructRational
  dip swap
  mul
  dip mul
  constructRational

negRational :: Rational : s :-> Rational : s
negRational = deconstructRational # neg # constructRational

mkRational_
  :: (forall s0. ErrInstr (("numerator" :! Integer, "denominator" :! Natural) : s0))
  -> Integer : Natural : s :-> Rational : s
mkRational_ onZeroDenominator = do
  swap
  dup
  ifEq0
    do
      toNamed #denominator
      swap
      toNamed #numerator
      pair
      onZeroDenominator
    do
      swap
      constructRational

type instance ErrorArg "zero_denominator" =
  ("numerator" :! Integer, "denominator" :! Natural)

instance CustomErrorHasDoc "zero_denominator" where
  customErrClass = ErrClassActionException
  customErrDocMdCause =
    "Attempt to construct rational with zero denominator."