{-# LANGUAGE FlexibleInstances #-}
module Data.Variant
        ( Variant (..)
        , flatten
        , toInteger
        , toDouble
        , toBool
        , toAList
        , (~==)
        , (~/=)
        , lookup
        , elem
        , keyExists
        , merge
        , scopeMerge
        , keys
        , values
        , vmap, vamap
        , wrapf, wrapfs
        , wrapf1, wrapfs1
        , call, callMaybe, callDef
        )
where

import Prelude hiding (toInteger, lookup, elem)
import Data.List hiding (lookup, elem)
import qualified Data.List as List
import Data.Maybe
import Safe
import Text.Printf
import Data.Monoid

data Variant = Null
             | Integer Integer
             | Double Double
             | String String
             | Bool Bool
             | List [Variant]
             | AList [(Variant, Variant)]
             | Function ([Variant] -> Variant)
             deriving (Show, Eq)

instance Show ([Variant] -> Variant)
    where show _ = "<<function>>"

instance Eq ([Variant] -> Variant)
    where (==) a b = False

instance Ord Variant
    where
        compare (String a) b = compare a $ flatten b
        compare a (String b) = compare (flatten a) b
        compare (Double a) b = compare a $ toDouble b
        compare a (Double b) = compare (toDouble a) b
        compare (Integer a) b = compare a $ toInteger b
        compare a (Integer b) = compare (toInteger a) b
        compare (Bool a) b = compare a $ toBool b
        compare a (Bool b) = compare (toBool a) b
        compare Null Null = EQ
        compare Null _ = LT
        compare _ Null = GT
        compare _ _ = EQ

instance Monoid Variant
    where
        mempty = Null
        mappend Null a = a
        mappend a Null = a
        mappend (List xs) (List ys) = List (xs ++ ys)
        mappend (AList xs) (AList ys) = AList (xs ++ ys)
        mappend (List xs) (AList ys) = List (xs ++ map snd ys)
        mappend (AList xs) (List ys) = List (map snd xs ++ ys)
        mappend a b = String (flatten a ++ flatten b)

flatten :: Variant -> String
flatten (String s) = s
flatten (Integer i) = show i
flatten (Double d) = cullFracPart . printf "%f" $ d
flatten (Bool True) = "1"
flatten (Bool False) = ""
flatten Null = ""
flatten (List xs) = concat . intersperse " " . map flatten $ xs
flatten (AList xs) = flatten . List . map snd $ xs
flatten (Function _) = "<<function>>"

cullFracPart :: String -> String
cullFracPart str = reverse $ dropWhile (`List.elem` ['0', '.']) $ reverse str

toMaybeInteger :: Variant -> Maybe Integer
toMaybeInteger (String s) = maybeRead s
toMaybeInteger (Integer i) = Just i
toMaybeInteger (Double d) = Just $ round d
toMaybeInteger (Bool True) = Just 1
toMaybeInteger (Bool False) = Just 0
toMaybeInteger _ = Nothing

toInteger :: Variant -> Integer
toInteger = fromMaybe 0 . toMaybeInteger

toMaybeDouble :: Variant -> Maybe Double
toMaybeDouble (String s) = maybeRead s
toMaybeDouble (Integer i) = Just $ fromIntegral i
toMaybeDouble (Double d) = Just d
toMaybeDouble (Bool True) = Just 1
toMaybeDouble (Bool False) = Just 0
toMaybeDouble _ = Nothing

toDouble :: Variant -> Double
toDouble = fromMaybe 0 . toMaybeDouble

toBool :: Variant -> Bool
toBool (Bool b) = b
toBool (Double d) = d /= 0
toBool a = toInteger a /= 0

toAList :: Variant -> [(Variant, Variant)]
toAList (AList xs) = xs
toAList (List xs) = zip (map Integer [0..]) xs
toAList _ = []

instance Num Variant where
    (+) = varAdd
    (-) = varSub
    (*) = varMul
    abs = varAbs
    signum = varSignum
    fromInteger = Integer

varAdd :: Variant -> Variant -> Variant
varAdd (Double a) b = Double $ a + toDouble b
varAdd a (Double b) = Double $ toDouble a + b
varAdd a b = Integer $ toInteger a + toInteger b

varSub :: Variant -> Variant -> Variant
varSub (Double a) b = Double $ a - toDouble b
varSub a (Double b) = Double $ toDouble a - b
varSub a b = Integer $ toInteger a - toInteger b

varMul :: Variant -> Variant -> Variant
varMul (Double a) b = Double $ a * toDouble b
varMul a (Double b) = Double $ toDouble a * b
varMul a b = Integer $ toInteger a * toInteger b

varAbs :: Variant -> Variant
varAbs (Integer i) = Integer (-i)
varAbs (Double i) = Double (-i)
varAbs b = b

varSignum :: Variant -> Variant
varSignum (Integer i) = Integer $ signum i
varSignum (Double i) = Double $ signum i
varSignum a = Integer 1

maybeRead :: Read a => String -> Maybe a
maybeRead s =
    let xs = reads s
    in if null xs then Nothing else (Just . fst . head) xs

(~==) :: Variant -> Variant -> Bool
(~==) a b = flatten a == flatten b

(~/=) a b = flatten a /= flatten b

lookup :: Variant -> Variant -> Variant
lookup key (List xs) =
    let index = fromIntegral . toInteger $ key
    in atDef Null xs index
lookup key (AList xs) =
    let mayVal = List.lookup key xs
    in fromMaybe Null mayVal
lookup _ _ = Null

keyExists :: Variant -> Variant -> Bool
keyExists key (List xs) =
    let index = fromIntegral . toInteger $ key
    in (index < length xs) && (index >= 0)
keyExists key (AList xs) =
    key `List.elem` map fst xs
keyExists _ _ = False

elem :: Variant -> Variant -> Variant
elem key (List xs) = Bool $ List.elem key xs
elem key (AList xs) = Bool $ List.elem key $ map snd xs
elem _ _ = Bool False

merge :: Variant -> Variant -> Variant
merge a b =
    let al = toAList a
        bl = toAList b
    in AList (al ++ bl)

-- Scope merge: First operand has precedence over second; if first argument is
-- scalar, then it becomes the new scope, otherwise both scopes are merged, the
-- first one taking precedence.
scopeMerge :: Variant -> Variant -> Variant
scopeMerge a@(AList xs) b = merge a b
scopeMerge a@(List xs) b = merge a b
scopeMerge Null b = b
scopeMerge a b = a

keys :: Variant -> [Variant]
keys v = map fst $ toAList v

values :: Variant -> [Variant]
values v = map snd $ toAList v

vmap :: (Variant -> a) -> Variant -> [a]
vmap f v = map f $ values v

vamap :: ((Variant, Variant) -> a) -> Variant -> [a]
vamap f v = map f $ toAList v

wrapfs :: (Variant -> [Variant] -> Variant) -> Variant -> Variant
wrapfs f s = Function $ f s

wrapf :: ([Variant] -> Variant) -> Variant
wrapf = Function

wrapfs1 :: (Variant -> Variant -> Variant) -> Variant -> Variant
wrapfs1 f s = wrapf (\(a:_) -> f s a)

wrapf1 :: (Variant -> Variant) -> Variant
wrapf1 f = wrapf (\(a:_) -> f a)

callMaybe :: Variant -> [Variant] -> Maybe Variant
callMaybe (Function f) args = Just $ f args
callMaybe _ _ = Nothing

callDef :: Variant -> [Variant] -> Variant -> Variant
callDef f args def = fromMaybe def $ callMaybe f args

call :: Variant -> [Variant] -> Variant
call f args = callDef f args Null