More efficient `toUnfoldable` without the `Ord` constraint

Keep a List of Maps as the unfold state, instead of a single Map.
Unioning Maps was linear time (making the entire function quadratic?).
Prepending to a List is constant time and doesn't require an `Ord k` instance.

-- | This module defines a type of maps as balanced 2-3 trees, based on
-- | <http://www.cs.princeton.edu/~dpw/courses/cos326-12/ass/2-3-trees.pdf>

module Data.Map
  ( Map
  , showTree
  , empty
  , isEmpty
  , singleton
  , checkValid
  , insert
  , lookup
  , lookupLE
  , lookupLT
  , lookupGE
  , lookupGT
  , findMin
  , findMax
  , fromFoldable
  , fromFoldableWith
  , toList
  , toUnfoldable
  , delete
  , pop
  , member
  , alter
  , update
  , keys
  , values
  , union
  , unionWith
  , unions
  , size
  , mapWithKey
  ) where

import Prelude

import Data.Foldable (foldl, foldMap, foldr, class Foldable)
import Data.List (List(..), (:), length, nub)
import Data.Maybe (Maybe(..), maybe, isJust, fromMaybe)
import Data.Monoid (class Monoid)
import Data.Traversable (traverse, class Traversable)
import Data.Tuple (Tuple(..), uncurry, snd)
import Data.Unfoldable (class Unfoldable, unfoldr)

import Partial.Unsafe (unsafePartial)

-- | `Map k v` represents maps from keys of type `k` to values of type `v`.
data Map k v
  = Leaf
  | Two (Map k v) k v (Map k v)
  | Three (Map k v) k v (Map k v) k v (Map k v)

instance eqMap :: (Eq k, Eq v) => Eq (Map k v) where
  eq m1 m2 = toList m1 == toList m2

instance ordMap :: (Ord k, Ord v) => Ord (Map k v) where
  compare m1 m2 = compare (toList m1) (toList m2)

instance showMap :: (Show k, Show v) => Show (Map k v) where
  show m = "(fromList " <> show (toList m) <> ")"

instance semigroupMap :: Ord k => Semigroup (Map k v) where
  append = union

instance monoidMap :: Ord k => Monoid (Map k v) where
  mempty = empty

instance functorMap :: Functor (Map k) where
  map _ Leaf = Leaf
  map f (Two left k v right) = Two (map f left) k (f v) (map f right)
  map f (Three left k1 v1 mid k2 v2 right) = Three (map f left) k1 (f v1) (map f mid) k2 (f v2) (map f right)

instance foldableMap :: Foldable (Map k) where
  foldl   f z m = foldl   f z (values m)
  foldr   f z m = foldr   f z (values m)
  foldMap f   m = foldMap f   (values m)

instance traversableMap :: Ord k => Traversable (Map k) where
  traverse f ms = foldr (\x acc -> union <$> x <*> acc) (pure empty) ((map (uncurry singleton)) <$> (traverse f <$> toList ms))
  sequence = traverse id

-- | Render a `Map` as a `String`
showTree :: forall k v. (Show k, Show v) => Map k v -> String
showTree Leaf = "Leaf"
showTree (Two left k v right) =
  "Two (" <> showTree left <>
  ") (" <> show k <>
  ") (" <> show v <>
  ") (" <> showTree right <> ")"
showTree (Three left k1 v1 mid k2 v2 right) =
  "Three (" <> showTree left <>
  ") (" <> show k1 <>
  ") (" <> show v1 <>
  ") (" <> showTree mid <>
  ") (" <> show k2 <>
  ") (" <> show v2 <>
  ") (" <> showTree right <> ")"

-- | An empty map
empty :: forall k v. Map k v
empty = Leaf

-- | Test if a map is empty
isEmpty :: forall k v. Map k v -> Boolean
isEmpty Leaf = true
isEmpty _ = false

-- | Create a map with one key/value pair
singleton :: forall k v. k -> v -> Map k v
singleton k v = Two Leaf k v Leaf

-- | Check whether the underlying tree satisfies the 2-3 invariant
-- |
-- | This function is provided for internal use.
checkValid :: forall k v. Map k v -> Boolean
checkValid tree = length (nub (allHeights tree)) == one
  where
  allHeights :: Map k v -> List Int
  allHeights Leaf = pure zero
  allHeights (Two left _ _ right) = map (\n -> n + one) (allHeights left <> allHeights right)
  allHeights (Three left _ _ mid _ _ right) = map (\n -> n + one) (allHeights left <> allHeights mid <> allHeights right)

-- | Lookup a value for the specified key
lookup :: forall k v. Ord k => k -> Map k v -> Maybe v
lookup = unsafePartial \k tree ->
  case tree of
    Leaf -> Nothing
    _ ->
      let comp :: k -> k -> Ordering
          comp = compare
      in case tree of
        Two left k1 v right ->
          case comp k k1 of
            EQ -> Just v
            LT -> lookup k left
            _  -> lookup k right
        Three left k1 v1 mid k2 v2 right ->
          case comp k k1 of
            EQ -> Just v1
            c1 ->
              case c1, comp k k2 of
                _ , EQ -> Just v2
                LT, _  -> lookup k left
                _ , GT -> lookup k right
                _ , _  -> lookup k mid

-- | Lookup a value for the specified key, or the greatest one less than it
lookupLE :: forall k v. (Ord k) => k -> Map k v -> Maybe { key :: k, value :: v }
lookupLE _ Leaf = Nothing
lookupLE k (Two left k1 v1 right) = case compare k k1 of
  EQ -> Just { key: k1, value: v1 }
  GT -> Just $ fromMaybe { key: k1, value: v1 } $ lookupLE k right
  LT -> lookupLE k left
lookupLE k (Three left k1 v1 mid k2 v2 right) = case compare k k2 of
  EQ -> Just { key: k2, value: v2 }
  GT -> Just $ fromMaybe { key: k2, value: v2 } $ lookupLE k right
  LT -> lookupLE k $ Two left k1 v1 mid

-- | Lookup a value for the greatest key less than the specified key
lookupLT :: forall k v. (Ord k) => k -> Map k v -> Maybe { key :: k, value :: v }
lookupLT _ Leaf = Nothing
lookupLT k (Two left k1 v1 right) = case compare k k1 of
  EQ -> findMax left
  GT -> Just $ fromMaybe { key: k1, value: v1 } $ lookupLT k right
  LT -> lookupLT k left
lookupLT k (Three left k1 v1 mid k2 v2 right) = case compare k k2 of
  EQ -> findMax $ Two left k1 v1 mid
  GT -> Just $ fromMaybe { key: k2, value: v2 } $ lookupLT k right
  LT -> lookupLT k $ Two left k1 v1 mid

-- | Lookup a value for the specified key, or the least one greater than it
lookupGE :: forall k v. (Ord k) => k -> Map k v -> Maybe { key :: k, value :: v }
lookupGE _ Leaf = Nothing
lookupGE k (Two left k1 v1 right) = case compare k k1 of
  EQ -> Just { key: k1, value: v1 }
  LT -> Just $ fromMaybe { key: k1, value: v1 } $ lookupGE k left
  GT -> lookupGE k right
lookupGE k (Three left k1 v1 mid k2 v2 right) = case compare k k1 of
  EQ -> Just { key: k1, value: v1 }
  LT -> Just $ fromMaybe { key: k1, value: v1 } $ lookupGE k left
  GT -> lookupGE k $ Two mid k2 v2 right

-- | Lookup a value for the least key greater than the specified key
lookupGT :: forall k v. (Ord k) => k -> Map k v -> Maybe { key :: k, value :: v }
lookupGT _ Leaf = Nothing
lookupGT k (Two left k1 v1 right) = case compare k k1 of
  EQ -> findMin right
  LT -> Just $ fromMaybe { key: k1, value: v1 } $ lookupGT k left
  GT -> lookupGT k right
lookupGT k (Three left k1 v1 mid k2 v2 right) = case compare k k1 of
  EQ -> findMin $ Two mid k2 v2 right
  LT -> Just $ fromMaybe { key: k1, value: v1 } $ lookupGT k left
  GT -> lookupGT k $ Two mid k2 v2 right

-- | Returns the pair with the greatest key
findMax :: forall k v. Map k v -> Maybe { key :: k, value :: v }
findMax Leaf = Nothing
findMax (Two _ k1 v1 right) = Just $ fromMaybe { key: k1, value: v1 } $ findMax right
findMax (Three _ _ _ _ k2 v2 right) = Just $ fromMaybe { key: k2, value: v2 } $ findMax right

-- | Returns the pair with the least key
findMin :: forall k v. Map k v -> Maybe { key :: k, value :: v }
findMin Leaf = Nothing
findMin (Two left k1 v1 _) = Just $ fromMaybe { key: k1, value: v1 } $ findMin left
findMin (Three left k1 v1 _ _ _ _) = Just $ fromMaybe { key: k1, value: v1 } $ findMin left

-- | Test if a key is a member of a map
member :: forall k v. Ord k => k -> Map k v -> Boolean
member k m = isJust (k `lookup` m)

data TreeContext k v
  = TwoLeft k v (Map k v)
  | TwoRight (Map k v) k v
  | ThreeLeft k v (Map k v) k v (Map k v)
  | ThreeMiddle (Map k v) k v k v (Map k v)
  | ThreeRight (Map k v) k v (Map k v) k v

fromZipper :: forall k v. Ord k => List (TreeContext k v) -> Map k v -> Map k v
fromZipper Nil tree = tree
fromZipper (Cons x ctx) tree =
  case x of
    TwoLeft k1 v1 right -> fromZipper ctx (Two tree k1 v1 right)
    TwoRight left k1 v1 -> fromZipper ctx (Two left k1 v1 tree)
    ThreeLeft k1 v1 mid k2 v2 right -> fromZipper ctx (Three tree k1 v1 mid k2 v2 right)
    ThreeMiddle left k1 v1 k2 v2 right -> fromZipper ctx (Three left k1 v1 tree k2 v2 right)
    ThreeRight left k1 v1 mid k2 v2 -> fromZipper ctx (Three left k1 v1 mid k2 v2 tree)

data KickUp k v = KickUp (Map k v) k v (Map k v)

-- | Insert a key/value pair into a map
insert :: forall k v. Ord k => k -> v -> Map k v -> Map k v
insert = down Nil
  where
  comp :: k -> k -> Ordering
  comp = compare

  down :: List (TreeContext k v) -> k -> v -> Map k v -> Map k v
  down ctx k v Leaf = up ctx (KickUp Leaf k v Leaf)
  down ctx k v (Two left k1 v1 right) =
    case comp k k1 of
      EQ -> fromZipper ctx (Two left k v right)
      LT -> down (Cons (TwoLeft k1 v1 right) ctx) k v left
      _  -> down (Cons (TwoRight left k1 v1) ctx) k v right
  down ctx k v (Three left k1 v1 mid k2 v2 right) =
    case comp k k1 of
      EQ -> fromZipper ctx (Three left k v mid k2 v2 right)
      c1 ->
        case c1, comp k k2 of
          _ , EQ -> fromZipper ctx (Three left k1 v1 mid k v right)
          LT, _  -> down (Cons (ThreeLeft k1 v1 mid k2 v2 right) ctx) k v left
          GT, LT -> down (Cons (ThreeMiddle left k1 v1 k2 v2 right) ctx) k v mid
          _ , _  -> down (Cons (ThreeRight left k1 v1 mid k2 v2) ctx) k v right

  up :: List (TreeContext k v) -> KickUp k v -> Map k v
  up Nil (KickUp left k v right) = Two left k v right
  up (Cons x ctx) kup =
    case x, kup of
      TwoLeft k1 v1 right, KickUp left k v mid -> fromZipper ctx (Three left k v mid k1 v1 right)
      TwoRight left k1 v1, KickUp mid k v right -> fromZipper ctx (Three left k1 v1 mid k v right)
      ThreeLeft k1 v1 c k2 v2 d, KickUp a k v b -> up ctx (KickUp (Two a k v b) k1 v1 (Two c k2 v2 d))
      ThreeMiddle a k1 v1 k2 v2 d, KickUp b k v c -> up ctx (KickUp (Two a k1 v1 b) k v (Two c k2 v2 d))
      ThreeRight a k1 v1 b k2 v2, KickUp c k v d -> up ctx (KickUp (Two a k1 v1 b) k2 v2 (Two c k v d))

-- | Delete a key and its corresponding value from a map.
delete :: forall k v. Ord k => k -> Map k v -> Map k v
delete k m = maybe m snd (pop k m)

-- | Delete a key and its corresponding value from a map, returning the value
-- | as well as the subsequent map.
pop :: forall k v. Ord k => k -> Map k v -> Maybe (Tuple v (Map k v))
pop = down Nil
  where
  comp :: k -> k -> Ordering
  comp = compare

  down :: List (TreeContext k v) -> k -> Map k v -> Maybe (Tuple v (Map k v))
  down = unsafePartial \ctx k m -> case m of
    Leaf -> Nothing
    Two left k1 v1 right ->
      case right, comp k k1 of
        Leaf, EQ -> Just (Tuple v1 (up ctx Leaf))
        _   , EQ -> let max = maxNode left
                     in Just (Tuple v1 (removeMaxNode (Cons (TwoLeft max.key max.value right) ctx) left))
        _   , LT -> down (Cons (TwoLeft k1 v1 right) ctx) k left
        _   , _  -> down (Cons (TwoRight left k1 v1) ctx) k right
    Three left k1 v1 mid k2 v2 right ->
      let leaves =
            case left, mid, right of
              Leaf, Leaf, Leaf -> true
              _   , _   , _    -> false
      in case leaves, comp k k1, comp k k2 of
        true, EQ, _  -> Just (Tuple v1 (fromZipper ctx (Two Leaf k2 v2 Leaf)))
        true, _ , EQ -> Just (Tuple v2 (fromZipper ctx (Two Leaf k1 v1 Leaf)))
        _   , EQ, _  -> let max = maxNode left
                         in Just (Tuple v1 (removeMaxNode (Cons (ThreeLeft max.key max.value mid k2 v2 right) ctx) left))
        _   , _ , EQ -> let max = maxNode mid
                         in Just (Tuple v2 (removeMaxNode (Cons (ThreeMiddle left k1 v1 max.key max.value right) ctx) mid))
        _   , LT, _  -> down (Cons (ThreeLeft k1 v1 mid k2 v2 right) ctx) k left
        _   , GT, LT -> down (Cons (ThreeMiddle left k1 v1 k2 v2 right) ctx) k mid
        _   , _ , _  -> down (Cons (ThreeRight left k1 v1 mid k2 v2) ctx) k right

  up :: List (TreeContext k v) -> Map k v -> Map k v
  up = unsafePartial \ctxs tree ->
    case ctxs of
      Nil -> tree
      Cons x ctx ->
        case x, tree of
          TwoLeft k1 v1 Leaf, Leaf -> fromZipper ctx (Two Leaf k1 v1 Leaf)
          TwoRight Leaf k1 v1, Leaf -> fromZipper ctx (Two Leaf k1 v1 Leaf)
          TwoLeft k1 v1 (Two m k2 v2 r), l -> up ctx (Three l k1 v1 m k2 v2 r)
          TwoRight (Two l k1 v1 m) k2 v2, r -> up ctx (Three l k1 v1 m k2 v2 r)
          TwoLeft k1 v1 (Three b k2 v2 c k3 v3 d), a -> fromZipper ctx (Two (Two a k1 v1 b) k2 v2 (Two c k3 v3 d))
          TwoRight (Three a k1 v1 b k2 v2 c) k3 v3, d -> fromZipper ctx (Two (Two a k1 v1 b) k2 v2 (Two c k3 v3 d))
          ThreeLeft k1 v1 Leaf k2 v2 Leaf, Leaf -> fromZipper ctx (Three Leaf k1 v1 Leaf k2 v2 Leaf)
          ThreeMiddle Leaf k1 v1 k2 v2 Leaf, Leaf -> fromZipper ctx (Three Leaf k1 v1 Leaf k2 v2 Leaf)
          ThreeRight Leaf k1 v1 Leaf k2 v2, Leaf -> fromZipper ctx (Three Leaf k1 v1 Leaf k2 v2 Leaf)
          ThreeLeft k1 v1 (Two b k2 v2 c) k3 v3 d, a -> fromZipper ctx (Two (Three a k1 v1 b k2 v2 c) k3 v3 d)
          ThreeMiddle (Two a k1 v1 b) k2 v2 k3 v3 d, c -> fromZipper ctx (Two (Three a k1 v1 b k2 v2 c) k3 v3 d)
          ThreeMiddle a k1 v1 k2 v2 (Two c k3 v3 d), b -> fromZipper ctx (Two a k1 v1 (Three b k2 v2 c k3 v3 d))
          ThreeRight a k1 v1 (Two b k2 v2 c) k3 v3, d -> fromZipper ctx (Two a k1 v1 (Three b k2 v2 c k3 v3 d))
          ThreeLeft k1 v1 (Three b k2 v2 c k3 v3 d) k4 v4 e, a -> fromZipper ctx (Three (Two a k1 v1 b) k2 v2 (Two c k3 v3 d) k4 v4 e)
          ThreeMiddle (Three a k1 v1 b k2 v2 c) k3 v3 k4 v4 e, d -> fromZipper ctx (Three (Two a k1 v1 b) k2 v2 (Two c k3 v3 d) k4 v4 e)
          ThreeMiddle a k1 v1 k2 v2 (Three c k3 v3 d k4 v4 e), b -> fromZipper ctx (Three a k1 v1 (Two b k2 v2 c) k3 v3 (Two d k4 v4 e))
          ThreeRight a k1 v1 (Three b k2 v2 c k3 v3 d) k4 v4, e -> fromZipper ctx (Three a k1 v1 (Two b k2 v2 c) k3 v3 (Two d k4 v4 e))

  maxNode :: Map k v -> { key :: k, value :: v }
  maxNode = unsafePartial \m -> case m of
    Two _ k v Leaf -> { key: k, value: v }
    Two _ _ _ right -> maxNode right
    Three _ _ _ _ k v Leaf -> { key: k, value: v }
    Three _ _ _ _ _ _ right -> maxNode right

  removeMaxNode :: List (TreeContext k v) -> Map k v -> Map k v
  removeMaxNode = unsafePartial \ctx m ->
    case m of
      Two Leaf _ _ Leaf -> up ctx Leaf
      Two left k v right -> removeMaxNode (Cons (TwoRight left k v) ctx) right
      Three Leaf k1 v1 Leaf _ _ Leaf -> up (Cons (TwoRight Leaf k1 v1) ctx) Leaf
      Three left k1 v1 mid k2 v2 right -> removeMaxNode (Cons (ThreeRight left k1 v1 mid k2 v2) ctx) right

-- | Insert the value, delete a value, or update a value for a key in a map
alter :: forall k v. Ord k => (Maybe v -> Maybe v) -> k -> Map k v -> Map k v
alter f k m = case f (k `lookup` m) of
  Nothing -> delete k m
  Just v -> insert k v m

-- | Update or delete the value for a key in a map
update :: forall k v. Ord k => (v -> Maybe v) -> k -> Map k v -> Map k v
update f k m = alter (maybe Nothing f) k m

-- | Convert any foldable collection of key/value pairs to a map.
-- | On key collision, later values take precedence over earlier ones.
fromFoldable :: forall f k v. (Ord k, Foldable f) => f (Tuple k v) -> Map k v
fromFoldable = foldl (\m (Tuple k v) -> insert k v m) empty

-- | Convert any foldable collection of key/value pairs to a map.
-- | On key collision, the values are configurably combined.
fromFoldableWith :: forall f k v. (Ord k, Foldable f) => (v -> v -> v) -> f (Tuple k v) -> Map k v
fromFoldableWith f = foldl (\m (Tuple k v) -> alter (combine v) k m) empty where
  combine v (Just v') = Just $ f v v'
  combine v Nothing = Just v

-- | Convert a map to a list of key/value pairs
toList :: forall k v. Map k v -> List (Tuple k v)
toList Leaf = Nil
toList (Two left k v right) = toList left <> Tuple k v : toList right
toList (Three left k1 v1 mid k2 v2 right) = toList left <> Tuple k1 v1 : toList mid <> Tuple k2 v2 : toList right

-- | Convert a map to an unfoldable structure of key/value pairs
toUnfoldable :: forall f k v. Unfoldable f => Map k v -> f (Tuple k v)
toUnfoldable m = unfoldr go (m : Nil) where
  go Nil = Nothing
  go (hd : tl) = case hd of
    Leaf -> go tl
    Two left k v right ->
      Just $ Tuple (Tuple k v) (left : right : tl)
    Three left k1 v1 mid k2 v2 right ->
      Just $ Tuple (Tuple k1 v1) ((singleton k2 v2) : left : mid : right : tl)

-- | Get a list of the keys contained in a map
keys :: forall k v. Map k v -> List k
keys Leaf = Nil
keys (Two left k _ right) = keys left <> pure k <> keys right
keys (Three left k1 _ mid k2 _ right) = keys left <> pure k1 <> keys mid <> pure k2 <> keys right

-- | Get a list of the values contained in a map
values :: forall k v. Map k v -> List v
values Leaf = Nil
values (Two left _ v right) = values left <> pure v <> values right
values (Three left _ v1 mid _ v2 right) = values left <> pure v1 <> values mid <> pure v2 <> values right

-- | Compute the union of two maps, using the specified function
-- | to combine values for duplicate keys.
unionWith :: forall k v. Ord k => (v -> v -> v) -> Map k v -> Map k v -> Map k v
unionWith f m1 m2 = foldl go m2 (toList m1)
  where
  go m (Tuple k v) = alter (Just <<< maybe v (f v)) k m

-- | Compute the union of two maps, preferring values from the first map in the case
-- | of duplicate keys
union :: forall k v. Ord k => Map k v -> Map k v -> Map k v
union = unionWith const

-- | Compute the union of a collection of maps
unions :: forall k v f. (Ord k, Foldable f) => f (Map k v) -> Map k v
unions = foldl union empty

-- | Calculate the number of key/value pairs in a map
size :: forall k v. Map k v -> Int
size = length <<< values

-- | Apply a function of two arguments to each key/value pair, producing a new map
mapWithKey :: forall k v v'. (k -> v -> v') -> Map k v -> Map k v'
mapWithKey _ Leaf = Leaf
mapWithKey f (Two left k v right) = Two (mapWithKey f left) k (f k v) (mapWithKey f right)
mapWithKey f (Three left k1 v1 mid k2 v2 right) = Three (mapWithKey f left) k1 (f k1 v1) (mapWithKey f mid) k2 (f k2 v2) (mapWithKey f right)

-- | This module defines a type of maps as balanced 2-3 trees, based on
-- | <http://www.cs.princeton.edu/~dpw/courses/cos326-12/ass/2-3-trees.pdf>

module Data.Map
  ( Map
  , showTree
  , empty
  , isEmpty
  , singleton
  , checkValid
  , insert
  , lookup
  , lookupLE
  , lookupLT
  , lookupGE
  , lookupGT
  , findMin
  , findMax
  , fromFoldable
  , fromFoldableWith
  , toList
  , toUnfoldable
  , delete
  , pop
  , member
  , alter
  , update
  , keys
  , values
  , union
  , unionWith
  , unions
  , size
  , mapWithKey
  ) where

import Prelude

import Data.Foldable (foldl, foldMap, foldr, class Foldable)
import Data.List (List(..), (:), length, nub)
import Data.Maybe (Maybe(..), maybe, isJust, fromMaybe)
import Data.Monoid (class Monoid)
import Data.Traversable (traverse, class Traversable)
import Data.Tuple (Tuple(..), uncurry, snd)
import Data.Unfoldable (class Unfoldable, unfoldr)

import Partial.Unsafe (unsafePartial)

-- | `Map k v` represents maps from keys of type `k` to values of type `v`.
data Map k v
  = Leaf
  | Two (Map k v) k v (Map k v)
  | Three (Map k v) k v (Map k v) k v (Map k v)

instance eqMap :: (Eq k, Eq v) => Eq (Map k v) where
  eq m1 m2 = toList m1 == toList m2

instance ordMap :: (Ord k, Ord v) => Ord (Map k v) where
  compare m1 m2 = compare (toList m1) (toList m2)

instance showMap :: (Show k, Show v) => Show (Map k v) where
  show m = "(fromList " <> show (toList m) <> ")"

instance semigroupMap :: Ord k => Semigroup (Map k v) where
  append = union

instance monoidMap :: Ord k => Monoid (Map k v) where
  mempty = empty

instance functorMap :: Functor (Map k) where
  map _ Leaf = Leaf
  map f (Two left k v right) = Two (map f left) k (f v) (map f right)
  map f (Three left k1 v1 mid k2 v2 right) = Three (map f left) k1 (f v1) (map f mid) k2 (f v2) (map f right)

instance foldableMap :: Foldable (Map k) where
  foldl   f z m = foldl   f z (values m)
  foldr   f z m = foldr   f z (values m)
  foldMap f   m = foldMap f   (values m)

instance traversableMap :: Ord k => Traversable (Map k) where
  traverse f ms = foldr (\x acc -> union <$> x <*> acc) (pure empty) ((map (uncurry singleton)) <$> (traverse f <$> toList ms))
  sequence = traverse id

-- | Render a `Map` as a `String`
showTree :: forall k v. (Show k, Show v) => Map k v -> String
showTree Leaf = "Leaf"
showTree (Two left k v right) =
  "Two (" <> showTree left <>
  ") (" <> show k <>
  ") (" <> show v <>
  ") (" <> showTree right <> ")"
showTree (Three left k1 v1 mid k2 v2 right) =
  "Three (" <> showTree left <>
  ") (" <> show k1 <>
  ") (" <> show v1 <>
  ") (" <> showTree mid <>
  ") (" <> show k2 <>
  ") (" <> show v2 <>
  ") (" <> showTree right <> ")"

-- | An empty map
empty :: forall k v. Map k v
empty = Leaf

-- | Test if a map is empty
isEmpty :: forall k v. Map k v -> Boolean
isEmpty Leaf = true
isEmpty _ = false

-- | Create a map with one key/value pair
singleton :: forall k v. k -> v -> Map k v
singleton k v = Two Leaf k v Leaf

-- | Check whether the underlying tree satisfies the 2-3 invariant
-- |
-- | This function is provided for internal use.
checkValid :: forall k v. Map k v -> Boolean
checkValid tree = length (nub (allHeights tree)) == one
  where
  allHeights :: Map k v -> List Int
  allHeights Leaf = pure zero
  allHeights (Two left _ _ right) = map (\n -> n + one) (allHeights left <> allHeights right)
  allHeights (Three left _ _ mid _ _ right) = map (\n -> n + one) (allHeights left <> allHeights mid <> allHeights right)

-- | Lookup a value for the specified key
lookup :: forall k v. Ord k => k -> Map k v -> Maybe v
lookup = unsafePartial \k tree ->
  case tree of
    Leaf -> Nothing
    _ ->
      let comp :: k -> k -> Ordering
          comp = compare
      in case tree of
        Two left k1 v right ->
          case comp k k1 of
            EQ -> Just v
            LT -> lookup k left
            _  -> lookup k right
        Three left k1 v1 mid k2 v2 right ->
          case comp k k1 of
            EQ -> Just v1
            c1 ->
              case c1, comp k k2 of
                _ , EQ -> Just v2
                LT, _  -> lookup k left
                _ , GT -> lookup k right
                _ , _  -> lookup k mid

-- | Lookup a value for the specified key, or the greatest one less than it
lookupLE :: forall k v. (Ord k) => k -> Map k v -> Maybe { key :: k, value :: v }
lookupLE _ Leaf = Nothing
lookupLE k (Two left k1 v1 right) = case compare k k1 of
  EQ -> Just { key: k1, value: v1 }
  GT -> Just $ fromMaybe { key: k1, value: v1 } $ lookupLE k right
  LT -> lookupLE k left
lookupLE k (Three left k1 v1 mid k2 v2 right) = case compare k k2 of
  EQ -> Just { key: k2, value: v2 }
  GT -> Just $ fromMaybe { key: k2, value: v2 } $ lookupLE k right
  LT -> lookupLE k $ Two left k1 v1 mid

-- | Lookup a value for the greatest key less than the specified key
lookupLT :: forall k v. (Ord k) => k -> Map k v -> Maybe { key :: k, value :: v }
lookupLT _ Leaf = Nothing
lookupLT k (Two left k1 v1 right) = case compare k k1 of
  EQ -> findMax left
  GT -> Just $ fromMaybe { key: k1, value: v1 } $ lookupLT k right
  LT -> lookupLT k left
lookupLT k (Three left k1 v1 mid k2 v2 right) = case compare k k2 of
  EQ -> findMax $ Two left k1 v1 mid
  GT -> Just $ fromMaybe { key: k2, value: v2 } $ lookupLT k right
  LT -> lookupLT k $ Two left k1 v1 mid

-- | Lookup a value for the specified key, or the least one greater than it
lookupGE :: forall k v. (Ord k) => k -> Map k v -> Maybe { key :: k, value :: v }
lookupGE _ Leaf = Nothing
lookupGE k (Two left k1 v1 right) = case compare k k1 of
  EQ -> Just { key: k1, value: v1 }
  LT -> Just $ fromMaybe { key: k1, value: v1 } $ lookupGE k left
  GT -> lookupGE k right
lookupGE k (Three left k1 v1 mid k2 v2 right) = case compare k k1 of
  EQ -> Just { key: k1, value: v1 }
  LT -> Just $ fromMaybe { key: k1, value: v1 } $ lookupGE k left
  GT -> lookupGE k $ Two mid k2 v2 right

-- | Lookup a value for the least key greater than the specified key
lookupGT :: forall k v. (Ord k) => k -> Map k v -> Maybe { key :: k, value :: v }
lookupGT _ Leaf = Nothing
lookupGT k (Two left k1 v1 right) = case compare k k1 of
  EQ -> findMin right
  LT -> Just $ fromMaybe { key: k1, value: v1 } $ lookupGT k left
  GT -> lookupGT k right
lookupGT k (Three left k1 v1 mid k2 v2 right) = case compare k k1 of
  EQ -> findMin $ Two mid k2 v2 right
  LT -> Just $ fromMaybe { key: k1, value: v1 } $ lookupGT k left
  GT -> lookupGT k $ Two mid k2 v2 right

-- | Returns the pair with the greatest key
findMax :: forall k v. Map k v -> Maybe { key :: k, value :: v }
findMax Leaf = Nothing
findMax (Two _ k1 v1 right) = Just $ fromMaybe { key: k1, value: v1 } $ findMax right
findMax (Three _ _ _ _ k2 v2 right) = Just $ fromMaybe { key: k2, value: v2 } $ findMax right

-- | Returns the pair with the least key
findMin :: forall k v. Map k v -> Maybe { key :: k, value :: v }
findMin Leaf = Nothing
findMin (Two left k1 v1 _) = Just $ fromMaybe { key: k1, value: v1 } $ findMin left
findMin (Three left k1 v1 _ _ _ _) = Just $ fromMaybe { key: k1, value: v1 } $ findMin left

-- | Test if a key is a member of a map
member :: forall k v. Ord k => k -> Map k v -> Boolean
member k m = isJust (k `lookup` m)

data TreeContext k v
  = TwoLeft k v (Map k v)
  | TwoRight (Map k v) k v
  | ThreeLeft k v (Map k v) k v (Map k v)
  | ThreeMiddle (Map k v) k v k v (Map k v)
  | ThreeRight (Map k v) k v (Map k v) k v

fromZipper :: forall k v. Ord k => List (TreeContext k v) -> Map k v -> Map k v
fromZipper Nil tree = tree
fromZipper (Cons x ctx) tree =
  case x of
    TwoLeft k1 v1 right -> fromZipper ctx (Two tree k1 v1 right)
    TwoRight left k1 v1 -> fromZipper ctx (Two left k1 v1 tree)
    ThreeLeft k1 v1 mid k2 v2 right -> fromZipper ctx (Three tree k1 v1 mid k2 v2 right)
    ThreeMiddle left k1 v1 k2 v2 right -> fromZipper ctx (Three left k1 v1 tree k2 v2 right)
    ThreeRight left k1 v1 mid k2 v2 -> fromZipper ctx (Three left k1 v1 mid k2 v2 tree)

data KickUp k v = KickUp (Map k v) k v (Map k v)

-- | Insert a key/value pair into a map
insert :: forall k v. Ord k => k -> v -> Map k v -> Map k v
insert = down Nil
  where
  comp :: k -> k -> Ordering
  comp = compare

  down :: List (TreeContext k v) -> k -> v -> Map k v -> Map k v
  down ctx k v Leaf = up ctx (KickUp Leaf k v Leaf)
  down ctx k v (Two left k1 v1 right) =
    case comp k k1 of
      EQ -> fromZipper ctx (Two left k v right)
      LT -> down (Cons (TwoLeft k1 v1 right) ctx) k v left
      _  -> down (Cons (TwoRight left k1 v1) ctx) k v right
  down ctx k v (Three left k1 v1 mid k2 v2 right) =
    case comp k k1 of
      EQ -> fromZipper ctx (Three left k v mid k2 v2 right)
      c1 ->
        case c1, comp k k2 of
          _ , EQ -> fromZipper ctx (Three left k1 v1 mid k v right)
          LT, _  -> down (Cons (ThreeLeft k1 v1 mid k2 v2 right) ctx) k v left
          GT, LT -> down (Cons (ThreeMiddle left k1 v1 k2 v2 right) ctx) k v mid
          _ , _  -> down (Cons (ThreeRight left k1 v1 mid k2 v2) ctx) k v right

  up :: List (TreeContext k v) -> KickUp k v -> Map k v
  up Nil (KickUp left k v right) = Two left k v right
  up (Cons x ctx) kup =
    case x, kup of
      TwoLeft k1 v1 right, KickUp left k v mid -> fromZipper ctx (Three left k v mid k1 v1 right)
      TwoRight left k1 v1, KickUp mid k v right -> fromZipper ctx (Three left k1 v1 mid k v right)
      ThreeLeft k1 v1 c k2 v2 d, KickUp a k v b -> up ctx (KickUp (Two a k v b) k1 v1 (Two c k2 v2 d))
      ThreeMiddle a k1 v1 k2 v2 d, KickUp b k v c -> up ctx (KickUp (Two a k1 v1 b) k v (Two c k2 v2 d))
      ThreeRight a k1 v1 b k2 v2, KickUp c k v d -> up ctx (KickUp (Two a k1 v1 b) k2 v2 (Two c k v d))

-- | Delete a key and its corresponding value from a map.
delete :: forall k v. Ord k => k -> Map k v -> Map k v
delete k m = maybe m snd (pop k m)

-- | Delete a key and its corresponding value from a map, returning the value
-- | as well as the subsequent map.
pop :: forall k v. Ord k => k -> Map k v -> Maybe (Tuple v (Map k v))
pop = down Nil
  where
  comp :: k -> k -> Ordering
  comp = compare

  down :: List (TreeContext k v) -> k -> Map k v -> Maybe (Tuple v (Map k v))
  down = unsafePartial \ctx k m -> case m of
    Leaf -> Nothing
    Two left k1 v1 right ->
      case right, comp k k1 of
        Leaf, EQ -> Just (Tuple v1 (up ctx Leaf))
        _   , EQ -> let max = maxNode left
                     in Just (Tuple v1 (removeMaxNode (Cons (TwoLeft max.key max.value right) ctx) left))
        _   , LT -> down (Cons (TwoLeft k1 v1 right) ctx) k left
        _   , _  -> down (Cons (TwoRight left k1 v1) ctx) k right
    Three left k1 v1 mid k2 v2 right ->
      let leaves =
            case left, mid, right of
              Leaf, Leaf, Leaf -> true
              _   , _   , _    -> false
      in case leaves, comp k k1, comp k k2 of
        true, EQ, _  -> Just (Tuple v1 (fromZipper ctx (Two Leaf k2 v2 Leaf)))
        true, _ , EQ -> Just (Tuple v2 (fromZipper ctx (Two Leaf k1 v1 Leaf)))
        _   , EQ, _  -> let max = maxNode left
                         in Just (Tuple v1 (removeMaxNode (Cons (ThreeLeft max.key max.value mid k2 v2 right) ctx) left))
        _   , _ , EQ -> let max = maxNode mid
                         in Just (Tuple v2 (removeMaxNode (Cons (ThreeMiddle left k1 v1 max.key max.value right) ctx) mid))
        _   , LT, _  -> down (Cons (ThreeLeft k1 v1 mid k2 v2 right) ctx) k left
        _   , GT, LT -> down (Cons (ThreeMiddle left k1 v1 k2 v2 right) ctx) k mid
        _   , _ , _  -> down (Cons (ThreeRight left k1 v1 mid k2 v2) ctx) k right

  up :: List (TreeContext k v) -> Map k v -> Map k v
  up = unsafePartial \ctxs tree ->
    case ctxs of
      Nil -> tree
      Cons x ctx ->
        case x, tree of
          TwoLeft k1 v1 Leaf, Leaf -> fromZipper ctx (Two Leaf k1 v1 Leaf)
          TwoRight Leaf k1 v1, Leaf -> fromZipper ctx (Two Leaf k1 v1 Leaf)
          TwoLeft k1 v1 (Two m k2 v2 r), l -> up ctx (Three l k1 v1 m k2 v2 r)
          TwoRight (Two l k1 v1 m) k2 v2, r -> up ctx (Three l k1 v1 m k2 v2 r)
          TwoLeft k1 v1 (Three b k2 v2 c k3 v3 d), a -> fromZipper ctx (Two (Two a k1 v1 b) k2 v2 (Two c k3 v3 d))
          TwoRight (Three a k1 v1 b k2 v2 c) k3 v3, d -> fromZipper ctx (Two (Two a k1 v1 b) k2 v2 (Two c k3 v3 d))
          ThreeLeft k1 v1 Leaf k2 v2 Leaf, Leaf -> fromZipper ctx (Three Leaf k1 v1 Leaf k2 v2 Leaf)
          ThreeMiddle Leaf k1 v1 k2 v2 Leaf, Leaf -> fromZipper ctx (Three Leaf k1 v1 Leaf k2 v2 Leaf)
          ThreeRight Leaf k1 v1 Leaf k2 v2, Leaf -> fromZipper ctx (Three Leaf k1 v1 Leaf k2 v2 Leaf)
          ThreeLeft k1 v1 (Two b k2 v2 c) k3 v3 d, a -> fromZipper ctx (Two (Three a k1 v1 b k2 v2 c) k3 v3 d)
          ThreeMiddle (Two a k1 v1 b) k2 v2 k3 v3 d, c -> fromZipper ctx (Two (Three a k1 v1 b k2 v2 c) k3 v3 d)
          ThreeMiddle a k1 v1 k2 v2 (Two c k3 v3 d), b -> fromZipper ctx (Two a k1 v1 (Three b k2 v2 c k3 v3 d))
          ThreeRight a k1 v1 (Two b k2 v2 c) k3 v3, d -> fromZipper ctx (Two a k1 v1 (Three b k2 v2 c k3 v3 d))
          ThreeLeft k1 v1 (Three b k2 v2 c k3 v3 d) k4 v4 e, a -> fromZipper ctx (Three (Two a k1 v1 b) k2 v2 (Two c k3 v3 d) k4 v4 e)
          ThreeMiddle (Three a k1 v1 b k2 v2 c) k3 v3 k4 v4 e, d -> fromZipper ctx (Three (Two a k1 v1 b) k2 v2 (Two c k3 v3 d) k4 v4 e)
          ThreeMiddle a k1 v1 k2 v2 (Three c k3 v3 d k4 v4 e), b -> fromZipper ctx (Three a k1 v1 (Two b k2 v2 c) k3 v3 (Two d k4 v4 e))
          ThreeRight a k1 v1 (Three b k2 v2 c k3 v3 d) k4 v4, e -> fromZipper ctx (Three a k1 v1 (Two b k2 v2 c) k3 v3 (Two d k4 v4 e))

  maxNode :: Map k v -> { key :: k, value :: v }
  maxNode = unsafePartial \m -> case m of
    Two _ k v Leaf -> { key: k, value: v }
    Two _ _ _ right -> maxNode right
    Three _ _ _ _ k v Leaf -> { key: k, value: v }
    Three _ _ _ _ _ _ right -> maxNode right

  removeMaxNode :: List (TreeContext k v) -> Map k v -> Map k v
  removeMaxNode = unsafePartial \ctx m ->
    case m of
      Two Leaf _ _ Leaf -> up ctx Leaf
      Two left k v right -> removeMaxNode (Cons (TwoRight left k v) ctx) right
      Three Leaf k1 v1 Leaf _ _ Leaf -> up (Cons (TwoRight Leaf k1 v1) ctx) Leaf
      Three left k1 v1 mid k2 v2 right -> removeMaxNode (Cons (ThreeRight left k1 v1 mid k2 v2) ctx) right

-- | Insert the value, delete a value, or update a value for a key in a map
alter :: forall k v. Ord k => (Maybe v -> Maybe v) -> k -> Map k v -> Map k v
alter f k m = case f (k `lookup` m) of
  Nothing -> delete k m
  Just v -> insert k v m

-- | Update or delete the value for a key in a map
update :: forall k v. Ord k => (v -> Maybe v) -> k -> Map k v -> Map k v
update f k m = alter (maybe Nothing f) k m

-- | Convert any foldable collection of key/value pairs to a map.
-- | On key collision, later values take precedence over earlier ones.
fromFoldable :: forall f k v. (Ord k, Foldable f) => f (Tuple k v) -> Map k v
fromFoldable = foldl (\m (Tuple k v) -> insert k v m) empty

-- | Convert any foldable collection of key/value pairs to a map.
-- | On key collision, the values are configurably combined.
fromFoldableWith :: forall f k v. (Ord k, Foldable f) => (v -> v -> v) -> f (Tuple k v) -> Map k v
fromFoldableWith f = foldl (\m (Tuple k v) -> alter (combine v) k m) empty where
  combine v (Just v') = Just $ f v v'
  combine v Nothing = Just v

-- | Convert a map to a list of key/value pairs
toList :: forall k v. Map k v -> List (Tuple k v)
toList Leaf = Nil
toList (Two left k v right) = toList left <> Tuple k v : toList right
toList (Three left k1 v1 mid k2 v2 right) = toList left <> Tuple k1 v1 : toList mid <> Tuple k2 v2 : toList right

-- | Convert a map to an unfoldable structure of key/value pairs
toUnfoldable :: forall f k v. (Ord k, Unfoldable f) => Map k v -> f (Tuple k v)
toUnfoldable = unfoldr go
  where
  go :: Map k v -> Maybe (Tuple (Tuple k v) (Map k v))
  go Leaf = Nothing
  go (Two left k v right) = Just $ Tuple (Tuple k v) (left <> right)
  go (Three left k1 v1 mid k2 v2 right) = Just $ Tuple (Tuple k1 v1) (insert k2 v2 (left <> mid <> right))

-- | Get a list of the keys contained in a map
keys :: forall k v. Map k v -> List k
keys Leaf = Nil
keys (Two left k _ right) = keys left <> pure k <> keys right
keys (Three left k1 _ mid k2 _ right) = keys left <> pure k1 <> keys mid <> pure k2 <> keys right

-- | Get a list of the values contained in a map
values :: forall k v. Map k v -> List v
values Leaf = Nil
values (Two left _ v right) = values left <> pure v <> values right
values (Three left _ v1 mid _ v2 right) = values left <> pure v1 <> values mid <> pure v2 <> values right

-- | Compute the union of two maps, using the specified function
-- | to combine values for duplicate keys.
unionWith :: forall k v. Ord k => (v -> v -> v) -> Map k v -> Map k v -> Map k v
unionWith f m1 m2 = foldl go m2 (toList m1)
  where
  go m (Tuple k v) = alter (Just <<< maybe v (f v)) k m

-- | Compute the union of two maps, preferring values from the first map in the case
-- | of duplicate keys
union :: forall k v. Ord k => Map k v -> Map k v -> Map k v
union = unionWith const

-- | Compute the union of a collection of maps
unions :: forall k v f. (Ord k, Foldable f) => f (Map k v) -> Map k v
unions = foldl union empty

-- | Calculate the number of key/value pairs in a map
size :: forall k v. Map k v -> Int
size = length <<< values

-- | Apply a function of two arguments to each key/value pair, producing a new map
mapWithKey :: forall k v v'. (k -> v -> v') -> Map k v -> Map k v'
mapWithKey _ Leaf = Leaf
mapWithKey f (Two left k v right) = Two (mapWithKey f left) k (f k v) (mapWithKey f right)
mapWithKey f (Three left k1 v1 mid k2 v2 right) = Three (mapWithKey f left) k1 (f k1 v1) (mapWithKey f mid) k2 (f k2 v2) (mapWithKey f right)


Relax constraint to `Bind m` for `bindWriterT`

-- | This module defines the writer monad transformer, `WriterT`.

module Control.Monad.Writer.Trans
  ( WriterT(..), runWriterT, execWriterT, mapWriterT
  , module Control.Monad.Trans.Class
  , module Control.Monad.Writer.Class
  ) where

import Prelude

import Control.Alt (class Alt, (<|>))
import Control.Alternative (class Alternative)
import Control.Monad.Cont.Class (class MonadCont, callCC)
import Control.Monad.Eff.Class (class MonadEff, liftEff)
import Control.Monad.Error.Class (class MonadError, catchError, throwError)
import Control.Monad.Reader.Class (class MonadAsk, class MonadReader, ask, local)
import Control.Monad.Rec.Class (class MonadRec, tailRecM, Step(..))
import Control.Monad.State.Class (class MonadState, state)
import Control.Monad.Trans.Class (class MonadTrans, lift)
import Control.Monad.Writer.Class (class MonadTell, tell, class MonadWriter, censor, listen, listens, pass)
import Control.MonadPlus (class MonadPlus)
import Control.MonadZero (class MonadZero)
import Control.Plus (class Plus, empty)

import Data.Monoid (class Monoid, mempty)
import Data.Newtype (class Newtype)
import Data.Tuple (Tuple(..), snd)

-- | The writer monad transformer.
-- |
-- | This monad transformer extends the base monad with a monoidal accumulator of
-- | type `w`.
-- |
-- | The `MonadWriter` type class describes the operations supported by this monad.
newtype WriterT w m a = WriterT (m (Tuple a w))

-- | Run a computation in the `WriterT` monad.
runWriterT :: forall w m a. WriterT w m a -> m (Tuple a w)
runWriterT (WriterT x) = x

-- | Run a computation in the `WriterT` monad, discarding the result.
execWriterT :: forall w m a. Functor m => WriterT w m a -> m w
execWriterT (WriterT m) = snd <$> m

-- | Change the accumulator and base monad types in a `WriterT` monad action.
mapWriterT :: forall w1 w2 m1 m2 a b. (m1 (Tuple a w1) -> m2 (Tuple b w2)) -> WriterT w1 m1 a -> WriterT w2 m2 b
mapWriterT f (WriterT m) = WriterT (f m)

derive instance newtypeWriterT :: Newtype (WriterT w m a) _

instance functorWriterT :: Functor m => Functor (WriterT w m) where
  map f = mapWriterT $ map \(Tuple a w) -> Tuple (f a) w

instance applyWriterT :: (Semigroup w, Apply m) => Apply (WriterT w m) where
  apply (WriterT f) (WriterT v) = WriterT
    let k (Tuple a w) (Tuple b w') = Tuple (a b) (w <> w')
    in k <$> f <*> v

instance applicativeWriterT :: (Monoid w, Applicative m) => Applicative (WriterT w m) where
  pure a = WriterT $ pure $ Tuple a mempty

instance altWriterT :: Alt m => Alt (WriterT w m) where
  alt (WriterT m) (WriterT n) = WriterT (m <|> n)

instance plusWriterT :: Plus m => Plus (WriterT w m) where
  empty = WriterT empty

instance alternativeWriterT :: (Monoid w, Alternative m) => Alternative (WriterT w m)

instance bindWriterT :: (Semigroup w, Bind m) => Bind (WriterT w m) where
  bind (WriterT m) k = WriterT $
    m >>= \(Tuple a w) ->
      case k a of
        WriterT wt ->
          map (\(Tuple b w') -> Tuple b (w <> w')) wt

instance monadWriterT :: (Monoid w, Monad m) => Monad (WriterT w m)

instance monadRecWriterT :: (Monoid w, MonadRec m) => MonadRec (WriterT w m) where
  tailRecM f a = WriterT $ tailRecM f' (Tuple a mempty)
    where
    f' (Tuple a w) =
      case f a of WriterT wt ->
        wt >>= \(Tuple m w1) ->
          pure case m of
            Loop a -> Loop (Tuple a (w <> w1))
            Done b -> Done (Tuple b (w <> w1))

instance monadZeroWriterT :: (Monoid w, MonadZero m) => MonadZero (WriterT w m)

instance monadPlusWriterT :: (Monoid w, MonadPlus m) => MonadPlus (WriterT w m)

instance monadTransWriterT :: Monoid w => MonadTrans (WriterT w) where
  lift m = WriterT do
    a <- m
    pure $ Tuple a mempty

instance monadEffWriter :: (Monoid w, MonadEff eff m) => MonadEff eff (WriterT w m) where
  liftEff = lift <<< liftEff

instance monadContWriterT :: (Monoid w, MonadCont m) => MonadCont (WriterT w m) where
  callCC f = WriterT $ callCC \c ->
    case f (\a -> WriterT $ c (Tuple a mempty)) of WriterT b -> b

instance monadErrorWriterT :: (Monoid w, MonadError e m) => MonadError e (WriterT w m) where
  throwError e = lift (throwError e)
  catchError (WriterT m) h = WriterT $ catchError m (\e -> case h e of WriterT a -> a)

instance monadAskWriterT :: (Monoid w, MonadAsk r m) => MonadAsk r (WriterT w m) where
  ask = lift ask

instance monadReaderWriterT :: (Monoid w, MonadReader r m) => MonadReader r (WriterT w m) where
  local f = mapWriterT (local f)

instance monadStateWriterT :: (Monoid w, MonadState s m) => MonadState s (WriterT w m) where
  state f = lift (state f)

instance monadTellWriterT :: (Monoid w, Monad m) => MonadTell w (WriterT w m) where
  tell = WriterT <<< pure <<< Tuple unit

instance monadWriterWriterT :: (Monoid w, Monad m) => MonadWriter w (WriterT w m) where
  listen (WriterT m) = WriterT do
    Tuple a w <- m
    pure $ Tuple (Tuple a w) w
  pass (WriterT m) = WriterT do
    Tuple (Tuple a f) w <- m
    pure $ Tuple a (f w)

-- | This module defines the writer monad transformer, `WriterT`.

module Control.Monad.Writer.Trans
  ( WriterT(..), runWriterT, execWriterT, mapWriterT
  , module Control.Monad.Trans.Class
  , module Control.Monad.Writer.Class
  ) where

import Prelude

import Control.Alt (class Alt, (<|>))
import Control.Alternative (class Alternative)
import Control.Monad.Cont.Class (class MonadCont, callCC)
import Control.Monad.Eff.Class (class MonadEff, liftEff)
import Control.Monad.Error.Class (class MonadError, catchError, throwError)
import Control.Monad.Reader.Class (class MonadAsk, class MonadReader, ask, local)
import Control.Monad.Rec.Class (class MonadRec, tailRecM, Step(..))
import Control.Monad.State.Class (class MonadState, state)
import Control.Monad.Trans.Class (class MonadTrans, lift)
import Control.Monad.Writer.Class (class MonadTell, tell, class MonadWriter, censor, listen, listens, pass)
import Control.MonadPlus (class MonadPlus)
import Control.MonadZero (class MonadZero)
import Control.Plus (class Plus, empty)

import Data.Monoid (class Monoid, mempty)
import Data.Newtype (class Newtype)
import Data.Tuple (Tuple(..), snd)

-- | The writer monad transformer.
-- |
-- | This monad transformer extends the base monad with a monoidal accumulator of
-- | type `w`.
-- |
-- | The `MonadWriter` type class describes the operations supported by this monad.
newtype WriterT w m a = WriterT (m (Tuple a w))

-- | Run a computation in the `WriterT` monad.
runWriterT :: forall w m a. WriterT w m a -> m (Tuple a w)
runWriterT (WriterT x) = x

-- | Run a computation in the `WriterT` monad, discarding the result.
execWriterT :: forall w m a. Functor m => WriterT w m a -> m w
execWriterT (WriterT m) = snd <$> m

-- | Change the accumulator and base monad types in a `WriterT` monad action.
mapWriterT :: forall w1 w2 m1 m2 a b. (m1 (Tuple a w1) -> m2 (Tuple b w2)) -> WriterT w1 m1 a -> WriterT w2 m2 b
mapWriterT f (WriterT m) = WriterT (f m)

derive instance newtypeWriterT :: Newtype (WriterT w m a) _

instance functorWriterT :: Functor m => Functor (WriterT w m) where
  map f = mapWriterT $ map \(Tuple a w) -> Tuple (f a) w

instance applyWriterT :: (Semigroup w, Apply m) => Apply (WriterT w m) where
  apply (WriterT f) (WriterT v) = WriterT
    let k (Tuple a w) (Tuple b w') = Tuple (a b) (w <> w')
    in k <$> f <*> v

instance applicativeWriterT :: (Monoid w, Applicative m) => Applicative (WriterT w m) where
  pure a = WriterT $ pure $ Tuple a mempty

instance altWriterT :: Alt m => Alt (WriterT w m) where
  alt (WriterT m) (WriterT n) = WriterT (m <|> n)

instance plusWriterT :: Plus m => Plus (WriterT w m) where
  empty = WriterT empty

instance alternativeWriterT :: (Monoid w, Alternative m) => Alternative (WriterT w m)

instance bindWriterT :: (Semigroup w, Monad m) => Bind (WriterT w m) where
  bind (WriterT m) k = WriterT $
    m >>= \(Tuple a w) ->
      case k a of WriterT wt ->
        wt >>= \(Tuple b w') ->
          pure $ Tuple b (w <> w')

instance monadWriterT :: (Monoid w, Monad m) => Monad (WriterT w m)

instance monadRecWriterT :: (Monoid w, MonadRec m) => MonadRec (WriterT w m) where
  tailRecM f a = WriterT $ tailRecM f' (Tuple a mempty)
    where
    f' (Tuple a w) =
      case f a of WriterT wt ->
        wt >>= \(Tuple m w1) ->
          pure case m of
            Loop a -> Loop (Tuple a (w <> w1))
            Done b -> Done (Tuple b (w <> w1))

instance monadZeroWriterT :: (Monoid w, MonadZero m) => MonadZero (WriterT w m)

instance monadPlusWriterT :: (Monoid w, MonadPlus m) => MonadPlus (WriterT w m)

instance monadTransWriterT :: Monoid w => MonadTrans (WriterT w) where
  lift m = WriterT do
    a <- m
    pure $ Tuple a mempty

instance monadEffWriter :: (Monoid w, MonadEff eff m) => MonadEff eff (WriterT w m) where
  liftEff = lift <<< liftEff

instance monadContWriterT :: (Monoid w, MonadCont m) => MonadCont (WriterT w m) where
  callCC f = WriterT $ callCC \c ->
    case f (\a -> WriterT $ c (Tuple a mempty)) of WriterT b -> b

instance monadErrorWriterT :: (Monoid w, MonadError e m) => MonadError e (WriterT w m) where
  throwError e = lift (throwError e)
  catchError (WriterT m) h = WriterT $ catchError m (\e -> case h e of WriterT a -> a)

instance monadAskWriterT :: (Monoid w, MonadAsk r m) => MonadAsk r (WriterT w m) where
  ask = lift ask

instance monadReaderWriterT :: (Monoid w, MonadReader r m) => MonadReader r (WriterT w m) where
  local f = mapWriterT (local f)

instance monadStateWriterT :: (Monoid w, MonadState s m) => MonadState s (WriterT w m) where
  state f = lift (state f)

instance monadTellWriterT :: (Monoid w, Monad m) => MonadTell w (WriterT w m) where
  tell = WriterT <<< pure <<< Tuple unit

instance monadWriterWriterT :: (Monoid w, Monad m) => MonadWriter w (WriterT w m) where
  listen (WriterT m) = WriterT do
    Tuple a w <- m
    pure $ Tuple (Tuple a w) w
  pass (WriterT m) = WriterT do
    Tuple (Tuple a f) w <- m
    pure $ Tuple a (f w)


bug fix & strict check

bug fix & strict check

{- | This module implements `DamageVector`:
   | an array of `Damage`s that can be applied to one fleet of `Ship`s.
   | most functions prefixed with `calc` in this module
   | produces either `LR DamageVector` or a single `DamageVector`.
   | which can then be converted `NormalDamageVector` or `CombinedDamageVector`
   | depending on what the corresponding data is standing for.
 -}
module KanColle.DamageAnalysis.DamageVector
  ( DamageVector, getDV, mkDV

  , NormalDamageVector
  , GCombinedDamageVector
  , CombinedDamageVector, CombinedDamageVectorAC
  , FleetRole(..)

  , calcKoukuDamage
  , calcKoukuDamageAC
  , calcKoukuDamageCombined
  , calcSupportAirAttackDamage
  , calcSupportAirAttackDamageAC
  , calcSupportHouraiDamage
  , calcSupportHouraiDamageAC
  , calcHougekiDamage
  , calcRaigekiDamage

  , calcRaigekiDamageAC
  , calcHougekiDamageAC

  , calcLandBasedKoukuDamage

  , toGCombined
  , toCombined
  , toCombinedAC

  , applyDamageVector
  , applyNormalDamageVector
  , applyCombinedDamageVector
  ) where

import Prelude
import Data.Monoid
import Data.Int as Int
import Data.Array as A
import Data.Array.ST as STA
import Data.Unfoldable hiding (fromMaybe)
import Control.Monad.Eff
import Control.Monad.ST as ST
import KanColle.Util
import KanColle.DamageAnalysis.Types
import KanColle.DamageAnalysis.Damage

import Data.Foreign
import Data.Foldable
import KanColle.KCAPI.Battle
import Data.Maybe
import Data.String (joinWith)

-- when flagship protection happens
-- and extra 0.1 will be added to that value
-- for our purpose we don't need to deal with that
-- so all damages can be normalized as Ints
normalizeDamage :: Number -> Int
normalizeDamage x
  | x >= 0.0 = Int.floor x
  | x == (-1.0) = 0
  | otherwise =
    let warning = "invalid damage number: " <> show x
    in traceWarn warning (\_ -> 0)

-- | `DamageVector` is a 6-element array of `Damage`s
newtype DamageVector = DV (Array Damage)

-- | fetches the array inside `DamageVector`
getDV :: DamageVector -> Array Damage
getDV (DV v) = v

-- | creates `DamageVector` from arrays.
-- | array has to be of length 6
mkDV :: Array Damage -> DamageVector
mkDV xs = if check
    then DV xs
    else throwWith "mkDV: array size should be 6"
  where
    check = A.length xs == 6

instance semigroupDamageVector :: Semigroup DamageVector where
  append (DV a) (DV b) = DV (A.zipWith (<>) a b)

instance monoidDamageVector :: Monoid DamageVector where
  mempty = DV (replicate 6 mempty)

-- | `DamageVector` for normal battles
type NormalDamageVector = NormalBattle DamageVector

-- | `DamageVector` for battles involving comined fleets
type GCombinedDamageVector = GCombinedBattle DamageVector
type CombinedDamageVector = CombinedBattle DamageVector
type CombinedDamageVectorAC = CombinedBattleAC DamageVector

debugShowDV :: DamageVector -> String
debugShowDV (DV xs) = joinWith "," (map (show <<< damageToInt) xs)

-- | get `DamageVector` from raw `fDam` and `eDam` fields
fromFDamAndEDam :: KoukuStage3 -> LR DamageVector
fromFDamAndEDam ks3 =
    { left: fromMaybe mempty ((convertFEDam >>> mkDV) <$> fDam)
    , right: fromMaybe mempty ((convertFEDam >>> mkDV) <$> eDam) }
  where
    fDam = getKoukuStage3FDam ks3
    eDam = getKoukuStage3EDam ks3

-- L: ally fleet, LL: ally main, LR: ally escort
-- R: enemy fleet: RL: enemy main, RR: enemy escort
fromFDamAndEDamAC :: forall a.
                  { api_fdam :: Array Number
                  , api_edam :: Array Number | a} -> LR (LR DamageVector)
fromFDamAndEDamAC v = (lrMap >>> lrMap) mkDV
    { left: allyDams
    , right: enemyDams }
  where
    longFDam = convertFEDam v.api_fdam
    allyDams = fleetSplit false longFDam

    longEDam = convertFEDam v.api_edam
    enemyDams = fleetSplit false longEDam

-- (internal use only)
-- an "-1" is put in front of both api_fdam and api_edam
-- we first drop that element and then convert
-- integers intos Damages
convertFEDam :: Array Number -> Array Damage
convertFEDam = unsafeArrTail >>> map (normalizeDamage >>> mkDamage)

-- | calculate damage from kouku (aerial) stages
calcKoukuDamage :: Kouku -> LR DamageVector
calcKoukuDamage kk = fromFDamAndEDam kk.api_stage3

-- calcKoukuDamageAC :: Kouku -> LR DamageVector
-- calcKoukuDamageAC kk = fromFDamAndEDam kk.api_stage3_combined

-- | calculate damage from ally vs. enemy main fleet kouku stages (for Abyssal Combined fleet)
calcKoukuDamageACEMain :: forall a . {api_stage3 :: KoukuStage3 | a} -> LR DamageVector
calcKoukuDamageACEMain kk = fromFDamAndEDam kk.api_stage3

-- | calculate damage from ally vs. enemy escort fleet kouku stages (for Abyssal Combined fleet)
calcKoukuDamageACEEscort :: forall a. {api_stage3_combined :: KoukuStage3 | a} -> LR DamageVector
calcKoukuDamageACEEscort kk = fromFDamAndEDam kk.api_stage3_combined

calcKoukuDamageAC :: forall a.
                  { api_stage3 :: KoukuStage3
                  , api_stage3_combined :: KoukuStage3
                  | a} -> LR (LR DamageVector)
calcKoukuDamageAC kk =
    { left:
      { left: resultEMain.left <> resultEEscort.left
      , right: mempty }
    , right:
      { left: resultEMain.right
      , right: resultEEscort.right } }
  where
    resultEMain = calcKoukuDamageACEMain kk
    resultEEscort = calcKoukuDamageACEEscort kk

-- | calculate damage from kouku (aerial) stages (combined fleet)
-- | note that only escort fleet is taking damage. so we just need DamageVector
calcKoukuDamageCombined :: Kouku -> DamageVector
calcKoukuDamageCombined kk = DV (convertFEDam (kk.api_stage3_combined.api_fdam))

-- | calculate land based kouku damages. whose api_stage3 only have an api_edam field
-- | (because only enemies are taking damage)
calcLandBasedKoukuDamage :: Kouku -> DamageVector
calcLandBasedKoukuDamage lbkk = case getKoukuStage3 lbkk  of
    Just stage3 -> DV (convertFEDam stage3.api_edam)
    Nothing -> mempty

-- | calculate damage from raigeki (torpedo) stages
calcRaigekiDamage :: Raigeki -> LR DamageVector
calcRaigekiDamage = fromFDamAndEDam

calcRaigekiDamageAC :: Raigeki -> LR (LR DamageVector)
calcRaigekiDamageAC = fromFDamAndEDamAC

-- | calculate damage from hougeki (shelling) stages
calcHougekiDamage :: Hougeki -> LR DamageVector
calcHougekiDamage h =
    if lengthCheck
      then let resultArr = runPure (STA.runSTArray resultDV)
           in lrMap DV (fleetSplit false resultArr)
      else throwWith "invalid: api_df_list and api_damage length mismatch"
  where
    cAI :: Foreign -> Array Int
    cAI = unsafeFromForeign
    cAN :: Foreign -> Array Number
    cAN = unsafeFromForeign

    lengthCheck = A.length eventTargets == A.length eventDamages

    eventTargets :: Array Int
    eventTargets = map (cAI >>> toOne) (unsafeArrTail h.api_df_list)
      where
        -- merge targets like [a,a] into a-1 (so that it's zero-indexed)
        -- there are 2 checks:
        --   elements of api_df_list should:
        --   * be non-empty
        --   * all elements of it should be the same
        -- the reason for minus 1 in the end is to make it zero-based
        toOne :: Array Int -> Int
        toOne xs = case A.uncons xs of
          Just {head: y, tail: ys} ->
            if all (\v -> v == y) ys || all (\v -> v == -1) ys
              then y - 1
              else throwWith "invalid: elements are different in api_df_list"
          Nothing -> throwWith "invalid: empty api_df_list element"

    eventDamages :: Array Damage
    eventDamages = map (cAN >>> convert) (unsafeArrTail h.api_damage)
      where
        convert :: Array Number -> Damage
        convert xs = mkDamage totalDmg
          where
            -- there are double attacks in the damage list
            -- because damecon triggers *AFTER* the attack actions are completed
            -- we need to sum damage values before turning it into real Damage
            totalDmg = sum (map normalizeDamage xs :: Array Int)

    accumulateDamage :: forall h r.
                        STA.STArray h Damage -- array reference
                     -> Int -> Damage -- target index and corresponding damage
                     -> Eff (st :: ST.ST h | r) Unit
    accumulateDamage arr targetInd damage = do
        dmg <- peekSTArrayUnsafe arr targetInd
        pokeSTArrayUnsafe arr targetInd (dmg <> damage :: Damage)
        pure unit
    resultDV :: forall h r . Eff (st :: ST.ST h | r) (STA.STArray h Damage)
    resultDV = do
        arr <- STA.thaw (replicate 12 mempty)
        A.zipWithA (accumulateDamage arr) eventTargets eventDamages
        pure arr

calcHougekiDamageAC :: Hougeki -> LR (LR DamageVector)
calcHougekiDamageAC h =
    if lengthCheck
      then let resultArr = runPure (STA.runSTArray resultDV)
           in { left:
                { left: DV (A.slice 0 6 resultArr)
                , right: DV (A.slice 6 12 resultArr) }
              , right:
                { left: DV (A.slice 12 18 resultArr)
                , right: DV (A.slice 18 24 resultArr) }
              }
      else throwWith "invalid: api_df_list / api_damage / api_at_eflag length mismatch"
  where
    cAI :: Foreign -> Array Int
    cAI = unsafeFromForeign
    cAN :: Foreign -> Array Number
    cAN = unsafeFromForeign

    dfList = unsafeArrTail h.api_df_list
    atEFlag :: Array Int
    atEFlag = unsafeArrTail h.api_at_eflag
    damageList = unsafeArrTail h.api_damage

    lengthCheck = A.length dfList == A.length atEFlag
               && A.length dfList == A.length damageList

    eventTargetsRaw :: Array Int
    eventTargetsRaw = map (cAI >>> toOne) (unsafeArrTail h.api_df_list)
      where
        -- merge targets like [a,a] into a-1 (so that it's zero-indexed)
        -- there are 2 checks:
        --   elements of api_df_list should:
        --   * be non-empty
        --   * all elements of it should be the same
        -- the reason for minus 1 in the end is to make it zero-based
        toOne :: Array Int -> Int
        toOne xs = case A.uncons xs of
          Just {head: y, tail: ys} ->
            if all (\v -> v == y) ys || all (\v -> v == -1) ys
              then y - 1
              else throwWith "invalid: elements are different in api_df_list"
          Nothing -> throwWith "invalid: empty api_df_list element"

    eventTargets :: Array Int
    eventTargets = A.zipWith updateTarget eventTargetsRaw atEFlag
      where
        updateTarget rawTarget eFlag = case eFlag of
            -- ally attacks in this turn, so the target is an enemy
            0 -> rawTarget + 12
            -- enemy attacks in this turn, so the target is an ally
            -- nothing to be doing in this case
            1 -> rawTarget
            _ -> throwWith "invalid api_at_eflag element"

    eventDamages :: Array Damage
    eventDamages = map (cAN >>> convert) (unsafeArrTail h.api_damage)
      where
        convert :: Array Number -> Damage
        convert xs = mkDamage totalDmg
          where
            -- there are double attacks in the damage list
            -- because damecon triggers *AFTER* the attack actions are completed
            -- we need to sum damage values before turning it into real Damage
            totalDmg = sum (map normalizeDamage xs :: Array Int)
    accumulateDamage :: forall h r.
                        STA.STArray h Damage -- array reference
                     -> Int -> Damage -- target index and corresponding damage
                     -> Eff (st :: ST.ST h | r) Unit
    accumulateDamage arr targetInd damage = do
        dmg <- peekSTArrayUnsafe arr targetInd
        pokeSTArrayUnsafe arr targetInd (dmg <> damage :: Damage)
        pure unit
    resultDV :: forall h r . Eff (st :: ST.ST h | r) (STA.STArray h Damage)
    resultDV = do
        arr <- STA.thaw (replicate 24 mempty)
        A.zipWithA (accumulateDamage arr) eventTargets eventDamages
        pure arr

-- | calculate damage from support airstrike stages.
-- | note that only enemy is taking damage so this results in
-- | a single `DamageVector`.
calcSupportAirAttackDamage :: SupportAirInfo -> DamageVector
calcSupportAirAttackDamage info = DV $ convertFEDam info.api_stage3.api_edam

calcSupportAirAttackDamageAC :: SupportAirInfo -> LR (LR DamageVector)
calcSupportAirAttackDamageAC info = calcKoukuDamageAC info

-- | calculate damage from support shelling stages.
-- | note that only enemy is taking damage so this results in
-- | a single `DamageVector`.
calcSupportHouraiDamage :: SupportHouraiInfo -> DamageVector
calcSupportHouraiDamage info = DV $ convertFEDam info.api_damage

-- TODO: to be verified, for now I just guess it's an array of raw length 1+12
calcSupportHouraiDamageAC :: SupportHouraiInfo -> LR (LR DamageVector)
calcSupportHouraiDamageAC info = { left: memptyLR, right: lrMap mkDV converted }
  where
    converted = fleetSplit false (convertFEDam info.api_damage)

-- | ally fleet's role in this battle
data FleetRole
  = FRMain -- ally main vs. enemy main
  | FREscort -- ally escort vs. enemy main
  | FRSupport -- ally support exped vs. enemy main
  | FRLandBased -- ally LBAS vs. enemy main
--  | FRMainEEscort -- ally main vs. enemy escort
--  | FRLandBasedEEscort -- ally LBAS vs. enemy escort

-- | `toCombined role dv` converts a `LR DamageVector` whose left part
-- | is playing role `role` into `CombinedDamageVector`
toGCombined :: FleetRole -> LR DamageVector -> GCombinedDamageVector
toGCombined r dv = case r of
    FRMain      ->
        { allyMain: dv.left, allyEscort: mempty
        , enemyMain: dv.right, enemyEscort: mempty }
    FREscort    ->
        { allyMain: mempty, allyEscort: dv.left
        , enemyMain: dv.right, enemyEscort: mempty }
    FRSupport   ->
        { allyMain: mempty, allyEscort: mempty
        , enemyMain: dv.right, enemyEscort: mempty }
    FRLandBased ->
        { allyMain: mempty, allyEscort: mempty
        , enemyMain: dv.right, enemyEscort: mempty }

toCombined :: FleetRole -> LR DamageVector -> CombinedDamageVector
toCombined r dv = toCombinedBattle (toGCombined r dv)

-- TODO: add check on dv.left.right, which should be all 0
toCombinedAC :: LR (LR DamageVector) -> CombinedDamageVectorAC
toCombinedAC dv =
    { main: dv.left.left
    , enemyMain: dv.right.left
    , enemyEscort: dv.right.right
    }

-- | apply a single `DamageVector` on a single fleet
applyDamageVector :: DamageVector -> FleetInfo Ship -> FleetInfo Ship
applyDamageVector dv fleet = A.zipWith combine (getDV dv) fleet
  where
    combine :: Damage -> Maybe Ship -> Maybe Ship
    combine dmg ms = applyDamage dmg <$> ms

-- | apply `NormalDamageVector` on a normal fleet of ships (including enemy ships)
applyNormalDamageVector :: NormalDamageVector -> NormalFleetInfo Ship -> NormalFleetInfo Ship
applyNormalDamageVector ndv fleet =
    dupAsNormalBattle applyDamageVector
      `appNormalBattle` ndv
      `appNormalBattle` fleet

-- | apply `CombinedDamageVector` on a combined fleet of ships (including enemy ships)
applyCombinedDamageVector :: CombinedDamageVector -> CombinedFleetInfo Ship -> CombinedFleetInfo Ship
applyCombinedDamageVector ndv fleet =
    dupAsCombinedBattle applyDamageVector
      `appCombinedBattle` ndv
      `appCombinedBattle` fleet

{- | This module implements `DamageVector`:
   | an array of `Damage`s that can be applied to one fleet of `Ship`s.
   | most functions prefixed with `calc` in this module
   | produces either `LR DamageVector` or a single `DamageVector`.
   | which can then be converted `NormalDamageVector` or `CombinedDamageVector`
   | depending on what the corresponding data is standing for.
 -}
module KanColle.DamageAnalysis.DamageVector
  ( DamageVector, getDV, mkDV

  , NormalDamageVector
  , GCombinedDamageVector
  , CombinedDamageVector, CombinedDamageVectorAC
  , FleetRole(..)

  , calcKoukuDamage
  , calcKoukuDamageAC
  , calcKoukuDamageCombined
  , calcSupportAirAttackDamage
  , calcSupportAirAttackDamageAC
  , calcSupportHouraiDamage
  , calcSupportHouraiDamageAC
  , calcHougekiDamage
  , calcRaigekiDamage

  , calcRaigekiDamageAC
  , calcHougekiDamageAC

  , calcLandBasedKoukuDamage

  , toGCombined
  , toCombined
  , toCombinedAC

  , applyDamageVector
  , applyNormalDamageVector
  , applyCombinedDamageVector
  ) where

import Prelude
import Data.Monoid
import Data.Int as Int
import Data.Array as A
import Data.Array.ST as STA
import Data.Unfoldable hiding (fromMaybe)
import Control.Monad.Eff
import Control.Monad.ST as ST
import KanColle.Util
import KanColle.DamageAnalysis.Types
import KanColle.DamageAnalysis.Damage

import Data.Foreign
import Data.Foldable
import KanColle.KCAPI.Battle
import Data.Maybe
import Data.String (joinWith)

-- when flagship protection happens
-- and extra 0.1 will be added to that value
-- for our purpose we don't need to deal with that
-- so all damages can be normalized as Ints
normalizeDamage :: Number -> Int
normalizeDamage x
  | x >= 0.0 = Int.floor x
  | x == (-1.0) = 0
  | otherwise =
    let warning = "invalid damage number: " <> show x
    in traceWarn warning (\_ -> 0)

-- | `DamageVector` is a 6-element array of `Damage`s
newtype DamageVector = DV (Array Damage)

-- | fetches the array inside `DamageVector`
getDV :: DamageVector -> Array Damage
getDV (DV v) = v

-- | creates `DamageVector` from arrays.
-- | array has to be of length 6
mkDV :: Array Damage -> DamageVector
mkDV xs = if check
    then DV xs
    else throwWith "mkDV: array size should be 6"
  where
    check = A.length xs == 6

instance semigroupDamageVector :: Semigroup DamageVector where
  append (DV a) (DV b) = DV (A.zipWith (<>) a b)

instance monoidDamageVector :: Monoid DamageVector where
  mempty = DV (replicate 6 mempty)

-- | `DamageVector` for normal battles
type NormalDamageVector = NormalBattle DamageVector

-- | `DamageVector` for battles involving comined fleets
type GCombinedDamageVector = GCombinedBattle DamageVector
type CombinedDamageVector = CombinedBattle DamageVector
type CombinedDamageVectorAC = CombinedBattleAC DamageVector

debugShowDV :: DamageVector -> String
debugShowDV (DV xs) = joinWith "," (map (show <<< damageToInt) xs)

-- TODO: abandon this and always use the safe version
-- | get `DamageVector` from raw `fDam` and `eDam` fields
fromFDamAndEDam :: KoukuStage3 -> LR DamageVector
fromFDamAndEDam ks3 =
    { left: DV (convertFEDam fDam)
    , right: DV (convertFEDam eDam) }
  where
    fDam = fromMaybe mempty (getKoukuStage3FDam ks3)
    eDam = fromMaybe mempty (getKoukuStage3EDam ks3)

-- L: ally fleet, LL: ally main, LR: ally escort
-- R: enemy fleet: RL: enemy main, RR: enemy escort
fromFDamAndEDamAC :: forall a.
                  { api_fdam :: Array Number
                  , api_edam :: Array Number | a} -> LR (LR DamageVector)
fromFDamAndEDamAC v = (lrMap >>> lrMap) DV
    { left: allyDams
    , right: enemyDams }
  where
    longFDam = convertFEDam v.api_fdam
    allyDams = fleetSplit false longFDam

    longEDam = convertFEDam v.api_edam
    enemyDams = fleetSplit false longEDam

-- (internal use only)
-- an "-1" is put in front of both api_fdam and api_edam
-- we first drop that element and then convert
-- integers intos Damages
convertFEDam :: Array Number -> Array Damage
convertFEDam = unsafeArrTail >>> map (normalizeDamage >>> mkDamage)

-- | calculate damage from kouku (aerial) stages
calcKoukuDamage :: Kouku -> LR DamageVector
calcKoukuDamage kk = fromFDamAndEDam kk.api_stage3

-- calcKoukuDamageAC :: Kouku -> LR DamageVector
-- calcKoukuDamageAC kk = fromFDamAndEDam kk.api_stage3_combined

-- | calculate damage from ally vs. enemy main fleet kouku stages (for Abyssal Combined fleet)
calcKoukuDamageACEMain :: forall a . {api_stage3 :: KoukuStage3 | a} -> LR DamageVector
calcKoukuDamageACEMain kk = fromFDamAndEDam kk.api_stage3

-- | calculate damage from ally vs. enemy escort fleet kouku stages (for Abyssal Combined fleet)
calcKoukuDamageACEEscort :: forall a. {api_stage3_combined :: KoukuStage3 | a} -> LR DamageVector
calcKoukuDamageACEEscort kk = fromFDamAndEDam kk.api_stage3_combined

calcKoukuDamageAC :: forall a.
                  { api_stage3 :: KoukuStage3
                  , api_stage3_combined :: KoukuStage3
                  | a} -> LR (LR DamageVector)
calcKoukuDamageAC kk =
    { left:
      { left: resultEMain.left <> resultEEscort.left
      , right: mempty }
    , right:
      { left: resultEMain.right
      , right: resultEEscort.right } }
  where
    resultEMain = calcKoukuDamageACEMain kk
    resultEEscort = calcKoukuDamageACEEscort kk

-- | calculate damage from kouku (aerial) stages (combined fleet)
-- | note that only escort fleet is taking damage. so we just need DamageVector
calcKoukuDamageCombined :: Kouku -> DamageVector
calcKoukuDamageCombined kk = DV (convertFEDam (kk.api_stage3_combined.api_fdam))

-- | calculate land based kouku damages. whose api_stage3 only have an api_edam field
-- | (because only enemies are taking damage)
calcLandBasedKoukuDamage :: Kouku -> DamageVector
calcLandBasedKoukuDamage lbkk = case getKoukuStage3 lbkk  of
    Just stage3 -> DV (convertFEDam stage3.api_edam)
    Nothing -> mempty

-- | calculate damage from raigeki (torpedo) stages
calcRaigekiDamage :: Raigeki -> LR DamageVector
calcRaigekiDamage = fromFDamAndEDam

calcRaigekiDamageAC :: Raigeki -> LR (LR DamageVector)
calcRaigekiDamageAC = fromFDamAndEDamAC

-- | calculate damage from hougeki (shelling) stages
calcHougekiDamage :: Hougeki -> LR DamageVector
calcHougekiDamage h =
    if lengthCheck
      then let resultArr = runPure (STA.runSTArray resultDV)
           in lrMap DV (fleetSplit false resultArr)
      else throwWith "invalid: api_df_list and api_damage length mismatch"
  where
    cAI :: Foreign -> Array Int
    cAI = unsafeFromForeign
    cAN :: Foreign -> Array Number
    cAN = unsafeFromForeign

    lengthCheck = A.length eventTargets == A.length eventDamages

    eventTargets :: Array Int
    eventTargets = map (cAI >>> toOne) (unsafeArrTail h.api_df_list)
      where
        -- merge targets like [a,a] into a-1 (so that it's zero-indexed)
        -- there are 2 checks:
        --   elements of api_df_list should:
        --   * be non-empty
        --   * all elements of it should be the same
        -- the reason for minus 1 in the end is to make it zero-based
        toOne :: Array Int -> Int
        toOne xs = case A.uncons xs of
          Just {head: y, tail: ys} ->
            if all (\v -> v == y) ys || all (\v -> v == -1) ys
              then y - 1
              else throwWith "invalid: elements are different in api_df_list"
          Nothing -> throwWith "invalid: empty api_df_list element"

    eventDamages :: Array Damage
    eventDamages = map (cAN >>> convert) (unsafeArrTail h.api_damage)
      where
        convert :: Array Number -> Damage
        convert xs = mkDamage totalDmg
          where
            -- there are double attacks in the damage list
            -- because damecon triggers *AFTER* the attack actions are completed
            -- we need to sum damage values before turning it into real Damage
            totalDmg = sum (map normalizeDamage xs :: Array Int)

    accumulateDamage :: forall h r.
                        STA.STArray h Damage -- array reference
                     -> Int -> Damage -- target index and corresponding damage
                     -> Eff (st :: ST.ST h | r) Unit
    accumulateDamage arr targetInd damage = do
        dmg <- peekSTArrayUnsafe arr targetInd
        pokeSTArrayUnsafe arr targetInd (dmg <> damage :: Damage)
        pure unit
    resultDV :: forall h r . Eff (st :: ST.ST h | r) (STA.STArray h Damage)
    resultDV = do
        arr <- STA.thaw (replicate 12 mempty)
        A.zipWithA (accumulateDamage arr) eventTargets eventDamages
        pure arr

calcHougekiDamageAC :: Hougeki -> LR (LR DamageVector)
calcHougekiDamageAC h =
    if lengthCheck
      then let resultArr = runPure (STA.runSTArray resultDV)
           in { left:
                { left: DV (A.slice 0 6 resultArr)
                , right: DV (A.slice 6 12 resultArr) }
              , right:
                { left: DV (A.slice 12 18 resultArr)
                , right: DV (A.slice 18 24 resultArr) }
              }
      else throwWith "invalid: api_df_list / api_damage / api_at_eflag length mismatch"
  where
    cAI :: Foreign -> Array Int
    cAI = unsafeFromForeign
    cAN :: Foreign -> Array Number
    cAN = unsafeFromForeign

    dfList = unsafeArrTail h.api_df_list
    atEFlag :: Array Int
    atEFlag = unsafeArrTail h.api_at_eflag
    damageList = unsafeArrTail h.api_damage

    lengthCheck = A.length dfList == A.length atEFlag
               && A.length dfList == A.length damageList

    eventTargetsRaw :: Array Int
    eventTargetsRaw = map (cAI >>> toOne) (unsafeArrTail h.api_df_list)
      where
        -- merge targets like [a,a] into a-1 (so that it's zero-indexed)
        -- there are 2 checks:
        --   elements of api_df_list should:
        --   * be non-empty
        --   * all elements of it should be the same
        -- the reason for minus 1 in the end is to make it zero-based
        toOne :: Array Int -> Int
        toOne xs = case A.uncons xs of
          Just {head: y, tail: ys} ->
            if all (\v -> v == y) ys || all (\v -> v == -1) ys
              then y - 1
              else throwWith "invalid: elements are different in api_df_list"
          Nothing -> throwWith "invalid: empty api_df_list element"

    eventTargets :: Array Int
    eventTargets = A.zipWith updateTarget eventTargetsRaw atEFlag
      where
        updateTarget rawTarget eFlag = case eFlag of
            -- ally attacks in this turn, so the target is an enemy
            0 -> rawTarget + 12
            -- enemy attacks in this turn, so the target is an ally
            -- nothing to be doing in this case
            1 -> rawTarget
            _ -> throwWith "invalid api_at_eflag element"

    eventDamages :: Array Damage
    eventDamages = map (cAN >>> convert) (unsafeArrTail h.api_damage)
      where
        convert :: Array Number -> Damage
        convert xs = mkDamage totalDmg
          where
            -- there are double attacks in the damage list
            -- because damecon triggers *AFTER* the attack actions are completed
            -- we need to sum damage values before turning it into real Damage
            totalDmg = sum (map normalizeDamage xs :: Array Int)
    accumulateDamage :: forall h r.
                        STA.STArray h Damage -- array reference
                     -> Int -> Damage -- target index and corresponding damage
                     -> Eff (st :: ST.ST h | r) Unit
    accumulateDamage arr targetInd damage = do
        dmg <- peekSTArrayUnsafe arr targetInd
        pokeSTArrayUnsafe arr targetInd (dmg <> damage :: Damage)
        pure unit
    resultDV :: forall h r . Eff (st :: ST.ST h | r) (STA.STArray h Damage)
    resultDV = do
        arr <- STA.thaw (replicate 24 mempty)
        A.zipWithA (accumulateDamage arr) eventTargets eventDamages
        pure arr

-- | calculate damage from support airstrike stages.
-- | note that only enemy is taking damage so this results in
-- | a single `DamageVector`.
calcSupportAirAttackDamage :: SupportAirInfo -> DamageVector
calcSupportAirAttackDamage info = DV $ convertFEDam info.api_stage3.api_edam

calcSupportAirAttackDamageAC :: SupportAirInfo -> LR (LR DamageVector)
calcSupportAirAttackDamageAC info = calcKoukuDamageAC info

-- | calculate damage from support shelling stages.
-- | note that only enemy is taking damage so this results in
-- | a single `DamageVector`.
calcSupportHouraiDamage :: SupportHouraiInfo -> DamageVector
calcSupportHouraiDamage info = DV $ convertFEDam info.api_damage

-- TODO: to be verified, for now I just guess it's an array of raw length 1+12
calcSupportHouraiDamageAC :: SupportHouraiInfo -> LR (LR DamageVector)
calcSupportHouraiDamageAC info = { left: memptyLR, right: lrMap mkDV converted }
  where
    converted = fleetSplit false (convertFEDam info.api_damage)

-- | ally fleet's role in this battle
data FleetRole
  = FRMain -- ally main vs. enemy main
  | FREscort -- ally escort vs. enemy main
  | FRSupport -- ally support exped vs. enemy main
  | FRLandBased -- ally LBAS vs. enemy main
--  | FRMainEEscort -- ally main vs. enemy escort
--  | FRLandBasedEEscort -- ally LBAS vs. enemy escort

-- | `toCombined role dv` converts a `LR DamageVector` whose left part
-- | is playing role `role` into `CombinedDamageVector`
toGCombined :: FleetRole -> LR DamageVector -> GCombinedDamageVector
toGCombined r dv = case r of
    FRMain      ->
        { allyMain: dv.left, allyEscort: mempty
        , enemyMain: dv.right, enemyEscort: mempty }
    FREscort    ->
        { allyMain: mempty, allyEscort: dv.left
        , enemyMain: dv.right, enemyEscort: mempty }
    FRSupport   ->
        { allyMain: mempty, allyEscort: mempty
        , enemyMain: dv.right, enemyEscort: mempty }
    FRLandBased ->
        { allyMain: mempty, allyEscort: mempty
        , enemyMain: dv.right, enemyEscort: mempty }

toCombined :: FleetRole -> LR DamageVector -> CombinedDamageVector
toCombined r dv = toCombinedBattle (toGCombined r dv)

-- TODO: add check on dv.left.right, which should be all 0
toCombinedAC :: LR (LR DamageVector) -> CombinedDamageVectorAC
toCombinedAC dv =
    { main: dv.left.left
    , enemyMain: dv.right.left
    , enemyEscort: dv.right.right
    }

-- | apply a single `DamageVector` on a single fleet
applyDamageVector :: DamageVector -> FleetInfo Ship -> FleetInfo Ship
applyDamageVector dv fleet = A.zipWith combine (getDV dv) fleet
  where
    combine :: Damage -> Maybe Ship -> Maybe Ship
    combine dmg ms = applyDamage dmg <$> ms

-- | apply `NormalDamageVector` on a normal fleet of ships (including enemy ships)
applyNormalDamageVector :: NormalDamageVector -> NormalFleetInfo Ship -> NormalFleetInfo Ship
applyNormalDamageVector ndv fleet =
    dupAsNormalBattle applyDamageVector
      `appNormalBattle` ndv
      `appNormalBattle` fleet

-- | apply `CombinedDamageVector` on a combined fleet of ships (including enemy ships)
applyCombinedDamageVector :: CombinedDamageVector -> CombinedFleetInfo Ship -> CombinedFleetInfo Ship
applyCombinedDamageVector ndv fleet =
    dupAsCombinedBattle applyDamageVector
      `appCombinedBattle` ndv
      `appCombinedBattle` fleet


Take advantage of record wildcards

Take advantage of record wildcards

module GulpPurescript.Options
  ( pscOptions
  , pscOptionsNoOutput
  , pscMakeOptions
  , pscDocsOptions
  ) where

import Control.Alt ((<|>))

import Data.Array (concat)
import Data.Either (Either(..), either)
import Data.Foreign (Foreign(), ForeignError(TypeMismatch), F())
import Data.Foreign.Class (IsForeign, read, readProp)
import Data.Foreign.NullOrUndefined (NullOrUndefined(..), runNullOrUndefined)
import Data.Maybe (Maybe(..), fromMaybe, maybe)
import Data.Tuple (Tuple())
import Data.Tuple.Nested (tuple2)

noPreludeOpt = "no-prelude"

noPreludeKey = "noPrelude"

noOptsOpt = "no-opts"

noOptsKey = "noOpts"

noMagicDoOpt = "no-magic-do"

noMagicDoKey = "noMagicDo"

noTcoOpt = "no-tco"

noTcoKey = "noTco"

verboseErrorsOpt = "verbose-errors"

verboseErrorsKey = "verboseErrors"

outputOpt = "output"

outputKey = outputOpt

browserNamespaceOpt = "browser-namespace"

browserNamespaceKey = "browserNamespace"

commentsOpt = "comments"

commentsKey = commentsOpt

noPrefixOpt = "no-prefix"

noPrefixKey = "noPrefix"

mainOpt = "main"

mainKey = mainOpt

moduleOpt = "module"

moduleKey = moduleOpt

codegenOpt = "codegen"

codegenKey = codegenOpt

externsOpt = "externs"

externsKey = externsOpt

formatOpt = "format"

formatKey = formatOpt

ffiOpt = "ffi"

ffiKey = ffiOpt

newtype Psc
  = Psc { noPrelude :: NullOrUndefined Boolean
        , noTco :: NullOrUndefined Boolean
        , noMagicDo :: NullOrUndefined Boolean
        , main :: NullOrUndefined (Either Boolean String)
        , noOpts :: NullOrUndefined Boolean
        , verboseErrors :: NullOrUndefined Boolean
        , comments :: NullOrUndefined Boolean
        , browserNamespace :: NullOrUndefined String
        , "module" :: NullOrUndefined [String]
        , codegen :: NullOrUndefined [String]
        , output :: NullOrUndefined String
        , externs :: NullOrUndefined String
        , noPrefix :: NullOrUndefined Boolean
        , ffi :: NullOrUndefined [String]
        }

newtype PscMake
  = PscMake { noPrelude :: NullOrUndefined Boolean
            , noOpts :: NullOrUndefined Boolean
            , noMagicDo :: NullOrUndefined Boolean
            , noTco :: NullOrUndefined Boolean
            , verboseErrors :: NullOrUndefined Boolean
            , comments :: NullOrUndefined Boolean
            , noPrefix :: NullOrUndefined Boolean
            , output :: NullOrUndefined String
            , ffi :: NullOrUndefined [String]
            }

newtype PscDocs
  = PscDocs { format :: NullOrUndefined Format }

data Format = Markdown | ETags | CTags

instance isForeignEither :: (IsForeign a, IsForeign b) => IsForeign (Either a b) where
  read a = (Left <$> read a :: F a) <|>
           (Right <$> read a :: F b)

instance isForeignPsc :: IsForeign Psc where
  read obj =
    Psc <$> ({ noPrelude: _
             , noTco: _
             , noMagicDo: _
             , main: _
             , noOpts: _
             , verboseErrors: _
             , comments: _
             , browserNamespace: _
             , "module": _
             , codegen: _
             , output: _
             , externs: _
             , noPrefix: _
             , ffi: _
             } <$> readProp noPreludeKey obj
               <*> readProp noTcoKey obj
               <*> readProp noMagicDoKey obj
               <*> readProp mainKey obj
               <*> readProp noOptsKey obj
               <*> readProp verboseErrorsKey obj
               <*> readProp commentsKey obj
               <*> readProp browserNamespaceKey obj
               <*> readProp moduleKey obj
               <*> readProp codegenKey obj
               <*> readProp outputKey obj
               <*> readProp externsKey obj
               <*> readProp noPrefixKey obj
               <*> readProp ffiKey obj)

instance isForeignPscMake :: IsForeign PscMake where
  read obj =
    PscMake <$> ({ output: _
                 , noPrelude: _
                 , noTco: _
                 , noMagicDo: _
                 , noOpts: _
                 , verboseErrors: _
                 , comments: _
                 , noPrefix: _
                 , ffi: _
                 } <$> readProp outputKey obj
                   <*> readProp noPreludeKey obj
                   <*> readProp noTcoKey obj
                   <*> readProp noMagicDoKey obj
                   <*> readProp noOptsKey obj
                   <*> readProp verboseErrorsKey obj
                   <*> readProp commentsKey obj
                   <*> readProp noPrefixKey obj
                   <*> readProp ffiKey obj)

instance isForeignPscDocs :: IsForeign PscDocs where
  read obj = PscDocs <<< { format: _ } <$> readProp formatKey obj

instance isForeignFormat :: IsForeign Format where
  read val = read val >>= (\a -> case a of
                                      "markdown" -> Right Markdown
                                      "etags" -> Right ETags
                                      "ctags" -> Right CTags
                                      a -> Left $ TypeMismatch "Format" a)

mkBoolean :: String -> NullOrUndefined Boolean -> [String]
mkBoolean key opt = maybe [] (\a -> if a then ["--" ++ key] else []) (runNullOrUndefined opt)

mkString :: String -> NullOrUndefined String -> [String]
mkString key opt = maybe [] (\a -> ["--" ++ key ++ "=" ++ a]) (runNullOrUndefined opt)

mkBooleanString :: String -> NullOrUndefined (Either Boolean String) -> [String]
mkBooleanString key opt = maybe [] (either (\a -> mkBoolean key (NullOrUndefined $ Just a))
                                           (\a -> mkString key (NullOrUndefined $ Just a)))
                                   (runNullOrUndefined opt)

mkStringArray :: String -> NullOrUndefined [String] -> [String]
mkStringArray key opt = concat $ mkString key <$> (NullOrUndefined <<< Just)
                                              <$> (fromMaybe [] $ runNullOrUndefined opt)

mkPathArray :: String -> NullOrUndefined [String] -> [String]
mkPathArray key opt = concat $ mkString key <$> (NullOrUndefined <<< Just)
                                            <$> (fromMaybe [] (runNullOrUndefined opt) >>= expandGlob)

foreign import expandGlob
  """
  var expandGlob = (function () {
    var glob = require("glob");
    return function (pattern) {
      return glob.sync(pattern);
    };
  }());
  """ :: String -> [String]

mkFormat :: String -> NullOrUndefined Format -> [String]
mkFormat key opt = mkString key (maybe j (\a -> case a of
                                                     Markdown -> i "markdown"
                                                     ETags -> i "etags"
                                                     CTags -> i "ctags") $ runNullOrUndefined opt)
  where i a = NullOrUndefined $ Just a
        j = NullOrUndefined Nothing

foldPscOptions :: Psc -> [String]
foldPscOptions (Psc a) = mkBoolean noPreludeOpt a.noPrelude <>
                         mkBoolean noTcoOpt a.noTco <>
                         mkBoolean noMagicDoOpt a.noMagicDo <>
                         mkBooleanString mainOpt a.main <>
                         mkBoolean noOptsOpt a.noOpts <>
                         mkBoolean verboseErrorsOpt a.verboseErrors <>
                         mkBoolean commentsOpt a.comments <>
                         mkString browserNamespaceOpt a.browserNamespace <>
                         mkStringArray moduleOpt a."module" <>
                         mkStringArray codegenOpt a.codegen <>
                         mkString outputOpt a.output <>
                         mkString externsOpt a.externs <>
                         mkBoolean noPrefixOpt a.noPrefix <>
                         mkPathArray ffiOpt a.ffi

pscOptions :: Foreign -> [String]
pscOptions opts = either (const []) foldPscOptions parsed
  where parsed = read opts :: F Psc

pscOptionsNoOutput :: Foreign -> Tuple (Maybe String) [String]
pscOptionsNoOutput opts = either (const $ tuple2 Nothing []) fold parsed
  where parsed = read opts :: F Psc
        fold (Psc a) = tuple2 (runNullOrUndefined a.output)
                              (foldPscOptions (Psc $ a { output = NullOrUndefined Nothing }))

pscMakeOptions :: Foreign -> [String]
pscMakeOptions opts = either (const []) fold parsed
  where parsed = read opts :: F PscMake
        fold (PscMake a) = mkString outputOpt a.output <>
                           mkBoolean noPreludeOpt a.noPrelude <>
                           mkBoolean noTcoOpt a.noTco <>
                           mkBoolean noMagicDoOpt a.noMagicDo <>
                           mkBoolean noOptsOpt a.noOpts <>
                           mkBoolean verboseErrorsOpt a.verboseErrors <>
                           mkBoolean commentsOpt a.comments <>
                           mkBoolean noPrefixOpt a.noPrefix <>
                           mkPathArray ffiOpt a.ffi

pscDocsOptions :: Foreign -> [String]
pscDocsOptions opts = either (const []) fold parsed
  where parsed = read opts :: F PscDocs
        fold (PscDocs a) = mkFormat formatOpt a.format

module GulpPurescript.Options
  ( pscOptions
  , pscOptionsNoOutput
  , pscMakeOptions
  , pscDocsOptions
  ) where

import Control.Alt ((<|>))

import Data.Array (concat)
import Data.Either (Either(..), either)
import Data.Foreign (Foreign(), ForeignError(TypeMismatch), F())
import Data.Foreign.Class (IsForeign, read, readProp)
import Data.Foreign.NullOrUndefined (NullOrUndefined(..), runNullOrUndefined)
import Data.Maybe (Maybe(..), fromMaybe, maybe)
import Data.Tuple (Tuple())
import Data.Tuple.Nested (tuple2)

noPreludeOpt = "no-prelude"

noPreludeKey = "noPrelude"

noOptsOpt = "no-opts"

noOptsKey = "noOpts"

noMagicDoOpt = "no-magic-do"

noMagicDoKey = "noMagicDo"

noTcoOpt = "no-tco"

noTcoKey = "noTco"

verboseErrorsOpt = "verbose-errors"

verboseErrorsKey = "verboseErrors"

outputOpt = "output"

outputKey = outputOpt

browserNamespaceOpt = "browser-namespace"

browserNamespaceKey = "browserNamespace"

commentsOpt = "comments"

commentsKey = commentsOpt

noPrefixOpt = "no-prefix"

noPrefixKey = "noPrefix"

mainOpt = "main"

mainKey = mainOpt

moduleOpt = "module"

moduleKey = moduleOpt

codegenOpt = "codegen"

codegenKey = codegenOpt

externsOpt = "externs"

externsKey = externsOpt

formatOpt = "format"

formatKey = formatOpt

ffiOpt = "ffi"

ffiKey = ffiOpt

newtype Psc
  = Psc { noPrelude :: NullOrUndefined Boolean
        , noTco :: NullOrUndefined Boolean
        , noMagicDo :: NullOrUndefined Boolean
        , main :: NullOrUndefined (Either Boolean String)
        , noOpts :: NullOrUndefined Boolean
        , verboseErrors :: NullOrUndefined Boolean
        , comments :: NullOrUndefined Boolean
        , browserNamespace :: NullOrUndefined String
        , "module" :: NullOrUndefined [String]
        , codegen :: NullOrUndefined [String]
        , output :: NullOrUndefined String
        , externs :: NullOrUndefined String
        , noPrefix :: NullOrUndefined Boolean
        , ffi :: NullOrUndefined [String]
        }

newtype PscMake
  = PscMake { noPrelude :: NullOrUndefined Boolean
            , noOpts :: NullOrUndefined Boolean
            , noMagicDo :: NullOrUndefined Boolean
            , noTco :: NullOrUndefined Boolean
            , verboseErrors :: NullOrUndefined Boolean
            , comments :: NullOrUndefined Boolean
            , noPrefix :: NullOrUndefined Boolean
            , output :: NullOrUndefined String
            , ffi :: NullOrUndefined [String]
            }

newtype PscDocs
  = PscDocs { format :: NullOrUndefined Format }

data Format = Markdown | ETags | CTags

instance isForeignEither :: (IsForeign a, IsForeign b) => IsForeign (Either a b) where
  read a = (Left <$> read a :: F a) <|>
           (Right <$> read a :: F b)

instance isForeignPsc :: IsForeign Psc where
  read obj =
    (\a b c d e f g h i j k l m o ->
    Psc { noPrelude: a
        , noTco: b
        , noMagicDo: c
        , main: d
        , noOpts: e
        , verboseErrors: f
        , comments: g
        , browserNamespace: h
        , "module": i
        , codegen: j
        , output: k
        , externs: l
        , noPrefix: m
        , ffi: o
        }) <$> readProp noPreludeKey obj
           <*> readProp noTcoKey obj
           <*> readProp noMagicDoKey obj
           <*> readProp mainKey obj
           <*> readProp noOptsKey obj
           <*> readProp verboseErrorsKey obj
           <*> readProp commentsKey obj
           <*> readProp browserNamespaceKey obj
           <*> readProp moduleKey obj
           <*> readProp codegenKey obj
           <*> readProp outputKey obj
           <*> readProp externsKey obj
           <*> readProp noPrefixKey obj
           <*> readProp ffiKey obj

instance isForeignPscMake :: IsForeign PscMake where
  read obj =
    (\a b c d e f g h i ->
    PscMake { output: a
            , noPrelude: b
            , noTco: c
            , noMagicDo: d
            , noOpts: e
            , verboseErrors: f
            , comments: g
            , noPrefix: h
            , ffi: i
            }) <$> readProp outputKey obj
               <*> readProp noPreludeKey obj
               <*> readProp noTcoKey obj
               <*> readProp noMagicDoKey obj
               <*> readProp noOptsKey obj
               <*> readProp verboseErrorsKey obj
               <*> readProp commentsKey obj
               <*> readProp noPrefixKey obj
               <*> readProp ffiKey obj

instance isForeignPscDocs :: IsForeign PscDocs where
  read obj = (\a -> PscDocs { format: a }) <$> readProp formatKey obj

instance isForeignFormat :: IsForeign Format where
  read val = read val >>= (\a -> case a of
                                      "markdown" -> Right Markdown
                                      "etags" -> Right ETags
                                      "ctags" -> Right CTags
                                      a -> Left $ TypeMismatch "Format" a)

mkBoolean :: String -> NullOrUndefined Boolean -> [String]
mkBoolean key opt = maybe [] (\a -> if a then ["--" ++ key] else []) (runNullOrUndefined opt)

mkString :: String -> NullOrUndefined String -> [String]
mkString key opt = maybe [] (\a -> ["--" ++ key ++ "=" ++ a]) (runNullOrUndefined opt)

mkBooleanString :: String -> NullOrUndefined (Either Boolean String) -> [String]
mkBooleanString key opt = maybe [] (either (\a -> mkBoolean key (NullOrUndefined $ Just a))
                                           (\a -> mkString key (NullOrUndefined $ Just a)))
                                   (runNullOrUndefined opt)

mkStringArray :: String -> NullOrUndefined [String] -> [String]
mkStringArray key opt = concat $ mkString key <$> (NullOrUndefined <<< Just)
                                              <$> (fromMaybe [] $ runNullOrUndefined opt)

mkPathArray :: String -> NullOrUndefined [String] -> [String]
mkPathArray key opt = concat $ mkString key <$> (NullOrUndefined <<< Just)
                                            <$> (fromMaybe [] (runNullOrUndefined opt) >>= expandGlob)

foreign import expandGlob
  """
  var expandGlob = (function () {
    var glob = require("glob");
    return function (pattern) {
      return glob.sync(pattern);
    };
  }());
  """ :: String -> [String]

mkFormat :: String -> NullOrUndefined Format -> [String]
mkFormat key opt = mkString key (maybe j (\a -> case a of
                                                     Markdown -> i "markdown"
                                                     ETags -> i "etags"
                                                     CTags -> i "ctags") $ runNullOrUndefined opt)
  where i a = NullOrUndefined $ Just a
        j = NullOrUndefined Nothing

foldPscOptions :: Psc -> [String]
foldPscOptions (Psc a) = mkBoolean noPreludeOpt a.noPrelude <>
                         mkBoolean noTcoOpt a.noTco <>
                         mkBoolean noMagicDoOpt a.noMagicDo <>
                         mkBooleanString mainOpt a.main <>
                         mkBoolean noOptsOpt a.noOpts <>
                         mkBoolean verboseErrorsOpt a.verboseErrors <>
                         mkBoolean commentsOpt a.comments <>
                         mkString browserNamespaceOpt a.browserNamespace <>
                         mkStringArray moduleOpt a."module" <>
                         mkStringArray codegenOpt a.codegen <>
                         mkString outputOpt a.output <>
                         mkString externsOpt a.externs <>
                         mkBoolean noPrefixOpt a.noPrefix <>
                         mkPathArray ffiOpt a.ffi

pscOptions :: Foreign -> [String]
pscOptions opts = either (const []) foldPscOptions parsed
  where parsed = read opts :: F Psc

pscOptionsNoOutput :: Foreign -> Tuple (Maybe String) [String]
pscOptionsNoOutput opts = either (const $ tuple2 Nothing []) fold parsed
  where parsed = read opts :: F Psc
        fold (Psc a) = tuple2 (runNullOrUndefined a.output)
                              (foldPscOptions (Psc $ a { output = NullOrUndefined Nothing }))

pscMakeOptions :: Foreign -> [String]
pscMakeOptions opts = either (const []) fold parsed
  where parsed = read opts :: F PscMake
        fold (PscMake a) = mkString outputOpt a.output <>
                           mkBoolean noPreludeOpt a.noPrelude <>
                           mkBoolean noTcoOpt a.noTco <>
                           mkBoolean noMagicDoOpt a.noMagicDo <>
                           mkBoolean noOptsOpt a.noOpts <>
                           mkBoolean verboseErrorsOpt a.verboseErrors <>
                           mkBoolean commentsOpt a.comments <>
                           mkBoolean noPrefixOpt a.noPrefix <>
                           mkPathArray ffiOpt a.ffi

pscDocsOptions :: Foreign -> [String]
pscDocsOptions opts = either (const []) fold parsed
  where parsed = read opts :: F PscDocs
        fold (PscDocs a) = mkFormat formatOpt a.format