\import Algebra.Meta
\import Algebra.Monoid
\import Algebra.Ordered
\import Arith.Fin.Order
\import Arith.Nat \using (pred \as pred')
\import Data.Bool
\import Data.Fin
\import Data.Or
\import Equiv
\import Function
\import Function.Meta ($)
\import Logic.Unique
\import Logic
\import Logic.Meta
\import Meta
\import Order.Biordered
\import Order.PartialOrder
\import Order.StrictOrder
\import Paths
\import Paths.Meta
\import Set
\import Set.Fin.KFin

\class FinSet \extends KFinSet, DecSet {
  | finEq : TruncP (QEquiv {Fin card} {E})
  | finSurj => \case finEq \with {
    | inP e => inP (e, e.isSurj)
  }

  \default decideEq x y => \case finEq \with {
    | inP e => \case (FinOrder _).decideEq (e.ret x) (e.ret y) \with {
      | yes q => yes ((symQEquiv e).isInj q)
      | no q => no \lam p => q (pmap e.ret p)
    }
  }

  | # => /=
  | nonEqualApart p => p
  | #-irreflexive p => p idp
  | #-symmetric => /=-sym
  | #-comparison x y z x/=z => \case decideEq x y \with {
    | yes x=y => byRight (\lam y=z => x/=z (x=y *> y=z))
    | no x/=y => byLeft x/=y
  }
  | tightness {x} {y} x//=y => \case decideEq x y \with {
    | yes x=y => x=y
    | no x/=y => absurd (x//=y x/=y)
  }

  \lemma choice {B : E -> \Type} (p : \Pi (x : E) -> TruncP (B x)) : TruncP (\Pi (x : E) -> B x)
    => \case finEq \with {
      | inP e => \case finiteAC (\lam j => p (e j)) \with {
        | inP h => inP \lam x => transport B (e.f_ret x) $ h (e.ret x)
      }
    }

  \lemma liftDepSurj {A B : E -> \Type} {f : \Pi {x : E} -> A x -> B x} (isSur : \Pi {x : E} (y : B x) -> ∃ (a : A x) (f a = y)) (g : \Pi (x : E) -> B x) : ∃ (g' : \Pi (x : E) -> A x) ∀ x (f (g' x) = g x)
    => TruncP.map (choice (\lam x => isSur (g x))) (\lam h => (\lam x => (h x)., \lam x => (h x).))

  \lemma liftSurj {A B : \Type} {f : A -> B} (isSur : IsSurj f) (g : E -> B) : ∃ (g' : E -> A) ∀ x (f (g' x) = g x)
    => liftDepSurj (\lam {_} => isSur) g

  \lemma splitSurj {A : \Type} {f : A -> E} (isSur : IsSurj f) : IsSplitSurj f
    => liftSurj isSur id
} \where {
    \use \level levelProp {A : \Set} (X Y : FinSet A) : X = Y
      => ext (\case X.finEq, Y.finEq \with {
        | inP Xcard=A, inP Ycard=A => FinCardBij (Xcard=A `transQEquiv` symQEquiv Ycard=A)
      })

  \lemma ||-finiteAC {n : Nat} {A : \Type} {B : Fin n -> \Type} (f : \Pi (i : Fin n) -> A || B i) : A || (\Pi (i : Fin n) -> B i) \elim n
    |  => byRight \case __
    | suc n => \case f , ||-finiteAC (\lam i => f (suc i)) \with {
      | byLeft a, _ => byLeft a
      | _, byLeft a => byLeft a
      | byRight b0, byRight bs => byRight \case \elim __ \with {
        |  => b0
        | suc j => bs j
      }
    }

  \lemma finiteAC {n : Nat} {A : Fin n -> \Type} (f : \Pi (i : Fin n) -> TruncP (A i)) : TruncP (\Pi (i : Fin n) -> A i)
    => \case ||-finiteAC {_} {Empty} (\lam i => \case f i \with {
      | inP a => byRight a
    }) \with {
      | byLeft ()
      | byRight r => inP r
    }

    \func searchFin {n : Nat} (A : Fin n -> \Prop) (d : \Pi (j : Fin n) -> Dec (A j))
      : Or (\Sigma (j0 : Fin n) (A j0) (\Pi (j : Fin n) -> A j -> (j0 : Nat) <= j)) (\Pi (j : Fin n) -> Not (A j)) \elim n
      |  => inr (\case __)
      | suc n => \case d , searchFin _ (\lam j => d (suc j)) \with {
        | yes e, _ => inl (, e, \lam j _ => zero<=_)
        | no e, inl (j0,Aj0,u) => inl (suc j0, Aj0, \case \elim __ \with {
          |  => \lam A0 => absurd (e A0)
          | suc j => \lam Aj+1 => suc<=suc (u j Aj+1)
        })
        | no e, inr f => inr \case \elim __ \with {
          |  => e
          | suc j => f j
        }
      }

  \func searchFin-equiv {B : \Type} {n : Nat} (e : QEquiv {Fin n} {B}) (A : B -> \Prop) (d : \Pi (b : B) -> Dec (A b))
    : Or (\Sigma (b : B) (A b)) (\Pi (b : B) -> Not (A b))
    => \case searchFin (\lam j => A (e j)) (\lam j => d (e j)) \with {
         | inl p => inl (e p., p.)
         | inr q => inr \lam b Ab => q (e.ret b) $ transportInv A (e.f_ret b) Ab
       }

  \lemma searchFin-unique {n : Nat} (A : Fin n -> \Prop) (d : \Pi (j : Fin n) -> Dec (A j)) (c : ∃ (j : Fin n) (A j))
      : Contr (\Sigma (j0 : Fin n) (A j0) (\Pi (j : Fin n) -> A j -> (j0 : Nat) <= j))
      => \case \elim c \with {
        | inP (j,Aj) => isProp=>isContr (\lam t s => ext $ nat_fin_= $ <=-antisymmetric (t. s. s.) (s. t. t.)) \case searchFin A d \with {
          | inl r => r
          | inr e => absurd (e j Aj)
        }
      }

    \func pred {n : Nat} (x : Fin (suc (suc n))) : Fin (suc n) \elim x
      |  => 
      | suc x => x

    \lemma suc-isInj {n : Nat} (x y : Fin n) (p : suc x = suc y) : x = y \elim n
      | suc n => pmap pred p

    \lemma FinCardBij {n m : Nat} (p : QEquiv {Fin n} {Fin m}) : n = m
      => <=-antisymmetric (FinCardInj p p.isInj) (FinCardInj p.ret (symQEquiv p).isInj)

    \lemma FinCardInj {n m : Nat} (f : Fin n -> Fin m) (f-inj : IsInj f) : n <= m \elim n, m
      | , _ => zero<=_
      | suc n,  => \case f 
      | suc n, suc m => suc<=suc (FinCardInj (\lam x => skip (f ) (f (suc x)) (\lam f_0=f_suc => \case f-inj f_0=f_suc))
                                             (\lam {x} {y} p => suc-isInj x y (f-inj (skip-isInj p))))

    \func skip {n : Nat} (x0 x : Fin (suc n)) (d : x0 /= x) : Fin n \elim n, x0, x
      | _, ,  => absurd (d idp)
      | _, , suc x => x
      | suc n, suc x0,  => 
      | suc n, suc x0, suc x => suc (skip x0 x (\lam x0=x => d (pmap (suc __) x0=x)))

    \lemma skip-isInj {n : Nat} {x0 x y : Fin (suc n)} {d : x0 /= x} {d' : x0 /= y} (p : skip x0 x d = skip x0 y d') : x = y \elim n, x0, x, y
      | _, , , _ => absurd (d idp)
      | _, , _,  => absurd (d' idp)
      | _, _, ,  => idp
      | suc n, , suc x, suc y => pmap (suc __) p
      | suc n, suc x0, , suc y => \case p
      | suc n, suc x0, suc x,  => \case p
      | suc n, suc x0, suc x, suc y => pmap (suc __) (skip-isInj (suc-isInj _ _ p))

    \lemma sface_skip {n : Nat} {x0 x : Fin (suc n)} {d : x0 /= x} : sface x0 (skip x0 x d) = x \elim n, x0, x
      | _, ,  => absurd (d idp)
      | suc n, , suc x => idp
      | suc n, suc x0,  => idp
      | suc n, suc x0, suc x => pmap fsuc sface_skip

    \lemma skip_sface {n : Nat} {k : Fin (suc n)} {i : Fin n} {d : k /= sface k i} : skip k (sface k i) d = i \elim n, k, i
      | suc n, , i => idp
      | suc n, suc k,  => idp
      | suc n, suc k, suc i => pmap fsuc skip_sface

    \open NatOrder

    \lemma skip_< {n : Nat} {x0 x x' : Fin (suc n)} {d : x0 /= x} {d' : x0 /= x'} (p : skip x0 x d < skip x0 x' d') : x < x' \elim n, x0, x, x', p
      | _, , , _, _ => absurd (d idp)
      | _, , _, , _ => absurd (d' idp)
      | suc n, , suc x, suc x', p => suc<suc p
      | suc n, suc x0, , suc x', p => zero<suc
      | suc n, suc x0, suc x, suc x', suc<suc p => suc<suc (skip_< p)

    \lemma <_skip {n : Nat} {x0 x x' : Fin (suc n)} {d : x0 /= x} {d' : x0 /= x'} (p : x < x') : skip x0 x d < skip x0 x' d' \elim n, x0, x, x', p
      | _, , , _, _ => absurd (d idp)
      | _, , _, , _ => absurd (d' idp)
      | suc n, , suc x, suc x', suc<suc p => p
      | suc n, suc x0, , suc x', p => zero<suc
      | suc n, suc x0, suc x, suc x', suc<suc p => suc<suc (<_skip p)

    \lemma skip-left {n : Nat} {x0 x : Fin (suc n)} {d : x0 /= x} (p : x < x0) : skip x0 x d = {Nat} x \elim n, x0, x, p
      | , , , ()
      | suc n, suc x0, , _ => idp
      | suc n, suc x0, suc x, suc<suc p => pmap suc (skip-left p)

    \lemma skip-right {n : Nat} {x0 x : Fin (suc n)} {d : x0 /= x} (p : x0 < x) : skip x0 x d = {Nat} pred' x \elim n, x0, x, p
      | , , , ()
      | suc n, , suc x, _ => idp
      | suc n, suc x0, , _ => idp
      | suc n, suc x0, suc x, suc<suc p => pmap suc (skip-right p) *> suc_pred (/=-sym $ NatSemiring.<_/= $ zero<=_ <∘r p)

    \lemma fromArray {A : \Set} (l : Array A) (p : \Pi (a : A) -> ∃ (i : Fin l.len) (l i = a)) (q : \Pi {i j : Fin l.len} -> l i = l j -> i = j) : FinSet A l.len \cowith
      | finEq => IsEquiv.fromInjSurj q p
  }

\lemma card_Equiv {A B : FinSet} (e : QEquiv {A} {B}) : A.card = B.card
  => \case A.finEq, B.finEq \with {
    | inP e1, inP e2 => FinSet.FinCardBij $ e1 `transQEquiv` e `transQEquiv` symQEquiv e2
  }

\lemma card_inj {A B : FinSet} (f : A -> B) (inj : IsInj f) : A.card <= B.card
  => \case A.finEq, B.finEq \with {
    | inP e1, inP e2 => FinSet.FinCardInj (\lam x => e2.ret (f (e1 x))) \lam p => e1.isInj $ inj $ (symQEquiv e2).isInj p
  }

\lemma FinSet-transport (A : FinSet) {B : \Set} (e : QEquiv {A} {B}) : FinSet B A.card \cowith
  | finEq => \case A.finEq \with {
    | inP e' => inP (transQEquiv e' e)
  }

\lemma transport_zero {n m : Nat} (p : n = m) : transport Fin (pmap suc p)  =  \elim p
  | idp => idp

\lemma transport_suc {n m : Nat} (p : n = m) (x : Fin n) : transport Fin (pmap suc p) (suc x) = suc (transport Fin p x) \elim p
  | idp => idp

\lemma fin_transport {n m : Nat} {j : Fin n} {p : n = m} : transport Fin p j = {Nat} j \elim p
  | idp => idp

\func fin-inc_<= {n m : Nat} (p : n <= m) (i : Fin n) : Fin m
  => toFin i (fin_< i NatSemiring.<∘l p)

\func fin-inc {n m : Nat} (i : Fin n) : Fin (n + m) \elim n, i
  | suc n,  => 
  | suc n, suc i => suc (fin-inc i)
  \where {
    \lemma char {n m : Nat} {i : Fin n} : fin-inc {n} {m} i = {Nat} i \elim n, i
      | suc n,  => idp
      | suc n, suc i => pmap suc char

    \lemma char_nat {n : Nat} (i : Nat) : fin-inc {suc i} {n} (i Nat.mod suc i) = {Nat} i
      => char *> mod_< id<suc
  }

\func fin-inc-right {n m : Nat} (i : Fin m) : Fin (n + m)
  => fin-inc_<= (<=_+ zero<=_ <=-refl) i
  \where {
    \lemma char_nat {n : Nat} (i : Nat) : fin-inc-right {n} {suc i} (i Nat.mod suc i) = {Nat} i
      => toFin=id *> mod_< id<suc
  }

\func fin-raise {n k : Nat} (i : Fin n) : Fin (k + n) \elim k
  |  => i
  | suc k => suc (fin-raise i)

\func sface {n : Nat} (k : Fin (suc n)) (i : Fin n) : Fin (suc n) \elim n, k, i
  | suc n, , i => suc i
  | suc n, suc k,  => 
  | suc n, suc k, suc i => suc (sface k i)

\lemma sface-skip {n : Nat} {k : Fin (suc n)} {i : Fin n} : sface k i /= k \elim n, k, i
  | suc n, , i => \case __
  | suc n, suc k,  => \case __
  | suc n, suc k, suc i => fsuc/= sface-skip

\lemma sface-inj {n : Nat} {k : Fin (suc n)} {i i' : Fin n} (p : sface k i = sface k i') : i = i' \elim n, k, i, i'
  | suc n, , i, i' => pmap (fpred i) p
  | suc n, suc k, ,  => idp
  | suc n, suc k, , suc i' => \case p
  | suc n, suc k, suc i,  => \case p
  | suc n, suc k, suc i, suc i' => pmap fsuc $ sface-inj $ pmap (fpred k) p

\instance EmptyFin : FinSet Empty
  | card => 
  | finEq => inP \new QEquiv {
    | f => \case __
    | ret => \case __
    | ret_f => \case __
    | f_sec => \case __
  }

\instance UnitFin : FinSet (\Sigma)
  | card => 
  | finEq => inP \new QEquiv {
    | f _ => ()
    | ret _ => 
    | ret_f () => idp
    | f_sec _ => idp
  }

\instance BoolFin : FinSet Bool 
  | finEq => inP equiv
  \where
    \func equiv : QEquiv {Fin } {Bool} \cowith
      | f => \case __ \with {
        |  => true
        |  => false
      }
      | ret => \case __ \with {
        | true => 
        | false => 
      }
      | ret_f => \case \elim __ \with {
        |  => idp
        |  => idp
      }
      | f_sec => cases __ idp

\instance FinFin (n : Nat) : FinSet (Fin n)
  | card => n
  | finEq => inP idEquiv
  | decideEq => decideEq

\func DecFin {P : \Prop} (d : Dec P) : FinSet P \cowith
  | card => Dec.rec (\lam _ => ) (\lam _ => ) d
  | finEq => inP $ \case \elim d \with {
    | yes e => \new QEquiv {
                 | f _ => e
                 | ret _ => 
                 | ret_f () => idp
                 | f_sec => prop-isProp e
               }
    | no q => \new QEquiv {
                | f => \lam ()
                | ret e => absurd (q e)
                | ret_f => \lam ()
                | f_sec e => absurd (q e)
              }
  }

\func FinDec {A : FinSet} (B : A -> Decide) : Decide (\Pi (a : A) -> B a) \cowith
  | decide => \case A.finEq \with {
    | inP e => \case fin-dec (\lam j => B (e j)) \with {
      | yes p => yes \lam a => transport (B __) (e.f_ret a) $ p (e.ret a)
      | no q => no \lam f => q \lam j => f (e j)
    }
  }
  \where {
    \func fin-dec {n : Nat} (B : Fin n -> Decide) : Dec (\Pi (j : Fin n) -> B j) \elim n
      |  => yes \case __
      | suc n => \case decide {B }, fin-dec (\lam j => B (suc j)) \with {
        | yes p, yes q => yes \case \elim __ \with {
          |  => p
          | suc j => q j
        }
        | _, no q => no \lam f => q \lam j => f (suc j)
        | no p, _ => no \lam f => p (f )
      }
  }

\instance OrFin (S T : FinSet) : FinSet (Or S T)
  | card => S.card + T.card
  | finEq => \case S.finEq, T.finEq \with {
    | inP Se, inP Te => inP $ transQEquiv (aux {S.card} {T.card}) (Or.Or_QEquiv Se Te)
  }
  \where {
    \func aux {n m : Nat} : QEquiv {Fin (n + m)} {Or (Fin n) (Fin m)} \cowith
      | f (i : Fin (n + m)) : Or (Fin n) (Fin m) \elim n, i {
        | , i => inr i
        | suc n,  => inl 
        | suc n, suc i => Or.map (f i) fsuc id
      }
      | ret (e : Or (Fin n) (Fin m)) : Fin (n + m) \with {
        | inl s => fin-inc s
        | inr t => fin-raise t
      }
      | ret_f (i : Fin (n + m)) : ret (f i) = i \elim n, i {
        | , i => idp
        | suc n,  => idp
        | suc n, suc i => cases (f i arg addPath) \with {
          | inl s, p => \have q => inv (pmap ret p) *> ret_f i \in \case \elim n, \elim i, \elim s, \elim q \with {
            | suc n, i, s, q => pmap fsuc q
          }
          | inr t, p => pmap (\lam i => suc i) (inv (pmap ret p) *> ret_f i)
        }
      }
      | f_sec => \case \elim __ \with {
        | inl s => ret_inl-lem s
        | inr t => ret_inr-lem t
      }

    \lemma ret_inl-lem {n m : Nat} (i : Fin n) : OrFin.aux.f {n} {m} (fin-inc i) = inl i \elim n, i
      | suc n,  => idp
      | suc n, suc i => rewrite ret_inl-lem idp

    \lemma ret_inr-lem {n m : Nat} (t : Fin m) : aux.f (fin-raise {m} {n} t) = inr t \elim n
      |  => idp
      | suc n => rewrite (ret_inr-lem {n} t) idp
  }

\func finSet {A : FinSet} => A

\lemma finiteDependentChoice {A : \Type} (R : A -> A -> \Type) (a0 : A) (step : \Pi (a : A) -> ∃ (b : A) (R a b)) (n : Nat)
  : ∃ (f : Fin (suc n) -> A) (f  = a0) (\Pi (i : Fin n) -> R (f i) (f (suc i)))
  \elim n
  |  => inP (\lam _ => a0, idp, \case __)
  | suc n => \case finiteDependentChoice R a0 step n \with {
    | inP (f, f0=a0, chain) => \case step (f (finLast n)) \with {
      | inP (b, Rfnb) => inP (extendSeq f b, extendSeq-0 f b *> f0=a0, extendChain R f b chain Rfnb)
    }
  }
  \where {
    \func extendSeq {A : \Type} {n : Nat} (f : Fin (suc n) -> A) (b : A) (i : Fin (suc (suc n))) : A \elim n, i
      | n,  => f 
      | ,  => b
      | suc n, suc i => extendSeq (\lam j => f (suc j)) b i

    \func extendSeq-0 {A : \Type} {n : Nat} (f : Fin (suc n) -> A) (b : A) : extendSeq f b  = f  \elim n
      |  => idp
      | suc n => idp

    \func extendSeq-last {A : \Type} {n : Nat} (f : Fin (suc n) -> A) (b : A) : extendSeq f b (finLast (suc n)) = b \elim n
      |  => idp
      | suc n => extendSeq-last (\lam j => f (suc j)) b

    \func extendSeq-agree {A : \Type} {n : Nat} (f : Fin (suc n) -> A) (b : A) (i : Fin (suc n))
      : extendSeq f b i = f i \elim n, i
      | ,  => idp
      | suc n,  => idp
      | suc n, suc i => extendSeq-agree (\lam j => f (suc j)) b i

    \func extendChain {A : \Type} (R : A -> A -> \Type) {n : Nat}
                      (f : Fin (suc n) -> A) (b : A)
                      (chain : \Pi (i : Fin n) -> R (f i) (f (suc i)))
                      (Rfnb : R (f (finLast n)) b)
                      (i : Fin (suc n)) : R (extendSeq f b i) (extendSeq f b (suc i)) \elim n, i
      | ,  => transport2 R (inv (extendSeq-0 f b)) (inv (extendSeq-last f b)) Rfnb
      | suc n,  =>
        \have | p1 : extendSeq f b  = f  => extendSeq-0 f b
              | p2 : extendSeq f b  = f  => extendSeq-agree f b (fsuc fzero)
        \in transport2 R (inv p1) (inv p2) (chain )
      | suc n, suc i => extendChain R (\lam j => f (suc j)) b (\lam j => chain (suc j)) Rfnb i
  }