module Problem14 where

open import Relation.Binary.PropositionalEquality
open import Data.Product
open import Data.Sum
open import Data.Nat
open import Data.Nat.Properties
open import Data.List
open import Data.Bool

open ≡-Reasoning

postulate A : Set
postulate _⋆_ : A → A → A
postulate ι : A
infixl 50 _⋆_
postulate ⋆-assoc : ∀ a b c → a ⋆ (b ⋆ c) ≡ a ⋆ b ⋆ c
postulate ⋆-identityʳ : ∀ y → y ⋆ ι ≡ y
postulate ⋆-identityˡ : ∀ y → ι ⋆ y ≡ y

_⋆⋆_ : A → ℕ → A
n ⋆⋆ zero = ι
n ⋆⋆ suc k = n ⋆ (n ⋆⋆ k)

⋆⋆-homoʳ : (x : A) (n k : ℕ) → x ⋆⋆ (n + k) ≡ (x ⋆⋆ n) ⋆ (x ⋆⋆ k)
⋆⋆-homoʳ x zero    k = sym (⋆-identityˡ (x ⋆⋆ k))
⋆⋆-homoʳ x (suc n) k = begin
  x ⋆ (x ⋆⋆ (n + k))        ≡⟨ cong (x ⋆_) (⋆⋆-homoʳ x n k) ⟩
  x ⋆ ((x ⋆⋆ n) ⋆ (x ⋆⋆ k)) ≡⟨ ⋆-assoc x (x ⋆⋆ n) (x ⋆⋆ k) ⟩
  x ⋆  (x ⋆⋆ n) ⋆ (x ⋆⋆ k)  ∎

⋆⋆-comm : (x : A) (k : ℕ) → (x ⋆⋆ k) ⋆ x ≡ x ⋆ (x ⋆⋆ k)
⋆⋆-comm x k = begin
  (x ⋆⋆ k) ⋆ x        ≡˘⟨ cong ((x ⋆⋆ k) ⋆_) (⋆-identityʳ x) ⟩
  (x ⋆⋆ k) ⋆ (x ⋆⋆ 1) ≡˘⟨ ⋆⋆-homoʳ x k 1 ⟩
  x ⋆⋆ (k + 1)        ≡⟨  cong (x ⋆⋆_) (+-comm k 1) ⟩
  x ⋆⋆ (1 + k)        ∎

fromBits : List Bool → ℕ
fromBits [] = 0
fromBits (false ∷ bs) = 2 * fromBits bs
fromBits (true ∷ bs) = 1 + 2 * fromBits bs

expBySquare : A → A → List Bool → A
expBySquare y x [] = y
expBySquare y x (false ∷ n) = expBySquare y (x ⋆ x) n
expBySquare y x (true ∷ n)  = expBySquare (y ⋆ x) (x ⋆ x) n

expBySquare-homoˡ :
  (x y z : A) (n : List Bool) →
  expBySquare (x ⋆ y) z n ≡ x ⋆ expBySquare y z n
expBySquare-homoˡ x y z []          = refl
expBySquare-homoˡ x y z (false ∷ n) = expBySquare-homoˡ x y (z ⋆ z) n
expBySquare-homoˡ x y z (true  ∷ n) = begin
  expBySquare (x ⋆  y ⋆ z)  (z ⋆ z) n ≡˘⟨ cong (λ ◌ → expBySquare ◌ (z ⋆ z) n) (⋆-assoc x y z) ⟩
  expBySquare (x ⋆ (y ⋆ z)) (z ⋆ z) n ≡⟨  expBySquare-homoˡ x (y ⋆ z) (z ⋆ z) n ⟩
  x  ⋆ expBySquare (y ⋆ z)  (z ⋆ z) n ∎

-- Special case of commutativity
expBySquare-comm :
  (x : A) (k : ℕ) (n : List Bool) →
  expBySquare ι (x ⋆⋆ k) n ⋆ x ≡ expBySquare x (x ⋆⋆ k) n
expBySquare-comm x k []          = ⋆-identityˡ x
expBySquare-comm x k (false ∷ n) = begin
  expBySquare ι ((x ⋆⋆ k) ⋆ (x ⋆⋆ k)) n ⋆ x ≡˘⟨ cong (λ ◌ → expBySquare ι ◌ n ⋆ x) (⋆⋆-homoʳ x k k) ⟩
  expBySquare ι (x ⋆⋆ (k + k))        n ⋆ x ≡⟨  expBySquare-comm x (k + k) n ⟩
  expBySquare x (x ⋆⋆ (k + k))        n     ≡⟨  cong (λ ◌ → expBySquare x ◌ n) (⋆⋆-homoʳ x k k) ⟩
  expBySquare x ((x ⋆⋆ k) ⋆ (x ⋆⋆ k)) n     ∎
expBySquare-comm x k (true  ∷ n) = begin
  expBySquare (ι ⋆ (x ⋆⋆ k)) ((x ⋆⋆ k) ⋆ (x ⋆⋆ k)) n ⋆ x  ≡⟨  cong (λ ◌ → expBySquare ◌ ((x ⋆⋆ k) ⋆ (x ⋆⋆ k)) n ⋆ x) (⋆-identityˡ (x ⋆⋆ k)) ⟩
  expBySquare (x ⋆⋆ k)       ((x ⋆⋆ k) ⋆ (x ⋆⋆ k)) n ⋆ x  ≡˘⟨ cong₂ (λ ◌ᵃ ◌ᵇ → expBySquare ◌ᵃ ◌ᵇ n ⋆ x) (⋆-identityʳ (x ⋆⋆ k)) (⋆⋆-homoʳ x k k) ⟩
  expBySquare ((x ⋆⋆ k) ⋆ ι) (x ⋆⋆ (k + k))        n ⋆ x  ≡⟨  cong (_⋆ x) (expBySquare-homoˡ (x ⋆⋆ k) ι (x ⋆⋆ (k + k)) n) ⟩
  (x ⋆⋆ k) ⋆   expBySquare ι (x ⋆⋆ (k + k))        n ⋆ x  ≡˘⟨ ⋆-assoc (x ⋆⋆ k) (expBySquare ι (x ⋆⋆ (k + k)) n) x ⟩
  (x ⋆⋆ k) ⋆  (expBySquare ι (x ⋆⋆ (k + k))        n ⋆ x) ≡⟨  cong ((x ⋆⋆ k) ⋆_) (expBySquare-comm x (k + k) n) ⟩
  (x ⋆⋆ k) ⋆   expBySquare x (x ⋆⋆ (k + k))        n      ≡˘⟨ expBySquare-homoˡ (x ⋆⋆ k) x (x ⋆⋆ (k + k)) n ⟩
  expBySquare ((x ⋆⋆ k) ⋆ x) (x ⋆⋆ (k + k))        n      ≡⟨  cong₂ (λ ◌ᵃ ◌ᵇ → expBySquare ◌ᵃ ◌ᵇ n) (⋆⋆-comm x k) (⋆⋆-homoʳ x k k) ⟩
  expBySquare (x ⋆ (x ⋆⋆ k)) ((x ⋆⋆ k) ⋆ (x ⋆⋆ k)) n      ∎

-- Special case of homomorphism because multiplication isn't commutative
expBySquare-homoʳ :
  (x y : A) (n : List Bool) →
  expBySquare x (y ⋆ y) n ≡ expBySquare x y n ⋆ expBySquare ι y n
expBySquare-homoʳ x y []          = sym (⋆-identityʳ x)
expBySquare-homoʳ x y (false ∷ n) = expBySquare-homoʳ x (y ⋆ y) n
expBySquare-homoʳ x y (true  ∷ n) = begin
  expBySquare (x ⋆ (y ⋆ y)) (y ⋆ y ⋆ (y ⋆ y)) n                      ≡⟨ expBySquare-homoʳ (x ⋆ (y ⋆ y)) (y ⋆ y) n ⟩
  expBySquare (x ⋆ (y ⋆ y)) (y ⋆ y) n ⋆ expBySquare ι (y ⋆ y) n      ≡⟨  cong (λ ◌ → expBySquare ◌ (y ⋆ y) n ⋆ expBySquare ι (y ⋆ y) n) (⋆-assoc x y y) ⟩
  expBySquare ((x ⋆ y) ⋆ y) (y ⋆ y) n ⋆ expBySquare ι (y ⋆ y) n      ≡⟨  cong (_⋆ expBySquare ι (y ⋆ y) n) (expBySquare-homoˡ (x ⋆ y) y (y ⋆ y) n) ⟩
  (x ⋆ y) ⋆ expBySquare y (y ⋆ y) n ⋆ expBySquare ι (y ⋆ y) n        ≡˘⟨ cong (λ ◌ → x ⋆ y ⋆ expBySquare y (y ⋆ ◌) n ⋆ expBySquare ι (y ⋆ y) n) (⋆-identityʳ y) ⟩
  (x ⋆ y) ⋆ expBySquare y (y ⋆⋆ 2) n ⋆ expBySquare ι (y ⋆ y) n       ≡˘⟨ cong (λ ◌ → x ⋆ y ⋆ ◌ ⋆ expBySquare ι (y ⋆ y) n) (expBySquare-comm y 2 n) ⟩
  (x ⋆ y) ⋆ (expBySquare ι (y ⋆⋆ 2) n ⋆ y) ⋆ expBySquare ι (y ⋆ y) n ≡⟨  cong (λ ◌ → x ⋆ y ⋆ (expBySquare ι (y ⋆ ◌) n ⋆ y) ⋆ expBySquare ι (y ⋆ y) n) (⋆-identityʳ y) ⟩
  (x ⋆ y) ⋆ (expBySquare ι (y ⋆ y) n ⋆ y) ⋆ expBySquare ι (y ⋆ y) n  ≡⟨  cong (_⋆ expBySquare ι (y ⋆ y) n) (⋆-assoc (x ⋆ y) (expBySquare ι (y ⋆ y) n) y) ⟩
  (x ⋆ y) ⋆ expBySquare ι (y ⋆ y) n ⋆ y ⋆ expBySquare ι (y ⋆ y) n    ≡˘⟨ cong (λ ◌ → ◌ ⋆ y ⋆ expBySquare ι (y ⋆ y) n) (expBySquare-homoˡ (x ⋆ y) ι (y ⋆ y) n) ⟩
  expBySquare ((x ⋆ y) ⋆ ι) (y ⋆ y) n ⋆ y ⋆ expBySquare ι (y ⋆ y) n  ≡⟨  cong (λ ◌ → expBySquare ◌ (y ⋆ y) n ⋆ y ⋆ expBySquare ι (y ⋆ y) n) (⋆-identityʳ (x ⋆ y)) ⟩
  expBySquare (x ⋆ y) (y ⋆ y) n ⋆ y ⋆ expBySquare ι (y ⋆ y) n        ≡˘⟨ ⋆-assoc (expBySquare (x ⋆ y) (y ⋆ y) n) y (expBySquare ι (y ⋆ y) n) ⟩
  expBySquare (x ⋆ y) (y ⋆ y) n ⋆ (y ⋆ expBySquare ι (y ⋆ y) n)      ≡˘⟨ cong (expBySquare (x ⋆ y) (y ⋆ y) n ⋆_) (expBySquare-homoˡ y ι (y ⋆ y) n) ⟩
  expBySquare (x ⋆ y) (y ⋆ y) n ⋆ expBySquare (y ⋆ ι) (y ⋆ y) n      ≡⟨  cong (λ ◌ → expBySquare (x ⋆ y) (y ⋆ y) n ⋆ expBySquare ◌ (y ⋆ y) n) (⋆-identityʳ y) ⟩
  expBySquare (x ⋆ y) (y ⋆ y) n ⋆ expBySquare y       (y ⋆ y) n      ≡˘⟨ cong (λ ◌ → expBySquare (x ⋆ y) (y ⋆ y) n ⋆ expBySquare ◌ (y ⋆ y) n) (⋆-identityˡ y) ⟩
  expBySquare (x ⋆ y) (y ⋆ y) n ⋆ expBySquare (ι ⋆ y) (y ⋆ y) n      ∎

_⋆⋆ᵇ_ : A → List Bool → A
x ⋆⋆ᵇ n = expBySquare ι x n

⋆⋆ᵇ-homoˡ : (x : A) (n : List Bool) → (x ⋆ x) ⋆⋆ᵇ n ≡ (x ⋆⋆ᵇ n) ⋆ (x ⋆⋆ᵇ n)
⋆⋆ᵇ-homoˡ x n = expBySquare-homoʳ ι x n

proof-lemma : ∀ x n → x ⋆⋆ᵇ n ≡ x ⋆⋆ fromBits n → (x ⋆ x) ⋆⋆ᵇ n ≡ x ⋆⋆ (2 * fromBits n)
proof-lemma x n prf = begin
  (x ⋆ x) ⋆⋆ᵇ n                         ≡⟨  ⋆⋆ᵇ-homoˡ x n ⟩
  (x ⋆⋆ᵇ n)         ⋆ (x ⋆⋆ᵇ n)         ≡⟨  cong₂ _⋆_ prf prf ⟩
  (x ⋆⋆ fromBits n) ⋆ (x ⋆⋆ fromBits n) ≡˘⟨ ⋆⋆-homoʳ x (fromBits n) (fromBits n) ⟩
  x ⋆⋆ (fromBits n + fromBits n)        ≡˘⟨ cong (λ ◌ → x ⋆⋆ (fromBits n + ◌)) (+-identityʳ (fromBits n)) ⟩
  x ⋆⋆ (2 * fromBits n)                 ∎

proof : ∀ n k → n ⋆⋆ᵇ k ≡ n ⋆⋆ (fromBits k)
proof x []          = refl
proof x (false ∷ n) = proof-lemma x n (proof x n)
proof x (true  ∷ n) = begin
  expBySquare (ι ⋆ x) (x ⋆ x) n ≡⟨  cong (λ ◌ → expBySquare ◌ (x ⋆ x) n) (⋆-identityˡ x) ⟩
  expBySquare x       (x ⋆ x) n ≡˘⟨ cong (λ ◌ → expBySquare ◌ (x ⋆ x) n) (⋆-identityʳ x) ⟩
  expBySquare (x ⋆ ι) (x ⋆ x) n ≡⟨  expBySquare-homoˡ x ι (x ⋆ x) n ⟩
  x ⋆ ((x ⋆ x) ⋆⋆ᵇ n)           ≡⟨  cong (x ⋆_) (proof-lemma x n (proof x n)) ⟩
  x ⋆ (x ⋆⋆ (2 * fromBits n))   ∎