{-# OPTIONS --safe  #-}

open import Level
open import Data.Nat using (ℕ)
open import Data.Maybe
open import Data.Sum
open import Data.List
open import Data.Vec using (Vec)
import Data.Vec as V
open import Data.Fin
open import Data.List.NonEmpty using (List⁺; snocView; _∷ʳ′_)
open import Relation.Binary.PropositionalEquality -- using (_≡_; sym; refl)
open import Function.Base
open import Data.Product

open import Tactic.Cong using (cong!)

open import Finite
open import Auxiliary-Definitions

module Mealy {ℓ : Level} (I O : Set ℓ) where

record Mealy-Machine-on {qℓ} (Q : Set qℓ) : Set (qℓ ⊔ ℓ) where
  -- Given a set of states Q
  -- a Mealy-Graph is what remains for an entire Mealy machine
  field
    q₀ : Q
    δ : Q → I → Q
    out : Q → I → O
    -- (λ is a reserved keyword, so we use a different name than in the L# paper)

  δ* : Q → List I → Q
  δ* q [] = q
  δ* q (x ∷ w) = δ* (δ q x) w

  δ*-++ : ∀ {q} v w → δ* q (v ++ w) ≡ δ* (δ* q v) w
  δ*-++ {q} [] w = refl
  δ*-++ {q} (x ∷ v) w = δ*-++ {δ q x} v w

  δ₀* : List I → Q
  δ₀* = δ* q₀

  δ₀*-snoc : ∀ (v : List I) i → δ₀* (v ∷ʳ i) ≡ δ (δ₀* v) i
  δ₀*-snoc v i = δ*-++ v (i ∷ [])

  ⟦_/_⟧* : Q → List I → List O
  ⟦ q / [] ⟧* = []
  ⟦ q / (x ∷ w) ⟧* =
    (out q x) ∷ ⟦ δ q x / w ⟧*

  ⟦_/_⟧ⁿ : Q → ∀ {n} → Vec I n → Vec O n
  ⟦ q / V.[] ⟧ⁿ = V.[]
  ⟦ q / (x V.∷ w) ⟧ⁿ =
    (out q x) V.∷ ⟦ q / w ⟧ⁿ

  -- The semantics of a mealy machine is equivalent
  -- to a map from non-empty input sequences to output symbols.
  ⟦_⟧+ : List⁺ I → O
  ⟦_⟧+ w with snocView w
  ... | (v ∷ʳ′ i) = out (δ* q₀ v) i

  ⟦_⟧* : List I → List O
  ⟦_⟧* w = ⟦ q₀ / w ⟧*

  ⟦_⟧ⁿ : ∀ {n} → Vec I n → Vec O n
  ⟦_⟧ⁿ w = ⟦ q₀ / w ⟧ⁿ

  reindex : ∀ {a : Level} {P : Set a} → (Q → P) → (P → Q) → Mealy-Machine-on P
  reindex f f⁻¹ =
    record { q₀ = f q₀ ; δ = λ p i → f (δ (f⁻¹ p) i) ; out = λ p i → out (f⁻¹ p) i }

module Equiv {a b c : Level}  (_≈_ : O → O → Set c) {Q₁ : Set a} {Q₂ : Set b} (M₁ : Mealy-Machine-on Q₁) (M₂ : Mealy-Machine-on Q₂) where
  private
    module M₁ = Mealy-Machine-on M₁
    module M₂ = Mealy-Machine-on M₂

  Mealy-Equivalent : Set _
  Mealy-Equivalent = ∀ (w : List⁺ I) → M₁.⟦ w ⟧+ ≈ M₂.⟦ w ⟧+

  record is-Mealy-morphism (h : Q₁ → Q₂) : Set (a ⊔ b ⊔ ℓ ⊔ c) where
    field
      preserves-q₀ : h M₁.q₀ ≡ M₂.q₀
      preserves-δ : ∀ q i → h (M₁.δ q i) ≡ M₂.δ (h q) i
      preserves-out : ∀ q i → M₁.out q i ≈ M₂.out (h q) i

    preserves-δ* : ∀ q w → h (M₁.δ* q w) ≡ M₂.δ* (h q) w
    preserves-δ* q [] = refl
    preserves-δ* q (x ∷ w) =
      begin
      h (M₁.δ* q (x ∷ w))     ≡⟨⟩
      h (M₁.δ* (M₁.δ q x) w)  ≡⟨ preserves-δ* _ w ⟩
      M₂.δ* (h (M₁.δ q x)) w  ≡⟨ cong! (preserves-δ q x) ⟩
      M₂.δ* (M₂.δ (h q) x) w  ≡⟨⟩
      M₂.δ* (h q) (x ∷ w)
      ∎
      where open ≡-Reasoning

    -- a proof that a morphism between Mealy machines
    -- preserves the semantics:
    equivalent-semantics : Mealy-Equivalent
    equivalent-semantics w with snocView w
    ... | (v ∷ʳ′ i) rewrite (sym preserves-q₀)
                    rewrite (sym (preserves-δ* M₁.q₀ v))
                    = preserves-out _ i

module Reindex {a b : Level} {Q₁ : Set a} {Q₂ : Set b}
               (M : Mealy-Machine-on Q₁)
               (f : Q₁ → Q₂)
               (g : Q₂ → Q₁)
               where
  private
    module M = Mealy-Machine-on M

  open Equiv {a} {b} _≡_ {Q₁} {Q₂}
  reindex-morphism : (∀ x → g (f x) ≡ x)
                    → is-Mealy-morphism M (Mealy-Machine-on.reindex M f g) f
  reindex-morphism splitting =
    record
    { preserves-q₀ = refl
    ; preserves-δ = λ q i → cong (λ X → f (δ X i)) (sym (splitting q))
    ; preserves-out = λ q i → cong (λ X → out X i) (sym (splitting q))
    }
    where open Mealy-Machine-on M

  reindex-correct : (∀ x → g (f x) ≡ x) → Mealy-Equivalent M (M.reindex f g)
  reindex-correct splitting =
    is-Mealy-morphism.equivalent-semantics
      M (M.reindex f g) (reindex-morphism splitting)


record Mealy-Machine : Set ℓ where
  field
    n : ℕ
  Q = Fin n
  field
    structure : Mealy-Machine-on Q

  open Mealy-Machine-on structure public