{-# OPTIONS --safe #-}

open import Data.Nat.Properties using (+-0-monoid; +-0-commutativeMonoid)
open import Data.List using (filter)
open import Data.Integer using () renaming (+_ to pos)
open import Data.Integer.Ext
open import Data.Nat.GeneralisedArithmetic using (iterate)

open import Agda.Builtin.FromNat

open import Ledger.Prelude hiding (iterate)
open import Ledger.Abstract
open import Ledger.Transaction

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

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


record RewardUpdate : Set where


  constructor ⟦_,_,_,_⟧ʳᵘ


  field
    Δt Δr Δf : 
    rs : Credential  Coin


record Snapshot : Set where
  field
    stake        : Credential  Coin
    delegations  : Credential  KeyHash
    -- poolParameters : KeyHash ⇀ PoolParam

record Snapshots : Set where
  field
    mark set go  : Snapshot
    feeSS        : Coin



record EpochState : Type where


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


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


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


instance
  unquoteDecl To-RewardUpdate To-Snapshot To-Snapshots To-EpochState To-NewEpochState = derive-To
    (   (quote RewardUpdate   , To-RewardUpdate)
       (quote Snapshot       , To-Snapshot)
       (quote Snapshots      , To-Snapshots)
       (quote EpochState     , To-EpochState)
     [ (quote NewEpochState  , To-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 RwdAddr using (stake)
open GovActionState using (returnAddr)


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


stakeDistr : UTxO  DState  PState  Snapshot
stakeDistr utxo stᵈ pState =  aggregate₊ (stakeRelation ᶠˢ) , stakeDelegs 
  where
    open DState stᵈ using (stakeDelegs; rewards)
    m = mapˢ  a  (a , cbalance (utxo ∣^' λ i  getStakeCred i  just a))) (dom rewards)
    stakeRelation = m  proj₁ rewards

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 (proj₁ 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 _⊢_⇀⦇_,SNAP⦈_ : LState  Snapshots    Snapshots  Type where
  SNAP : let open LState lstate; open UTxOState utxoSt; open CertState certState
             stake = stakeDistr utxo dState pState
    in
    lstate   mark , set , go , feeSS  ⇀⦇ tt ,SNAP⦈  stake , mark , set , fees 

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


  EPOCH : let


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



      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