(.module:
  [library
   [lux #*
    [abstract
     equivalence
     [monad (#+ Monad do)]
     ["." order (#+ Order)]]
    [control
     ["." maybe]]
    [data
     ["p" product]
     [collection
      ["." list ("#\." monoid fold)]]]
    [macro
     ["." code]
     ["." template]]
    [math
     [number
      ["n" nat]]]]])

(def: error_message
  "Invariant violation")

(type: Color
  #Red
  #Black)

(type: (Node k v)
  {#color Color
   #key k
   #value v
   #left (Maybe (Node k v))
   #right (Maybe (Node k v))})

(template [<create> <color>]
  [(def: (<create> key value left right)
     (All [k v] (-> k v (Maybe (Node k v)) (Maybe (Node k v)) (Node k v)))
     {#color <color>
      #key key
      #value value
      #left left
      #right right})]

  [red   #Red]
  [black #Black]
  )

(type: .public (Dictionary k v)
  {#&order (Order k)
   #root (Maybe (Node k v))})

(def: .public (empty order)
  (All [k v] (-> (Order k) (Dictionary k v)))
  {#&order order
   #root #.None})

... TODO: Doing inneficient access of Order functions due to compiler bug.
... TODO: Must improve it as soon as bug is fixed.
(def: .public (value key dict)
  (All [k v] (-> k (Dictionary k v) (Maybe v)))
  (let [... (^open "_\.") (value@ #&order dict)
        ]
    (loop [node (value@ #root dict)]
      (case node
        #.None
        #.None

        (#.Some node)
        (let [node_key (value@ #key node)]
          (cond (\ dict = node_key key)
                ... (_\= node_key key)
                (#.Some (value@ #value node))

                (\ dict < node_key key)
                ... (_\< node_key key)
                (recur (value@ #left node))

                ... (_\> (value@ #key node) key)
                (recur (value@ #right node))))
        ))))

... TODO: Doing inneficient access of Order functions due to compiler bug.
... TODO: Must improve it as soon as bug is fixed.
(def: .public (key? dict key)
  (All [k v] (-> (Dictionary k v) k Bit))
  (let [... (^open "_\.") (value@ #&order dict)
        ]
    (loop [node (value@ #root dict)]
      (case node
        #.None
        #0

        (#.Some node)
        (let [node_key (value@ #key node)]
          (or (\ dict = node_key key)
              ... (_\= node_key key)
              (if (\ dict < node_key key)
                ... (_\< node_key key)
                (recur (value@ #left node))
                (recur (value@ #right node)))))))))

(template [<name> <side>]
  [(def: .public (<name> dict)
     {#.doc (example (~~ (template.text ["Yields value under the " <name> "imum key."])))}
     (All [k v] (-> (Dictionary k v) (Maybe v)))
     (case (value@ #root dict)
       #.None
       #.None

       (#.Some node)
       (loop [node node]
         (case (value@ <side> node)
           #.None
           (#.Some (value@ #value node))

           (#.Some side)
           (recur side)))))]

  [min #left]
  [max #right]
  )

(def: .public (size dict)
  (All [k v] (-> (Dictionary k v) Nat))
  (loop [node (value@ #root dict)]
    (case node
      #.None
      0

      (#.Some node)
      (++ (n.+ (recur (value@ #left node))
               (recur (value@ #right node)))))))

(def: .public empty?
  (All [k v] (-> (Dictionary k v) Bit))
  (|>> ..size (n.= 0)))

(template [<name> <other_color> <self_color> <no_change>]
  [(def: (<name> self)
     (All [k v] (-> (Node k v) (Node k v)))
     (case (value@ #color self)
       <other_color>
       (with@ #color <self_color> self)

       <self_color>
       <no_change>
       ))]

  [blackened #Red   #Black self]
  [reddened  #Black #Red   (panic! error_message)]
  )

(def: (with_left addition center)
  (All [k v] (-> (Node k v) (Node k v) (Node k v)))
  (case (value@ #color center)
    #Red
    (red (value@ #key center)
         (value@ #value center)
         (#.Some addition)
         (value@ #right center))
    
    #Black
    (with_expansions
      [<default_behavior> (as_is (black (value@ #key center)
                                        (value@ #value center)
                                        (#.Some addition)
                                        (value@ #right center)))]
      (case (value@ #color addition)
        #Red
        (case (value@ #left addition)
          (^multi (#.Some left)
                  {(value@ #color left) #Red})
          (red (value@ #key addition)
               (value@ #value addition)
               (#.Some (blackened left))
               (#.Some (black (value@ #key center)
                              (value@ #value center)
                              (value@ #right addition)
                              (value@ #right center))))

          _
          (case (value@ #right addition)
            (^multi (#.Some right)
                    {(value@ #color right) #Red})
            (red (value@ #key right)
                 (value@ #value right)
                 (#.Some (black (value@ #key addition)
                                (value@ #value addition)
                                (value@ #left addition)
                                (value@ #left right)))
                 (#.Some (black (value@ #key center)
                                (value@ #value center)
                                (value@ #right right)
                                (value@ #right center))))

            _
            <default_behavior>))
        
        #Black
        <default_behavior>))))

(def: (with_right addition center)
  (All [k v] (-> (Node k v) (Node k v) (Node k v)))
  (case (value@ #color center)
    #Red
    (red (value@ #key center)
         (value@ #value center)
         (value@ #left center)
         (#.Some addition))
    
    #Black
    (with_expansions
      [<default_behavior> (as_is (black (value@ #key center)
                                        (value@ #value center)
                                        (value@ #left center)
                                        (#.Some addition)))]
      (case (value@ #color addition)
        #Red
        (case (value@ #right addition)
          (^multi (#.Some right)
                  {(value@ #color right) #Red})
          (red (value@ #key addition)
               (value@ #value addition)
               (#.Some (black (value@ #key center)
                              (value@ #value center)
                              (value@ #left center)
                              (value@ #left addition)))
               (#.Some (blackened right)))

          _
          (case (value@ #left addition)
            (^multi (#.Some left)
                    {(value@ #color left) #Red})
            (red (value@ #key left)
                 (value@ #value left)
                 (#.Some (black (value@ #key center)
                                (value@ #value center)
                                (value@ #left center)
                                (value@ #left left)))
                 (#.Some (black (value@ #key addition)
                                (value@ #value addition)
                                (value@ #right left)
                                (value@ #right addition))))

            _
            <default_behavior>))
        
        #Black
        <default_behavior>))))

(def: .public (has key value dict)
  (All [k v] (-> k v (Dictionary k v) (Dictionary k v)))
  (let [(^open "_\.") (value@ #&order dict)
        root' (loop [?root (value@ #root dict)]
                (case ?root
                  #.None
                  (#.Some (red key value #.None #.None))

                  (#.Some root)
                  (let [reference (value@ #key root)]
                    (`` (cond (~~ (template [<comp> <tag> <add>]
                                    [(<comp> reference key)
                                     (let [side_root (value@ <tag> root)
                                           outcome (recur side_root)]
                                       (if (same? side_root outcome)
                                         ?root
                                         (#.Some (<add> (maybe.trusted outcome)
                                                        root))))]

                                    [_\<                           #left  ..with_left]
                                    [(order.> (value@ #&order dict)) #right ..with_right]
                                    ))

                              ... (_\= reference key)
                              (#.Some (with@ #value value root))
                              )))
                  ))]
    (with@ #root root' dict)))

(def: (left_balanced key value ?left ?right)
  (All [k v] (-> k v (Maybe (Node k v)) (Maybe (Node k v)) (Node k v)))
  (case ?left
    (^multi (#.Some left)
            {(value@ #color left) #Red}
            {(value@ #left left) (#.Some left>>left)}
            {(value@ #color left>>left) #Red})
    (red (value@ #key left)
         (value@ #value left)
         (#.Some (blackened left>>left))
         (#.Some (black key value (value@ #right left) ?right)))

    (^multi (#.Some left)
            {(value@ #color left) #Red}
            {(value@ #right left) (#.Some left>>right)}
            {(value@ #color left>>right) #Red})
    (red (value@ #key left>>right)
         (value@ #value left>>right)
         (#.Some (black (value@ #key left)
                        (value@ #value left)
                        (value@ #left left)
                        (value@ #left left>>right)))
         (#.Some (black key value
                        (value@ #right left>>right)
                        ?right)))

    _
    (black key value ?left ?right)))

(def: (right_balanced key value ?left ?right)
  (All [k v] (-> k v (Maybe (Node k v)) (Maybe (Node k v)) (Node k v)))
  (case ?right
    (^multi (#.Some right)
            {(value@ #color right) #Red}
            {(value@ #right right) (#.Some right>>right)}
            {(value@ #color right>>right) #Red})
    (red (value@ #key right)
         (value@ #value right)
         (#.Some (black key value ?left (value@ #left right)))
         (#.Some (blackened right>>right)))

    (^multi (#.Some right)
            {(value@ #color right) #Red}
            {(value@ #left right) (#.Some right>>left)}
            {(value@ #color right>>left) #Red})
    (red (value@ #key right>>left)
         (value@ #value right>>left)
         (#.Some (black key value ?left (value@ #left right>>left)))
         (#.Some (black (value@ #key right)
                        (value@ #value right)
                        (value@ #right right>>left)
                        (value@ #right right))))

    _
    (black key value ?left ?right)))

(def: (without_left key value ?left ?right)
  (All [k v] (-> k v (Maybe (Node k v)) (Maybe (Node k v)) (Node k v)))
  (case ?left
    (^multi (#.Some left)
            {(value@ #color left) #Red})
    (red key value (#.Some (blackened left)) ?right)

    _
    (case ?right
      (^multi (#.Some right)
              {(value@ #color right) #Black})
      (right_balanced key value ?left (#.Some (reddened right)))

      (^multi (#.Some right)
              {(value@ #color right) #Red}
              {(value@ #left right) (#.Some right>>left)}
              {(value@ #color right>>left) #Black})
      (red (value@ #key right>>left)
           (value@ #value right>>left)
           (#.Some (black key value ?left (value@ #left right>>left)))
           (#.Some (right_balanced (value@ #key right)
                                   (value@ #value right)
                                   (value@ #right right>>left)
                                   (\ maybe.functor map reddened (value@ #right right)))))

      _
      (panic! error_message))
    ))

(def: (without_right key value ?left ?right)
  (All [k v] (-> k v (Maybe (Node k v)) (Maybe (Node k v)) (Node k v)))
  (case ?right
    (^multi (#.Some right)
            {(value@ #color right) #Red})
    (red key value ?left (#.Some (blackened right)))

    _
    (case ?left
      (^multi (#.Some left)
              {(value@ #color left) #Black})
      (left_balanced key value (#.Some (reddened left)) ?right)

      (^multi (#.Some left)
              {(value@ #color left) #Red}
              {(value@ #right left) (#.Some left>>right)}
              {(value@ #color left>>right) #Black})
      (red (value@ #key left>>right)
           (value@ #value left>>right)
           (#.Some (left_balanced (value@ #key left)
                                  (value@ #value left)
                                  (\ maybe.functor map reddened (value@ #left left))
                                  (value@ #left left>>right)))
           (#.Some (black key value (value@ #right left>>right) ?right)))

      _
      (panic! error_message)
      )))

(def: (prepended ?left ?right)
  (All [k v] (-> (Maybe (Node k v)) (Maybe (Node k v)) (Maybe (Node k v))))
  (case [?left ?right]
    [#.None _]
    ?right

    [_ #.None]
    ?left

    [(#.Some left) (#.Some right)]
    (case [(value@ #color left) (value@ #color right)]
      [#Red #Red]
      (do maybe.monad
        [fused (prepended (value@ #right left) (value@ #right right))]
        (case (value@ #color fused)
          #Red
          (in (red (value@ #key fused)
                   (value@ #value fused)
                   (#.Some (red (value@ #key left)
                                (value@ #value left)
                                (value@ #left left)
                                (value@ #left fused)))
                   (#.Some (red (value@ #key right)
                                (value@ #value right)
                                (value@ #right fused)
                                (value@ #right right)))))

          #Black
          (in (red (value@ #key left)
                   (value@ #value left)
                   (value@ #left left)
                   (#.Some (red (value@ #key right)
                                (value@ #value right)
                                (#.Some fused)
                                (value@ #right right)))))))

      [#Red #Black]
      (#.Some (red (value@ #key left)
                   (value@ #value left)
                   (value@ #left left)
                   (prepended (value@ #right left)
                              ?right)))
      
      [#Black #Red]
      (#.Some (red (value@ #key right)
                   (value@ #value right)
                   (prepended ?left
                              (value@ #left right))
                   (value@ #right right)))

      [#Black #Black]
      (do maybe.monad
        [fused (prepended (value@ #right left) (value@ #left right))]
        (case (value@ #color fused)
          #Red
          (in (red (value@ #key fused)
                   (value@ #value fused)
                   (#.Some (black (value@ #key left)
                                  (value@ #value left)
                                  (value@ #left left)
                                  (value@ #left fused)))
                   (#.Some (black (value@ #key right)
                                  (value@ #value right)
                                  (value@ #right fused)
                                  (value@ #right right)))))
          
          #Black
          (in (without_left (value@ #key left)
                            (value@ #value left)
                            (value@ #left left)
                            (#.Some (black (value@ #key right)
                                           (value@ #value right)
                                           (#.Some fused)
                                           (value@ #right right)))))
          ))
      )

    _
    (undefined)))

(def: .public (lacks key dict)
  (All [k v] (-> k (Dictionary k v) (Dictionary k v)))
  (let [(^open "_\.") (value@ #&order dict)
        [?root found?] (loop [?root (value@ #root dict)]
                         (case ?root
                           #.None
                           [#.None #0]

                           (#.Some root)
                           (let [root_key (value@ #key root)
                                 root_val (value@ #value root)]
                             (if (_\= root_key key)
                               [(prepended (value@ #left root)
                                           (value@ #right root))
                                #1]
                               (let [go_left? (_\< root_key key)]
                                 (case (recur (if go_left?
                                                (value@ #left root)
                                                (value@ #right root)))
                                   [#.None #0]
                                   [#.None #0]

                                   [side_outcome _]
                                   (if go_left?
                                     (case (value@ #left root)
                                       (^multi (#.Some left)
                                               {(value@ #color left) #Black})
                                       [(#.Some (without_left root_key root_val side_outcome (value@ #right root)))
                                        #0]

                                       _
                                       [(#.Some (red root_key root_val side_outcome (value@ #right root)))
                                        #0])
                                     (case (value@ #right root)
                                       (^multi (#.Some right)
                                               {(value@ #color right) #Black})
                                       [(#.Some (without_right root_key root_val (value@ #left root) side_outcome))
                                        #0]

                                       _
                                       [(#.Some (red root_key root_val (value@ #left root) side_outcome))
                                        #0])
                                     )))
                               ))
                           ))]
    (case ?root
      #.None
      (if found?
        (with@ #root ?root dict)
        dict)

      (#.Some root)
      (with@ #root (#.Some (blackened root)) dict)
      )))

(def: .public (revised key transform dict)
  (All [k v] (-> k (-> v v) (Dictionary k v) (Dictionary k v)))
  (case (..value key dict)
    (#.Some old)
    (..has key (transform old) dict)

    #.None
    dict))

(def: .public (of_list order list)
  (All [k v] (-> (Order k) (List [k v]) (Dictionary k v)))
  (list\fold (function (_ [key value] dict)
               (has key value dict))
             (empty order)
             list))

(template [<name> <type> <output>]
  [(def: .public (<name> dict)
     (All [k v] (-> (Dictionary k v) (List <type>)))
     (loop [node (value@ #root dict)]
       (case node
         #.None
         (list)

         (#.Some node')
         ($_ list\compose
             (recur (value@ #left node'))
             (list <output>)
             (recur (value@ #right node'))))))]

  [entries [k v] [(value@ #key node') (value@ #value node')]]
  [keys    k     (value@ #key node')]
  [values  v     (value@ #value node')]
  )

(implementation: .public (equivalence (^open ",\."))
  (All [k v] (-> (Equivalence v) (Equivalence (Dictionary k v))))
  
  (def: (= reference sample)
    (let [(^open "/\.") (value@ #&order reference)]
      (loop [entriesR (entries reference)
             entriesS (entries sample)]
        (case [entriesR entriesS]
          [#.End #.End]
          #1

          [(#.Item [keyR valueR] entriesR') (#.Item [keyS valueS] entriesS')]
          (and (/\= keyR keyS)
               (,\= valueR valueS)
               (recur entriesR' entriesS'))

          _
          #0)))))