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

open import Algebra using (Semiring)
open import Level using (Level; 0ℓ; _⊔_)

module Data.Matrix.Semiadditive {c ℓ : Level} (R : Semiring c ℓ) where

module R = Semiring 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.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.Semiadditive using (Semiadditive)
open import Data.Matrix.Category R 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.Transform R 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 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 : ℕ

inj₁ : (M : Matrix A C) (N : Matrix B C) → (M ∥ N) · (I ≑ 𝟎) ≋ M
inj₁ {A} {C} M N = begin
    (M ∥ N) · (I ≑ 𝟎)   ≈⟨ ∥-·-≑ M N I 𝟎 ⟩
    (M · I) [+] (N · 𝟎) ≈⟨ [+]-cong ·-Iʳ (·-𝟎ʳ N) ⟩
    M [+] 𝟎             ≈⟨ [+]-𝟎ʳ M ⟩
    M ∎
  where
    open ≈-Reasoning (Matrixₛ A C)

inj₂ : (M : Matrix A C) (N : Matrix B C) → (M ∥ N) · (𝟎 ≑ I) ≋ N
inj₂ {A} {C} {B} M N = begin
    (M ∥ N) · (𝟎 ≑ I)   ≈⟨ ∥-·-≑ M N 𝟎 I ⟩
    (M · 𝟎) [+] (N · I) ≈⟨ [+]-cong (·-𝟎ʳ M) ·-Iʳ ⟩
    𝟎 [+] N             ≈⟨ [+]-𝟎ˡ N ⟩
    N ∎
  where
    open ≈-Reasoning (Matrixₛ B C)

opaque

  unfolding Matrix _∷ᵥ_

  split-∥ : (A : ℕ) (M : Matrix (A ℕ.+ B) C) → Σ[ M₁ ∈ Matrix A C ] Σ[ M₂ ∈ Matrix B C ] M₁ ∥ M₂ ≡ M
  split-∥ zero M = []ᵥ , M , []ᵥ-∥ M
  split-∥ (suc A) M′
    rewrite ≡.sym (head-∷-tailₕ M′)
    using M₀ ← headₕ M′
    using M ← tailₕ M′
    with split-∥ A M
  ... | M₁ , M₂ , M₁∥M₂≡M = M₀ ∷ₕ M₁ , M₂ , (begin
      (M₀ ∷ₕ M₁) ∥ M₂ ≡⟨ ∷ₕ-∥ M₀ M₁ M₂ ⟨
      M₀ ∷ₕ M₁ ∥ M₂   ≡⟨ ≡.cong (M₀ ∷ₕ_) M₁∥M₂≡M ⟩
      M₀ ∷ₕ M ∎)
    where
      open ≡-Reasoning

  split-≑ : (B : ℕ) (M : Matrix A (B ℕ.+ C)) → Σ[ M₁ ∈ Matrix A B ] Σ[ M₂ ∈ Matrix A C ] M₁ ≑ M₂ ≡ M
  split-≑ zero M = []ₕ , M , []ₕ-≑ M
  split-≑ (suc B) (M₀ ∷ M) with split-≑ B M
  ... | M₁ , M₂ , M₁≑M₂≡M = M₀ ∷ᵥ M₁ , M₂ , (begin
      (M₀ ∷ᵥ M₁) ≑ M₂ ≡⟨ ∷ᵥ-≑ M₀ M₁ M₂ ⟨
      M₀ ∷ᵥ M₁ ≑ M₂   ≡⟨ ≡.cong (M₀ ∷ᵥ_) M₁≑M₂≡M ⟩
      M₀ ∷ᵥ M ∎)
    where
      open ≡-Reasoning

∥-uniq
    : (H : Matrix (A ℕ.+ B) C)
      (M : Matrix A C)
      (N : Matrix B C)
    → H · (I ≑ 𝟎) ≋ M
    → H · (𝟎 ≑ I) ≋ N
    → M ∥ N ≋ H
∥-uniq {A} {B} {C} H M N eq₁ eq₂
  with (H₁ , H₂ , H₁∥H₂≡H) ← split-∥ A H
  rewrite ≡.sym H₁∥H₂≡H = begin
    M ∥ N                                         ≈⟨ ∥-cong eq₁ eq₂ ⟨
    (H₁ ∥ H₂) · (I {A} ≑ 𝟎) ∥ (H₁ ∥ H₂) · (𝟎 ≑ I) ≈⟨ ∥-cong (inj₁ H₁ H₂) (inj₂ H₁ H₂) ⟩
    (H₁ ∥ H₂) ∎
  where
    open ≈-Reasoning (Matrixₛ (A ℕ.+ B) C)

proj₁ : (M : Matrix A B) (N : Matrix A C) → (I ∥ 𝟎) · (M ≑ N) ≋ M
proj₁ {A} {B} M N = begin
    (I ∥ 𝟎) · (M ≑ N)   ≈⟨ ∥-·-≑ I 𝟎 M N ⟩
    (I · M) [+] (𝟎 · N) ≈⟨ [+]-cong ·-Iˡ (·-𝟎ˡ N) ⟩
    M [+] 𝟎             ≈⟨ [+]-𝟎ʳ M ⟩
    M ∎
  where
    open ≈-Reasoning (Matrixₛ A B)

proj₂ : (M : Matrix A B) (N : Matrix A C) → (𝟎 ∥ I) · (M ≑ N) ≋ N
proj₂ {A} {_} {C} M N = begin
    (𝟎 ∥ I) · (M ≑ N)   ≈⟨ ∥-·-≑ 𝟎 I M N ⟩
    (𝟎 · M) [+] (I · N) ≈⟨ [+]-cong (·-𝟎ˡ M) ·-Iˡ ⟩
    𝟎 [+] N             ≈⟨ [+]-𝟎ˡ N ⟩
    N ∎
  where
    open ≈-Reasoning (Matrixₛ A C)

≑-uniq
    : (H : Matrix A (B ℕ.+ C))
      (M : Matrix A B)
      (N : Matrix A C)
    → (I ∥ 𝟎) · H ≋ M
    → (𝟎 ∥ I) · H ≋ N
    → M ≑ N ≋ H
≑-uniq {A} {B} {C} H M N eq₁ eq₂
  with (H₁ , H₂ , H₁≑H₂≡H) ← split-≑ B H
  rewrite ≡.sym H₁≑H₂≡H = begin
      M ≑ N                                         ≈⟨ ≑-cong eq₁ eq₂ ⟨
      (I {B} ∥ 𝟎) · (H₁ ≑ H₂) ≑ (𝟎 ∥ I) · (H₁ ≑ H₂) ≈⟨ ≑-cong (proj₁ H₁ H₂) (proj₂ H₁ H₂) ⟩
      H₁ ≑ H₂                                       ∎
  where
    open ≈-Reasoning (Matrixₛ A (B ℕ.+ C))

isCoproduct : IsCoproduct Mat (I {A} ≑ 𝟎) (𝟎 ≑ I {B})
isCoproduct {A} {B} = record
    { [_,_] = _∥_
    ; inject₁ = λ {a} {b} {c} → inj₁ b c
    ; inject₂ = λ {a} {b} {c} → inj₂ b c
    ; unique = λ eq₁ eq₂ → ∥-uniq _ _ _ eq₁ eq₂
    }

isProduct : IsProduct Mat (I {A} ∥ 𝟎) (𝟎 ∥ I {B})
isProduct {A} {B} = record
    { ⟨_,_⟩ = _≑_
    ; project₁ = λ {a} {b} {c} → proj₁ b c
    ; project₂ = λ {a} {b} {c} → proj₂ b c
    ; unique = λ eq₁ eq₂ → ≑-uniq _ _ _ eq₁ eq₂
    }

opaque

  unfolding Matrix

  π₁∘i₁ : (I {A} ∥ 𝟎 {B}) · (I ≑ 𝟎) ≋ I
  π₁∘i₁ {A} = begin
      (I ∥ 𝟎) · (I ≑ 𝟎)   ≈⟨ ∥-·-≑ I 𝟎 I 𝟎 ⟩
      (I · I) [+] (𝟎 · 𝟎) ≈⟨ [+]-cong ·-Iˡ (·-𝟎ˡ 𝟎) ⟩
      I [+] 𝟎             ≈⟨ [+]-𝟎ʳ I ⟩
      I                   ∎
    where
      open ≈-Reasoning (Matrixₛ A A)

  π₂∘i₂ : (𝟎 {A} {B} ∥ I) · (𝟎 ≑ I) ≋ I
  π₂∘i₂ {A} {B} = begin
      (𝟎 ∥ I) · (𝟎 ≑ I)   ≈⟨ ∥-·-≑ 𝟎 I 𝟎 I ⟩
      (𝟎 · 𝟎) [+] (I · I) ≈⟨ [+]-cong (·-𝟎ˡ 𝟎) ·-Iˡ  ⟩
      𝟎 [+] I             ≈⟨ [+]-𝟎ˡ I ⟩
      I                   ∎
    where
      open ≈-Reasoning (Matrixₛ B B)

  π₁∘i₂ : (I {A} ∥ 𝟎 {B}) · (𝟎 ≑ I) ≋ 𝟎 {B} {A}
  π₁∘i₂ {A} {B} = begin
      (I ∥ 𝟎) · (𝟎 ≑ I)   ≈⟨ ∥-·-≑ I 𝟎 𝟎 I ⟩
      (I · 𝟎) [+] (𝟎 · I) ≈⟨ [+]-cong (·-𝟎ʳ I) (·-𝟎ˡ I) ⟩
      𝟎 [+] 𝟎             ≈⟨ [+]-𝟎ʳ 𝟎 ⟩
      𝟎                   ∎
    where
      open ≈-Reasoning (Matrixₛ B A)

  π₂∘i₁ : (𝟎 {A} {B} ∥ I) · (I ≑ 𝟎) ≋ 𝟎 {A} {B}
  π₂∘i₁ {A} {B} = begin
      (𝟎 ∥ I) · (I ≑ 𝟎)   ≈⟨ ∥-·-≑ 𝟎 I I 𝟎 ⟩
      (𝟎 · I) [+] (I · 𝟎) ≈⟨ [+]-cong (·-𝟎ˡ I) (·-𝟎ʳ I) ⟩
      𝟎 [+] 𝟎             ≈⟨ [+]-𝟎ʳ 𝟎 ⟩
      𝟎                   ∎
    where
      open ≈-Reasoning (Matrixₛ A B)

  permute
      : (I ≑ 𝟎 {A} {B}) · (I ∥ 𝟎 {B} {A}) · (𝟎 {B} {A} ≑ I) · (𝟎 {A} {B} ∥ I)
      ≋ (𝟎 {B} {A} ≑ I) · (𝟎 {A} {B} ∥ I) · (I ≑ 𝟎 {A} {B}) · (I ∥ 𝟎 {B} {A})
  permute {A} {B} = begin
      (I ≑ 𝟎) · (I ∥ 𝟎 {B} {A}) · (𝟎 {B} {A} ≑ I) · (𝟎 {A} {B} ∥ I)   ≈⟨ ·-resp-≋ ≋.refl ·-assoc ⟨
      (I ≑ 𝟎) · ((I ∥ 𝟎 {B} {A}) · (𝟎 {B} {A} ≑ I)) · (𝟎 {A} {B} ∥ I) ≈⟨ ·-resp-≋ ≋.refl (·-resp-≋ π₁∘i₂ ≋.refl) ⟩
      (I ≑ 𝟎) · 𝟎 {B} {A} · (𝟎 {A} {B} ∥ I)                           ≈⟨ ·-resp-≋ ≋.refl (·-𝟎ˡ (𝟎 ∥ I)) ⟩
      (I ≑ 𝟎 {A} {B}) · 𝟎                                             ≈⟨ ·-𝟎ʳ (I ≑ 𝟎) ⟩
      𝟎                                                               ≈⟨ ·-𝟎ʳ (𝟎 ≑ I) ⟨
      (𝟎 {B} {A} ≑ I) · 𝟎                                             ≈⟨ ·-resp-≋ ≋.refl (·-𝟎ˡ (I ∥ 𝟎)) ⟨
      (𝟎 ≑ I) · 𝟎 {A} {B} · (I ∥ 𝟎 {B} {A})                           ≈⟨ ·-resp-≋ ≋.refl (·-resp-≋ π₂∘i₁ ≋.refl) ⟨
      (𝟎 {B} {A} ≑ I) · ((𝟎 ∥ I) · (I ≑ 𝟎 {A} {B})) · (I ∥ 𝟎 {B} {A}) ≈⟨ ·-resp-≋ ≋.refl ·-assoc ⟩
      (𝟎 {B} {A} ≑ I) · (𝟎 {A} {B} ∥ I) · (I ≑ 𝟎 {A} {B}) · (I ∥ 𝟎)   ∎
    where
      open ≈-Reasoning (Matrixₛ (A ℕ.+ B) (A ℕ.+ B))

biproduct : Biproduct Mat A B
biproduct {A} {B} = record
    { A⊕B = A ℕ.+ B
    ; π₁ = I ∥ 𝟎
    ; π₂ = 𝟎 ∥ I
    ; i₁ = I ≑ 𝟎
    ; i₂ = 𝟎 ≑ I
    ; isBiproduct = record
        { isCoproduct = isCoproduct
        ; isProduct = isProduct
        ; π₁∘i₁≈id = π₁∘i₁
        ; π₂∘i₂≈id = π₂∘i₂
        ; permute = permute
        }
    }

opaque

  unfolding _≋_

  ¡-unique : (E : Matrix 0 B) → []ᵥ ≋ E
  ¡-unique E = ≋.reflexive (≡.sym ([]ᵥ-! E))

  !-unique : (E : Matrix A 0) → []ₕ ≋ E
  !-unique E = ≋.reflexive (≡.sym ([]ₕ-! E))

isInitial : IsInitial Mat 0
isInitial = record
    { ¡ = []ᵥ
    ; ¡-unique = ¡-unique
    }

isTerminal : IsTerminal Mat 0
isTerminal = record
    { ! = []ₕ
    ; !-unique = !-unique
    }

zeroObj : Zero Mat
zeroObj = record
    { 𝟘 = 0
    ; isZero = record
        { isInitial = isInitial
        ; isTerminal = isTerminal
        }
    }

Mat-Semiadditive : Semiadditive Mat
Mat-Semiadditive = record
    { zero = zeroObj
    ; biproducts = record
        { biproduct = biproduct
        }
    }

open Semiadditive Mat-Semiadditive using (cartesian)

Mat-CC : CartesianCategory 0ℓ c (c ⊔ ℓ)
Mat-CC = record
    { U = Mat
    ; cartesian = cartesian
    }