load "nat.quox"
load "maybe.quox"
load "bool.quox"

namespace vec {

def0 Vec : ℕ → ★ → ★ =
  λ n A ⇒
  caseω n return ★ of {
    zero           ⇒ {nil};
    succ _, 0.Tail ⇒ A × Tail
  }

def drop-nil-dep : 0.(A : ★) → 0.(P : Vec 0 A → ★) →
                   (xs : Vec 0 A) → P 'nil → P xs =
  λ A P xs p ⇒ case xs return xs' ⇒ P xs' of { 'nil ⇒ p }

def drop-nil : 0.(A B : ★) → Vec 0 A → B → B =
  λ A B ⇒ drop-nil-dep A (λ _ ⇒ B)

def match-dep :
    0.(A : ★) → 0.(P : (n : ℕ) → Vec n A → ★) →
    ω.(P 0 'nil) →
    ω.((n : ℕ) → (x : A) → (xs : Vec n A) → P (succ n) (x, xs)) →
    (n : ℕ) → (xs : Vec n A) → P n xs =
  λ A P pn pc n ⇒
    case n return n' ⇒ (xs : Vec n' A) → P n' xs of {
      0 ⇒ λ nil ⇒ drop-nil-dep A (P 0) nil pn;
      succ len ⇒ λ cons ⇒
        case cons return cons' ⇒ P (succ len) cons' of {
          (first, rest) ⇒ pc len first rest
        }
    }

def elim : 0.(A : ★) → 0.(P : (n : ℕ) → Vec n A → ★) →
           P 0 'nil →
           ω.((x : A) → 0.(n : ℕ) → 0.(xs : Vec n A) →
              P n xs → P (succ n) (x, xs)) →
           (n : ℕ) → (xs : Vec n A) → P n xs =
  λ A P pn pc n ⇒
    case n return n' ⇒ (xs' : Vec n' A) → P n' xs' of {
      zero ⇒ λ nil ⇒
        case nil return nil' ⇒ P 0 nil' of { 'nil ⇒ pn };
      succ n, ih ⇒ λ cons ⇒
        case cons return cons' ⇒ P (succ n) cons' of {
          (first, rest) ⇒ pc first n rest (ih rest)
        }
    }

def elim2 : 0.(A B : ★) → 0.(P : (n : ℕ) → Vec n A → Vec n B → ★) →
            P 0 'nil 'nil →
            ω.((x : A) → (y : B) → 0.(n : ℕ) →
               0.(xs : Vec n A) → 0.(ys : Vec n B) →
               P n xs ys → P (succ n) (x, xs) (y, ys)) →
            (n : ℕ) → (xs : Vec n A) → (ys : Vec n B) → P n xs ys =
  λ A B P pn pc n ⇒
    case n return n' ⇒ (xs : Vec n' A) → (ys : Vec n' B) → P n' xs ys of {
      zero ⇒ λ nila nilb ⇒
        drop-nil-dep A (λ n ⇒ P 0 n nilb) nila
          (drop-nil-dep B (λ n ⇒ P 0 'nil n) nilb pn);
      succ n, ih ⇒ λ consa consb ⇒
        case consa return consa' ⇒ P (succ n) consa' consb of { (a, as) ⇒
          case consb return consb' ⇒ P (succ n) (a, as) consb' of { (b, bs) ⇒
            pc a b n as bs (ih as bs)
          }
        }
    }

-- haha gross
def elimω : 0.(A : ★) → 0.(P : (n : ℕ) → Vec n A → ★) →
            ω.(P 0 'nil) →
            ω.(ω.(x : A) → 0.(n : ℕ) → 0.(xs : Vec n A) →
               ω.(P n xs) → P (succ n) (x, xs)) →
            ω.(n : ℕ) → ω.(xs : Vec n A) → P n xs =
  λ A P pn pc n ⇒
    caseω n return n' ⇒ ω.(xs' : Vec n' A) → P n' xs' of {
      zero ⇒ λ nil ⇒
        caseω nil return nil' ⇒ P 0 nil' of { 'nil ⇒ pn };
      succ n, ω.ih ⇒ λ cons ⇒
        caseω cons return cons' ⇒ P (succ n) cons' of {
          (first, rest) ⇒ pc first n rest (ih rest)
        }
    }

def elimω2 : 0.(A B : ★) → 0.(P : (n : ℕ) → Vec n A → Vec n B → ★) →
             ω.(P 0 'nil 'nil) →
             ω.(ω.(x : A) → ω.(y : B) → 0.(n : ℕ) →
                0.(xs : Vec n A) → 0.(ys : Vec n B) →
                ω.(P n xs ys) → P (succ n) (x, xs) (y, ys)) →
             ω.(n : ℕ) → ω.(xs : Vec n A) → ω.(ys : Vec n B) → P n xs ys =
  λ A B P pn pc n ⇒
    caseω n return n' ⇒ ω.(xs : Vec n' A) → ω.(ys : Vec n' B) → P n' xs ys of {
      zero ⇒ λ nila nilb ⇒
        drop-nil-dep A (λ n ⇒ P 0 n nilb) nila
          (drop-nil-dep B (λ n ⇒ P 0 'nil n) nilb pn);
      succ n, ω.ih ⇒ λ consa consb ⇒
        caseω consa return consa' ⇒ P (succ n) consa' consb of { (a, as) ⇒
          caseω consb return consb' ⇒ P (succ n) (a, as) consb' of { (b, bs) ⇒
            pc a b n as bs (ih as bs)
          }
        }
    }

def zip-with : 0.(A B C : ★) → ω.(A → B → C) →
               (n : ℕ) → Vec n A → Vec n B → Vec n C =
  λ A B C f ⇒
    elim2 A B (λ n _ _ ⇒ Vec n C) 'nil (λ a b _ _ _ abs ⇒ (f a b, abs))

def zip-withω : 0.(A B C : ★) → ω.(ω.A → ω.B → C) →
                ω.(n : ℕ) → ω.(Vec n A) → ω.(Vec n B) → Vec n C =
  λ A B C f ⇒
    elimω2 A B (λ n _ _ ⇒ Vec n C) 'nil (λ a b _ _ _ abs ⇒ (f a b, abs))

#[compile-scheme "(lambda% (n xs) xs)"]
def up : 0.(A : ★) → (n : ℕ) → Vec n A → Vec¹ n A =
  λ A n ⇒
    case n return n' ⇒ Vec n' A → Vec¹ n' A of {
      zero ⇒ λ xs ⇒
        case xs return Vec¹ 0 A of { 'nil ⇒ 'nil };
      succ n', f' ⇒ λ xs ⇒
        case xs return Vec¹ (succ n') A of {
          (first, rest) ⇒ (first, f' rest)
        }
    }

}

def0 Vec = vec.Vec


namespace list {

def0 List : ★ → ★ =
  λ A ⇒ (len : ℕ) × Vec len A

def Nil : 0.(A : ★) → List A =
  λ A ⇒ (0, 'nil)

def Cons : 0.(A : ★) → A → List A → List A =
  λ A x xs ⇒ case xs return List A of { (len, elems) ⇒ (succ len, x, elems) }

def single : 0.(A : ★) → A → List A =
  λ A x ⇒ Cons A x (Nil A)

def elim : 0.(A : ★) → 0.(P : List A → ★) →
           P (Nil A) →
           ω.((x : A) → 0.(xs : List A) → P xs → P (Cons A x xs)) →
           (xs : List A) → P xs =
  λ A P pn pc xs ⇒
    case xs return xs' ⇒ P xs' of { (len, elems) ⇒
      vec.elim A (λ n xs ⇒ P (n, xs))
        pn (λ x n xs ih ⇒ pc x (n, xs) ih)
        len elems
    }

def elimω : 0.(A : ★) → 0.(P : List A → ★) →
            ω.(P (Nil A)) →
            ω.(ω.(x : A) → 0.(xs : List A) → ω.(P xs) → P (Cons A x xs)) →
            ω.(xs : List A) → P xs =
  λ A P pn pc xs ⇒
    caseω xs return xs' ⇒ P xs' of { (len, elems) ⇒
      vec.elimω A (λ n xs ⇒ P (n, xs))
        pn (λ x n xs ih ⇒ pc x (n, xs) ih)
        len elems
    }

def match-dep :
    0.(A : ★) → 0.(P : List A → ★) →
    ω.(P (Nil A)) → ω.((x : A) → (xs : List A) → P (Cons A x xs)) →
    (xs : List A) → P xs =
  λ A P pn pc xs ⇒
    case xs return xs' ⇒ P xs' of {
      (len, elems) ⇒
        vec.match-dep A (λ n xs ⇒ P (n, xs)) pn (λ n x xs ⇒ pc x (n, xs))
          len elems
    }

def match : 0.(A B : ★) → ω.B → ω.(A → List A → B) → List A → B =
  λ A B ⇒ match-dep A (λ _ ⇒ B)


-- [fixme] List A <: List¹ A should be automatic, imo
#[compile-scheme "(lambda (xs) xs)"]
def up : 0.(A : ★) → List A → List¹ A =
  λ A xs ⇒
  case xs return List¹ A of { (len, elems) ⇒
    case nat.dup! len return List¹ A of { [p] ⇒
      caseω p return List¹ A of { (lenω, eq0) ⇒
        case eq0 return List¹ A of { [eq] ⇒
          (lenω, vec.up A lenω (coe (𝑖 ⇒ Vec (eq @𝑖) A) @1 @0 elems))
        }
      }
    }
  }

def foldr : 0.(A B : ★) → B → ω.(A → B → B) → List A → B =
  λ A B z f xs ⇒ elim A (λ _ ⇒ B) z (λ x _ y ⇒ f x y) xs

def foldl : 0.(A B : ★) → B → ω.(B → A → B) → List A → B =
  λ A B z f xs ⇒
    foldr A (B → B) (λ b ⇒ b) (λ a g b ⇒ g (f b a)) xs z

def map : 0.(A B : ★) → ω.(A → B) → List A → List B =
  λ A B f ⇒ foldr A (List B) (Nil B) (λ x ys ⇒ Cons B (f x) ys)


-- ugh
def foldrω : 0.(A B : ★) → ω.B → ω.(ω.A → ω.B → B) → ω.(List A) → B =
  λ A B z f xs ⇒ elimω A (λ _ ⇒ B) z (λ x _ y ⇒ f x y) xs

def foldlω : 0.(A B : ★) → ω.B → ω.(ω.B → ω.A → B) → ω.(List A) → B =
  λ A B z f xs ⇒
    foldrω A (ω.B → B) (λ b ⇒ b) (λ a g b ⇒ g (f b a)) xs z

def mapω : 0.(A B : ★) → ω.(ω.A → B) → ω.(List A) → List B =
  λ A B f ⇒ foldrω A (List B) (Nil B) (λ x ys ⇒ Cons B (f x) ys)


def0 All : (A : ★) → (P : A → ★) → List A → ★ =
  λ A P xs ⇒ foldr¹ A ★ True (λ x ps ⇒ P x × ps) (up A xs)

def append : 0.(A : ★) → List A → List A → List A =
  λ A xs ys ⇒ foldr A (List A) ys (Cons A) xs

def reverse : 0.(A : ★) → List A → List A =
  λ A ⇒ foldl A (List A) (Nil A) (λ xs x ⇒ Cons A x xs)


def find : 0.(A : ★) → ω.(ω.A → Bool) → ω.(List A) → Maybe A =
  λ A p ⇒
    foldlω A (Maybe A) (Nothing A) (λ m x ⇒ maybe.or A m (maybe.check A p x))

def cons-first : 0.(A : ★) → ω.A → List (List A) → List (List A) =
  λ A x ⇒
    match (List A) (List (List A))
      (single (List A) (single A x))
      (λ xs xss ⇒ Cons (List A) (Cons A x xs) xss)

def split : 0.(A : ★) → ω.(ω.A → Bool) → ω.(List A) → List (List A) =
  λ A p ⇒
    foldrω A (List (List A))
      (Nil (List A))
      (λ x xss ⇒ bool.if (List (List A)) (p x)
          (Cons (List A) (Nil A) xss)
          (cons-first A x xss))

def break : 0.(A : ★) → ω.(ω.A → Bool) → ω.(List A) → List A × List A =
  λ A p xs ⇒
  let0 Lst = List A; Lst2 = (Lst × Lst) ∷ ★; State = Either Lst Lst2 in
  letω LeftS = Left Lst Lst2; RightS = Right Lst Lst2 in
  letω res =
    foldlω A State
      (LeftS (Nil A))
      (λ acc x ⇒
        either.foldω Lst Lst2 State
          (λ xs ⇒ bool.if State (p x)
            (RightS (xs, list.single A x))
            (LeftS  (Cons A x xs)))
          (λ xsys ⇒
            RightS (fst xsys, Cons A x (snd xsys))) acc)
      xs ∷ State in
  letω res =
    either.fold Lst Lst2 Lst2 (λ xs ⇒ (Nil A, xs)) (λ xsys ⇒ xsys) res in
  (reverse A (fst res), reverse A (snd res))

def uncons : 0.(A : ★) → List A → Maybe (A × List A) =
  λ A ⇒
    match A (Maybe (A × List A))
      (Nothing (A × List A))
      (λ x xs ⇒ Just (A × List A) (x, xs))

def head : 0.(A : ★) → ω.(List A) → Maybe A =
  λ A xs ⇒ maybe.mapω (A × List A) A (λ xxs ⇒ fst xxs) (uncons A xs)

def tail : 0.(A : ★) → ω.(List A) → Maybe (List A) =
  λ A xs ⇒ maybe.mapω (A × List A) (List A) (λ xxs ⇒ snd xxs) (uncons A xs)

def tail-or-nil : 0.(A : ★) → ω.(List A) → List A =
  λ A xs ⇒ maybe.fold (List A) (List A) (Nil A) (λ xs ⇒ xs) (tail A xs)

def slip : 0.(A : ★) → List A × List A → List A × List A =
  λ A xsys ⇒
    case xsys return List A × List A of { (xs, ys) ⇒
      maybe.fold (A × List A) (List A → List A × List A)
        (λ     xs ⇒ (xs, Nil A))
        (λ yys xs ⇒
          case yys return List A × List A of { (y, ys) ⇒ (Cons A y xs, ys) })
        (uncons A ys) xs
    }

def split-at' : 0.(A : ★) → ℕ → List A → List A × List A =
  λ A n xs ⇒
    (case n return List A × List A → List A × List A of {
      0         ⇒ λ xsys ⇒ xsys;
      succ _, f ⇒ λ xsys ⇒ f (slip A xsys)
    }) (Nil A, xs)

def split-at : 0.(A : ★) → ℕ → List A → List A × List A =
  λ A n xs ⇒
    case split-at' A n xs return List A × List A of {
      (xs', ys) ⇒ (reverse A xs', ys)
    }

def filter : 0.(A : ★) → ω.(ω.A → Bool) → ω.(List A) → List A =
  λ A p ⇒
    foldrω A (List A)
      (Nil A)
      (λ x xs ⇒ bool.if (List A) (p x) (Cons A x xs) xs)

def length : 0.(A : ★) → ω.(List A) → ℕ =
  λ A xs ⇒ fst xs


def0 ZipWithFailureVec : (m n : ℕ) → (A B : ★) → Vec m A → Vec n B → ★ =
  λ m n A B xs ys ⇒ Sing (Vec m A) xs × Sing (Vec n B) ys × [0. Not (m ≡ n : ℕ)]

def0 ZipWithFailure : (A B : ★) → List A → List B → ★ =
  λ A B xs ys ⇒ ZipWithFailureVec (fst xs) (fst ys) A B (snd xs) (snd ys)

{-
-- unfinished
def zip-with : 0.(A B C : ★) → ω.(A → B → C) →
               (xs : List A) → (ys : List B) →
               Either (ZipWithFailure A B xs ys) (List C) =
  λ A B C f xs' ys' ⇒
  let0 Ret = Either (ZipWithFailure A B xs' ys') (List C) in
  case xs' return Ret of { (m', xs) ⇒
  case ys' return Ret of { (n', ys) ⇒
  case nat.dup! m' return Ret of { [m!] ⇒
  let ω.m = fst m!; 0.mm' = get0 (m ≡ m' : ℕ) (snd m!) in
  case nat.dup! n' return Ret of { [n!] ⇒
  let ω.n = fst n!; 0.nn' = get0 (n ≡ n' : ℕ) (snd n!) in
  let1 xs = coe (𝑖 ⇒ Vec (mm' @𝑖) A) @1 @0 xs ∷ Vec m A in
  let1 ys = coe (𝑖 ⇒ Vec (nn' @𝑖) B) @1 @0 ys ∷ Vec n B in
    dec.elim (m ≡ n : ℕ) Ret
      (λ mn ⇒
        let xs = coe (𝑖 ⇒ Vec (mn @𝑖) A) xs ∷ Vec n A in
        Right (ZipWithFailure A B xs' ys') (List C)
          (n, vec.zip-with A B C n xs ys))
      (λ nmn ⇒
        Left (ZipWithFailure A B xs' ys') (List C)
          (?, ?, [nmn]) -- <---------------------
      (nat.eq? m n)
  }}}}
-}


def zip-withω : 0.(A B C : ★) → ω.(ω.A → ω.B → C) →
                ω.(xs : List A) → ω.(ys : List B) →
                Either (ZipWithFailure A B xs ys) (List C) =
  λ A B C f xs' ys' ⇒
    let0 Err = ZipWithFailure A B xs' ys';
         Ret = Either Err (List C) in
    letω m = fst xs'; xs = snd xs';
         n = fst ys'; ys = snd ys' in
    dec.elim (m ≡ n : ℕ) (λ _ ⇒ Ret)
      (λ mn ⇒
        letω xs = coe (𝑖 ⇒ Vec (mn @𝑖) A) xs in
        Right Err (List C) (n, vec.zip-withω A B C f n xs ys))
      (λ nmn ⇒
        Left Err (List C) (sing (Vec m A) xs, sing (Vec n B) ys, [nmn]))
      (nat.eq? m n)

def zip-with-uneven :
    0.(A B C : ★) → ω.(ω.A → ω.B → C) → ω.(List A) → ω.(List B) → List C =
  λ A B C f xs ys ⇒
    caseω nat.min (fst xs) (fst ys)
    return ω.(List A) → ω.(List B) → List C of {
      0             ⇒ λ _ _ ⇒ Nil C;
      succ _, ω.fih ⇒ λ xs ys ⇒
        maybe.foldω (A × List A) (List C) (Nil C)
          (λ xxs ⇒ maybe.foldω (B × List B) (List C) (Nil C)
            (λ yys ⇒ Cons C (f (fst xxs) (fst yys)) (fih (snd xxs) (snd yys)))
            (list.uncons B ys))
          (list.uncons A xs)
    } xs ys


{-
-- unfinished
def zip-with : 0.(A B C : ★) → ω.(A → B → C) →
               (xs : List A) → (ys : List B) →
               Either (Sing (List A) xs × Sing (List B) ys ×
                       Not (length A xs ≡ length B ys : ℕ))
                      (List C) =
  λ A B C f xs' ys' ⇒
    let0 Err = (Sing (List A) xs' × Sing (List B) ys' ×
                Not (length A xs' ≡ length B ys' : ℕ)) ∷ ★;
         Ret = Either Err (List C) in
    case xs' return Ret of { (m', xs) ⇒
      case ys' return Ret of { (n', ys) ⇒
        case nat.dup! m' return Ret of { [msing] ⇒
          case nat.dup! n' return Ret of { [nsing] ⇒
            let1 m = fst msing; n = fst nsing in
            let0 mm' = get0 (m ≡ m' : ℕ) (snd msing);
                 nn' = get0 (n ≡ n' : ℕ) (snd nsing) in
            dec.elim (m ≡ n : ℕ) (λ _ ⇒ Ret)
              (λ mn ⇒
                let0 m'n = trans ℕ m' m n (sym ℕ m m' mm') mn ∷ m' ≡ n : ℕ in
                let1 xs = coe (𝑖 ⇒ Vec (m'n @𝑖) A) xs ∷ Vec n A;
                     ys = coe (𝑖 ⇒ Vec (nn' @𝑖) B) @1 @0 ys ∷ Vec n B in
                Right Err (List C) (n, vec.zip-with A B C f n xs ys))
              (λ nmn ⇒
                let xs =
                  ((m, coe (𝑖 ⇒ Vec (mm' @𝑖) A) @1 @0 xs),
                   [δ 𝑗 ⇒ (mm' @𝑗, coe (𝑖 ⇒ Vec (mm' @𝑖) A) @1 @𝑗 xs)])
                  ∷ Sing (List A) xs' in
                  -- sing (List A) (m, coe (𝑖 ⇒ Vec (mm' @𝑖) A) @1 @0 xs);
                let ys = sing (List B) (n, coe (𝑖 ⇒ Vec (nn' @𝑖) B) @1 @0 ys) in
                Left Err (List C) (xs, ys, nmn))
          }
        }
      }
    }
-}


postulate0 SchemeList : ★ → ★

#[compile-scheme
  "(lambda (list) (cons (length list) (fold-right cons 'nil list)))"]
postulate from-scheme : 0.(A : ★) → SchemeList A → List A

#[compile-scheme
  "(lambda (list)
    (let loop [(acc '()) (list (cdr list))]
      (if (pair? list)
        (loop (cons (car list) acc) (cdr list))
        (reverse acc))))"]
postulate to-scheme : 0.(A : ★) → List A → SchemeList A

}

def0 List = list.List