{-# OPTIONS --without-K --safe #-}

open import Algebra.Bundles using (CommutativeSemiring)
open import Level using (Level; 0ℓ; _⊔_)

module Data.Matrix.SemiadditiveDagger {c ℓ : Level} (R : CommutativeSemiring c ℓ) where

module R = CommutativeSemiring R

import Data.Nat as ℕ
import Data.Nat.Properties as ℕ-Props
import Data.Vec.Relation.Binary.Pointwise.Inductive as PW
import Relation.Binary.Reasoning.Setoid as ≈-Reasoning

open import Categories.Category.Cartesian.Bundle using (CartesianCategory)
open import Categories.Category.Cocartesian using (Cocartesian)
open import Categories.Category.Dagger using (HasDagger)
open import Categories.Object.Biproduct using (Biproduct)
open import Categories.Object.Coproduct using (IsCoproduct)
open import Categories.Object.Initial using (IsInitial)
open import Categories.Object.Product using (IsProduct)
open import Categories.Object.Terminal using (IsTerminal)
open import Categories.Object.Zero using (Zero)
open import Category.Dagger.Semiadditive using (SemiadditiveDagger)
open import Category.Semiadditive using (Semiadditive)
open import Data.Matrix.Category R.semiring using (Mat; _·_; ≑-·; ·-Iˡ; ·-Iʳ; ·-𝟎ˡ; ·-𝟎ʳ; ·-∥; ∥-·-≑; ·-resp-≋; ·-assoc)
open import Data.Matrix.Core R.setoid using (Matrix; Matrixₛ; _≋_; module ≋; ∥-cong; ≑-cong; ᵀ-cong)
open import Data.Matrix.Monoid R.+-monoid using (𝟎; 𝟎ᵀ; 𝟎≑𝟎; 𝟎∥𝟎; _[+]_; [+]-cong; [+]-𝟎ˡ; [+]-𝟎ʳ)
open import Data.Matrix.Raw using (_ᵀ; _ᵀᵀ; mapRows; []ᵥ; []ᵥ-∥; []ₕ; []ₕ-!; []ₕ-≑; _∷ᵥ_; _∷ₕ_; ∷ᵥ-ᵀ; _∥_; _≑_; ∷ₕ-ᵀ; ∷ₕ-≑; []ᵥ-ᵀ; head-∷-tailₕ; headₕ; tailₕ; ∷ₕ-∥; ∷ᵥ-≑; []ᵥ-!)
open import Data.Matrix.Semiadditive R.semiring using (Mat-Semiadditive)
open import Data.Matrix.Transform R.semiring using (I; Iᵀ; [_]_; _[_]; -[-]ᵀ; [-]--cong; [-]-[]ᵥ; [⟨⟩]-[]ₕ)
open import Data.Nat using (ℕ)
open import Data.Product using (_,_; Σ-syntax)
open import Data.Vec using (Vec; map; replicate; _++_)
open import Data.Vec.Properties using (map-cong; map-const)
open import Data.Vector.Bisemimodule R.semiring using (_∙_ ; ∙-cong)
open import Data.Vector.Core R.setoid using (Vector; Vectorₛ; module ≊; _≊_)
open import Data.Vector.Monoid R.+-monoid using () renaming (⟨ε⟩ to ⟨0⟩)
open import Data.Vector.Raw using (⟨⟩)
open import Data.Vector.Vec using (replicate-++)
open import Function using (_∘_)
open import Relation.Binary.PropositionalEquality as ≡ using (_≡_; module ≡-Reasoning)

open R
open Vec
open ℕ.ℕ

private
  variable
    A B C D E F : ℕ

opaque
  unfolding _∙_
  ∙-comm : (V W : Vector A) → V ∙ W ≈ W ∙ V
  ∙-comm [] [] = refl
  ∙-comm (x ∷ V) (w ∷ W) = +-cong (*-comm x w) (∙-comm V W)

opaque
  unfolding _[_] [_]_ _ᵀ []ᵥ _∷ₕ_ _≋_ _∷ᵥ_
  [-]-ᵀ : (M : Matrix A B) (V : Vector A) →  M [ V ] ≊ [ V ] (M ᵀ)
  [-]-ᵀ [] V = ≊.sym (≊.reflexive ([-]-[]ᵥ V))
  [-]-ᵀ (M₀ ∷ M) V = begin
      M₀ ∙ V ∷ map (_∙ V) M         ≈⟨ ∙-comm M₀ V PW.∷ (PW.map⁺ (λ {x} ≊V → trans (∙-comm x V) (∙-cong ≊.refl ≊V)) ≋.refl) ⟩
      V ∙ M₀ ∷ map (V ∙_) M         ≡⟨⟩
      map (V ∙_) (M₀ ∷ᵥ M)          ≡⟨ ≡.cong (map (V ∙_) ∘ (M₀ ∷ᵥ_)) (M ᵀᵀ) ⟨
      map (V ∙_) (M₀ ∷ᵥ M ᵀ ᵀ)      ≡⟨ ≡.cong (map (V ∙_)) (∷ₕ-ᵀ M₀ (M ᵀ)) ⟨
      map (V ∙_) ((M₀ ∷ₕ (M ᵀ)) ᵀ)  ∎
    where
      open ≈-Reasoning (Vectorₛ _)

opaque
  unfolding []ᵥ mapRows _∷ₕ_ _∷ᵥ_ _ᵀ _≋_
  ·-ᵀ
      : {A B C : ℕ}
        (M : Matrix A B)
        (N : Matrix B C)
      → (N · M) ᵀ ≋ M ᵀ · N ᵀ
  ·-ᵀ {A} {B} {zero} M [] = begin
      []ᵥ                   ≡⟨ map-const (M ᵀ) ⟨⟩ ⟨
      map (λ _ → ⟨⟩) (M ᵀ)  ≡⟨ map-cong [-]-[]ᵥ (M ᵀ) ⟨
      map ([_] []ᵥ) (M ᵀ)   ∎
    where
      open ≈-Reasoning (Matrixₛ 0 A)
  ·-ᵀ {A} {B} {suc C} M (N₀ ∷ N) = begin
      map ([_] M) (N₀ ∷ N) ᵀ        ≡⟨ -[-]ᵀ (N₀ ∷ N) M ⟨
      map ((N₀ ∷ N) [_]) (M ᵀ)      ≈⟨ PW.map⁺ (λ {V} ≋V → ≊.trans ([-]-ᵀ (N₀ ∷ N) V) ([-]--cong {A = (N₀ ∷ᵥ N) ᵀ} ≋V ≋.refl)) ≋.refl ⟩
      map ([_] ((N₀ ∷ N) ᵀ)) (M ᵀ)  ≡⟨ map-cong (λ V → ≡.cong ([ V ]_) (∷ᵥ-ᵀ N₀ N)) (M ᵀ) ⟩
      map ([_] (N₀ ∷ₕ N ᵀ)) (M ᵀ)   ∎
    where
      open ≈-Reasoning (Matrixₛ (suc C) A)

opaque
  unfolding _≋_
  ᵀ-involutive : (M : Matrix A B) → (M ᵀ) ᵀ ≋ M
  ᵀ-involutive M = ≋.reflexive (M ᵀᵀ)

opaque
  unfolding Matrix _∥_ _ᵀ _≑_ _∷ₕ_
  ∥-ᵀ : (M : Matrix A C) (N : Matrix B C) → (M ∥ N) ᵀ ≡ M ᵀ ≑ N ᵀ
  ∥-ᵀ {A} {zero} {B} [] [] = ≡.sym (replicate-++ A B [])
  ∥-ᵀ (M₀ ∷ M) (N₀ ∷ N) = begin
      (M₀ ++ N₀) ∷ₕ ((M ∥ N) ᵀ) ≡⟨ ≡.cong ((M₀ ++ N₀) ∷ₕ_) (∥-ᵀ M N) ⟩
      (M₀ ++ N₀) ∷ₕ (M ᵀ ≑ N ᵀ) ≡⟨ ∷ₕ-≑ M₀ N₀ (M ᵀ) (N ᵀ) ⟩
      (M₀ ∷ₕ M ᵀ) ≑ (N₀ ∷ₕ N ᵀ) ∎
    where
      open ≡-Reasoning

≑-ᵀ : (M : Matrix A B) (N : Matrix A C) → (M ≑ N) ᵀ ≡ M ᵀ ∥ N ᵀ
≑-ᵀ M N = begin
    (M ≑ N) ᵀ           ≡⟨ ≡.cong₂ (λ h₁ h₂ → (h₁ ≑ h₂) ᵀ) (M ᵀᵀ) (N ᵀᵀ) ⟨
    (M ᵀ ᵀ ≑ N ᵀ ᵀ ) ᵀ  ≡⟨ ≡.cong (_ᵀ) (∥-ᵀ (M ᵀ) (N ᵀ)) ⟨
    (M ᵀ ∥ N ᵀ ) ᵀ ᵀ    ≡⟨ (M ᵀ ∥ N ᵀ ) ᵀᵀ ⟩
    M ᵀ ∥ N ᵀ           ∎
  where
    open ≡-Reasoning

[I∥𝟎]ᵀ : (I ∥ 𝟎 {B} {A}) ᵀ ≋ I ≑ 𝟎
[I∥𝟎]ᵀ {B} {A} = begin
    (I ∥ 𝟎) ᵀ ≡⟨ ∥-ᵀ I 𝟎 ⟩
    I ᵀ ≑ 𝟎 ᵀ ≡⟨ ≡.cong₂ _≑_ Iᵀ 𝟎ᵀ ⟩
    I ≑ 𝟎     ∎
  where
    open ≈-Reasoning (Matrixₛ A (A ℕ.+ B))

[𝟎∥I]ᵀ : (𝟎 {A} {B} ∥ I) ᵀ ≋ 𝟎 ≑ I
[𝟎∥I]ᵀ {A} {B} = begin
    (𝟎 ∥ I) ᵀ ≡⟨ ∥-ᵀ 𝟎 I ⟩
    𝟎 ᵀ ≑ I ᵀ ≡⟨ ≡.cong₂ _≑_ 𝟎ᵀ Iᵀ ⟩
    𝟎 ≑ I     ∎
  where
    open ≈-Reasoning (Matrixₛ B (A ℕ.+ B))

Mat-HasDagger : HasDagger Mat
Mat-HasDagger = record
    { _† = λ M → M ᵀ
    ; †-identity = ≋.reflexive Iᵀ
    ; †-homomorphism = λ {f = f} {g} → ·-ᵀ f g
    ; †-resp-≈ = ᵀ-cong
    ; †-involutive = ᵀ-involutive
    }

Mat-SemiadditiveDagger : SemiadditiveDagger Mat
Mat-SemiadditiveDagger = record
    { semiadditive = Mat-Semiadditive
    ; dagger = Mat-HasDagger
    ; π₁† = [I∥𝟎]ᵀ
    ; π₂† = [𝟎∥I]ᵀ
    ; ⟨⟩-† = λ {f = M} {N} → ≋.reflexive (≑-ᵀ M N)
    }