module Trees where

open import Data.Nat
open import Data.Sum
open import Data.Product
open import Relation.Binary.PropositionalEquality
open import Data.Bool
open import Data.Empty
open import Data.Unit

data Flag : Set where
  F₁ : Flag
  F₂ : Flag

data Tree A : Set where
  Leaf    :                                    Tree A
  Branch  : (m : Flag)(a : A)(t₁ t₂ : Tree A) → Tree A

insertFullTree : ∀ {A} → A → Tree A → Tree A
insertFullTree x Leaf              = Branch F₂ x Leaf Leaf 
insertFullTree x (Branch _ a t₁ t₂) = Branch F₁ a (insertFullTree x t₁) t₂

data Maybe A : Set where
  just : A → Maybe A
  nothing : Maybe A

data FullTree A : ℕ → Tree A → Set where
  full-leaf   : FullTree A 0 Leaf 
  full-branch : {a : A}{t₁ t₂ : Tree A}{n : ℕ}
              → FullTree A n t₁
              → FullTree A n t₂
              → FullTree A (1 + n) (Branch F₂ a t₁ t₂)

data BalancedTree A : ℕ → Tree A → Set where
  bal-leaf    : BalancedTree A 0 Leaf
  bal-branch₁ : {a : A}{t₁ t₂ : Tree A}{n : ℕ}
              → BalancedTree A (1 + n) t₁
              → FullTree A n t₂
              → BalancedTree A (2 + n) (Branch F₁ a t₁ t₂) 
  bal-branch₂ : {a : A}{t₁ t₂ : Tree A}{n : ℕ}
              → FullTree A n t₁
              → BalancedTree A n t₂
              → BalancedTree A (1 + n) (Branch F₂ a t₁ t₂)

preserve : {A : Set}{a : A}{n : ℕ}{t : Tree A}
           → FullTree A n t
           → BalancedTree A (1 + n) (insertFullTree a t)
preserve full-leaf            = bal-branch₂ full-leaf bal-leaf
preserve (full-branch y y') = bal-branch₁ (preserve y) y'

data X (A : Set) (n : ℕ) : Maybe (Tree A) → Set where
  bal-x : ∀ {t} → BalancedTree A n t → X A n (just t)
  ful-x : ∀ {t} → FullTree A n t     → X A n nothing

-- just if insert was successful, nothing if it failed due to being full
insertTree : ∀ {A} → A → Tree A → Maybe (Tree A)
insertTree _ Leaf = nothing
insertTree x (Branch F₁ a t₁ t₂) with insertTree x t₁
... | just t₁′ = just (Branch F₁ a t₁′ t₂)
... | nothing = just (Branch F₂ a t₁ (insertFullTree x t₂))
insertTree x (Branch F₂ a t₁ t₂) with insertTree x t₂
... | just t₂′ = just (Branch F₂ a t₁ t₂′)
... | nothing = nothing

maybe-different : ∀ {A : Set} {a : A} {b : Maybe A} {c : Set}
                → just a ≡ b → nothing ≡ b → c
maybe-different refl ()

nested-failure : ∀ {A : Set}{a a' : A}{t₁ t₂}
      → nothing ≡ insertTree a t₂
      → nothing ≡ (insertTree a (Branch F₂ a' t₁ t₂))
nested-failure {_}{a}{_}{_}{t₂} p with insertTree a t₂
... | nothing = refl
nested-failure () | just k

insert-full-fails : {A : Set} {a : A} {n : ℕ} (t : Tree A) → FullTree A n t
                  → nothing ≡ insertTree a t
insert-full-fails .Leaf full-leaf = refl
insert-full-fails {A}{a} {suc n}
   .(Branch F₂ a' t₁ t₂) (full-branch {a'} {t₁} {t₂} y y')
   with inspect (insertTree a t₂)
... | nothing with-≡ eq = nested-failure eq
... | just m with-≡ eq = maybe-different eq (insert-full-fails {A}{a} t₂ y')