{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeFamilies #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}

module Balance.Test where

import Control.Arrow ((***))
import Data.Data ()
import Data.Either (Either(..), either, isRight)
import Data.Eq (Eq(..))
import Data.Function (($), (.), id, const)
import Data.Functor ((<$>))
import Data.Map.Strict (Map)
import Data.Maybe (Maybe(..))
import Data.Ord (Ord(..))
import Data.String (String)
import Data.Text (Text)
import Data.Tuple (snd)
import Prelude (Integer, error)
import Test.Tasty
import Test.Tasty.HUnit
import Text.Show (Show(..))
import qualified Data.List as List
import qualified Data.Map.Strict as Map
import qualified Data.NonNull as NonNull
import qualified Data.Text as Text

import qualified Data.TreeMap.Strict as TreeMap
import qualified Hcompta as H
import qualified Hcompta.Lib.Strict as Strict

elide :: String -> String
elide s | List.length s > 42 = List.take 42 s List.++ ['…']
elide s = s

type Unit = Text
type Quantity = Integer
amounts :: (H.Addable q, Ord u) => [(u, q)] -> Map u q
amounts = Map.fromListWith (H.+)
pol :: H.Polarizable t => (forall q. q -> (Unit, q)) -> (t -> (Unit, H.Polarized t))
pol u = u . H.polarize
usd :: t -> (Unit, t)
usd = ("$",)
eur :: t -> (Unit, t)
eur = ("€",)
gbp :: t -> (Unit, t)
gbp = ("£",)

type Account = Text
account :: Text -> TreeMap.Path Text
account txt = case go txt of
	 [] -> error "account"
	 a:as -> NonNull.ncons a as
	where
	sep = '/'
	go t =
		case Text.uncons t of
		 Nothing -> []
		 Just (x, xs) ->
			if x == sep
			then go xs
			else
				let (chunk, rest) = Text.break (== sep) t in
				chunk:go rest

type Posting = (Account, [(Unit, Quantity)])
instance H.To (TreeMap.Path a) (TreeMap.Path a) where to = id

tests :: TestTree
tests = testGroup "Balance"
 [ testGroup "balanceOf" $
	let (==>) (input::[Posting])
		 ( byAccount::[(Account, [(Unit, H.Polarized Quantity)])]
		 , byUnit::[(Unit, (H.Polarized Quantity, [Account]))] ) =
		testCase (elide $ show input) $
		List.foldl (H.+=) H.zero
		 ((account *** Map.map H.polarize . amounts) <$> input) @?=
		H.Balance
		 { H.balByAccount = TreeMap.from_List const $
			(account *** amounts) <$> byAccount
		 , H.balByUnit = Map.fromList $
			(id *** (\(q, as) -> H.SumByUnit q $ Map.fromList $ ((,()) . account <$> as))) <$> byUnit
		 }
	in
	 [ [ ("/A", [usd 1]) ] ==> (,)
		[ ("/A", [pol usd 1]) ]
		[ usd (H.polarize 1, ["/A"]) ]
	 , [ ("/A", [usd 1])
		 , ("/A", [usd (-1)]) ] ==> (,)
		[ ("/A", [usd $ H.PolBoth (-1) 1]) ]
		[ usd (H.PolBoth (-1) 1, ["/A"]) ]
	 , [ ("/A", [usd 1])
		 , ("/A", [eur (-1)]) ] ==> (,)
		[ ("/A", [ pol usd 1
		         , pol eur (-1) ]) ]
		[ usd (H.polarize 1, ["/A"])
		, eur (H.polarize (-1), ["/A"]) ]
	 , [ ("/A", [usd 1])
		 , ("/B", [usd (-1)]) ] ==> (,)
		[ ("/A", [pol usd 1])
		, ("/B", [pol usd (-1)]) ]
		[ usd (H.PolBoth (-1) 1, ["/A", "/B"]) ]
	 , [ ("/A", [usd 1])
		 , ("/B", [eur (-1)]) ] ==> (,)
		[ ("/A", [pol usd 1])
		, ("/B", [pol eur (-1)]) ]
		[ usd (H.polarize 1, ["/A"])
		, eur (H.polarize (-1), ["/B"]) ]
	 , [ ("/A", [usd 1])
		 , ("/B", [usd 1]) ] ==> (,)
		[ ("/A", [pol usd 1])
		, ("/B", [pol usd 1]) ]
		[ usd (H.polarize 2, ["/A", "/B"]) ]
	 , [ ("/A", [usd 1, eur 2])
		 , ("/A", [usd (-1), eur (-2)]) ] ==> (,)
		[ ("/A", [ usd $ H.PolBoth (-1) 1
		         , eur $ H.PolBoth (-2) 2 ]) ]
		[ usd (H.PolBoth (-1) 1, ["/A"])
		, eur (H.PolBoth (-2) 2, ["/A"]) ]
	 , [ ("/A", [usd 1, eur 2, gbp 3])
		 , ("/B", [usd (-1), eur (-2), gbp (-3)]) ] ==> (,)
		[ ("/A", [ pol usd 1
		         , pol eur 2
		         , pol gbp 3 ])
		, ("/B", [ pol usd (-1)
		         , pol eur (-2)
		         , pol gbp (-3) ]) ]
		[ usd (H.PolBoth (-1) 1, ["/A", "/B"])
		, eur (H.PolBoth (-2) 2, ["/A", "/B"])
		, gbp (H.PolBoth (-3) 3, ["/A", "/B"]) ]
	 ]
 , testGroup "unionBal" $
	let (==>)
		 (input::
		 [( [(Account, [(Unit, H.Polarized Quantity)])]
		  , [(Unit, (H.Polarized Quantity, [Account]))] )])
		 (expected::
		 ( [(Account, [(Unit, H.Polarized Quantity)])]
		 , [(Unit, (H.Polarized Quantity, [Account]))] )) =
		testCase (elide $ show input) $
		List.foldl (H.+) H.zero (bal <$> input) @?= bal expected
		where
		bal (byAccount, byUnit) =
			H.Balance
			 { H.balByAccount = TreeMap.from_List const $
				(account *** amounts) <$> byAccount
			 , H.balByUnit = Map.fromList $
				(id *** (\(q, as) -> H.SumByUnit q $ Map.fromList $ ((,()) . account <$> as))) <$> byUnit
			 }
	in
	 [ [ (,)
			[ ("/A", [pol usd 1]) ]
			[ usd (H.polarize 1, ["/A"]) ]
		 , (,)
			[ ("/A", [pol usd 1]) ]
			[ usd (H.polarize 1, ["/A"]) ]
		 ] ==> (,)
			[ ("/A", [pol usd 2]) ]
			[ usd (H.polarize 2, ["/A"]) ]
	 , [ (,)
			[ ("/A", [pol usd 1]) ]
			[ usd (H.polarize 1, ["/A"]) ]
		 , (,)
			[ ("/B", [pol usd 1]) ]
			[ usd (H.polarize 1, ["/B"]) ]
		 ] ==> (,)
			[ ("/A", [pol usd 1])
			, ("/B", [pol usd 1]) ]
			[ usd (H.polarize 2, ["/A", "/B"]) ]
	 , [ (,)
			[ ("/A", [pol usd 1]) ]
			[ usd (H.polarize 1, ["/A"]) ]
		 , (,)
			[ ("/B", [pol eur 1]) ]
			[ eur (H.polarize 1, ["/B"]) ]
		 ] ==> (,)
			[ ("/A", [pol usd 1])
			, ("/B", [pol eur 1]) ]
			[ usd (H.polarize 1, ["/A"])
			, eur (H.polarize 1, ["/B"]) ]
	 , [ (,)
			[ ("/A", [pol usd 1, pol eur 2]) ]
			[ usd (H.polarize 1, ["/A"])
			, eur (H.polarize 2, ["/A"]) ]
		 , (,)
			[ ("/B", [pol usd (-1), pol eur (-2)]) ]
			[ usd (H.polarize (-1), ["/B"])
			, eur (H.polarize (-2), ["/B"]) ]
		 ] ==> (,)
			[ ("/A", [pol usd 1   , pol eur 2])
			, ("/B", [pol usd (-1), pol eur (-2)]) ]
			[ usd (H.PolBoth (-1) 1, ["/A", "/B"])
			, eur (H.PolBoth (-2) 2, ["/A", "/B"]) ]
	 , [ (,)
			[ ("/A", [pol usd 1, pol eur 2]) ]
			[ usd (H.polarize 1, ["/A"])
			, eur (H.polarize 2, ["/A"]) ]
		 , (,)
			[ ("/B", [pol usd (-1), pol eur (-2)]) ]
			[ usd (H.polarize (-1), ["/B"])
			, eur (H.polarize (-2), ["/B"]) ]
		 , (,)
			[ ("/C", [pol gbp 3]) ]
			[ gbp (H.polarize 3, ["/C"]) ]
		 ] ==> (,)
			[ ("/A", [pol usd 1   , pol eur 2])
			, ("/B", [pol usd (-1), pol eur (-2)])
			, ("/C", [pol gbp 3]) ]
			[ usd (H.PolBoth (-1) 1, ["/A", "/B"])
			, eur (H.PolBoth (-2) 2, ["/A", "/B"])
			, gbp (H.polarize 3, ["/C"]) ]
	 ]
 , testGroup "clusiveBalByAccount" $
	let (==>)
		 (input::[(Account, [(Unit, H.Polarized Quantity)])])
		 (expected::[(Account, ( [(Unit, H.Polarized Quantity)]
		                       , [(Unit, H.Polarized Quantity)] ))]) =
		testCase (elide $ show input) $
		H.clusiveBalByAccount (bal input) @?= exbal expected
		where
		bal byAccount =
			TreeMap.from_List const $
			(account *** amounts) <$> byAccount
		exbal byAccount =
			TreeMap.from_List const $
			(account *** (\(e, i) -> Strict.Clusive (amounts e) (amounts i) )) <$> byAccount
	in
	 [ [] ==> []
	 , [ ("/A", [pol usd 1]) ] ==>
		[ ("/A", (,) [pol usd 1] [pol usd 1]) ]
	 , [ ("/A/A", [pol usd 1]) ] ==>
		[ ("/A",   (,) []          [pol usd 1])
		, ("/A/A", (,) [pol usd 1] [pol usd 1]) ]
	 , [ ("/A/B", [pol usd 1]) ] ==>
		[ ("/A",   (,) []          [pol usd 1])
		, ("/A/B", (,) [pol usd 1] [pol usd 1]) ]
	 , [ ("/A/B/C", [pol usd 1]) ] ==>
		[ ("/A",     (,) []          [pol usd 1])
		, ("/A/B",   (,) []          [pol usd 1])
		, ("/A/B/C", (,) [pol usd 1] [pol usd 1]) ]
	 , [ ("/A/B",   [pol usd (-1)])
		 , ("/A/B/C", [pol usd 1])
		 ] ==>
		[ ("/A",     (,) []             [usd $ H.PolBoth (-1) 1])
		, ("/A/B",   (,) [pol usd (-1)] [usd $ H.PolBoth (-1) 1])
		, ("/A/B/C", (,) [pol usd 1]    [pol usd 1]) ]
	 , [ ("/A/B",   [pol usd 1])
		 , ("/A/B/C", [pol usd 1])
		 , ("/A/B/D", [pol usd 1])
		 ] ==>
		[ ("/A",     (,) []          [pol usd 3])
		, ("/A/B",   (,) [pol usd 1] [pol usd 3])
		, ("/A/B/C", (,) [pol usd 1] [pol usd 1])
		, ("/A/B/D", (,) [pol usd 1] [pol usd 1]) ]
	 , [ ("/A/B", [pol usd 1])
		 , ("/A/C", [pol usd 1])
		 , ("/D/B", [pol usd 1])
		 ] ==>
		[ ("/A",   (,) []          [pol usd 2])
		, ("/A/B", (,) [pol usd 1] [pol usd 1])
		, ("/A/C", (,) [pol usd 1] [pol usd 1])
		, ("/D",   (,) []          [pol usd 1])
		, ("/D/B", (,) [pol usd 1] [pol usd 1]) ]
	 ]
 , testGroup "deviationByUnit" $
	let (==>) (input::[Posting])
		 (expected::[(Unit, (H.Polarized Quantity, [Account]))]) =
		testCase (elide $ show input) $
		H.deviationByUnit
		 H.Balance
			 { H.balByAccount
			 , H.balByUnit
			 } @?=
		H.DeviationByUnit balDev
		where
		balByAccount =
			TreeMap.from_List const $
			(account *** (H.polarize <$>) . amounts) <$> input
		balByUnit = H.sum balByAccount
		balDev =
			Map.fromList $
			(id *** (\(q, as) ->
				H.SumByUnit q $ Map.fromList $
				((,()) . account <$> as))
			) <$> expected
	in
	 [ [] ==> []
	 , [ ("/A", [usd 1]) ] ==>
		[usd (H.polarize 1, [])]
	 , [ ("/A", [usd 1])
		 , ("/B", []) ] ==>
		[usd (H.polarize 1, ["/B"])]
	 , [ ("/A",   [usd 1])
		 , ("/A/B", []) ] ==>
		[usd (H.polarize 1, ["/A/B"])]
	 , [ ("/A",   [usd 1])
		 , ("/A/B", [usd (-1)]) ] ==>
		[]
	 , [ ("/A",   [usd 1])
		 , ("/A/B", [eur (-1)]) ] ==>
		[ usd (H.polarize 1, ["/A/B"])
		, eur (H.polarize (-1), ["/A"])
		]
	 , [ ("/A", [usd 1])
		 , ("/B", [usd 1]) ] ==>
		[usd (H.polarize 2, [])]
	 , [ ("/A", [usd 1])
		 , ("/B", [eur 2]) ] ==>
		[ usd (H.polarize 1, ["/B"])
		, eur (H.polarize 2, ["/A"])
		]
	 , [ ("/A", [usd 1, eur 2])
		 , ("/B", [usd (-1), eur 2]) ] ==>
		[ eur (H.polarize 4, []) ]
	 , [ ("/A", [usd 1, eur 2])
		 , ("/B", [usd (-1), eur 2])
		 , ("/C", [gbp 3]) ] ==>
		[ eur (H.polarize 4, ["/C"])
		, gbp (H.polarize 3, ["/A", "/B"]) ]
	 ]
 , testGroup "equilibrium" $
	let (==>) (input::[Posting])
		 (expected::Either [(Unit, (H.Polarized Quantity, [Account]))] [Posting]) =
		testCase (elide $ show input) $
		let o = snd (H.equilibrium $ postings input) in
		let e = either (Left . units) (Right . postings) expected in
		let is = H.isEquilibriumInferrable $
			H.deviationByUnit
			 H.Balance
				 { H.balByAccount
				 , H.balByUnit
				 } in
		(isRight o, o) @?= (is, e)
		where
		postings
		 :: [Posting]
		 -> Map (TreeMap.Path Text)
		        [( TreeMap.Path Text
		         , Map Unit (H.Polarized Quantity) )]
		postings =
			Map.fromList .
			((\(acct, amts) ->
				let a = account acct in (a,)
				[(a, H.polarize <$> amounts amts)] ) <$>)
		units
		 :: [(Unit, (H.Polarized Quantity, [Account]))]
		 -> [(Unit, H.SumByUnit (TreeMap.Path Text) (H.Polarized Quantity))]
		units =
			((id *** (\(q, as) ->
				H.SumByUnit q $ Map.fromList $
				((,()) . account <$> as))) <$>)
		balByAccount =
			TreeMap.from_List const $
			(account *** (H.polarize <$>) . amounts) <$> input
		balByUnit = H.sum balByAccount
	in
	 [ [] ==> Right []
	 , [ ("/A", [usd 1]) ] ==> Left
		 [ usd (H.polarize 1, []) ]
	 , [ ("/A", [usd 1])
		 , ("/B", []) ] ==> Right
		 [ ("/A", [usd 1])
		 , ("/B", [usd (-1)]) ]
	 , [ ("/A",   [usd 1])
		 , ("/A/B", []) ] ==> Right
		 [ ("/A",   [usd 1])
		 , ("/A/B", [usd (-1)]) ]
	 , [ ("/A",   [usd 1])
		 , ("/A/B", [usd (-1)]) ] ==> Right
		 [ ("/A",   [usd 1])
		 , ("/A/B", [usd (-1)]) ]
	 , [ ("/A",   [usd 1])
		 , ("/A/B", [eur (-1)]) ] ==> Right
		 [ ("/A",   [usd 1, eur 1])
		 , ("/A/B", [usd (-1), eur (-1)]) ]
	 , [ ("/A", [usd 1])
		 , ("/B", [usd 1]) ] ==> Left
		 [ usd (H.polarize 2, []) ]
	 , [ ("/A", [usd 1])
		 , ("/B", [eur 2]) ] ==> Right
		 [ ("/A", [usd 1, eur (-2)])
		 , ("/B", [usd (-1), eur 2]) ]
	 , [ ("/A", [usd 1, eur 2])
		 , ("/B", [usd (-1), eur 2]) ] ==> Left
		 [ eur (H.polarize 4, []) ]
	 , [ ("/A", [usd 1, eur 2])
		 , ("/B", [usd (-1), eur 2])
		 , ("/C", [gbp 3]) ] ==> Left
		 [ gbp (H.polarize 3, ["/A", "/B"]) ]
	 ]
 ]