module Problem15 where

open import Data.Vec as Vec
open import Data.Fin
open import Data.Nat using (ℕ; zero; suc)
open import Data.Product hiding (map)
open import Relation.Binary.PropositionalEquality
open import Data.Vec.Properties
open import Relation.Nullary
open import Function

open ≡-Reasoning

variable n : ℕ
variable A B : Set

-- Useful properties not in stdlib 1.7.3

map-tabulate : (g : A → B) (f : Fin n → A) → map g (tabulate f) ≡ tabulate (g ∘′ f)
map-tabulate g f = begin
  map g (tabulate f)       ≡⟨  cong (map g) (tabulate-allFin f) ⟩
  map g (map f (allFin _)) ≡˘⟨ map-∘ g f (allFin _) ⟩
  map (g ∘′ f) (allFin _)  ≡˘⟨ tabulate-allFin (g ∘′ f) ⟩
  tabulate (g ∘′ f)        ∎

Sur : ∀(v : Vec (Fin n) n) → Set
Sur {n} v = ∀(x : Fin n) → ∃[ i ] lookup v i ≡ x

Inj : ∀(v : Vec A n ) → Set
Inj {_}{n} v = (a b : Fin n) → lookup v a ≡ lookup v b → a ≡ b

record Perm (n : ℕ) : Set where
    constructor P
    field indices : Vec (Fin n) n
    field surjective : Sur indices
    field injective  : Inj indices

    _⟨$⟩_ : Fin n → Fin n
    _⟨$⟩_ = lookup indices

open Perm

-- Agda's proof irrelevance is too janky to work here
-- And we'd need functional extensionality + UIP to prove this directly
postulate cong-Perm : ∀ (p q : Perm n) → Perm.indices p ≡ Perm.indices q →  p ≡ q

permute : Perm n → Vec A n → Vec A n
permute (P p _ _) v = map (lookup v) p

_⊡_ : Perm n → Perm n → Perm n
(p ⊡ r) .indices    = tabulate ((p ⟨$⟩_) ∘ (r ⟨$⟩_))
(p ⊡ r) .surjective x .proj₁ = r .surjective (p .surjective x .proj₁) .proj₁
(p ⊡ r) .surjective x .proj₂ =
  let (i , prf₁) = p .surjective x in
  let (j , prf₂) = r .surjective i in
  begin
    lookup (tabulate ((p ⟨$⟩_) ∘ (r ⟨$⟩_))) j ≡⟨ lookup∘tabulate _ j ⟩
    p ⟨$⟩ (r ⟨$⟩ j)                           ≡⟨ cong (p ⟨$⟩_) prf₂ ⟩
    p ⟨$⟩ i                                   ≡⟨ prf₁ ⟩
    x                                         ∎
(p ⊡ r) .injective x y prf  =
  r .injective x y $′
  p .injective (r ⟨$⟩ x) (r ⟨$⟩ y) $′
  begin
    p ⟨$⟩ (r ⟨$⟩ x)                           ≡˘⟨ lookup∘tabulate _ x ⟩
    lookup (tabulate ((p ⟨$⟩_) ∘ (r ⟨$⟩_))) x ≡⟨  prf ⟩
    lookup (tabulate ((p ⟨$⟩_) ∘ (r ⟨$⟩_))) y ≡⟨  lookup∘tabulate _ y ⟩
    p ⟨$⟩ (r ⟨$⟩ y)                           ∎

infixl 20 _⊡_

composition : (v : Vec A n)(p q : Perm n) → permute (p ⊡ q) v ≡ permute q (permute p v)
composition v p q = begin
  map (lookup v) ((p ⊡ q) .indices)                           ≡⟨⟩
  map (lookup v) (tabulate ((p ⟨$⟩_) ∘′ (q ⟨$⟩_)))            ≡⟨  map-tabulate (lookup v) ((p ⟨$⟩_) ∘′ (q ⟨$⟩_)) ⟩
  tabulate (lookup v ∘′ (p ⟨$⟩_) ∘′ (q ⟨$⟩_))                 ≡⟨⟩
  tabulate (lookup v ∘′ (p ⟨$⟩_) ∘′ lookup (q .indices))      ≡˘⟨ tabulate-cong (λ i → lookup-map i (lookup v ∘′ (p ⟨$⟩_)) (q .indices)) ⟩
  tabulate (lookup (map (lookup v ∘′ (p ⟨$⟩_)) (q .indices))) ≡⟨  tabulate∘lookup (map (lookup v ∘′ (p ⟨$⟩_)) (q .indices)) ⟩
  map (lookup v ∘′ (p ⟨$⟩_)) (q .indices)                     ≡⟨⟩
  map (lookup v ∘′ lookup (p .indices)) (q .indices)          ≡˘⟨  map-cong (λ i → lookup-map i (lookup v) (p .indices)) (q .indices) ⟩
  map (lookup (map (lookup v) (p .indices))) (q .indices)     ∎

assoc : ∀ (p q r : Perm n) → p ⊡ (q ⊡ r) ≡ p ⊡ q ⊡ r
assoc p q r =
  cong-Perm (p ⊡ (q ⊡ r)) (p ⊡ q ⊡ r) $′
  begin
    tabulate ((p ⟨$⟩_) ∘′ ((q ⊡ r) ⟨$⟩_))       ≡⟨  tabulate-cong (λ i → cong (p ⟨$⟩_) (lookup∘tabulate ((q ⟨$⟩_) ∘′ (r ⟨$⟩_)) i)) ⟩
    tabulate ((p ⟨$⟩_) ∘′ (q ⟨$⟩_) ∘′ (r ⟨$⟩_)) ≡˘⟨ tabulate-cong (λ i → lookup∘tabulate ((p ⟨$⟩_) ∘′ (q ⟨$⟩_)) (r ⟨$⟩ i)) ⟩
    tabulate (((p ⊡ q) ⟨$⟩_) ∘′ (r ⟨$⟩_))       ∎

ι : Perm n
ι {n} .indices = allFin n
ι {n} .surjective x = x , lookup-allFin x
ι {n} .injective x y prf = begin
  x                   ≡˘⟨ lookup-allFin x ⟩
  lookup (allFin n) x ≡⟨  prf ⟩
  lookup (allFin n) y ≡⟨  lookup-allFin y ⟩
  y                   ∎


identityˡ : ∀(p : Perm n) → ι ⊡ p ≡ p
identityˡ p =
  cong-Perm (ι ⊡ p) p $′
  begin
    tabulate (lookup (allFin _) ∘′ (p ⟨$⟩_)) ≡⟨ tabulate-cong (λ i → lookup-allFin (p ⟨$⟩ i)) ⟩
    tabulate (p ⟨$⟩_)                        ≡⟨ tabulate∘lookup (p .indices) ⟩
    p .indices                               ∎


identityʳ : ∀(p : Perm n) → p ⊡ ι ≡ p
identityʳ {n} p =
  cong-Perm (p ⊡ ι) p $′
  begin
    tabulate ((p ⟨$⟩_) ∘′ lookup (allFin n)) ≡⟨ tabulate-cong (λ i → cong (p ⟨$⟩_) (lookup-allFin i)) ⟩
    tabulate (p ⟨$⟩_)                        ≡⟨ tabulate∘lookup (p .indices) ⟩
    p .indices                               ∎


_⁻¹ : Perm n → Perm n
(p ⁻¹) .indices = tabulate (proj₁ ∘ p .surjective)
(p ⁻¹) .surjective x .proj₁ = p ⟨$⟩ x
(p ⁻¹) .surjective x .proj₂ =
  p .injective (lookup (tabulate (proj₁ ∘ p .surjective)) (p ⟨$⟩ x)) x $′
  begin
    p ⟨$⟩ lookup (tabulate (proj₁ ∘ p .surjective)) (p ⟨$⟩ x) ≡⟨ cong (p ⟨$⟩_) (lookup∘tabulate (proj₁ ∘ p .surjective) (p ⟨$⟩ x)) ⟩
    p ⟨$⟩ p .surjective (p ⟨$⟩ x) .proj₁                      ≡⟨ p .surjective (p ⟨$⟩ x) .proj₂ ⟩
    p ⟨$⟩ x                                                   ∎
(p ⁻¹) .injective x y prf =
  begin
    x                                                 ≡˘⟨ p .surjective x .proj₂ ⟩
    p ⟨$⟩ p .surjective x .proj₁                      ≡˘⟨ cong (p ⟨$⟩_) (lookup∘tabulate (proj₁ ∘ p .surjective) x) ⟩
    p ⟨$⟩ lookup (tabulate (proj₁ ∘ p .surjective)) x ≡⟨  cong (p ⟨$⟩_) prf ⟩
    p ⟨$⟩ lookup (tabulate (proj₁ ∘ p .surjective)) y ≡⟨  cong (p ⟨$⟩_) (lookup∘tabulate (proj₁ ∘ p .surjective) y) ⟩
    p ⟨$⟩ p .surjective y .proj₁                      ≡⟨  p .surjective y .proj₂ ⟩
    y                                                 ∎

infixl 30 _⁻¹


inverseˡ : ∀(p : Perm n) → p ⁻¹ ⊡ p ≡ ι
inverseˡ p =
  cong-Perm (p ⁻¹ ⊡ p) ι $′
  begin
    tabulate ((p ⁻¹ ⟨$⟩_) ∘′ (p ⟨$⟩_))          ≡⟨ tabulate-cong (λ i → lookup∘tabulate (proj₁ ∘ p .surjective) (p ⟨$⟩ i)) ⟩
    tabulate (proj₁ ∘ p .surjective ∘ (p ⟨$⟩_)) ≡⟨ tabulate-cong (λ i → p .injective _ i (p .surjective (p ⟨$⟩ i) .proj₂)) ⟩
    tabulate id                                 ∎

inverseʳ : ∀(p : Perm n) → p ⊡ p ⁻¹ ≡ ι
inverseʳ p =
  cong-Perm (p ⊡ p ⁻¹) ι $′
  begin
    tabulate ((p ⟨$⟩_) ∘′ (p ⁻¹ ⟨$⟩_))          ≡⟨ tabulate-cong (λ i → cong (p ⟨$⟩_) (lookup∘tabulate (proj₁ ∘ p .surjective) i)) ⟩
    tabulate ((p ⟨$⟩_) ∘ proj₁ ∘ p .surjective) ≡⟨ tabulate-cong (λ i → p .surjective i .proj₂) ⟩
    tabulate id                                 ∎