-- The algebra of relations: domains and ranges, restrictions, converses,
-- and composition
-- Part of Cumulus, the Grace17 mathematics library.
-- Copyright (c) 2026 Grant Passmore.
-- Released under the Apache License 2.0. See LICENSE in the repository root.
-- Authors: Grant Passmore

module Grace.Cumulus.Relation.Algebra
import Set.Function.Image

-- Domains and ranges

theorem domain.contains (x y r : Set)
  requires h : relates(r, x, y)
  shows x ∈ domain(r)
proof
  have member : (x, y) ∈ r := h
  apply domain.membership.backward
  constructor
  . unfold coordinateBound
    apply set.union.backward
    witness {x, y}
    constructor
    . apply set.union.backward
      witness (x, y)
      constructor
      . exact member
      . simp [orderedPair, pair.membership]
    . simp [pair.membership]
  . witness y
    exact member

theorem domain.membership_iff (x r : Set)
  shows x ∈ domain(r) ↔ ∃ y, relates(r, x, y)
proof
  constructor
  . intro h
    simp [domain.membership] at h
    obtain ⟨y, hy⟩ := h.right
    witness y
    exact hy
  . intro h
    obtain ⟨y, hy⟩ := h
    exact domain.contains with hy

theorem range.membership_iff (y r : Set)
  shows y ∈ range(r) ↔ ∃ x, relates(r, x, y)
proof
  constructor
  . intro h
    simp [range.membership] at h
    obtain ⟨x, hx⟩ := h.right
    witness x
    exact hx
  . intro h
    obtain ⟨x, hx⟩ := h
    exact range.contains at x, y, r with hx

theorem domain.subset (r A B : Set)
  requires relation : RelationBetween(r, A, B)
  shows domain(r) ⊆ A
proof
  intros x hx
  obtain ⟨y, hy⟩ := domain.membership_iff.forward with hx
  have bounds := RelationBetween.bounds with relation, hy
  exact bounds.left

theorem range.subset (r A B : Set)
  requires relation : RelationBetween(r, A, B)
  shows range(r) ⊆ B
proof
  intros y hy
  obtain ⟨x, hx⟩ := range.membership_iff.forward with hy
  have bounds := RelationBetween.bounds with relation, hx
  exact bounds.right

theorem domain.monotone (r s : Set)
  requires h : r ⊆ s
  shows domain(r) ⊆ domain(s)
proof
  intros x hx
  obtain ⟨y, hy⟩ := domain.membership_iff.forward with hx
  have inS : (x, y) ∈ s := h with (x, y), hy
  exact domain.contains at x, y, s with inS

theorem range.monotone (r s : Set)
  requires h : r ⊆ s
  shows range(r) ⊆ range(s)
proof
  intros y hy
  obtain ⟨x, hx⟩ := range.membership_iff.forward with hy
  have inS : (x, y) ∈ s := h with (x, y), hx
  exact range.contains with inS

theorem domain.empty
  shows domain(∅) = ∅
proof
  apply Subset.empty_iff.forward
  intros x hx
  obtain ⟨y, hy⟩ := domain.membership_iff.forward with hx
  have member : (x, y) ∈ ∅ := hy
  exfalso
  exact set.empty with member

theorem domain.union2 (r s : Set)
  shows domain(r ∪ s) = domain(r) ∪ domain(s)
proof
  ext x
  constructor
  . intro hx
    obtain ⟨y, hy⟩ := domain.membership_iff.forward with hx
    have member : (x, y) ∈ r ∪ s := hy
    simp [union2.membership] at member
    simp [union2.membership]
    cases member with hr | hs
    . left
      exact domain.contains at x, y, r with hr
    . right
      exact domain.contains at x, y, s with hs
  . intro hx
    simp [union2.membership] at hx
    have leftSubset := union2.subset_left at r, s
    have rightSubset := union2.subset_right at r, s
    cases hx with hr | hs
    . exact domain.monotone with leftSubset, x, hr
    . exact domain.monotone with rightSubset, x, hs

theorem range.union2 (r s : Set)
  shows range(r ∪ s) = range(r) ∪ range(s)
proof
  ext y
  constructor
  . intro hy
    obtain ⟨x, hx⟩ := range.membership_iff.forward with hy
    have member : (x, y) ∈ r ∪ s := hx
    simp [union2.membership] at member
    simp [union2.membership]
    cases member with hr | hs
    . left
      exact range.contains with hr
    . right
      exact range.contains with hs
  . intro hy
    simp [union2.membership] at hy
    have leftSubset := union2.subset_left at r, s
    have rightSubset := union2.subset_right at r, s
    cases hy with hr | hs
    . exact range.monotone with leftSubset, y, hr
    . exact range.monotone with rightSubset, y, hs

theorem RelationBetween.domain_range (r A B : Set)
  requires relation : RelationBetween(r, A, B)
  shows r ⊆ domain(r) × range(r)
proof
  have included : r ⊆ A × B := relation
  intros p hp
  have shape : p ∈ A × B := included with p, hp
  simp [product.membership] at shape
  obtain ⟨bound, x, hx, y, hy, equation⟩ := shape
  rw equation at hp
  rw equation
  have inDomain := domain.contains at x, y, r with hp
  have inRange := range.contains with hp
  exact product.pair with inDomain, inRange

theorem FunctionBetween.domain (f A B : Set)
  requires function : f : A → B
  shows domain(f) = A
proof
  apply Subset.antisymmetric
  . exact domain.subset with function.left
  . intros x hx
    have related := FunctionBetween.relates_apply with function, hx
    exact domain.contains with related

theorem FunctionBetween.range (f A B : Set)
  requires function : f : A → B
  shows range(f) = functionImage(f, A, B)
proof
  ext y
  constructor
  . intro hy
    obtain ⟨x, hx⟩ := range.membership_iff.forward with hy
    have bounds := RelationBetween.bounds with function.left, hx
    have value := FunctionBetween.apply_value with function, bounds.left, hx
    simp [functionImage.membership]
    constructor
    . exact bounds.right
    . witness x
      constructor
      . exact bounds.left
      . symm
        exact value
  . intro hy
    simp [functionImage.membership] at hy
    obtain ⟨inB, x, hx, shape⟩ := hy
    have related := FunctionBetween.relates_apply with function, hx
    rw shape
    exact range.contains at x, f(x), f with related

-- Restrictions

theorem restrict.subset (r A : Set)
  shows restrict(r, A) ⊆ r
proof
  intros p hp
  simp [restrict.membership] at hp
  exact hp.left

theorem restrict.relation (r A B S : Set)
  requires relation : RelationBetween(r, A, B)
  shows RelationBetween(restrict(r, S), S ∩ A, B)
proof
  intros p hp
  simp [restrict.membership] at hp
  obtain ⟨hr, x, hx, y, equation⟩ := hp
  rw equation at hr
  rw equation
  have bounds := RelationBetween.bounds at r, A, B, x, y with relation, hr
  apply product.pair
  . simp [intersection.membership]
  . exact bounds.right

theorem restrict.whole (r A B : Set)
  requires relation : RelationBetween(r, A, B)
  shows restrict(r, A) = r
proof
  have included : r ⊆ A × B := relation
  apply Subset.antisymmetric
  . exact restrict.subset at r, A
  . intros p hp
    have shape : p ∈ A × B := included with p, hp
    simp [product.membership] at shape
    obtain ⟨bound, x, hx, y, hy, equation⟩ := shape
    simp [restrict.membership]
    constructor
    . exact hp
    . witness x
      constructor
      . exact hx
      . witness y
        exact equation

theorem restrict.empty (r : Set)
  shows restrict(r, ∅) = ∅
proof
  apply Subset.empty_iff.forward
  intros p hp
  simp [restrict.membership] at hp

theorem restrict.intersection (r S T : Set)
  shows restrict(restrict(r, S), T) = restrict(r, S ∩ T)
proof
  ext p
  constructor
  . intro hp
    simp [restrict.membership] at hp
    obtain ⟨inner, x, hxT, y, equation⟩ := hp
    simp [restrict.membership] at inner
    obtain ⟨hr, z, hzS, w, equationZ⟩ := inner
    rw equation at equationZ
    have coordinates := orderedPair.injective with equationZ
    simp [restrict.membership]
    constructor
    . exact hr
    . witness x
      constructor
      . simp [intersection.membership]
      . witness y
        exact equation
  . intro hp
    simp [restrict.membership] at hp
    obtain ⟨hr, x, hx, y, equation⟩ := hp
    simp [intersection.membership] at hx
    simp [restrict.membership]
    constructor
    . constructor
      . exact hr
      . exact ⟨x, hx.left, y, equation⟩
    . exact ⟨x, hx.right, y, equation⟩

theorem restrict.monotone (r S T : Set)
  requires h : S ⊆ T
  shows restrict(r, S) ⊆ restrict(r, T)
proof
  intros p hp
  simp [restrict.membership] at hp
  obtain ⟨hr, x, hx, y, equation⟩ := hp
  simp [restrict.membership]
  constructor
  . exact hr
  . exact ⟨x, (h with x, hx), y, equation⟩

theorem restrict.union2 (r S T : Set)
  shows restrict(r, S ∪ T) = restrict(r, S) ∪ restrict(r, T)
proof
  ext p
  constructor
  . intro hp
    simp [restrict.membership] at hp
    obtain ⟨hr, x, hx, y, equation⟩ := hp
    simp [union2.membership] at hx
    simp [union2.membership]
    cases hx with hS | hT
    . left
      simp [restrict.membership]
      exact ⟨hr, x, hS, y, equation⟩
    . right
      simp [restrict.membership]
      exact ⟨hr, x, hT, y, equation⟩
  . intro hp
    simp [union2.membership] at hp
    have leftSubset := union2.subset_left at S, T
    have rightSubset := union2.subset_right at S, T
    cases hp with hS | hT
    . exact restrict.monotone at r, S, S ∪ T with leftSubset, p, hS
    . exact restrict.monotone at r, T, S ∪ T with rightSubset, p, hT

theorem restrict.domain (r S : Set)
  shows domain(restrict(r, S)) = domain(r) ∩ S
proof
  ext x
  constructor
  . intro hx
    obtain ⟨y, hy⟩ := domain.membership_iff.forward with hx
    have related := restrict.relates.forward with hy
    simp [intersection.membership]
    constructor
    . exact domain.contains with related.right
    . exact related.left
  . intro hx
    simp [intersection.membership] at hx
    obtain ⟨y, hy⟩ := domain.membership_iff.forward with hx.left
    have related := restrict.relates.backward at r, S, x, y with ⟨hx.right, hy⟩
    exact domain.contains with related

-- Converses

theorem converse.converse (r A B : Set)
  requires relation : RelationBetween(r, A, B)
  shows converse(converse(r, A, B), B, A) = r
proof
  ext p
  constructor
  . intro hp
    simp [converse.membership] at hp
    obtain ⟨bound, y, hy, x, hx, equation, related⟩ := hp
    have member : (y, x) ∈ converse(r, A, B) := related
    simp [converse.membership] at member
    obtain ⟨boundInner, x2, hx2, y2, hy2, equationInner, relatedInner⟩ := member
    have coordinates := orderedPair.injective with equationInner
    rw equation
    rwa [← coordinates.left, ← coordinates.right] at relatedInner
  . intro hp
    have included : r ⊆ A × B := relation
    have shape : p ∈ A × B := included with p, hp
    simp [product.membership] at shape
    obtain ⟨bound, x, hx, y, hy, equation⟩ := shape
    rw equation at hp
    rw equation
    have once := converse.relates at r, A, B, x, y with hx, hy, hp
    exact converse.relates with hy, hx, once

theorem converse.monotone (r s A B : Set)
  requires h : r ⊆ s
  shows converse(r, A, B) ⊆ converse(s, A, B)
proof
  intros p hp
  simp [converse.membership] at hp
  obtain ⟨bound, x, hx, y, hy, equation, related⟩ := hp
  rw equation
  have inS : (x, y) ∈ s := h with (x, y), related
  exact converse.relates at s, A, B, x, y with hx, hy, inS

theorem converse.domain (r A B : Set)
  requires relation : RelationBetween(r, A, B)
  shows domain(converse(r, A, B)) = range(r)
proof
  ext y
  constructor
  . intro hy
    obtain ⟨x, hx⟩ := domain.membership_iff.forward with hy
    have member : (y, x) ∈ converse(r, A, B) := hx
    simp [converse.membership] at member
    obtain ⟨bound, x2, hx2, y2, hy2, equation, related⟩ := member
    have coordinates := orderedPair.injective with equation
    rw [coordinates.left]
    exact range.contains at x2, y2, r with related
  . intro hy
    obtain ⟨x, hx⟩ := range.membership_iff.forward with hy
    have bounds := RelationBetween.bounds with relation, hx
    have flipped := converse.relates with bounds.left, bounds.right, hx
    exact domain.contains with flipped

theorem converse.range (r A B : Set)
  requires relation : RelationBetween(r, A, B)
  shows range(converse(r, A, B)) = domain(r)
proof
  ext x
  constructor
  . intro hx
    obtain ⟨y, hy⟩ := range.membership_iff.forward with hx
    have member : (y, x) ∈ converse(r, A, B) := hy
    simp [converse.membership] at member
    obtain ⟨bound, x2, hx2, y2, hy2, equation, related⟩ := member
    have coordinates := orderedPair.injective with equation
    rw [coordinates.right]
    exact domain.contains with related
  . intro hx
    obtain ⟨y, hy⟩ := domain.membership_iff.forward with hx
    have bounds := RelationBetween.bounds with relation, hy
    have flipped := converse.relates with bounds.left, bounds.right, hy
    exact range.contains at y, x, converse(r, A, B) with flipped

-- Composition

theorem compose.associative (h g f A B C D : Set)
  requires fFunction : f : A → B
  requires gFunction : g : B → C
  requires hFunction : h : C → D
  shows compose(h, compose(g, f, A, B, C), A, C, D) = compose(compose(h, g, B, C, D), f, A, B, D)
proof
  have inner := compose.function with fFunction, gFunction
  have outer := compose.function with gFunction, hFunction
  apply FunctionBetween.ext with (compose.function with inner, hFunction), (compose.function with fFunction, outer)
  intros x hx
  have valueB := FunctionBetween.apply_closed with fFunction, hx
  have leftStep := compose.apply with inner, hFunction, hx
  have innerStep := compose.apply with fFunction, gFunction, hx
  have rightStep := compose.apply with fFunction, outer, hx
  have outerStep := compose.apply with gFunction, hFunction, valueB
  rw [leftStep, innerStep, rightStep, outerStep]
  rfl

theorem compose.restrict (g f A B C S : Set)
  requires fFunction : f : A → B
  requires gFunction : g : B → C
  requires hS : S ⊆ A
  shows restrict(compose(g, f, A, B, C), S) = compose(g, restrict(f, S), S, B, C)
proof
  have composed := compose.function with fFunction, gFunction
  have closed : ∀ x ∈ S, f(x) ∈ B
    intros x hx
    exact FunctionBetween.apply_closed with fFunction, (hS with x, hx)
  have restricted := restrict.function with fFunction, hS, closed
  have closedComposed : ∀ x ∈ S, compose(g, f, A, B, C)(x) ∈ C
    intros x hx
    exact FunctionBetween.apply_closed with composed, (hS with x, hx)
  have restrictedComposed := restrict.function with composed, hS, closedComposed
  have composedRestricted := compose.function with restricted, gFunction
  apply FunctionBetween.ext with restrictedComposed, composedRestricted
  intros x hx
  have hxA : x ∈ A := hS with x, hx
  have leftStep := restrict.apply with composed, hS, closedComposed, hx
  have innerStep := compose.apply with fFunction, gFunction, hxA
  have rightStep := compose.apply with restricted, gFunction, hx
  have restrictedStep := restrict.apply with fFunction, hS, closed, hx
  rw [leftStep, innerStep, rightStep, restrictedStep]
  rfl