{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

module Budget (
    -- * Types
    Budget (..),
    Category (..),
    SomeBudget (..),

    -- * Constructors
    calculate,
    calculateSomeBudget,

    -- * Re-exports
    module Money,
    module Rule,
    module Data.Pos,
) where

import Data.Kind
import Data.List
import Data.Pos
import Data.Proxy

import Control.Monad.State

import Money
import Rule (Rule)
import Rule as Rule

import GHC.TypeLits

{- | Budget type using a set distribution between the categories Needs, Wants and Savings by
| using the 50/30/20 rule.
-}
data Budget (c :: Symbol) (b :: (Nat, Nat, Nat)) = Budget
    { needs :: PosMoney c
    , wants :: PosMoney c
    , savings :: PosMoney c
    }
    deriving (Show)

data SomeBudget where
    SomeBudget :: SSymbol c -> Rule '(n, w, s) -> Budget c '(n, w, s) -> SomeBudget

instance (KnownNat x, KnownNat y, KnownNat z) => IntoRule (Budget c '(x, y, z)) '(x, y, z) where
    toRule _ = MkRule (SNat @x) (SNat @y) (SNat @z)

type Income (c :: Symbol) = PosMoney c

data Category = Needs | Wants | Savings deriving (Show)

fromBudget :: SSymbol c -> Rule '(n, w, s) -> Budget c '(n, w, s) -> SomeBudget
fromBudget = SomeBudget

fromKnownBudget :: forall c n w s. (KnownSymbol c, KnownNat n, KnownNat w, KnownNat s) => Budget c '(n, w, s) -> SomeBudget
fromKnownBudget = SomeBudget symbolSing rule
  where
    rule = MkRule (SNat @n) (SNat @w) (SNat @s)

calculate :: Rule '(n, w, s) -> Income c -> Budget c '(n, w, s)
calculate rule x = Budget needs wants savings
  where
    income = fromIntegral . unP . unPosMoney $ x
    (n, w, s) = Rule.fromRule rule
    basisList = [fromIntegral n, fromIntegral w, fromIntegral s]
    (needs : wants : savings : _) = map (PosMoney . unsafePos "") . distribute basisList $ income

calculateSomeBudget ::
    forall n w s.
    (KnownNat n, KnownNat w, KnownNat s) =>
    Rule '(n, w, s)
    -> SomePosMoney
    -> SomeBudget
calculateSomeBudget rule (SomePosMoney c x) = SomeBudget c rule (calculate rule x)

distribute :: [Double] -> Int -> [Int]
distribute xs target = flip evalState 0 . mapM (\x -> distributeState x weight target) $ xs
  where
    weight = sum xs

distributeState :: Double -> Double -> Int -> State Double Int
distributeState basis weight target = do
    old <- get
    let r = (basis / weight * fromIntegral target) + old
    put r
    pure $ round r - round old