{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections #-}
{-# 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, flip)
import Data.Functor ((<$>))
import Data.Maybe (Maybe(..))
import Data.Monoid (Monoid(..))
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.Map u q
amounts = Map.fromListWith H.quantity_add
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 -> H.Balance_Account 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.Get (H.Balance_Account a) (H.Balance_Account a) where get = id

tests :: TestTree
tests = testGroup "Balance"
 [ testGroup "balance_cons" $
	let (==>) (input::[Posting])
		 ( byAccount::[(Account, [(Unit, H.Polarized Quantity)])]
		 , byUnit::[(Unit, (H.Polarized Quantity, [Account]))] ) =
		testCase (elide $ show input) $
		List.foldl (flip H.balance_cons) H.balance_empty
		 ((account *** Map.map H.polarize . amounts) <$> input) @?=
		H.Balance
		 { H.balByAccount = TreeMap.from_List const $
			(account *** H.SumByAccount . amounts) <$> byAccount
		 , H.balByUnit = 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.Polarized_Both (-1) 1]) ]
		[ usd (H.Polarized_Both (-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.Polarized_Both (-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.Polarized_Both (-1) 1
		         , eur $ H.Polarized_Both (-2) 2 ]) ]
		[ usd (H.Polarized_Both (-1) 1, ["/A"])
		, eur (H.Polarized_Both (-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.Polarized_Both (-1) 1, ["/A", "/B"])
		, eur (H.Polarized_Both (-2) 2, ["/A", "/B"])
		, gbp (H.Polarized_Both (-3) 3, ["/A", "/B"]) ]
	 ]
 , testGroup "balance_union" $
	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.balance_union H.balance_empty (bal <$> input) @?= bal expected
		where
		bal (byAccount, byUnit) =
			H.Balance
			 { H.balByAccount = TreeMap.from_List const $
				(account *** H.SumByAccount . amounts) <$> byAccount
			 , H.balByUnit = 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.Polarized_Both (-1) 1, ["/A", "/B"])
			, eur (H.Polarized_Both (-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.Polarized_Both (-1) 1, ["/A", "/B"])
			, eur (H.Polarized_Both (-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 *** H.SumByAccount . amounts) <$> byAccount
		exbal byAccount =
			TreeMap.from_List const $
			(account *** (\(e, i) -> Strict.Clusive
			 (H.SumByAccount (amounts e))
			 (H.SumByAccount (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.Polarized_Both (-1) 1])
		, ("/A/B",   (,) [pol usd (-1)] [usd $ H.Polarized_Both (-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.SumByAccount . (H.polarize <$>) . amounts) <$> input
		balByUnit = H.balByUnit_of_BalByAccount balByAccount mempty
		balDev =
			H.BalByUnit $ 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.is_equilibrium_inferrable $
			H.deviationByUnit
			 H.Balance
				 { H.balByAccount
				 , H.balByUnit
				 } in
		(isRight o, o) @?= (is, e)
		where
		postings
		 :: [Posting]
		 -> Map.Map (H.Balance_Account Text)
		            [( H.Balance_Account Text
		             , Map.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 (H.Balance_Account Text) (H.Polarized Quantity))]
		units =
			((id *** (\(q, as) ->
				H.SumByUnit q $ Map.fromList $
				((,()) . account <$> as))) <$>)
		balByAccount =
			TreeMap.from_List const $
			(account *** H.SumByAccount . (H.polarize <$>) . amounts) <$> input
		balByUnit = H.balByUnit_of_BalByAccount balByAccount mempty
	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"]) ]
	 ]
 ]