module Motzkin where

open import Data.Nat using (ℕ; zero; suc; _+_; _≤?_; z≤n)
open import Data.Nat.Properties using (≤-totalOrder)

open import Data.Empty using (⊥-elim)
open import Data.Fin using (Fin; zero; suc; toℕ; raise)
open import Data.List using (List; []; _∷_; length; filter; cartesianProductWith; allFin; map)
open import Data.List.Membership.Propositional using (_∈_)
open import Data.List.Relation.Binary.Pointwise as Pointwise using (Pointwise; []; _∷_)
open import Data.List.Relation.Binary.Suffix as Suffix using (base; next)
open import Data.List.Relation.Unary.All as All using (All; []; _∷_)
open import Data.List.Relation.Unary.All.Properties as AllP using (all-filter)
open import Data.List.Relation.Unary.Any as Any using (Any; here; there)
open import Data.List.Relation.Unary.Any.Properties as AnyP using (lookup-result)
open import Data.List.Relation.Unary.Linked using ([]; [-]; _∷_)
open import Data.List.Relation.Unary.Sorted.TotalOrder ≤-totalOrder using (Sorted; sorted?)
open import Data.Product using (_×_; _,_; ∃₂; proj₁; proj₂; map₁)
open import Data.Product.Properties as Prod
open import Data.Sum using (_⊎_; [_,_])
open import Data.Vec as Vec using (Vec; []; _∷_; _++_; replicate; updateAt; toList)
open import Data.Vec.Properties using (map-id)
open import Function using (_∘_; case_of_; id)
open import Function.Equivalence using (_⇔_; equivalence)
open import Relation.Binary.Core using (Rel)
open import Relation.Binary.PropositionalEquality using (_≡_; refl; sym; cong; subst; setoid)
open import Relation.Nullary using (Dec; yes; no; ¬_)
import Relation.Nullary.Decidable as Dec
open import Relation.Nullary.Implication using (_→-dec_)
open import Relation.Unary using (Pred; Decidable)

_≋_ : ∀ {n} → Rel (List (Fin n)) _
_≋_ = Pointwise _≡_

Suffix : ∀ {n} → Rel (List (Fin n)) _
Suffix = Suffix.Suffix _≡_

counts : ∀ {n} → (xs : List (Fin n)) → Vec ℕ n
counts []       = replicate 0
counts (x ∷ xs) = updateAt x suc (counts xs)

sorted-0s : ∀ {n} → (Sorted ∘ toList {n = n} ∘ replicate) 0
sorted-0s {zero}        = []
sorted-0s {suc zero}    = [-]
sorted-0s {suc (suc n)} = z≤n ∷ sorted-0s

≋⇒≡ : ∀ {n} {xs ys : List (Fin n)} → xs ≋ ys → xs ≡ ys
≋⇒≡ [] = refl
≋⇒≡ (x≡y ∷ xs≋ys) = subst (λ _ → _ ≡ _ ∷ _) x≡y (cong (_ ∷_) (≋⇒≡ xs≋ys))

P : ∀ {n} → Pred (List (Fin n)) _
P xs = ∀ ys → Suffix ys xs → (Sorted ∘ toList ∘ counts) ys

data Motzkin {n} : List (Fin n) → Set where
  base : Motzkin []
  next : ∀ {x xs} → (Sorted ∘ toList ∘ counts) (x ∷ xs) → Motzkin xs → Motzkin (x ∷ xs)

P⇔Motzkin : ∀ {n} xs → P {n} xs ⇔ Motzkin xs
P⇔Motzkin xs = equivalence (⇒ xs) ⇐
  where
  ⇒ : ∀ xs → P xs → Motzkin xs
  ⇒ [] pxs = base
  ⇒ (x ∷ xs) pxs = next (pxs (x ∷ xs) (Suffix.refl refl)) (⇒ xs (λ ys ys<xs → pxs ys (next x ys<xs)))

  ⇐ : ∀ {xs} → Motzkin xs → P xs
  ⇐ base [] (base []) = sorted-0s
  ⇐ (next counts↗ pxs) ys (base xs≋ys) = subst (Sorted ∘ toList ∘ counts) (sym (≋⇒≡ xs≋ys)) counts↗
  ⇐ (next counts↗ pxs) ys (next _ ys<xs) = ⇐ pxs ys ys<xs

extendUpdate : ∀ {n m} → Fin n → Vec (Fin n) m × Vec ℕ n → Vec (Fin n) (suc m) × Vec ℕ n
extendUpdate x (xs , counts) = x ∷ xs , updateAt x suc counts

extendUpdateCount : ∀ {n m x xs-counts} →
                    proj₂ xs-counts ≡ (counts ∘ toList ∘ proj₁) xs-counts →
                    (proj₂ ∘ extendUpdate {n} {m} x) xs-counts ≡ (counts ∘ toList ∘ proj₁ ∘ extendUpdate x) xs-counts
extendUpdateCount refl = refl

generate : ∀ n m → List (Vec (Fin n) m × Vec ℕ n)
generate n zero    = ([] , replicate 0) ∷ []
generate n (suc m) =
  filter (sorted? _≤?_ ∘ toList ∘ proj₂)
         (cartesianProductWith extendUpdate
                               (allFin n)
                               (generate n m))

generate⇒counts : ∀ n m → All (λ (xs , cs) → cs ≡ counts (toList xs)) (generate n m)
generate⇒counts n zero = refl ∷ []
generate⇒counts n (suc m) =
  AllP.filter⁺
    (sorted? _≤?_ ∘ toList ∘ proj₂)
    (AllP.cartesianProductWith⁺
      (setoid (Fin n))
      (setoid (Vec (Fin n) m × Vec ℕ n))
      extendUpdate
      (allFin n)
      (generate n m)
      (λ {x} {y} px py → extendUpdateCount {n} {m} {x} {y} (All.lookup (generate⇒counts n m) py)))

generate⇒counts↗ : ∀ n m → All (Sorted ∘ toList ∘ counts ∘ toList ∘ proj₁) (generate n m)
generate⇒counts↗ n zero = sorted-0s ∷ []
generate⇒counts↗ n (suc m) =
  All.zipWith
    (λ (p₂≡cs , p₂↗) → subst (Sorted ∘ toList) p₂≡cs p₂↗)
    ( generate⇒counts n (suc m)
    , all-filter (sorted? _≤?_ ∘ toList ∘ proj₂)
                 (cartesianProductWith extendUpdate (allFin n) (generate n m)))

generate⇒next : ∀ n m → All (λ pair → ∃₂ λ y ys → proj₁ pair ≡ y ∷ (proj₁ ys) × ys ∈ generate n m) (generate n (suc m))
generate⇒next n m =
  AllP.filter⁺
    (sorted? _≤?_ ∘ toList ∘ proj₂)
    (AllP.cartesianProductWith⁺
      (setoid (Fin n))
      (setoid (Vec (Fin n) m × Vec ℕ n))
      extendUpdate
      (allFin n)
      (generate n m)
      λ px py → _ , _ , refl , py)

generate⇒Motzkin : ∀ n m → All (Motzkin ∘ toList ∘ proj₁) (generate n m)
generate⇒Motzkin n zero    = base ∷ []
generate⇒Motzkin n (suc m) =
  All.zipWith
    (λ (counts↗ , y , ys , xs≡y∷ys , pys) →
      subst
        (Motzkin ∘ toList)
        (sym xs≡y∷ys)
        (next
          (subst (Sorted ∘ toList ∘ counts ∘ toList) xs≡y∷ys counts↗)
          (All.lookup (generate⇒Motzkin n m) pys)))
    (generate⇒counts↗ n (suc m) , generate⇒next n m)

Motzkin⇒generate : ∀ {n xs} → Motzkin xs → (xs , counts xs) ∈ map (map₁ toList) (generate n (length xs))
Motzkin⇒generate base = here refl
Motzkin⇒generate (next {x} {xs} counts↗ pxs) =
  AnyP.map⁺
    ([ id
     , (λ ¬counts↗ → ⊥-elim (¬counts↗ (subst (Sorted ∘ toList) (,-injectiveʳ (lookup-result cpw)) counts↗)))
     ] (AnyP.filter⁺ (sorted? _≤?_ ∘ toList ∘ proj₂) cpw))
  where
  cpw =
    AnyP.cartesianProductWith⁺
      extendUpdate
      (λ { refl eq → cong (λ (xs , cs) → x ∷ xs , updateAt x suc cs) eq})
      (AnyP.tabulate⁺ {P = x ≡_} x refl)
      (AnyP.map⁻ (Motzkin⇒generate pxs))