{-# OPTIONS --safe #-}

module Problem24 where

open import Data.Bool using (true; false)
open import Data.Empty using (⊥-elim)
open import Data.List
open import Data.List.Properties
open import Data.List.Relation.Unary.All hiding (toList)
open import Data.Nat
open import Data.Nat.Properties
open import Data.Product
open import Data.Product.Properties
open import Relation.Binary.PropositionalEquality hiding ([_])
open import Relation.Nullary

open ≡-Reasoning

data Tree : Set where
    Leaf : Tree
    Branch : Tree → Tree → Tree

toList : ℕ → Tree → List ℕ
toList d Leaf = [ d ]
toList d (Branch t₁ t₂) = toList (suc d) t₁ ++ toList (suc d) t₂

-- We write a parser for lists into forests with a labelled root depth, from
-- right to left.
--
-- The empty list is an empty forest.
--
-- For a given natural `n` and forest `ts`, we try and merge the labelled tree
-- `(n , Leaf)` into the forest.
--
-- Merging a tree `(n , t)` into a forest `ts` is again done by induction.

LabelledTree : Set
LabelledTree = ℕ × Tree

Forest : Set
Forest = List LabelledTree

addTree : LabelledTree → Forest → Forest
addTree (n , t) []              = [ (n , t) ]
addTree (n , t) ((k , t′) ∷ ts) with n ≟ k
... | yes p = addTree (pred n , Branch t t′) ts
... | no  p = (n , t) ∷ (k , t′) ∷ ts

parseList : List ℕ → Forest
parseList []       = []
parseList (x ∷ xs) = addTree (x , Leaf) (parseList xs)

-- Properties ------------------------------------------------------------------

private
  variable
    k d : ℕ
    xs ys : List ℕ
    t t′ : Tree
    ts : Forest

-- A forest blocks at `d` if trying to add any `1+n+d` tree fails.

data Blocking (d : ℕ) : (ts : Forest) → Set where
  empty : Blocking d []
  cons  : {ts : Forest} → k ≤ d → Blocking d ((k , t) ∷ ts)

-- Adding a tree at a lower depth is blocking, even if it is merged into existing trees.

blocking-addTree :
  (d≥k : d ≥ k) (t : Tree) (ts : Forest) → Blocking d (addTree (k , t) ts)
blocking-addTree         d≥k t []              = cons d≥k
blocking-addTree {k = k} d≥k t ((n , t′) ∷ ts) with k ≟ n
... | yes refl = blocking-addTree (≤-trans pred[n]≤n d≥k) (Branch t t′) ts
... | no  p    = cons d≥k

-- Blocking is closed for higher depths

blocking-suc : Blocking d ts → Blocking (suc d) ts
blocking-suc empty      = empty
blocking-suc (cons k≤d) = cons (≤-step k≤d)

-- The fundamental theorem of blocking:
-- if a forest blocks at `d`, adding a tree with depth `1+d` conses the tree
-- onto the forest.

addTree-blocking[1+n] :
  (d : ℕ) (t : Tree) {ts : Forest} → Blocking d ts →
  addTree (suc d , t) ts ≡ (suc d , t) ∷ ts
addTree-blocking[1+n] d t empty          = refl
addTree-blocking[1+n] d t (cons {k} k≤d) with suc d ≟ k
... | yes 1+d≡k = ⊥-elim (≤⇒≯ k≤d (≤-reflexive 1+d≡k))
... | no  _     = refl

-- Adding two trees of depth `1+n` onto a forest blocking at `n` is equivalent
-- to adding their composite.

addTree-branch′ :
  (n : ℕ) (t₁ t₂ : Tree) (ts : Forest) →
  addTree (suc n , t₁) ((suc n , t₂) ∷ ts) ≡ addTree (n , Branch t₁ t₂) ts
addTree-branch′ n t₁ t₂ ts with suc n ≟ suc n
... | yes _ = refl
... | no  p = ⊥-elim (p refl)

addTree-branch :
  (n : ℕ) (t₁ t₂ : Tree) {ts : Forest} → Blocking n ts →
  addTree (suc n , t₁) (addTree (suc n , t₂) ts) ≡ addTree (n , Branch t₁ t₂) ts
addTree-branch n t₁ t₂ {ts} block = begin
  addTree (suc n , t₁) (addTree (suc n , t₂) ts) ≡⟨ cong (addTree (suc n , t₁)) (addTree-blocking[1+n] n t₂ block) ⟩
  addTree (suc n , t₁) ((suc n , t₂) ∷ ts)       ≡⟨ addTree-branch′ n t₁ t₂ ts ⟩
  addTree (n , Branch t₁ t₂) ts                  ∎

-- Main result of parse: parsing inverts conversion to lists, as long as the
-- tail blocks.
--
-- We prove a blocking lemma at the same time. Using the correctness proof
-- significantly reduces the amount of work.

blocking-parseList :
  (d : ℕ) (t : Tree) {xs : List ℕ} → Blocking d (parseList xs) →
  Blocking d (parseList (toList d t ++ xs))
parseList-++ :
  (d : ℕ) (t : Tree) {xs : List ℕ} → Blocking d (parseList xs) →
  parseList (toList d t ++ xs) ≡ addTree (d , t) (parseList xs)

blocking-parseList d t {xs} block =
  subst (Blocking d)
    (sym (parseList-++ d t block))
    (blocking-addTree ≤-refl t (parseList xs))

parseList-++ d Leaf                block = refl
parseList-++ d (Branch t₁ t₂) {xs} block = begin
  parseList ((toList (suc d) t₁ ++ toList (suc d) t₂) ++ xs) ≡⟨ cong parseList (++-assoc (toList (suc d) t₁) (toList (suc d) t₂) xs) ⟩
  parseList (toList (suc d) t₁ ++ toList (suc d) t₂ ++ xs)   ≡⟨ parseList-++ (suc d) t₁ (blocking-parseList (suc d) t₂ (blocking-suc block)) ⟩
  addTree (suc d , t₁) (parseList (toList (suc d) t₂ ++ xs)) ≡⟨ cong (addTree (suc d , t₁)) (parseList-++ (suc d) t₂ (blocking-suc block)) ⟩
  addTree (suc d , t₁) (addTree (suc d , t₂) (parseList xs)) ≡⟨ addTree-branch d t₁ t₂ block ⟩
  addTree (d , Branch t₁ t₂) (parseList xs)                  ∎

-- Proof that parse is an inverse of `toList` (modulo some wrapping)

parse-correct : (d : ℕ) (t : Tree) → parseList (toList d t) ≡ [ d , t ]
parse-correct d t = begin
  parseList (toList d t)       ≡˘⟨ cong parseList (++-identityʳ (toList d t)) ⟩
  parseList (toList d t ++ []) ≡⟨  parseList-++ d t empty ⟩
  addTree (d , t) []           ≡⟨⟩
  [ d , t ]                    ∎

-- Goal

toList-injective : ∀{n} t t′ → toList n t ≡ toList n t′ → t ≡ t′
toList-injective {n} t t′ eq = ,-injectiveʳ (∷-injectiveˡ (begin
  [ n , t ]               ≡˘⟨ parse-correct n t ⟩
  parseList (toList n t)  ≡⟨  cong parseList eq ⟩
  parseList (toList n t′) ≡⟨  parse-correct n t′ ⟩
  [ n , t′ ]              ∎))