-- Images and preimages of functions and their algebra
-- 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.Set.Function.Image
import Function
import Algebra

-- The preimage of Y under a function with domain A.
def functionPreimage (f A Y : Set) : Set := { x ∈ A | f(x) ∈ Y }

theorem functionPreimage.subset (f A Y : Set)
  shows functionPreimage(f, A, Y) ⊆ A
proof
  intros x hx
  simp [functionPreimage.membership] at hx
  exact hx.left

theorem functionPreimage.of_apply (f A Y x : Set)
  requires hx : x ∈ A
  requires value : f(x) ∈ Y
  shows x ∈ functionPreimage(f, A, Y)
proof
  simp [functionPreimage.membership]

-- Images

theorem functionImage.empty (f B : Set)
  shows functionImage(f, ∅, B) = ∅
proof
  ext y
  constructor
  . intro hy
    simp [functionImage.membership] at hy
  . intro hy
    exfalso
    simpa using hy

theorem functionImage.monotone (f A B S T : Set)
  requires h : S ⊆ T
  shows functionImage(f, S, B) ⊆ functionImage(f, T, B)
proof
  intros y hy
  simp [functionImage.membership] at hy
  obtain ⟨inB, x, hx, shape⟩ := hy
  simp [functionImage.membership]
  constructor
  . exact inB
  . refine ⟨x, (h with x, hx), shape⟩

theorem functionImage.union2 (f A B S T : Set)
  shows functionImage(f, S ∪ T, B) = functionImage(f, S, B) ∪ functionImage(f, T, B)
proof
  ext y
  constructor
  . intro hy
    simp [functionImage.membership] at hy
    obtain ⟨inB, x, hx, shape⟩ := hy
    simp [union2.membership] at hx
    simp [union2.membership]
    cases hx with hS | hT
    . left
      simp [functionImage.membership]
      exact ⟨inB, x, hS, shape⟩
    . right
      simp [functionImage.membership]
      exact ⟨inB, x, hT, shape⟩
  . intro hy
    simp [union2.membership] at hy
    cases hy with hS | hT
    . exact functionImage.monotone at f, A, B, S, S ∪ T with (union2.subset_left at S, T), y, hS
    . exact functionImage.monotone at f, A, B, T, S ∪ T with (union2.subset_right at S, T), y, hT

theorem functionImage.intersection_subset (f A B S T : Set)
  shows functionImage(f, S ∩ T, B) ⊆ functionImage(f, S, B) ∩ functionImage(f, T, B)
proof
  apply Subset.intersection_iff.backward
  constructor
  . exact functionImage.monotone at f, A, B, S ∩ T, S with (intersection.subset_left at S, T)
  . exact functionImage.monotone at f, A, B, S ∩ T, T with (intersection.subset_right at S, T)

theorem Injective.apply_cancel (f A B x y : Set)
  requires injective : InjectiveBetween(f, A, B)
  requires hx : x ∈ A
  requires hy : y ∈ A
  requires same : f(x) = f(y)
  shows x = y
proof
  have function := injective.left
  have valueB := FunctionBetween.apply_closed with function, hx
  have hxz := FunctionBetween.relates_apply with function, hx
  have hyz := FunctionBetween.relates_apply with function, hy
  rw ← same at hyz
  exact Injective.cancel with injective, hx, hy, valueB, hxz, hyz

theorem functionImage.intersection_of_injective (f A B S T : Set)
  requires injective : InjectiveBetween(f, A, B)
  requires hS : S ⊆ A
  requires hT : T ⊆ A
  shows functionImage(f, S ∩ T, B) = functionImage(f, S, B) ∩ functionImage(f, T, B)
proof
  apply Subset.antisymmetric
  . exact functionImage.intersection_subset at f, A, B, S, T
  . intros y hy
    simp [intersection.membership] at hy
    obtain ⟨inS, inT⟩ := hy
    simp [functionImage.membership] at inS
    simp [functionImage.membership] at inT
    obtain ⟨inB, x, hx, shapeX⟩ := inS
    obtain ⟨inB2, z, hz, shapeZ⟩ := inT
    have hxA : x ∈ A := hS with x, hx
    have hzA : z ∈ A := hT with z, hz
    have agree : f(x) = f(z)
      rw [← shapeX, ← shapeZ]
      rfl
    have same := Injective.apply_cancel with injective, hxA, hzA, agree
    simp [functionImage.membership]
    constructor
    . exact inB
    . witness x
      constructor
      . simp [intersection.membership]
      . exact shapeX

theorem functionImage.difference_subset (f A B S T : Set)
  shows functionImage(f, S, B) ∖ functionImage(f, T, B) ⊆ functionImage(f, S ∖ T, B)
proof
  intros y hy
  simp [difference.membership] at hy
  obtain ⟨inS, notT⟩ := hy
  simp [functionImage.membership] at inS
  obtain ⟨inB, x, hx, shape⟩ := inS
  simp [functionImage.membership]
  constructor
  . exact inB
  . witness x
    constructor
    . simp [difference.membership]
      constructor
      . exact hx
      . intro inT
        apply notT
        simp [functionImage.membership]
        exact ⟨inB, x, inT, shape⟩
    . exact shape

theorem functionImage.subset_iff (f A B S Y : Set)
  shows functionImage(f, S, B) ⊆ Y ↔ ∀ x ∈ S, f(x) ∈ B → f(x) ∈ Y
proof
  constructor
  . intros h x hx valueB
    have member : f(x) ∈ functionImage(f, S, B) by simp [functionImage.membership]
    exact h with f(x), member
  . intros h y hy
    simp [functionImage.membership] at hy
    obtain ⟨inB, x, hx, shape⟩ := hy
    rw shape
    rw shape at inB
    exact h with x, hx, inB

-- Preimages

theorem functionPreimage.empty (f A : Set)
  shows functionPreimage(f, A, ∅) = ∅
proof
  ext x
  simp [functionPreimage.membership]

theorem functionPreimage.codomain (f A B : Set)
  requires function : f : A → B
  shows functionPreimage(f, A, B) = A
proof
  ext x
  constructor
  . intro hx
    exact functionPreimage.subset at f, A, B with x, hx
  . intro hx
    exact functionPreimage.of_apply with hx, (FunctionBetween.apply_closed with function, hx)

theorem functionPreimage.monotone (f A Y Z : Set)
  requires h : Y ⊆ Z
  shows functionPreimage(f, A, Y) ⊆ functionPreimage(f, A, Z)
proof
  intros x hx
  simp [functionPreimage.membership] at hx
  simp [functionPreimage.membership]

theorem functionPreimage.union2 (f A Y Z : Set)
  shows functionPreimage(f, A, Y ∪ Z) = functionPreimage(f, A, Y) ∪ functionPreimage(f, A, Z)
proof
  ext x
  simp [functionPreimage.membership, union2.membership]

theorem functionPreimage.intersection (f A Y Z : Set)
  shows functionPreimage(f, A, Y ∩ Z) = functionPreimage(f, A, Y) ∩ functionPreimage(f, A, Z)
proof
  ext x
  simp [functionPreimage.membership, intersection.membership]

theorem functionPreimage.difference (f A Y Z : Set)
  shows functionPreimage(f, A, Y ∖ Z) = functionPreimage(f, A, Y) ∖ functionPreimage(f, A, Z)
proof
  ext x
  simp [functionPreimage.membership, difference.membership]

theorem functionPreimage.complement (f A B Y : Set)
  requires function : f : A → B
  shows functionPreimage(f, A, B ∖ Y) = A ∖ functionPreimage(f, A, Y)
proof
  have whole := functionPreimage.codomain with function
  rw [functionPreimage.difference, whole]
  rfl

-- Images and preimages together

theorem functionImage.preimage_subset (f A B Y : Set)
  shows functionImage(f, functionPreimage(f, A, Y), B) ⊆ Y
proof
  intros y hy
  simp [functionImage.membership] at hy
  obtain ⟨inB, x, hx, shape⟩ := hy
  simp [functionPreimage.membership] at hx
  rw shape
  exact hx.right

theorem functionImage.preimage_of_subset (f A B Y : Set)
  requires function : f : A → B
  requires h : Y ⊆ functionImage(f, A, B)
  shows functionImage(f, functionPreimage(f, A, Y), B) = Y
proof
  apply Subset.antisymmetric
  . exact functionImage.preimage_subset at f, A, B, Y
  . intros y hy
    have inImage : y ∈ functionImage(f, A, B) := h with y, hy
    simp [functionImage.membership] at inImage
    obtain ⟨inB, x, hx, shape⟩ := inImage
    simp [functionImage.membership]
    constructor
    . exact inB
    . witness x
      constructor
      . simp [functionPreimage.membership]
      . exact shape

theorem functionPreimage.image_superset (f A B S : Set)
  requires function : f : A → B
  requires hS : S ⊆ A
  shows S ⊆ functionPreimage(f, A, functionImage(f, S, B))
proof
  intros x hx
  simp [functionPreimage.membership]
  constructor
  . exact hS with x, hx
  . have value := FunctionBetween.apply_closed with function, (hS with x, hx)
    simp [functionImage.membership]

theorem functionPreimage.image_of_injective (f A B S : Set)
  requires injective : InjectiveBetween(f, A, B)
  requires hS : S ⊆ A
  shows functionPreimage(f, A, functionImage(f, S, B)) = S
proof
  have function := injective.left
  apply Subset.antisymmetric
  . intros x hx
    simp [functionPreimage.membership] at hx
    obtain ⟨hxA, inImage⟩ := hx
    simp [functionImage.membership] at inImage
    obtain ⟨inB, z, hz, shape⟩ := inImage
    have same := Injective.apply_cancel with injective, hxA, (hS with z, hz), shape
    rwa same
  . exact functionPreimage.image_superset with function, hS

theorem functionImage.compose (g f A B C S : Set)
  requires fFunction : f : A → B
  requires gFunction : g : B → C
  requires hS : S ⊆ A
  shows functionImage(compose(g, f, A, B, C), S, C) = functionImage(g, functionImage(f, S, B), C)
proof
  have composed := compose.function with fFunction, gFunction
  ext y
  constructor
  . intro hy
    simp [functionImage.membership] at hy
    obtain ⟨inC, x, hx, shape⟩ := hy
    have hxA : x ∈ A := hS with x, hx
    have step := compose.apply with fFunction, gFunction, hxA
    simp [functionImage.membership]
    constructor
    . exact inC
    . witness f(x)
      constructor
      . have value := FunctionBetween.apply_closed with fFunction, hxA
        simp [functionImage.membership]
        exact ⟨value, x, hx, rfl⟩
      . rw [shape, step]
        rfl
  . intro hy
    simp [functionImage.membership] at hy
    obtain ⟨inC, z, hz, shape⟩ := hy
    simp [functionImage.membership] at hz
    obtain ⟨inB, x, hx, shapeZ⟩ := hz
    have hxA : x ∈ A := hS with x, hx
    have step := compose.apply with fFunction, gFunction, hxA
    simp [functionImage.membership]
    constructor
    . exact inC
    . witness x
      constructor
      . exact hx
      . rw [shape, shapeZ, step]
        rfl

theorem functionPreimage.compose (g f A B C Z : Set)
  requires fFunction : f : A → B
  requires gFunction : g : B → C
  shows functionPreimage(compose(g, f, A, B, C), A, Z) = functionPreimage(f, A, functionPreimage(g, B, Z))
proof
  ext x
  constructor
  . intro hx
    simp [functionPreimage.membership] at hx
    obtain ⟨hxA, value⟩ := hx
    have step := compose.apply with fFunction, gFunction, hxA
    rw step at value
    have inB := FunctionBetween.apply_closed with fFunction, hxA
    simp [functionPreimage.membership]
  . intro hx
    simp [functionPreimage.membership] at hx
    obtain ⟨hxA, inner⟩ := hx
    simp [functionPreimage.membership] at inner
    have step := compose.apply with fFunction, gFunction, hxA
    simp [functionPreimage.membership]

theorem functionImage.identity_of_subset (A S : Set)
  requires hS : S ⊆ A
  shows functionImage(identity(A), S, A) = S
proof
  ext y
  constructor
  . intro hy
    simp [functionImage.membership] at hy
    obtain ⟨inA, x, hx, shape⟩ := hy
    have fixed := identity.apply with (hS with x, hx)
    rwa [shape, fixed]
  . intro hy
    simp [functionImage.membership]
    constructor
    . exact hS with y, hy
    . witness y
      constructor
      . exact hy
      . symm
        exact identity.apply with (hS with y, hy)

-- An injection is a bijection onto its image.
theorem Injective.corestrict (f A B : Set)
  requires injective : InjectiveBetween(f, A, B)
  shows InjectiveBetween(f, A, functionImage(f, A, B))
proof
  have function := Injective.function with injective
  constructor
  . exact functionImage.corestrict with function
  . intros x hx y hy z hz hxz hyz
    have hzB := functionImage.subset at f, A, B with z, hz
    exact injective.right with x, hx, y, hy, z, hzB, hxz, hyz

-- Restricting an injection keeps it injective, with the same image.
theorem restrict.injective (f A B S : Set)
  requires injective : InjectiveBetween(f, A, B)
  requires subset : S ⊆ A
  shows InjectiveBetween(restrict(f, S), S, B)
proof
  have function := Injective.function with injective
  have closed : ∀ x ∈ S, f(x) ∈ B
    intros x hx
    exact FunctionBetween.apply_closed with function, (subset with x, hx)
  have restricted := restrict.function with function, subset, closed
  constructor
  . exact restricted
  . intros x hx y hy z hz hxz hyz
    have valueX := FunctionBetween.apply_value with restricted, hx, hxz
    have valueY := FunctionBetween.apply_value with restricted, hy, hyz
    have stepX := restrict.apply with function, subset, closed, hx
    have stepY := restrict.apply with function, subset, closed, hy
    have same : f(x) = f(y) by cc
    exact Injective.apply_cancel with injective, (subset with x, hx), (subset with y, hy), same

theorem functionImage.restrict (f A B S : Set)
  requires function : f : A → B
  requires subset : S ⊆ A
  shows functionImage(restrict(f, S), S, B) = functionImage(f, S, B)
proof
  have closed : ∀ x ∈ S, f(x) ∈ B
    intros x hx
    exact FunctionBetween.apply_closed with function, (subset with x, hx)
  ext y
  constructor
  . intro hy
    simp [functionImage.membership] at hy
    obtain ⟨inB, x, hx, shape⟩ := hy
    have step := restrict.apply with function, subset, closed, hx
    rw step at shape
    simp [functionImage.membership]
    constructor
    . exact inB
    . refine ⟨x, hx, shape⟩
  . intro hy
    simp [functionImage.membership] at hy
    obtain ⟨inB, x, hx, shape⟩ := hy
    have step := restrict.apply with function, subset, closed, hx
    rw ← step at shape
    simp [functionImage.membership]
    constructor
    . exact inB
    . refine ⟨x, hx, shape⟩

theorem functionImage.insert (f A B x S : Set)
  requires function : f : A → B
  requires hx : x ∈ A
  shows functionImage(f, insert(x, S), B) = insert(f(x), functionImage(f, S, B))
proof
  have value := FunctionBetween.apply_closed with function, hx
  ext y
  constructor
  . intro hy
    simp [functionImage.membership] at hy
    obtain ⟨inB, z, hz, shape⟩ := hy
    have split := insert.membership.forward with hz
    apply insert.membership.backward
    cases split with same | inS
    . left
      rw [shape, same]
      rfl
    . right
      simp [functionImage.membership]
      constructor
      . exact inB
      . refine ⟨z, inS, shape⟩
  . intro hy
    have split := insert.membership.forward with hy
    cases split with same | inImage
    . have inserted : x ∈ insert(x, S)
        apply insert.membership.backward
        left
        rfl
      rw same
      simp [functionImage.membership]
      constructor
      . exact value
      . witness x
        constructor
        . exact inserted
        . rfl
    . simp [functionImage.membership] at inImage
      obtain ⟨inB, z, hz, shape⟩ := inImage
      have inserted : z ∈ insert(x, S)
        apply insert.membership.backward
        right
        exact hz
      simp [functionImage.membership]
      constructor
      . exact inB
      . refine ⟨z, inserted, shape⟩