{-# OPTIONS --safe #-}

open import Data.Nat.Properties using (+-0-monoid; +-0-commutativeMonoid)
open import Data.Integer using () renaming (+_ to pos)
open import Data.Nat.GeneralisedArithmetic using (iterate)
open import Data.Rational using (; floor; _*_; _÷_; _/_)
open import Data.Rational.Literals using (number; fromℤ)
import      Data.Rational as  renaming (_⊓_ to min)

open import Agda.Builtin.FromNat

open import Ledger.Prelude hiding (iterate; _/_; _*_)
open Filter using (filter)
open import Ledger.Conway.Abstract
open import Ledger.Conway.Transaction
open import Ledger.Conway.Types.Numeric.UnitInterval

open Number number renaming (fromNat to fromℕ)

module Ledger.Conway.Epoch
  (txs : _) (open TransactionStructure txs)
  (abs : AbstractFunctions txs) (open AbstractFunctions abs)
  where

open import Ledger.Conway.Certs govStructure
open import Ledger.Conway.Enact govStructure
open import Ledger.Conway.Gov txs
open import Ledger.Conway.Ledger txs abs
open import Ledger.Conway.Ratify txs
open import Ledger.Conway.Rewards txs abs
open import Ledger.Conway.Utxo txs abs




record EpochState : Type where


  constructor ⟦_,_,_,_,_⟧ᵉ'


  field
    acnt       : Acnt
    ss         : Snapshots
    ls         : LState
    es         : EnactState
    fut        : RatifyState



record HasEpochState {a} (A : Type a) : Type a where
  field EpochStateOf : A  EpochState
open HasEpochState ⦃...⦄ public

instance
  HasLState-EpochState : HasLState EpochState
  HasLState-EpochState .LStateOf = EpochState.ls

  HasEnactState-EpochState : HasEnactState EpochState
  HasEnactState-EpochState .EnactStateOf = EpochState.es

  HasDeposits-EpochState : HasDeposits EpochState
  HasDeposits-EpochState .DepositsOf = DepositsOf  LStateOf

  Hastreasury-EpochState : Hastreasury EpochState
  Hastreasury-EpochState .treasuryOf = Acnt.treasury  EpochState.acnt



record NewEpochState : Type where
  field
    lastEpoch   : Epoch
    epochState  : EpochState
    ru          : Maybe RewardUpdate


record HasNewEpochState {a} (A : Type a) : Type a where
  field NewEpochStateOf : A  NewEpochState
open HasNewEpochState ⦃...⦄ public
record HasLastEpoch     {a} (A : Type a) : Type a where field LastEpochOf     : A  Epoch
open HasLastEpoch     ⦃...⦄ public

instance
  HasLastEpoch-NewEpochState : HasLastEpoch NewEpochState
  HasLastEpoch-NewEpochState .LastEpochOf = NewEpochState.lastEpoch

  HasEpochState-NewEpochState : HasEpochState NewEpochState
  HasEpochState-NewEpochState .EpochStateOf = NewEpochState.epochState

  Hastreasury-NewEpochState : Hastreasury NewEpochState
  Hastreasury-NewEpochState .treasuryOf = treasuryOf  EpochStateOf

  HasLState-NewEpochState : HasLState NewEpochState
  HasLState-NewEpochState .LStateOf = LStateOf  EpochStateOf

  HasGovState-NewEpochState : HasGovState NewEpochState
  HasGovState-NewEpochState .GovStateOf = GovStateOf  LStateOf

  HasCertState-NewEpochState : HasCertState NewEpochState
  HasCertState-NewEpochState .CertStateOf = CertStateOf  LStateOf

  HasDReps-NewEpochState : HasDReps NewEpochState 
  HasDReps-NewEpochState .DRepsOf = DRepsOf  CertStateOf

  HasRewards-NewEpochState : HasRewards NewEpochState
  HasRewards-NewEpochState .RewardsOf = RewardsOf  CertStateOf

  unquoteDecl HasCast-RewardUpdate HasCast-EpochState HasCast-NewEpochState = derive-HasCast
    (   (quote RewardUpdate   , HasCast-RewardUpdate)
       (quote EpochState     , HasCast-EpochState)
     [ (quote NewEpochState  , HasCast-NewEpochState)])

instance _ = +-0-monoid; _ = +-0-commutativeMonoid

toRwdAddr : Credential  RwdAddr
toRwdAddr x = record { net = NetworkId ; stake = x }

getStakeCred : TxOut  Maybe Credential
getStakeCred (a , _ , _ , _) = stakeCred a

open GovActionState using (returnAddr)



createRUpd :   BlocksMade  EpochState  Coin  RewardUpdate
createRUpd slotsPerEpoch b es total
  =  Δt₁ , 0 - Δr₁ + Δr₂ , 0 - feeSS , rs 
  where
    prevPp      = PParamsOf (es .EpochState.es)
    reserves    = es .EpochState.acnt .Acnt.reserves
    pstakego    = es .EpochState.ss .Snapshots.go
    feeSS       = es .EpochState.ss .Snapshots.feeSS
    stake       = pstakego .Snapshot.stake
    delegs      = pstakego .Snapshot.delegations
    poolParams  = pstakego .Snapshot.poolParameters

    blocksMade = ∑[ m  b ] m

    rho = fromUnitInterval (prevPp .PParams.monetaryExpansion)
    η = fromℕ blocksMade ÷₀ (fromℕ slotsPerEpoch * ActiveSlotCoeff)
    Δr₁ = floor (ℚ.min 1 η * rho * fromℕ reserves)

    rewardPot = pos feeSS + Δr₁
    tau = fromUnitInterval (prevPp .PParams.treasuryCut)
    Δt₁ = floor (tau * fromℤ rewardPot)
    R = posPart (rewardPot - Δt₁)
    circulation = total - reserves

    rs = reward prevPp b R poolParams stake delegs circulation
    Δr₂ = R - ∑[ c  rs ] c



applyRUpd : RewardUpdate  EpochState  EpochState
applyRUpd  Δt , Δr , Δf , rs ⟧ʳᵘ
    treasury , reserves ⟧ᵃ
  , ss
  ,   utxo , fees , deposits , donations ⟧ᵘ
    , govSt
    ,   voteDelegs , stakeDelegs , rewards ⟧ᵈ , pState , gState ⟧ᶜˢ ⟧ˡ
  , es
  , fut
  ⟧ᵉ' =
    posPart (pos treasury + Δt + pos unregRU')
    , posPart (pos reserves + Δr) 
  , ss
  ,   utxo , posPart (pos fees + Δf) , deposits , donations 
    , govSt
    ,   voteDelegs , stakeDelegs , rewards ∪⁺ regRU  , pState , gState  
  , es
  , fut 
  where
    regRU     = rs  dom rewards
    unregRU   = rs  dom rewards 
    unregRU'  = ∑[ x  unregRU ] x

getOrphans : EnactState  GovState  GovState
getOrphans es govSt = proj₁ $ iterate step ([] , govSt) (length govSt)
  where
    step : GovState × GovState  GovState × GovState
    step (orps , govSt) =
      let
        isOrphan? a prev = ¬? (hasParent? es govSt a prev)
        (orps' , govSt') = partition
           (_ , record {action = a ; prevAction = prev})  isOrphan? (a .gaType) prev) govSt
      in
        (orps ++ orps' , govSt')


open RwdAddr using (stake)


gaDepositStake : GovState  Deposits  Credential  Coin
gaDepositStake govSt ds = aggregateBy
  (mapˢ  (gaid , addr)  (gaid , addr) , stake addr) govSt')
  (mapFromPartialFun  (gaid , _)  lookupᵐ? ds (GovActionDeposit gaid)) govSt')
  where govSt' = mapˢ (map₂ returnAddr) (fromList govSt)



opaque


  mkStakeDistrs : Snapshot  GovState  Deposits  (Credential  VDeleg)  StakeDistrs
  mkStakeDistrs ss govSt ds delegations .StakeDistrs.stakeDistr =
    aggregateBy  delegations  (Snapshot.stake ss ∪⁺ gaDepositStake govSt ds)


private variable
  nes nes' : NewEpochState
  e lastEpoch : Epoch
  fut fut' : RatifyState
  eps eps' eps'' : EpochState
  ls : LState
  acnt : Acnt
  es₀ : EnactState
  mark set go : Snapshot
  feeSS : Coin
  lstate : LState
  ss ss' : Snapshots
  ru : RewardUpdate
  mru : Maybe RewardUpdate


data _⊢_⇀⦇_,EPOCH⦈_ :   EpochState  Epoch  EpochState  Type where


  EPOCH : let


      open LState ls
      open CertState certState
      open RatifyState fut hiding (es)
      open UTxOState
      open PState; open DState; open GState
      open Acnt; open EnactState; open GovActionState


      esW               = RatifyState.es fut
      es                = record esW { withdrawals =  }
      tmpGovSt          = filter  x  proj₁ x  mapˢ proj₁ removed) govSt
      orphans           = fromList (getOrphans es tmpGovSt)
      removed'          = removed  orphans
      removedGovActions = flip concatMapˢ removed' λ (gaid , gaSt) 
        mapˢ (returnAddr gaSt ,_) ((utxoSt .deposits   GovActionDeposit gaid ) ˢ)
      govActionReturns = aggregate₊ (mapˢ  (a , _ , d)  a , d) removedGovActions ᶠˢ)

      trWithdrawals   = esW .withdrawals
      totWithdrawals  = ∑[ x  trWithdrawals ] x

      retired    = (pState .retiring) ⁻¹ e
      payout     = govActionReturns ∪⁺ trWithdrawals
      refunds    = pullbackMap payout toRwdAddr (dom (dState .rewards))
      unclaimed  = getCoin payout - getCoin refunds

      govSt' = filter  x  proj₁ x  mapˢ proj₁ removed') govSt

      dState' =  dState .voteDelegs , dState .stakeDelegs ,  dState .rewards ∪⁺ refunds 

      pState' =  pState .pools  retired  , pState .retiring  retired  

      gState' =  (if null govSt' then mapValues (1 +_) (gState .dreps) else (gState .dreps))
                , gState .ccHotKeys  ccCreds (es .cc) 

      certState' : CertState
      certState' =  dState' , pState' , gState' 

      utxoSt' =  utxoSt .utxo , utxoSt .fees , utxoSt .deposits  mapˢ (proj₁  proj₂) removedGovActions  , 0 

      acnt' = record acnt
        { treasury  = acnt .treasury  totWithdrawals + utxoSt .donations + unclaimed }
    in
    record { currentEpoch = e
           ; stakeDistrs = mkStakeDistrs  (Snapshots.mark ss') govSt'
                                          (utxoSt' .deposits) (voteDelegs dState)
           ; treasury = acnt .treasury ; GState gState
           ; pools = pState .pools ; delegatees = dState .voteDelegs }
          es ,  , false  ⇀⦇ govSt' ,RATIFIES⦈ fut'
       ls  ss ⇀⦇ tt ,SNAP⦈ ss'
    ────────────────────────────────
    _   acnt , ss , ls , es₀ , fut  ⇀⦇ e ,EPOCH⦈
         acnt' , ss' ,  utxoSt' , govSt' , certState'  , es , fut' 


data


  _⊢_⇀⦇_,NEWEPOCH⦈_ :   NewEpochState  Epoch  NewEpochState  Type


  where


  NEWEPOCH-New : let
      eps' = applyRUpd ru eps
    in
     e  lastEpoch + 1
     _  eps' ⇀⦇ e ,EPOCH⦈ eps''
      ────────────────────────────────
      _   lastEpoch , eps , just ru  ⇀⦇ e ,NEWEPOCH⦈  e , eps'' , nothing 

  NEWEPOCH-Not-New :
     e  lastEpoch + 1
      ────────────────────────────────
      _   lastEpoch , eps , mru  ⇀⦇ e ,NEWEPOCH⦈  lastEpoch , eps , mru 

  NEWEPOCH-No-Reward-Update :
     e  lastEpoch + 1
     _  eps ⇀⦇ e ,EPOCH⦈ eps'
      ────────────────────────────────
      _   lastEpoch , eps , nothing  ⇀⦇ e ,NEWEPOCH⦈  e , eps' , nothing