/-
Copyright (c) 2020 Yury G. Kudryashov. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Author: Yury G. Kudryashov
-/
import Mathlib.PrePort
import Mathlib.Lean3Lib.init.default
import Mathlib.linear_algebra.affine_space.midpoint
import Mathlib.algebra.module.ordered
import Mathlib.tactic.field_simp
import Mathlib.PostPort

universes u_1 u_2 u_3

namespace Mathlib

/-!
# Ordered modules as affine spaces

In this file we define the slope of a function `f : k → PE` taking values in an affine space over
`k` and prove some theorems about `slope` and `line_map` in the case when `PE` is an ordered
semimodule over `k`. The `slope` function naturally appears in the Mean Value Theorem, and in the
proof of the fact that a function with nonnegative second derivative on an interval is convex on
this interval. In the third part of this file we prove inequalities that will be used in
`analysis.convex.basic` to link convexity of a function on an interval to monotonicity of the slope,
see section docstring below for details.

## Implementation notes

We do not introduce the notion of ordered affine spaces (yet?). Instead, we prove various theorems
for an ordered semimodule interpreted as an affine space.

## Tags

affine space, ordered semimodule, slope
-/

/-!
### Definition of `slope` and basic properties

In this section we define `slope f a b` and prove some properties that do not require order on the
codomain.  -/

/-- `slope f a b = (b - a)⁻¹ • (f b -ᵥ f a)` is the slope of a function `f` on the interval
`[a, b]`. Note that `slope f a a = 0`, not the derivative of `f` at `a`. -/
def slope {k : Type u_1} {E : Type u_2} {PE : Type u_3} [field k] [add_comm_group E] [semimodule k E] [add_torsor E PE] (f : k → PE) (a : k) (b : k) : E :=
  b - a⁻¹ • (f b -ᵥ f a)

theorem slope_def_field {k : Type u_1} [field k] (f : k → k) (a : k) (b : k) : slope f a b = (f b - f a) / (b - a) :=
  Eq.symm div_eq_inv_mul

@[simp] theorem slope_same {k : Type u_1} {E : Type u_2} {PE : Type u_3} [field k] [add_comm_group E] [semimodule k E] [add_torsor E PE] (f : k → PE) (a : k) : slope f a a = 0 := sorry

theorem eq_of_slope_eq_zero {k : Type u_1} {E : Type u_2} {PE : Type u_3} [field k] [add_comm_group E] [semimodule k E] [add_torsor E PE] {f : k → PE} {a : k} {b : k} (h : slope f a b = 0) : f a = f b := sorry

theorem slope_comm {k : Type u_1} {E : Type u_2} {PE : Type u_3} [field k] [add_comm_group E] [semimodule k E] [add_torsor E PE] (f : k → PE) (a : k) (b : k) : slope f a b = slope f b a := sorry

/-- `slope f a c` is a linear combination of `slope f a b` and `slope f b c`. This version
explicitly provides coefficients. If `a ≠ c`, then the sum of the coefficients is `1`, so it is
actually an affine combination, see `line_map_slope_slope_sub_div_sub`. -/
theorem sub_div_sub_smul_slope_add_sub_div_sub_smul_slope {k : Type u_1} {E : Type u_2} {PE : Type u_3} [field k] [add_comm_group E] [semimodule k E] [add_torsor E PE] (f : k → PE) (a : k) (b : k) (c : k) : ((b - a) / (c - a)) • slope f a b + ((c - b) / (c - a)) • slope f b c = slope f a c := sorry

/-- `slope f a c` is an affine combination of `slope f a b` and `slope f b c`. This version uses
`line_map` to express this property. -/
theorem line_map_slope_slope_sub_div_sub {k : Type u_1} {E : Type u_2} {PE : Type u_3} [field k] [add_comm_group E] [semimodule k E] [add_torsor E PE] (f : k → PE) (a : k) (b : k) (c : k) (h : a ≠ c) : coe_fn (affine_map.line_map (slope f a b) (slope f b c)) ((c - b) / (c - a)) = slope f a c := sorry

/-- `slope f a b` is an affine combination of `slope f a (line_map a b r)` and
`slope f (line_map a b r) b`. We use `line_map` to express this property. -/
theorem line_map_slope_line_map_slope_line_map {k : Type u_1} {E : Type u_2} {PE : Type u_3} [field k] [add_comm_group E] [semimodule k E] [add_torsor E PE] (f : k → PE) (a : k) (b : k) (r : k) : coe_fn
    (affine_map.line_map (slope f (coe_fn (affine_map.line_map a b) r) b)
      (slope f a (coe_fn (affine_map.line_map a b) r)))
    r =
  slope f a b := sorry

/-!
### Monotonicity of `line_map`

In this section we prove that `line_map a b r` is monotone (strictly or not) in its arguments if
other arguments belong to specific domains.
-/

theorem line_map_mono_left {k : Type u_1} {E : Type u_2} [ordered_ring k] [ordered_add_comm_group E] [semimodule k E] [ordered_semimodule k E] {a : E} {a' : E} {b : E} {r : k} (ha : a ≤ a') (hr : r ≤ 1) : coe_fn (affine_map.line_map a b) r ≤ coe_fn (affine_map.line_map a' b) r := sorry

theorem line_map_strict_mono_left {k : Type u_1} {E : Type u_2} [ordered_ring k] [ordered_add_comm_group E] [semimodule k E] [ordered_semimodule k E] {a : E} {a' : E} {b : E} {r : k} (ha : a < a') (hr : r < 1) : coe_fn (affine_map.line_map a b) r < coe_fn (affine_map.line_map a' b) r := sorry

theorem line_map_mono_right {k : Type u_1} {E : Type u_2} [ordered_ring k] [ordered_add_comm_group E] [semimodule k E] [ordered_semimodule k E] {a : E} {b : E} {b' : E} {r : k} (hb : b ≤ b') (hr : 0 ≤ r) : coe_fn (affine_map.line_map a b) r ≤ coe_fn (affine_map.line_map a b') r := sorry

theorem line_map_strict_mono_right {k : Type u_1} {E : Type u_2} [ordered_ring k] [ordered_add_comm_group E] [semimodule k E] [ordered_semimodule k E] {a : E} {b : E} {b' : E} {r : k} (hb : b < b') (hr : 0 < r) : coe_fn (affine_map.line_map a b) r < coe_fn (affine_map.line_map a b') r := sorry

theorem line_map_mono_endpoints {k : Type u_1} {E : Type u_2} [ordered_ring k] [ordered_add_comm_group E] [semimodule k E] [ordered_semimodule k E] {a : E} {a' : E} {b : E} {b' : E} {r : k} (ha : a ≤ a') (hb : b ≤ b') (h₀ : 0 ≤ r) (h₁ : r ≤ 1) : coe_fn (affine_map.line_map a b) r ≤ coe_fn (affine_map.line_map a' b') r :=
  has_le.le.trans (line_map_mono_left ha h₁) (line_map_mono_right hb h₀)

theorem line_map_strict_mono_endpoints {k : Type u_1} {E : Type u_2} [ordered_ring k] [ordered_add_comm_group E] [semimodule k E] [ordered_semimodule k E] {a : E} {a' : E} {b : E} {b' : E} {r : k} (ha : a < a') (hb : b < b') (h₀ : 0 ≤ r) (h₁ : r ≤ 1) : coe_fn (affine_map.line_map a b) r < coe_fn (affine_map.line_map a' b') r := sorry

theorem line_map_lt_line_map_iff_of_lt {k : Type u_1} {E : Type u_2} [ordered_ring k] [ordered_add_comm_group E] [semimodule k E] [ordered_semimodule k E] {a : E} {b : E} {r : k} {r' : k} (h : r < r') : coe_fn (affine_map.line_map a b) r < coe_fn (affine_map.line_map a b) r' ↔ a < b := sorry

theorem left_lt_line_map_iff_lt {k : Type u_1} {E : Type u_2} [ordered_ring k] [ordered_add_comm_group E] [semimodule k E] [ordered_semimodule k E] {a : E} {b : E} {r : k} (h : 0 < r) : a < coe_fn (affine_map.line_map a b) r ↔ a < b := sorry

theorem line_map_lt_left_iff_lt {k : Type u_1} {E : Type u_2} [ordered_ring k] [ordered_add_comm_group E] [semimodule k E] [ordered_semimodule k E] {a : E} {b : E} {r : k} (h : 0 < r) : coe_fn (affine_map.line_map a b) r < a ↔ b < a :=
  left_lt_line_map_iff_lt h

theorem line_map_lt_right_iff_lt {k : Type u_1} {E : Type u_2} [ordered_ring k] [ordered_add_comm_group E] [semimodule k E] [ordered_semimodule k E] {a : E} {b : E} {r : k} (h : r < 1) : coe_fn (affine_map.line_map a b) r < b ↔ a < b := sorry

theorem right_lt_line_map_iff_lt {k : Type u_1} {E : Type u_2} [ordered_ring k] [ordered_add_comm_group E] [semimodule k E] [ordered_semimodule k E] {a : E} {b : E} {r : k} (h : r < 1) : b < coe_fn (affine_map.line_map a b) r ↔ b < a :=
  line_map_lt_right_iff_lt h

theorem line_map_le_line_map_iff_of_lt {k : Type u_1} {E : Type u_2} [linear_ordered_field k] [ordered_add_comm_group E] [semimodule k E] [ordered_semimodule k E] {a : E} {b : E} {r : k} {r' : k} (h : r < r') : coe_fn (affine_map.line_map a b) r ≤ coe_fn (affine_map.line_map a b) r' ↔ a ≤ b := sorry

theorem left_le_line_map_iff_le {k : Type u_1} {E : Type u_2} [linear_ordered_field k] [ordered_add_comm_group E] [semimodule k E] [ordered_semimodule k E] {a : E} {b : E} {r : k} (h : 0 < r) : a ≤ coe_fn (affine_map.line_map a b) r ↔ a ≤ b := sorry

@[simp] theorem left_le_midpoint {k : Type u_1} {E : Type u_2} [linear_ordered_field k] [ordered_add_comm_group E] [semimodule k E] [ordered_semimodule k E] {a : E} {b : E} : a ≤ midpoint k a b ↔ a ≤ b :=
  left_le_line_map_iff_le (iff.mpr inv_pos zero_lt_two)

theorem line_map_le_left_iff_le {k : Type u_1} {E : Type u_2} [linear_ordered_field k] [ordered_add_comm_group E] [semimodule k E] [ordered_semimodule k E] {a : E} {b : E} {r : k} (h : 0 < r) : coe_fn (affine_map.line_map a b) r ≤ a ↔ b ≤ a :=
  left_le_line_map_iff_le h

@[simp] theorem midpoint_le_left {k : Type u_1} {E : Type u_2} [linear_ordered_field k] [ordered_add_comm_group E] [semimodule k E] [ordered_semimodule k E] {a : E} {b : E} : midpoint k a b ≤ a ↔ b ≤ a :=
  line_map_le_left_iff_le (iff.mpr inv_pos zero_lt_two)

theorem line_map_le_right_iff_le {k : Type u_1} {E : Type u_2} [linear_ordered_field k] [ordered_add_comm_group E] [semimodule k E] [ordered_semimodule k E] {a : E} {b : E} {r : k} (h : r < 1) : coe_fn (affine_map.line_map a b) r ≤ b ↔ a ≤ b := sorry

@[simp] theorem midpoint_le_right {k : Type u_1} {E : Type u_2} [linear_ordered_field k] [ordered_add_comm_group E] [semimodule k E] [ordered_semimodule k E] {a : E} {b : E} : midpoint k a b ≤ b ↔ a ≤ b :=
  line_map_le_right_iff_le (inv_lt_one one_lt_two)

theorem right_le_line_map_iff_le {k : Type u_1} {E : Type u_2} [linear_ordered_field k] [ordered_add_comm_group E] [semimodule k E] [ordered_semimodule k E] {a : E} {b : E} {r : k} (h : r < 1) : b ≤ coe_fn (affine_map.line_map a b) r ↔ b ≤ a :=
  line_map_le_right_iff_le h

@[simp] theorem right_le_midpoint {k : Type u_1} {E : Type u_2} [linear_ordered_field k] [ordered_add_comm_group E] [semimodule k E] [ordered_semimodule k E] {a : E} {b : E} : b ≤ midpoint k a b ↔ b ≤ a :=
  right_le_line_map_iff_le (inv_lt_one one_lt_two)

/-!
### Convexity and slope

Given an interval `[a, b]` and a point `c ∈ (a, b)`, `c = line_map a b r`, there are a few ways to
say that the point `(c, f c)` is above/below the segment `[(a, f a), (b, f b)]`:

* compare `f c` to `line_map (f a) (f b) r`;
* compare `slope f a c` to `slope `f a b`;
* compare `slope f c b` to `slope f a b`;
* compare `slope f a c` to `slope f c b`.

In this section we prove equivalence of these four approaches. In order to make the statements more
readable, we introduce local notation `c = line_map a b r`. Then we prove lemmas like

```
lemma map_le_line_map_iff_slope_le_slope_left (h : 0 < r * (b - a)) :
  f c ≤ line_map (f a) (f b) r ↔ slope f a c ≤ slope f a b :=
```

For each inequality between `f c` and `line_map (f a) (f b) r` we provide 3 lemmas:

* `*_left` relates it to an inequality on `slope f a c` and `slope f a b`;
* `*_right` relates it to an inequality on `slope f a b` and `slope f c b`;
* no-suffix version relates it to an inequality on `slope f a c` and `slope f c b`.

Later these inequalities will be used in to restate `convex_on` in terms of monotonicity of the
slope.
-/

/-- Given `c = line_map a b r`, `a < c`, the point `(c, f c)` is non-strictly below the
segment `[(a, f a), (b, f b)]` if and only if `slope f a c ≤ slope f a b`. -/
theorem map_le_line_map_iff_slope_le_slope_left {k : Type u_1} {E : Type u_2} [linear_ordered_field k] [ordered_add_comm_group E] [semimodule k E] [ordered_semimodule k E] {f : k → E} {a : k} {b : k} {r : k} (h : 0 < r * (b - a)) : f (coe_fn (affine_map.line_map a b) r) ≤ coe_fn (affine_map.line_map (f a) (f b)) r ↔
  slope f a (coe_fn (affine_map.line_map a b) r) ≤ slope f a b := sorry

/-- Given `c = line_map a b r`, `a < c`, the point `(c, f c)` is non-strictly above the
segment `[(a, f a), (b, f b)]` if and only if `slope f a b ≤ slope f a c`. -/
theorem line_map_le_map_iff_slope_le_slope_left {k : Type u_1} {E : Type u_2} [linear_ordered_field k] [ordered_add_comm_group E] [semimodule k E] [ordered_semimodule k E] {f : k → E} {a : k} {b : k} {r : k} (h : 0 < r * (b - a)) : coe_fn (affine_map.line_map (f a) (f b)) r ≤ f (coe_fn (affine_map.line_map a b) r) ↔
  slope f a b ≤ slope f a (coe_fn (affine_map.line_map a b) r) :=
  map_le_line_map_iff_slope_le_slope_left h

/-- Given `c = line_map a b r`, `a < c`, the point `(c, f c)` is strictly below the
segment `[(a, f a), (b, f b)]` if and only if `slope f a c < slope f a b`. -/
theorem map_lt_line_map_iff_slope_lt_slope_left {k : Type u_1} {E : Type u_2} [linear_ordered_field k] [ordered_add_comm_group E] [semimodule k E] [ordered_semimodule k E] {f : k → E} {a : k} {b : k} {r : k} (h : 0 < r * (b - a)) : f (coe_fn (affine_map.line_map a b) r) < coe_fn (affine_map.line_map (f a) (f b)) r ↔
  slope f a (coe_fn (affine_map.line_map a b) r) < slope f a b :=
  lt_iff_lt_of_le_iff_le' (line_map_le_map_iff_slope_le_slope_left h) (map_le_line_map_iff_slope_le_slope_left h)

/-- Given `c = line_map a b r`, `a < c`, the point `(c, f c)` is strictly above the
segment `[(a, f a), (b, f b)]` if and only if `slope f a b < slope f a c`. -/
theorem line_map_lt_map_iff_slope_lt_slope_left {k : Type u_1} {E : Type u_2} [linear_ordered_field k] [ordered_add_comm_group E] [semimodule k E] [ordered_semimodule k E] {f : k → E} {a : k} {b : k} {r : k} (h : 0 < r * (b - a)) : coe_fn (affine_map.line_map (f a) (f b)) r < f (coe_fn (affine_map.line_map a b) r) ↔
  slope f a b < slope f a (coe_fn (affine_map.line_map a b) r) :=
  map_lt_line_map_iff_slope_lt_slope_left h

/-- Given `c = line_map a b r`, `c < b`, the point `(c, f c)` is non-strictly below the
segment `[(a, f a), (b, f b)]` if and only if `slope f a b ≤ slope f c b`. -/
theorem map_le_line_map_iff_slope_le_slope_right {k : Type u_1} {E : Type u_2} [linear_ordered_field k] [ordered_add_comm_group E] [semimodule k E] [ordered_semimodule k E] {f : k → E} {a : k} {b : k} {r : k} (h : 0 < (1 - r) * (b - a)) : f (coe_fn (affine_map.line_map a b) r) ≤ coe_fn (affine_map.line_map (f a) (f b)) r ↔
  slope f a b ≤ slope f (coe_fn (affine_map.line_map a b) r) b := sorry

/-- Given `c = line_map a b r`, `c < b`, the point `(c, f c)` is non-strictly above the
segment `[(a, f a), (b, f b)]` if and only if `slope f c b ≤ slope f a b`. -/
theorem line_map_le_map_iff_slope_le_slope_right {k : Type u_1} {E : Type u_2} [linear_ordered_field k] [ordered_add_comm_group E] [semimodule k E] [ordered_semimodule k E] {f : k → E} {a : k} {b : k} {r : k} (h : 0 < (1 - r) * (b - a)) : coe_fn (affine_map.line_map (f a) (f b)) r ≤ f (coe_fn (affine_map.line_map a b) r) ↔
  slope f (coe_fn (affine_map.line_map a b) r) b ≤ slope f a b :=
  map_le_line_map_iff_slope_le_slope_right h

/-- Given `c = line_map a b r`, `c < b`, the point `(c, f c)` is strictly below the
segment `[(a, f a), (b, f b)]` if and only if `slope f a b < slope f c b`. -/
theorem map_lt_line_map_iff_slope_lt_slope_right {k : Type u_1} {E : Type u_2} [linear_ordered_field k] [ordered_add_comm_group E] [semimodule k E] [ordered_semimodule k E] {f : k → E} {a : k} {b : k} {r : k} (h : 0 < (1 - r) * (b - a)) : f (coe_fn (affine_map.line_map a b) r) < coe_fn (affine_map.line_map (f a) (f b)) r ↔
  slope f a b < slope f (coe_fn (affine_map.line_map a b) r) b :=
  lt_iff_lt_of_le_iff_le' (line_map_le_map_iff_slope_le_slope_right h) (map_le_line_map_iff_slope_le_slope_right h)

/-- Given `c = line_map a b r`, `c < b`, the point `(c, f c)` is strictly above the
segment `[(a, f a), (b, f b)]` if and only if `slope f c b < slope f a b`. -/
theorem line_map_lt_map_iff_slope_lt_slope_right {k : Type u_1} {E : Type u_2} [linear_ordered_field k] [ordered_add_comm_group E] [semimodule k E] [ordered_semimodule k E] {f : k → E} {a : k} {b : k} {r : k} (h : 0 < (1 - r) * (b - a)) : coe_fn (affine_map.line_map (f a) (f b)) r < f (coe_fn (affine_map.line_map a b) r) ↔
  slope f (coe_fn (affine_map.line_map a b) r) b < slope f a b :=
  map_lt_line_map_iff_slope_lt_slope_right h

/-- Given `c = line_map a b r`, `a < c < b`, the point `(c, f c)` is non-strictly below the
segment `[(a, f a), (b, f b)]` if and only if `slope f a c ≤ slope f c b`. -/
theorem map_le_line_map_iff_slope_le_slope {k : Type u_1} {E : Type u_2} [linear_ordered_field k] [ordered_add_comm_group E] [semimodule k E] [ordered_semimodule k E] {f : k → E} {a : k} {b : k} {r : k} (hab : a < b) (h₀ : 0 < r) (h₁ : r < 1) : f (coe_fn (affine_map.line_map a b) r) ≤ coe_fn (affine_map.line_map (f a) (f b)) r ↔
  slope f a (coe_fn (affine_map.line_map a b) r) ≤ slope f (coe_fn (affine_map.line_map a b) r) b := sorry

/-- Given `c = line_map a b r`, `a < c < b`, the point `(c, f c)` is non-strictly above the
segment `[(a, f a), (b, f b)]` if and only if `slope f c b ≤ slope f a c`. -/
theorem line_map_le_map_iff_slope_le_slope {k : Type u_1} {E : Type u_2} [linear_ordered_field k] [ordered_add_comm_group E] [semimodule k E] [ordered_semimodule k E] {f : k → E} {a : k} {b : k} {r : k} (hab : a < b) (h₀ : 0 < r) (h₁ : r < 1) : coe_fn (affine_map.line_map (f a) (f b)) r ≤ f (coe_fn (affine_map.line_map a b) r) ↔
  slope f (coe_fn (affine_map.line_map a b) r) b ≤ slope f a (coe_fn (affine_map.line_map a b) r) :=
  map_le_line_map_iff_slope_le_slope hab h₀ h₁

/-- Given `c = line_map a b r`, `a < c < b`, the point `(c, f c)` is strictly below the
segment `[(a, f a), (b, f b)]` if and only if `slope f a c < slope f c b`. -/
theorem map_lt_line_map_iff_slope_lt_slope {k : Type u_1} {E : Type u_2} [linear_ordered_field k] [ordered_add_comm_group E] [semimodule k E] [ordered_semimodule k E] {f : k → E} {a : k} {b : k} {r : k} (hab : a < b) (h₀ : 0 < r) (h₁ : r < 1) : f (coe_fn (affine_map.line_map a b) r) < coe_fn (affine_map.line_map (f a) (f b)) r ↔
  slope f a (coe_fn (affine_map.line_map a b) r) < slope f (coe_fn (affine_map.line_map a b) r) b :=
  lt_iff_lt_of_le_iff_le' (line_map_le_map_iff_slope_le_slope hab h₀ h₁) (map_le_line_map_iff_slope_le_slope hab h₀ h₁)

/-- Given `c = line_map a b r`, `a < c < b`, the point `(c, f c)` is strictly above the
segment `[(a, f a), (b, f b)]` if and only if `slope f c b < slope f a c`. -/
theorem line_map_lt_map_iff_slope_lt_slope {k : Type u_1} {E : Type u_2} [linear_ordered_field k] [ordered_add_comm_group E] [semimodule k E] [ordered_semimodule k E] {f : k → E} {a : k} {b : k} {r : k} (hab : a < b) (h₀ : 0 < r) (h₁ : r < 1) : coe_fn (affine_map.line_map (f a) (f b)) r < f (coe_fn (affine_map.line_map a b) r) ↔
  slope f (coe_fn (affine_map.line_map a b) r) b < slope f a (coe_fn (affine_map.line_map a b) r) :=
  map_lt_line_map_iff_slope_lt_slope hab h₀ h₁


/-
Copyright (c) 2017 Johannes Hölzl. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Johannes Hölzl, Mario Carneiro, Kevin Buzzard, Yury Kudryashov, Eric Wieser
-/
import linear_algebra.span
import order.partial_sups
import algebra.algebra.prod

/-! ### Products of modules

> THIS FILE IS SYNCHRONIZED WITH MATHLIB4.
> Any changes to this file require a corresponding PR to mathlib4.

This file defines constructors for linear maps whose domains or codomains are products.

It contains theorems relating these to each other, as well as to `submodule.prod`, `submodule.map`,
`submodule.comap`, `linear_map.range`, and `linear_map.ker`.

## Main definitions

- products in the domain:
  - `linear_map.fst`
  - `linear_map.snd`
  - `linear_map.coprod`
  - `linear_map.prod_ext`
- products in the codomain:
  - `linear_map.inl`
  - `linear_map.inr`
  - `linear_map.prod`
- products in both domain and codomain:
  - `linear_map.prod_map`
  - `linear_equiv.prod_map`
  - `linear_equiv.skew_prod`
-/

universes u v w x y z u' v' w' y'
variables {R : Type u} {K : Type u'} {M : Type v} {V : Type v'} {M₂ : Type w} {V₂ : Type w'}
variables {M₃ : Type y} {V₃ : Type y'} {M₄ : Type z} {ι : Type x}
variables {M₅ M₆ : Type*}

section prod

namespace linear_map

variables (S : Type*) [semiring R] [semiring S]
variables [add_comm_monoid M] [add_comm_monoid M₂] [add_comm_monoid M₃] [add_comm_monoid M₄]
variables [add_comm_monoid M₅] [add_comm_monoid M₆]
variables [module R M] [module R M₂] [module R M₃] [module R M₄]
variables [module R M₅] [module R M₆]
variables (f : M →ₗ[R] M₂)

section
variables (R M M₂)

/-- The first projection of a product is a linear map. -/
def fst : M × M₂ →ₗ[R] M := { to_fun := prod.fst, map_add' := λ x y, rfl, map_smul' := λ x y, rfl }

/-- The second projection of a product is a linear map. -/
def snd : M × M₂ →ₗ[R] M₂ := { to_fun := prod.snd, map_add' := λ x y, rfl, map_smul' := λ x y, rfl }
end

@[simp] theorem fst_apply (x : M × M₂) : fst R M M₂ x = x.1 := rfl
@[simp] theorem snd_apply (x : M × M₂) : snd R M M₂ x = x.2 := rfl

theorem fst_surjective : function.surjective (fst R M M₂) := λ x, ⟨(x, 0), rfl⟩
theorem snd_surjective : function.surjective (snd R M M₂) := λ x, ⟨(0, x), rfl⟩

/-- The prod of two linear maps is a linear map. -/
@[simps] def prod (f : M →ₗ[R] M₂) (g : M →ₗ[R] M₃) : (M →ₗ[R] M₂ × M₃) :=
{ to_fun    := pi.prod f g,
  map_add'  := λ x y, by simp only [pi.prod, prod.mk_add_mk, map_add],
  map_smul' := λ c x, by simp only [pi.prod, prod.smul_mk, map_smul, ring_hom.id_apply] }

lemma coe_prod (f : M →ₗ[R] M₂) (g : M →ₗ[R] M₃) : ⇑(f.prod g) = pi.prod f g := rfl

@[simp] theorem fst_prod (f : M →ₗ[R] M₂) (g : M →ₗ[R] M₃) :
  (fst R M₂ M₃).comp (prod f g) = f := by ext; refl

@[simp] theorem snd_prod (f : M →ₗ[R] M₂) (g : M →ₗ[R] M₃) :
  (snd R M₂ M₃).comp (prod f g) = g := by ext; refl

@[simp] theorem pair_fst_snd : prod (fst R M M₂) (snd R M M₂) = linear_map.id :=
fun_like.coe_injective pi.prod_fst_snd

/-- Taking the product of two maps with the same domain is equivalent to taking the product of
their codomains.

See note [bundled maps over different rings] for why separate `R` and `S` semirings are used. -/
@[simps] def prod_equiv
  [module S M₂] [module S M₃] [smul_comm_class R S M₂] [smul_comm_class R S M₃] :
  ((M →ₗ[R] M₂) × (M →ₗ[R] M₃)) ≃ₗ[S] (M →ₗ[R] M₂ × M₃) :=
{ to_fun := λ f, f.1.prod f.2,
  inv_fun := λ f, ((fst _ _ _).comp f, (snd _ _ _).comp f),
  left_inv := λ f, by ext; refl,
  right_inv := λ f, by ext; refl,
  map_add' := λ a b, rfl,
  map_smul' := λ r a, rfl }

section
variables (R M M₂)

/-- The left injection into a product is a linear map. -/
def inl : M →ₗ[R] M × M₂ := prod linear_map.id 0

/-- The right injection into a product is a linear map. -/
def inr : M₂ →ₗ[R] M × M₂ := prod 0 linear_map.id

theorem range_inl : range (inl R M M₂) = ker (snd R M M₂) :=
begin
  ext x,
  simp only [mem_ker, mem_range],
  split,
  { rintros ⟨y, rfl⟩, refl },
  { intro h, exact ⟨x.fst, prod.ext rfl h.symm⟩ }
end

theorem ker_snd : ker (snd R M M₂) = range (inl R M M₂) :=
eq.symm $ range_inl R M M₂

theorem range_inr : range (inr R M M₂) = ker (fst R M M₂) :=
begin
  ext x,
  simp only [mem_ker, mem_range],
  split,
  { rintros ⟨y, rfl⟩, refl },
  { intro h, exact ⟨x.snd, prod.ext h.symm rfl⟩ }
end

theorem ker_fst : ker (fst R M M₂) = range (inr R M M₂) :=
eq.symm $ range_inr R M M₂

end

@[simp] theorem coe_inl : (inl R M M₂ : M → M × M₂) = λ x, (x, 0) := rfl
theorem inl_apply (x : M) : inl R M M₂ x = (x, 0) := rfl

@[simp] theorem coe_inr : (inr R M M₂ : M₂ → M × M₂) = prod.mk 0 := rfl
theorem inr_apply (x : M₂) : inr R M M₂ x = (0, x) := rfl

theorem inl_eq_prod : inl R M M₂ = prod linear_map.id 0 := rfl

theorem inr_eq_prod : inr R M M₂ = prod 0 linear_map.id := rfl

theorem inl_injective : function.injective (inl R M M₂) :=
λ _, by simp

/-- The coprod function `λ x : M × M₂, f x.1 + g x.2` is a linear map. -/
def coprod (f : M →ₗ[R] M₃) (g : M₂ →ₗ[R] M₃) : M × M₂ →ₗ[R] M₃ :=
f.comp (fst _ _ _) + g.comp (snd _ _ _)

@[simp] theorem coprod_apply (f : M →ₗ[R] M₃) (g : M₂ →ₗ[R] M₃) (x : M × M₂) :
  coprod f g x = f x.1 + g x.2 := rfl

@[simp] theorem coprod_inl (f : M →ₗ[R] M₃) (g : M₂ →ₗ[R] M₃) :
  (coprod f g).comp (inl R M M₂) = f :=
by ext; simp only [map_zero, add_zero, coprod_apply, inl_apply, comp_apply]

@[simp] theorem coprod_inr (f : M →ₗ[R] M₃) (g : M₂ →ₗ[R] M₃) :
  (coprod f g).comp (inr R M M₂) = g :=
by ext; simp only [map_zero, coprod_apply, inr_apply, zero_add, comp_apply]

@[simp] theorem coprod_inl_inr : coprod (inl R M M₂) (inr R M M₂) = linear_map.id :=
by ext; simp only [prod.mk_add_mk, add_zero, id_apply, coprod_apply,
  inl_apply, inr_apply, zero_add]

theorem comp_coprod (f : M₃ →ₗ[R] M₄) (g₁ : M →ₗ[R] M₃) (g₂ : M₂ →ₗ[R] M₃) :
  f.comp (g₁.coprod g₂) = (f.comp g₁).coprod (f.comp g₂) :=
ext $ λ x, f.map_add (g₁ x.1) (g₂ x.2)

theorem fst_eq_coprod : fst R M M₂ = coprod linear_map.id 0 := by ext; simp

theorem snd_eq_coprod : snd R M M₂ = coprod 0 linear_map.id := by ext; simp

@[simp] theorem coprod_comp_prod (f : M₂ →ₗ[R] M₄) (g : M₃ →ₗ[R] M₄)
  (f' : M →ₗ[R] M₂) (g' : M →ₗ[R] M₃) :
  (f.coprod g).comp (f'.prod g') = f.comp f' + g.comp g' :=
rfl

@[simp]
lemma coprod_map_prod (f : M →ₗ[R] M₃) (g : M₂ →ₗ[R] M₃) (S : submodule R M)
  (S' : submodule R M₂) :
  (submodule.prod S S').map (linear_map.coprod f g) = S.map f ⊔ S'.map g :=
set_like.coe_injective $ begin
  simp only [linear_map.coprod_apply, submodule.coe_sup, submodule.map_coe],
  rw [←set.image2_add, set.image2_image_left, set.image2_image_right],
  exact set.image_prod (λ m m₂, f m + g m₂),
end

/-- Taking the product of two maps with the same codomain is equivalent to taking the product of
their domains.

See note [bundled maps over different rings] for why separate `R` and `S` semirings are used. -/
@[simps] def coprod_equiv [module S M₃] [smul_comm_class R S M₃] :
  ((M →ₗ[R] M₃) × (M₂ →ₗ[R] M₃)) ≃ₗ[S] (M × M₂ →ₗ[R] M₃) :=
{ to_fun := λ f, f.1.coprod f.2,
  inv_fun := λ f, (f.comp (inl _ _ _), f.comp (inr _ _ _)),
  left_inv := λ f, by simp only [prod.mk.eta, coprod_inl, coprod_inr],
  right_inv := λ f, by simp only [←comp_coprod, comp_id, coprod_inl_inr],
  map_add' := λ a b,
    by { ext, simp only [prod.snd_add, add_apply, coprod_apply, prod.fst_add, add_add_add_comm] },
  map_smul' := λ r a,
    by { dsimp, ext, simp only [smul_add, smul_apply, prod.smul_snd, prod.smul_fst,
                                coprod_apply] } }

theorem prod_ext_iff {f g : M × M₂ →ₗ[R] M₃} :
  f = g ↔ f.comp (inl _ _ _) = g.comp (inl _ _ _) ∧ f.comp (inr _ _ _) = g.comp (inr _ _ _) :=
(coprod_equiv ℕ).symm.injective.eq_iff.symm.trans prod.ext_iff

/--
Split equality of linear maps from a product into linear maps over each component, to allow `ext`
to apply lemmas specific to `M →ₗ M₃` and `M₂ →ₗ M₃`.

See note [partially-applied ext lemmas]. -/
@[ext] theorem prod_ext {f g : M × M₂ →ₗ[R] M₃}
  (hl : f.comp (inl _ _ _) = g.comp (inl _ _ _))
  (hr : f.comp (inr _ _ _) = g.comp (inr _ _ _)) :
  f = g :=
prod_ext_iff.2 ⟨hl, hr⟩

/-- `prod.map` of two linear maps. -/
def prod_map (f : M →ₗ[R] M₃) (g : M₂ →ₗ[R] M₄) : (M × M₂) →ₗ[R] (M₃ × M₄) :=
(f.comp (fst R M M₂)).prod (g.comp (snd R M M₂))

lemma coe_prod_map (f : M →ₗ[R] M₃) (g : M₂ →ₗ[R] M₄) :
  ⇑(f.prod_map g) = prod.map f g := rfl

@[simp] theorem prod_map_apply (f : M →ₗ[R] M₃) (g : M₂ →ₗ[R] M₄) (x) :
  f.prod_map g x = (f x.1, g x.2) := rfl

lemma prod_map_comap_prod (f : M →ₗ[R] M₂) (g : M₃ →ₗ[R] M₄) (S : submodule R M₂)
  (S' : submodule R M₄) :
  (submodule.prod S S').comap (linear_map.prod_map f g) = (S.comap f).prod (S'.comap g) :=
set_like.coe_injective $ set.preimage_prod_map_prod f g _ _

lemma ker_prod_map (f : M →ₗ[R] M₂) (g : M₃ →ₗ[R] M₄) :
  (linear_map.prod_map f g).ker = submodule.prod f.ker g.ker :=
begin
  dsimp only [ker],
  rw [←prod_map_comap_prod, submodule.prod_bot],
end

@[simp]
lemma prod_map_id : (id : M →ₗ[R] M).prod_map (id : M₂ →ₗ[R] M₂) = id :=
linear_map.ext $ λ _, prod.mk.eta

@[simp]
lemma prod_map_one : (1 : M →ₗ[R] M).prod_map (1 : M₂ →ₗ[R] M₂) = 1 :=
linear_map.ext $ λ _, prod.mk.eta

lemma prod_map_comp (f₁₂ : M →ₗ[R] M₂) (f₂₃ : M₂ →ₗ[R] M₃) (g₁₂ : M₄ →ₗ[R] M₅) (g₂₃ : M₅ →ₗ[R] M₆) :
  f₂₃.prod_map g₂₃ ∘ₗ f₁₂.prod_map g₁₂ = (f₂₃ ∘ₗ f₁₂).prod_map (g₂₃ ∘ₗ g₁₂) := rfl

lemma prod_map_mul (f₁₂ : M →ₗ[R] M) (f₂₃ : M →ₗ[R] M) (g₁₂ : M₂ →ₗ[R] M₂) (g₂₃ : M₂ →ₗ[R] M₂) :
  f₂₃.prod_map g₂₃ * f₁₂.prod_map g₁₂ = (f₂₃ * f₁₂).prod_map (g₂₃ * g₁₂) := rfl

lemma prod_map_add (f₁ : M →ₗ[R] M₃) (f₂ : M →ₗ[R] M₃) (g₁ : M₂ →ₗ[R] M₄) (g₂ : M₂ →ₗ[R] M₄) :
  (f₁ + f₂).prod_map (g₁ + g₂) = f₁.prod_map g₁ + f₂.prod_map g₂ := rfl

@[simp] lemma prod_map_zero :
  (0 : M →ₗ[R] M₂).prod_map (0 : M₃ →ₗ[R] M₄) = 0 := rfl

@[simp] lemma prod_map_smul
 [module S M₃] [module S M₄] [smul_comm_class R S M₃] [smul_comm_class R S M₄]
 (s : S) (f : M →ₗ[R] M₃) (g : M₂ →ₗ[R] M₄) : prod_map (s • f) (s • g) = s • prod_map f g := rfl

variables (R M M₂ M₃ M₄)

/-- `linear_map.prod_map` as a `linear_map` -/
@[simps]
def prod_map_linear
 [module S M₃] [module S M₄] [smul_comm_class R S M₃] [smul_comm_class R S M₄] :
 ((M →ₗ[R] M₃) × (M₂ →ₗ[R] M₄)) →ₗ[S] ((M × M₂) →ₗ[R] (M₃ × M₄)) :=
{ to_fun := λ f, prod_map f.1 f.2,
  map_add' := λ _ _, rfl,
  map_smul' := λ _ _, rfl}

/-- `linear_map.prod_map` as a `ring_hom` -/
@[simps]
def prod_map_ring_hom : (M →ₗ[R] M) × (M₂ →ₗ[R] M₂) →+* ((M × M₂) →ₗ[R] (M × M₂)) :=
{ to_fun := λ f, prod_map f.1 f.2,
  map_one' := prod_map_one,
  map_zero' := rfl,
  map_add' := λ _ _, rfl,
  map_mul' := λ _ _, rfl }

variables {R M M₂ M₃ M₄}

section map_mul

variables {A : Type*} [non_unital_non_assoc_semiring A] [module R A]
variables {B : Type*} [non_unital_non_assoc_semiring B] [module R B]

lemma inl_map_mul (a₁ a₂ : A) : linear_map.inl R A B (a₁ * a₂) =
  linear_map.inl R A B a₁ * linear_map.inl R A B a₂ :=
prod.ext rfl (by simp)

lemma inr_map_mul (b₁ b₂ : B) : linear_map.inr R A B (b₁ * b₂) =
  linear_map.inr R A B b₁ * linear_map.inr R A B b₂ :=
prod.ext (by simp) rfl

end map_mul

end linear_map

end prod

namespace linear_map

variables (R M M₂)

variables [comm_semiring R]
variables [add_comm_monoid M] [add_comm_monoid M₂]
variables [module R M] [module R M₂]

/-- `linear_map.prod_map` as an `algebra_hom` -/
@[simps]
def prod_map_alg_hom : (module.End R M) × (module.End R M₂) →ₐ[R] module.End R (M × M₂) :=
{ commutes' := λ _, rfl, ..prod_map_ring_hom R M M₂ }

end linear_map

namespace linear_map
open submodule

variables [semiring R]
  [add_comm_monoid M] [add_comm_monoid M₂] [add_comm_monoid M₃] [add_comm_monoid M₄]
  [module R M] [module R M₂] [module R M₃] [module R M₄]

lemma range_coprod (f : M →ₗ[R] M₃) (g : M₂ →ₗ[R] M₃) :
  (f.coprod g).range = f.range ⊔ g.range :=
submodule.ext $ λ x, by simp [mem_sup]

lemma is_compl_range_inl_inr : is_compl (inl R M M₂).range (inr R M M₂).range :=
begin
  split,
  { rw disjoint_def,
    rintros ⟨_, _⟩ ⟨x, hx⟩ ⟨y, hy⟩,
    simp only [prod.ext_iff, inl_apply, inr_apply, mem_bot] at hx hy ⊢,
    exact ⟨hy.1.symm, hx.2.symm⟩ },
  { rw codisjoint_iff_le_sup,
    rintros ⟨x, y⟩ -,
    simp only [mem_sup, mem_range, exists_prop],
    refine ⟨(x, 0), ⟨x, rfl⟩, (0, y), ⟨y, rfl⟩, _⟩,
    simp }
end

lemma sup_range_inl_inr : (inl R M M₂).range ⊔ (inr R M M₂).range = ⊤ :=
is_compl.sup_eq_top is_compl_range_inl_inr

lemma disjoint_inl_inr : disjoint (inl R M M₂).range (inr R M M₂).range :=
by simp [disjoint_def, @eq_comm M 0, @eq_comm M₂ 0] {contextual := tt}; intros; refl
theorem map_coprod_prod (f : M →ₗ[R] M₃) (g : M₂ →ₗ[R] M₃)
  (p : submodule R M) (q : submodule R M₂) :
  map (coprod f g) (p.prod q) = map f p ⊔ map g q :=
begin
  refine le_antisymm _ (sup_le (map_le_iff_le_comap.2 _) (map_le_iff_le_comap.2 _)),
  { rw set_like.le_def, rintro _ ⟨x, ⟨h₁, h₂⟩, rfl⟩,
    exact mem_sup.2 ⟨_, ⟨_, h₁, rfl⟩, _, ⟨_, h₂, rfl⟩, rfl⟩ },
  { exact λ x hx, ⟨(x, 0), by simp [hx]⟩ },
  { exact λ x hx, ⟨(0, x), by simp [hx]⟩ }
end

theorem comap_prod_prod (f : M →ₗ[R] M₂) (g : M →ₗ[R] M₃)
  (p : submodule R M₂) (q : submodule R M₃) :
  comap (prod f g) (p.prod q) = comap f p ⊓ comap g q :=
submodule.ext $ λ x, iff.rfl

theorem prod_eq_inf_comap (p : submodule R M) (q : submodule R M₂) :
  p.prod q = p.comap (linear_map.fst R M M₂) ⊓ q.comap (linear_map.snd R M M₂) :=
submodule.ext $ λ x, iff.rfl

theorem prod_eq_sup_map (p : submodule R M) (q : submodule R M₂) :
  p.prod q = p.map (linear_map.inl R M M₂) ⊔ q.map (linear_map.inr R M M₂) :=
by rw [← map_coprod_prod, coprod_inl_inr, map_id]

lemma span_inl_union_inr {s : set M} {t : set M₂} :
  span R (inl R M  M₂ '' s ∪ inr R M M₂ '' t) = (span R s).prod (span R t) :=
by rw [span_union, prod_eq_sup_map, ← span_image, ← span_image]

@[simp] lemma ker_prod (f : M →ₗ[R] M₂) (g : M →ₗ[R] M₃) :
  ker (prod f g) = ker f ⊓ ker g :=
by rw [ker, ← prod_bot, comap_prod_prod]; refl

lemma range_prod_le (f : M →ₗ[R] M₂) (g : M →ₗ[R] M₃) :
  range (prod f g) ≤ (range f).prod (range g) :=
begin
  simp only [set_like.le_def, prod_apply, mem_range, set_like.mem_coe, mem_prod,
    exists_imp_distrib],
  rintro _ x rfl,
  exact ⟨⟨x, rfl⟩, ⟨x, rfl⟩⟩
end

lemma ker_prod_ker_le_ker_coprod {M₂ : Type*} [add_comm_group M₂] [module R M₂]
  {M₃ : Type*} [add_comm_group M₃] [module R M₃]
  (f : M →ₗ[R] M₃) (g : M₂ →ₗ[R] M₃) :
  (ker f).prod (ker g) ≤ ker (f.coprod g) :=
by { rintros ⟨y, z⟩, simp {contextual := tt} }

lemma ker_coprod_of_disjoint_range {M₂ : Type*} [add_comm_group M₂] [module R M₂]
  {M₃ : Type*} [add_comm_group M₃] [module R M₃]
  (f : M →ₗ[R] M₃) (g : M₂ →ₗ[R] M₃) (hd : disjoint f.range g.range) :
  ker (f.coprod g) = (ker f).prod (ker g) :=
begin
  apply le_antisymm _ (ker_prod_ker_le_ker_coprod f g),
  rintros ⟨y, z⟩ h,
  simp only [mem_ker, mem_prod, coprod_apply] at h ⊢,
  have : f y ∈ f.range ⊓ g.range,
  { simp only [true_and, mem_range, mem_inf, exists_apply_eq_apply],
    use -z,
    rwa [eq_comm, map_neg, ← sub_eq_zero, sub_neg_eq_add] },
  rw [hd.eq_bot, mem_bot] at this,
  rw [this] at h,
  simpa [this] using h,
end

end linear_map

namespace submodule
open linear_map

variables [semiring R]
variables [add_comm_monoid M] [add_comm_monoid M₂]
variables [module R M] [module R M₂]

lemma sup_eq_range (p q : submodule R M) : p ⊔ q = (p.subtype.coprod q.subtype).range :=
submodule.ext $ λ x, by simp [submodule.mem_sup, set_like.exists]

variables (p : submodule R M) (q : submodule R M₂)

@[simp] theorem map_inl : p.map (inl R M M₂) = prod p ⊥ :=
by { ext ⟨x, y⟩, simp only [and.left_comm, eq_comm, mem_map, prod.mk.inj_iff, inl_apply, mem_bot,
  exists_eq_left', mem_prod] }

@[simp] theorem map_inr : q.map (inr R M M₂) = prod ⊥ q :=
by ext ⟨x, y⟩; simp [and.left_comm, eq_comm]

@[simp] theorem comap_fst : p.comap (fst R M M₂) = prod p ⊤ :=
by ext ⟨x, y⟩; simp

@[simp] theorem comap_snd : q.comap (snd R M M₂) = prod ⊤ q :=
by ext ⟨x, y⟩; simp

@[simp] theorem prod_comap_inl : (prod p q).comap (inl R M M₂) = p := by ext; simp

@[simp] theorem prod_comap_inr : (prod p q).comap (inr R M M₂) = q := by ext; simp

@[simp] theorem prod_map_fst : (prod p q).map (fst R M M₂) = p :=
by ext x; simp [(⟨0, zero_mem _⟩ : ∃ x, x ∈ q)]

@[simp] theorem prod_map_snd : (prod p q).map (snd R M M₂) = q :=
by ext x; simp [(⟨0, zero_mem _⟩ : ∃ x, x ∈ p)]

@[simp] theorem ker_inl : (inl R M M₂).ker = ⊥ :=
by rw [ker, ← prod_bot, prod_comap_inl]

@[simp] theorem ker_inr : (inr R M M₂).ker = ⊥ :=
by rw [ker, ← prod_bot, prod_comap_inr]

@[simp] theorem range_fst : (fst R M M₂).range = ⊤ :=
by rw [range_eq_map, ← prod_top, prod_map_fst]

@[simp] theorem range_snd : (snd R M M₂).range = ⊤ :=
by rw [range_eq_map, ← prod_top, prod_map_snd]

variables (R M M₂)

/-- `M` as a submodule of `M × N`. -/
def fst : submodule R (M × M₂) := (⊥ : submodule R M₂).comap (linear_map.snd R M M₂)

/-- `M` as a submodule of `M × N` is isomorphic to `M`. -/
@[simps] def fst_equiv : submodule.fst R M M₂ ≃ₗ[R] M :=
{ to_fun := λ x, x.1.1,
  inv_fun := λ m, ⟨⟨m, 0⟩, by tidy⟩,
  map_add' := by simp,
  map_smul' := by simp,
  left_inv := by tidy,
  right_inv := by tidy, }

lemma fst_map_fst : (submodule.fst R M M₂).map (linear_map.fst R M M₂) = ⊤ :=
by tidy
lemma fst_map_snd : (submodule.fst R M M₂).map (linear_map.snd R M M₂) = ⊥ :=
by { tidy, exact 0, }

/-- `N` as a submodule of `M × N`. -/
def snd : submodule R (M × M₂) := (⊥ : submodule R M).comap (linear_map.fst R M M₂)

/-- `N` as a submodule of `M × N` is isomorphic to `N`. -/
@[simps] def snd_equiv : submodule.snd R M M₂ ≃ₗ[R] M₂ :=
{ to_fun := λ x, x.1.2,
  inv_fun := λ n, ⟨⟨0, n⟩, by tidy⟩,
  map_add' := by simp,
  map_smul' := by simp,
  left_inv := by tidy,
  right_inv := by tidy, }

lemma snd_map_fst : (submodule.snd R M M₂).map (linear_map.fst R M M₂) = ⊥ :=
by { tidy, exact 0, }
lemma snd_map_snd : (submodule.snd R M M₂).map (linear_map.snd R M M₂) = ⊤ :=
by tidy

lemma fst_sup_snd : submodule.fst R M M₂ ⊔ submodule.snd R M M₂ = ⊤ :=
begin
  rw eq_top_iff,
  rintro ⟨m, n⟩ -,
  rw [show (m, n) = (m, 0) + (0, n), by simp],
  apply submodule.add_mem (submodule.fst R M M₂ ⊔ submodule.snd R M M₂),
  { exact submodule.mem_sup_left (submodule.mem_comap.mpr (by simp)), },
  { exact submodule.mem_sup_right (submodule.mem_comap.mpr (by simp)), },
end

lemma fst_inf_snd : submodule.fst R M M₂ ⊓ submodule.snd R M M₂ = ⊥ := by tidy

lemma le_prod_iff {p₁ : submodule R M} {p₂ : submodule R M₂} {q : submodule R (M × M₂)} :
  q ≤ p₁.prod p₂ ↔ map (linear_map.fst R M M₂) q ≤ p₁ ∧ map (linear_map.snd R M M₂) q ≤ p₂ :=
begin
  split,
  { intros h,
    split,
    { rintros x ⟨⟨y1,y2⟩, ⟨hy1,rfl⟩⟩, exact (h hy1).1 },
    { rintros x ⟨⟨y1,y2⟩, ⟨hy1,rfl⟩⟩, exact (h hy1).2 }, },
  { rintros ⟨hH, hK⟩ ⟨x1, x2⟩ h, exact ⟨hH ⟨_ , h, rfl⟩, hK ⟨ _, h, rfl⟩⟩, }
end

lemma prod_le_iff {p₁ : submodule R M} {p₂ : submodule R M₂} {q : submodule R (M × M₂)} :
  p₁.prod p₂ ≤ q ↔ map (linear_map.inl R M M₂) p₁ ≤ q ∧ map (linear_map.inr R M M₂) p₂ ≤ q :=
begin
  split,
  { intros h,
    split,
    { rintros _ ⟨x, hx, rfl⟩, apply h, exact ⟨hx, zero_mem p₂⟩, },
    { rintros _ ⟨x, hx, rfl⟩, apply h, exact ⟨zero_mem p₁, hx⟩, }, },
  { rintros ⟨hH, hK⟩ ⟨x1, x2⟩ ⟨h1, h2⟩,
    have h1' : (linear_map.inl R _ _) x1 ∈ q, { apply hH, simpa using h1, },
    have h2' : (linear_map.inr R _ _) x2 ∈ q, { apply hK, simpa using h2, },
    simpa using add_mem h1' h2', }
end

lemma prod_eq_bot_iff {p₁ : submodule R M} {p₂ : submodule R M₂} :
  p₁.prod p₂ = ⊥ ↔ p₁ = ⊥ ∧ p₂ = ⊥ :=
by simp only [eq_bot_iff, prod_le_iff, (gc_map_comap _).le_iff_le, comap_bot, ker_inl, ker_inr]

lemma prod_eq_top_iff {p₁ : submodule R M} {p₂ : submodule R M₂} :
  p₁.prod p₂ = ⊤ ↔ p₁ = ⊤ ∧ p₂ = ⊤ :=
by simp only [eq_top_iff, le_prod_iff, ← (gc_map_comap _).le_iff_le, map_top, range_fst, range_snd]

end submodule

namespace linear_equiv

/-- Product of modules is commutative up to linear isomorphism. -/
@[simps apply]
def prod_comm (R M N : Type*) [semiring R] [add_comm_monoid M] [add_comm_monoid N]
  [module R M] [module R N] : (M × N) ≃ₗ[R] (N × M) :=
{ to_fun := prod.swap,
  map_smul' := λ r ⟨m, n⟩, rfl,
  ..add_equiv.prod_comm }

section

variables [semiring R]
variables [add_comm_monoid M] [add_comm_monoid M₂] [add_comm_monoid M₃] [add_comm_monoid M₄]
variables {module_M : module R M} {module_M₂ : module R M₂}
variables {module_M₃ : module R M₃} {module_M₄ : module R M₄}
variables (e₁ : M ≃ₗ[R] M₂) (e₂ : M₃ ≃ₗ[R] M₄)

/-- Product of linear equivalences; the maps come from `equiv.prod_congr`. -/
protected def prod :
  (M × M₃) ≃ₗ[R] (M₂ × M₄) :=
{ map_smul' := λ c x, prod.ext (e₁.map_smulₛₗ c _) (e₂.map_smulₛₗ c _),
  .. e₁.to_add_equiv.prod_congr e₂.to_add_equiv }

lemma prod_symm : (e₁.prod e₂).symm = e₁.symm.prod e₂.symm := rfl

@[simp] lemma prod_apply (p) :
  e₁.prod e₂ p = (e₁ p.1, e₂ p.2) := rfl

@[simp, norm_cast] lemma coe_prod :
  (e₁.prod e₂ : (M × M₃) →ₗ[R] (M₂ × M₄)) = (e₁ : M →ₗ[R] M₂).prod_map (e₂ : M₃ →ₗ[R] M₄) := rfl

end

section
variables [semiring R]
variables [add_comm_monoid M] [add_comm_monoid M₂] [add_comm_monoid M₃] [add_comm_group M₄]
variables {module_M : module R M} {module_M₂ : module R M₂}
variables {module_M₃ : module R M₃} {module_M₄ : module R M₄}
variables (e₁ : M ≃ₗ[R] M₂) (e₂ : M₃ ≃ₗ[R] M₄)

/-- Equivalence given by a block lower diagonal matrix. `e₁` and `e₂` are diagonal square blocks,
  and `f` is a rectangular block below the diagonal. -/
protected def skew_prod (f : M →ₗ[R] M₄) :
  (M × M₃) ≃ₗ[R] M₂ × M₄ :=
{ inv_fun := λ p : M₂ × M₄, (e₁.symm p.1, e₂.symm (p.2 - f (e₁.symm p.1))),
  left_inv := λ p, by simp,
  right_inv := λ p, by simp,
  .. ((e₁ : M →ₗ[R] M₂).comp (linear_map.fst R M M₃)).prod
    ((e₂ : M₃ →ₗ[R] M₄).comp (linear_map.snd R M M₃) +
      f.comp (linear_map.fst R M M₃)) }

@[simp] lemma skew_prod_apply (f : M →ₗ[R] M₄) (x) :
  e₁.skew_prod e₂ f x = (e₁ x.1, e₂ x.2 + f x.1) := rfl

@[simp] lemma skew_prod_symm_apply (f : M →ₗ[R] M₄) (x) :
  (e₁.skew_prod e₂ f).symm x = (e₁.symm x.1, e₂.symm (x.2 - f (e₁.symm x.1))) := rfl

end
end linear_equiv

namespace linear_map
open submodule

variables [ring R]
variables [add_comm_group M] [add_comm_group M₂] [add_comm_group M₃]
variables [module R M] [module R M₂] [module R M₃]

/-- If the union of the kernels `ker f` and `ker g` spans the domain, then the range of
`prod f g` is equal to the product of `range f` and `range g`. -/
lemma range_prod_eq {f : M →ₗ[R] M₂} {g : M →ₗ[R] M₃} (h : ker f ⊔ ker g = ⊤) :
  range (prod f g) = (range f).prod (range g) :=
begin
  refine le_antisymm (f.range_prod_le g) _,
  simp only [set_like.le_def, prod_apply, mem_range, set_like.mem_coe, mem_prod, exists_imp_distrib,
    and_imp, prod.forall, pi.prod],
  rintros _ _ x rfl y rfl,
  simp only [prod.mk.inj_iff, ← sub_mem_ker_iff],
  have : y - x ∈ ker f ⊔ ker g, { simp only [h, mem_top] },
  rcases mem_sup.1 this with ⟨x', hx', y', hy', H⟩,
  refine ⟨x' + x, _, _⟩,
  { simp only [mem_ker.mp hx', map_add, zero_add]},
  { simp [←eq_sub_iff_add_eq.1 H, map_add, add_left_inj, self_eq_add_right, mem_ker.mp hy'] }
end

end linear_map

namespace linear_map
/-!
## Tunnels and tailings

Some preliminary work for establishing the strong rank condition for noetherian rings.

Given a morphism `f : M × N →ₗ[R] M` which is `i : injective f`,
we can find an infinite decreasing `tunnel f i n` of copies of `M` inside `M`,
and sitting beside these, an infinite sequence of copies of `N`.

We picturesquely name these as `tailing f i n` for each individual copy of `N`,
and `tailings f i n` for the supremum of the first `n+1` copies:
they are the pieces left behind, sitting inside the tunnel.

By construction, each `tailing f i (n+1)` is disjoint from `tailings f i n`;
later, when we assume `M` is noetherian, this implies that `N` must be trivial,
and establishes the strong rank condition for any left-noetherian ring.
-/
section tunnel

-- (This doesn't work over a semiring: we need to use that `submodule R M` is a modular lattice,
-- which requires cancellation.)
variables [ring R]
variables {N : Type*} [add_comm_group M] [module R M] [add_comm_group N] [module R N]

open function

/-- An auxiliary construction for `tunnel`.
The composition of `f`, followed by the isomorphism back to `K`,
followed by the inclusion of this submodule back into `M`. -/
def tunnel_aux (f : M × N →ₗ[R] M) (Kφ : Σ K : submodule R M, K ≃ₗ[R] M) :
  M × N →ₗ[R] M :=
(Kφ.1.subtype.comp Kφ.2.symm.to_linear_map).comp f

lemma tunnel_aux_injective
  (f : M × N →ₗ[R] M) (i : injective f) (Kφ : Σ K : submodule R M, K ≃ₗ[R] M) :
  injective (tunnel_aux f Kφ) :=
(subtype.val_injective.comp Kφ.2.symm.injective).comp i

noncomputable theory

/-- Auxiliary definition for `tunnel`. -/
-- Even though we have `noncomputable theory`,
-- we get an error without another `noncomputable` here.
noncomputable def tunnel' (f : M × N →ₗ[R] M) (i : injective f) :
  ℕ → Σ (K : submodule R M), K ≃ₗ[R] M
| 0 := ⟨⊤, linear_equiv.of_top ⊤ rfl⟩
| (n+1) :=
⟨(submodule.fst R M N).map (tunnel_aux f (tunnel' n)),
  ((submodule.fst R M N).equiv_map_of_injective _ (tunnel_aux_injective f i (tunnel' n))).symm.trans
    (submodule.fst_equiv R M N)⟩

/--
Give an injective map `f : M × N →ₗ[R] M` we can find a nested sequence of submodules
all isomorphic to `M`.
-/
def tunnel (f : M × N →ₗ[R] M) (i : injective f) : ℕ →o (submodule R M)ᵒᵈ :=
⟨λ n, order_dual.to_dual (tunnel' f i n).1, monotone_nat_of_le_succ (λ n, begin
    dsimp [tunnel', tunnel_aux],
    rw [submodule.map_comp, submodule.map_comp],
    apply submodule.map_subtype_le,
  end)⟩

/--
Give an injective map `f : M × N →ₗ[R] M` we can find a sequence of submodules
all isomorphic to `N`.
-/
def tailing (f : M × N →ₗ[R] M) (i : injective f) (n : ℕ) : submodule R M :=
(submodule.snd R M N).map (tunnel_aux f (tunnel' f i n))

/-- Each `tailing f i n` is a copy of `N`. -/
def tailing_linear_equiv (f : M × N →ₗ[R] M) (i : injective f) (n : ℕ) : tailing f i n ≃ₗ[R] N :=
((submodule.snd R M N).equiv_map_of_injective _
  (tunnel_aux_injective f i (tunnel' f i n))).symm.trans (submodule.snd_equiv R M N)

lemma tailing_le_tunnel (f : M × N →ₗ[R] M) (i : injective f) (n : ℕ) :
  tailing f i n ≤ (tunnel f i n).of_dual :=
begin
  dsimp [tailing, tunnel_aux],
  rw [submodule.map_comp, submodule.map_comp],
  apply submodule.map_subtype_le,
end

lemma tailing_disjoint_tunnel_succ (f : M × N →ₗ[R] M) (i : injective f) (n : ℕ) :
  disjoint (tailing f i n) (tunnel f i (n+1)).of_dual :=
begin
  rw disjoint_iff,
  dsimp [tailing, tunnel, tunnel'],
  rw [submodule.map_inf_eq_map_inf_comap,
    submodule.comap_map_eq_of_injective (tunnel_aux_injective _ i _), inf_comm,
    submodule.fst_inf_snd, submodule.map_bot],
end

lemma tailing_sup_tunnel_succ_le_tunnel (f : M × N →ₗ[R] M) (i : injective f) (n : ℕ) :
  tailing f i n ⊔ (tunnel f i (n+1)).of_dual ≤ (tunnel f i n).of_dual :=
begin
  dsimp [tailing, tunnel, tunnel', tunnel_aux],
  rw [←submodule.map_sup, sup_comm, submodule.fst_sup_snd, submodule.map_comp, submodule.map_comp],
  apply submodule.map_subtype_le,
end

/-- The supremum of all the copies of `N` found inside the tunnel. -/
def tailings (f : M × N →ₗ[R] M) (i : injective f) : ℕ → submodule R M :=
partial_sups (tailing f i)

@[simp] lemma tailings_zero (f : M × N →ₗ[R] M) (i : injective f) :
  tailings f i 0 = tailing f i 0 :=
by simp [tailings]

@[simp] lemma tailings_succ (f : M × N →ₗ[R] M) (i : injective f) (n : ℕ) :
  tailings f i (n+1) = tailings f i n ⊔ tailing f i (n+1) :=
by simp [tailings]

lemma tailings_disjoint_tunnel (f : M × N →ₗ[R] M) (i : injective f) (n : ℕ) :
  disjoint (tailings f i n) (tunnel f i (n+1)).of_dual :=
begin
  induction n with n ih,
  { simp only [tailings_zero],
    apply tailing_disjoint_tunnel_succ, },
  { simp only [tailings_succ],
    refine disjoint.disjoint_sup_left_of_disjoint_sup_right _ _,
    apply tailing_disjoint_tunnel_succ,
    apply disjoint.mono_right _ ih,
    apply tailing_sup_tunnel_succ_le_tunnel, },
end

lemma tailings_disjoint_tailing (f : M × N →ₗ[R] M) (i : injective f) (n : ℕ) :
  disjoint (tailings f i n) (tailing f i (n+1)) :=
disjoint.mono_right (tailing_le_tunnel f i _) (tailings_disjoint_tunnel f i _)

end tunnel

section graph

variables [semiring R] [add_comm_monoid M] [add_comm_monoid M₂]
  [add_comm_group M₃] [add_comm_group M₄] [module R M] [module R M₂]
  [module R M₃] [module R M₄] (f : M →ₗ[R] M₂) (g : M₃ →ₗ[R] M₄)

/-- Graph of a linear map. -/
def graph : submodule R (M × M₂) :=
{ carrier := {p | p.2 = f p.1},
  add_mem' := λ a b (ha : _ = _) (hb : _ = _),
  begin
    change _ + _ = f (_ + _),
    rw [map_add, ha, hb]
  end,
  zero_mem' := eq.symm (map_zero f),
  smul_mem' := λ c x (hx : _ = _),
  begin
    change _ • _ = f (_ • _),
    rw [map_smul, hx]
  end }

@[simp] lemma mem_graph_iff (x : M × M₂) : x ∈ f.graph ↔ x.2 = f x.1 := iff.rfl

lemma graph_eq_ker_coprod : g.graph = ((-g).coprod linear_map.id).ker :=
begin
  ext x,
  change _ = _ ↔ -(g x.1) + x.2 = _,
  rw [add_comm, add_neg_eq_zero]
end

lemma graph_eq_range_prod : f.graph = (linear_map.id.prod f).range :=
begin
  ext x,
  exact ⟨λ hx, ⟨x.1, prod.ext rfl hx.symm⟩, λ ⟨u, hu⟩, hu ▸ rfl⟩
end

end graph

end linear_map


import Mathlib.Tactic.NormNum
import Mathlib.Tactic.Linarith
import Mathlib.Data.Nat.Basic
import Mathlib.Data.Int.Basic

namespace LC

-- pure simply typed lambda calculus
structure Tensor1d (α : Type) where
  size : Nat
  val :  Nat → α

def Tensor1d.empty [Inhabited α] : Tensor1d α where
  size := 0
  val := fun _ => default

def Tensor1d.extract (t: Tensor1d α) (l: Nat) (len: Nat) : Tensor1d α where
  size := len
  val := fun ix => t.val (l + ix)

def Tensor1d.map (f : α → α) (t : Tensor1d α) : Tensor1d α where
  size := t.size
  val := fun ix => f (t.val ix)

def Tensor1d.fill (t: Tensor1d α) (v: α) : Tensor1d α where
  size := t.size
  val := fun _ix => v

-- insert a slice into a tensor.
def Tensor1d.insertslice (t: Tensor1d α) (ix: Nat) (slice : Tensor1d α) : Tensor1d α where
  size := t.size + slice.size
  val := fun i =>
    if i < ix then t.val i
    else if i < ix + slice.size then slice.val (i - ix)
    else t.val (i - slice.size)

-- | TODO: implement fold
def Tensor1d.fold_rec (n: Nat) (arr: Nat → α) (f: β → α → β) (seed: β): β :=
  match n with
  | 0 => seed
  | n + 1 => f (Tensor1d.fold_rec n arr f seed) (arr n)

def Tensor1d.fold (f : β → α → β)  (seed : β) (t : Tensor1d α) : β :=
  Tensor1d.fold_rec t.size t.val f seed

structure Tensor2d (α : Type) where
  size : Nat × Nat
  val :  Nat × Nat → α

def Tensor2d.transpose (t: Tensor2d α) : Tensor2d α where
  size := (t.size.2, t.size.1)
  val := fun ix => t.val (ix.2, ix.1)

-- theorem 1: extract (map) = map extract

theorem Tensor1d.extract_map (t: Tensor1d α):
  (t.extract left len).map f = (t.map f).extract left len := by {
    simp[Tensor1d.extract, Tensor1d.map]
}

-- theorem 2: extract (fill v) = fill (extract v)

theorem Tensor1d.extract_fill (t: Tensor1d α):
  (t.extract left len).fill v = (t.fill v).extract left len := by {
    simp[Tensor1d.extract, Tensor1d.fill]
}

-- theorem 3 : map fusion -- map (f ∘ g) = map f ∘ map g
theorem Tensor1d.map_fusion (t: Tensor1d α):
  (t.map (g ∘ f)) = (t.map f).map g := by {
    simp[Tensor1d.map]
}

-- for loop
def scf.for.loop (f : Nat → β → β) (n n_minus_i: Nat) (acc: β) : β :=
  let i := n - n_minus_i
  match n_minus_i with
    | 0 => acc
    | n_minus_i' + 1 =>
      scf.for.loop f n n_minus_i' (f i acc)

def scf.for (n: Nat) (f: Nat → β → β) (seed: β) : β :=
  let i := 0
  scf.for.loop f n (n - i) seed

-- theorem 1 : for peeling at beginning
theorem scf.for.peel_begin (n : Nat) (f : Nat → β → β) (seed : β) :
  scf.for.loop f (n + 1) n (f 0 seed) = scf.for.loop f (n + 1) (n + 1) seed := by {
    simp[scf.for.loop]
  }

-- theorem 2 : for peeling at ending
theorem scf.for.peel_end (n : Nat) (f : Nat → β → β) (seed : β) :
  scf.for.loop f (n + 1) 0 (f n seed) = f n (scf.for.loop f n 0 seed) := by {
    simp[scf.for.loop]
  }

-- theorem 3: for fusion: if computations commute, then they can be fused.
-- TODO:
theorem scf.for.fusion (n : Nat) (f g : Nat → β → β)  (seed : β)
  (COMMUTE : ∀ (ix : ℕ)  (v : β),  f ix (g ix v) = g ix (f ix v)) :
  scf.for.loop f n n (scf.for.loop g n n seed) =
  scf.for.loop (fun i acc => f i (g i acc)) n n seed := by {
    induction n;
    case zero => {
      simp[loop];
    }
    case succ n' IH => {
      simp[loop];
      sorry
    }
  }

theorem scf.for.zero_n (f: Nat → β → β) (seed : β) :
  scf.for 0 f seed = seed := by {
    simp[scf.for, loop]
  }

  def scf.for.one_n (f: Nat → β → β) (seed : β) :
  scf.for 1 f seed = f 0 seed := by {
    simp[scf.for, loop]
  }

-- theorem 3 : arbitrary for peeling
theorem scf.for.peel_add (n m : Nat) (f : Nat → β → β) (seed : β)  :
  scf.for.loop f (n + m) ((n + m) - n) (scf.for.loop f n (n - 0) seed) = scf.for.loop f (n + m) (n + m - 0) seed := by {
    simp[scf.for.loop]
    revert m;
    induction n;
    case zero => {
      simp[loop]
    }
    case succ n' IH => {
      intros m;
      simp[loop];
      sorry
    }
  }

-- theorem 4 : tiling
-- proof obligation for chris :)
theorem Tensor1d.tile [Inhabited α] (t : Tensor1d α) (SIZE :4 ∣ t.size) (f : α → α):
  t.map f = scf.for (t.size / 4) (fun i acc =>
    let tile := t.extract (i * 4) 4
    let mapped_tile := tile.map f
    let out := acc.insertslice (i * 4) mapped_tile
    out) (Tensor1d.empty) := by {
    cases t;
    case mk size val =>
    simp at SIZE ⊢;
    have : { n : Nat //  size = n * 4 } := by {
      norm_num at SIZE
      -- have ⟨x, y⟩ := SIZE
      sorry
    }
    have ⟨n, N⟩ := this
    rw[N];
    revert size
    induction n
    case zero => {
      simp[scf.for, scf.for.loop]
      sorry
    }
    case succ n IH => {
      sorry
    }
}

-- transpose is an involution

theorem Tensor2d.transpose_involutive (t: Tensor2d α):
  (t.transpose).transpose = t := by {
    simp[Tensor2d.transpose]
}

inductive Val where
| int : Int → Val
| unit : Val
| nat : Nat → Val
| bool : Bool → Val
| tensor1d : Tensor1d Int → Val
| tensor2d : Tensor2d Int → Val
| pair : Val → Val → Val
| triple : Val → Val → Val → Val
| inl : Val → Val
| inr : Val → Val
deriving Inhabited

def Val.int! : Val → Int
| .int i => i
| _ => default

def Val.nat! : Val → Nat
| .nat i => i
| _ => default

def Val.bool! : Val → Bool
| .bool i => i
| _ => default

abbrev Var := Int

abbrev Env (α: Type) := Var → α

def Env.empty {α : Type} [Inhabited α]: Env α := fun _ => default
notation "∅" =>  Env.empty

def Env.set (e: Env α) (var: Var) (val: α) :=
  fun needle => if needle == var then val else e needle
notation e "[" var " := " val "]" => Env.set e var val

-- RHS of an assignment
inductive SSAIndex : Type
| STMT
| EXPR
| TERMINATOR
| REGION

-- NOTE: multiple regions can be converted into a single region by tagging the
-- input appropriately with inl/inr.
inductive SSA (Op: Type): SSAIndex → Type where
| assign (lhs: Var) (rhs: SSA Op .EXPR) (rest: SSA Op .STMT) : SSA Op .STMT
| nop : SSA Op .STMT
| ret (above : SSA Op .STMT) (v: Var): SSA Op .TERMINATOR
| pair (fst snd : Var) : SSA Op .EXPR
| op (o : Op) (arg: Var) (rgn: SSA Op .REGION) : SSA Op .EXPR
| const (k: Val) : SSA Op .EXPR
| rgn (arg: Var) (body: SSA Op .TERMINATOR) : SSA Op .REGION
| rgn0 : SSA Op .REGION
| rgnvar (v: Var) : SSA Op .REGION
| var (v: Var) : SSA Op .EXPR

abbrev Expr (Op: Type) := SSA Op .EXPR
abbrev Stmt (Op: Type) := SSA Op .STMT

class UserSemantics (Op: Type) where
  eval: (o: Op) → (arg: Val) → (rgn: Val → Val) → Val

def SSAIndex.eval : SSAIndex → Type
| .STMT => Env Val
| .TERMINATOR => Val
| .EXPR => Val
| .REGION => Val -> Val

def SSA.eval [S : UserSemantics Op] (e: Env Val) (re: Env (Val → Val)) : SSA Op k → k.eval
| .assign lhs rhs rest =>
  rest.eval (e.set lhs (rhs.eval e re)) re
| .nop => e
| .ret above v => (above.eval e re) v
| .pair fst snd => (e fst).pair (e snd)
| .const v => v
| .op o arg r => S.eval o (e arg) (r.eval Env.empty re)
| .var v => e v
| .rgnvar v => re v
| .rgn0 => id
| .rgn arg body => fun val => body.eval (e.set arg val) re

namespace ArithScfLinalg

inductive op
| add
| sub
| mul
| run
| for_
| if_
| fold1d -- fold
| map1d
| extract
| fill
| transpose
deriving DecidableEq

instance : UserSemantics op where
  eval
  | .add, .pair (.int x) (.int y), _ => .int (x + y)
  | .sub, .pair (.int x) (.int y), _ => .int (x - y)
  | .run, v, r => r v
  | .if_, (.bool cond), r => if cond then r (.inl .unit) else r (.inr .unit)
  | .for_, (.pair (.nat n) (.int seed)), r =>
      .int <| scf.for n (fun ix acc => (r (.pair (.int ix) (.int acc))).int!) seed
  | .map1d, (.tensor1d t), r => .tensor1d <| t.map fun v => (r (.int v)).int!
  | .fold1d, (.pair (.tensor1d t) (.int seed)), r =>
      .int <| t.fold (fun acc v => (r (.pair (.int acc) (.int v))).int!) seed
  | .extract, (.triple (.tensor1d t) (.nat l) (.nat len)), _ =>
      .tensor1d <| t.extract l len
  | _, _, _ => default

-- TODO: port Hacker's delight examples.

end ArithScfLinalg

end LC


import Lean
import Init.Classical

import SciLean.Core.Differential
import SciLean.Core.Adjoint
import SciLean.Core.HasAdjDiff

namespace SciLean

variable {α β γ : Type}
variable {X Y Z : Type} [SemiHilbert X] [SemiHilbert Y] [SemiHilbert Z]
variable {Y₁ Y₂ : Type} [SemiHilbert Y₁] [SemiHilbert Y₂]
variable {ι : Type} [Enumtype ι]

-- noncomputable
-- def adjointDifferential (f : X → Y) (x : X) (dy' : Y) : X := (∂ f x)† dy'

-- @[default_instance]
-- instance (f : X → Y) : PartialDagger f (adjointDifferential f) := ⟨⟩

-- Someting wrong here :(
-- noncomputable
-- def Smooth.adjointDifferential {X Y} [Hilbert X] [Hilbert Y] (f : X ⟿ Y) : X⟿Y⊸X := λ x ⟿ λ dy ⊸ adjoint (∂ f x) dy

-- @[default_instance]
-- instance (f : X → Y) : PartialDagger f (adjointDifferential f) := ⟨⟩

-- Question: Should there be `𝒯[y] Y` or `𝒯[f x] Y`?
-- Maybe return `(y:Y)×(𝒯[y] Y → 𝒯[x] X)×(f x = y)` but there is a problem with `Sigma` vs `PSigma`
-- noncomputable
-- def reverseDifferential (f : X → Y) (x : X) : Y×(Y→X) := (f x, λ dy => ∂† f x dy)

-- instance (priority:=low) (f : X → Y) : ReverseDifferential f (reverseDifferential f) := ⟨⟩

-- noncomputable
-- abbrev gradient (f : X → ℝ) (x : X) : X := ∂† f x 1

-- @[default_instance]
-- instance (f : X → ℝ) : Nabla f (gradient f) := ⟨⟩

-- noncomputable
-- abbrev Smooth.gradient (f : X ⟿ ℝ) : X⟿X := SmoothMap.mk (λ x => adjoint (λ dx => ∂ f x dx) 1) sorry_proof

-- instance (f : X ⟿ ℝ) : Nabla f (Smooth.gradient f) := ⟨⟩

-- Notation
-- ∇ s, f s         --> ∇ λ s => f s
-- ∇ s : ℝ, f s     --> ∇ λ s : ℝ => f s
-- ∇ s := t, f s    --> (∇ λ s => f s) t
syntax "∇" diffBinder "," term:66 : term
syntax "∇" "(" diffBinder ")" "," term:66 : term
macro_rules
| `(∇ $x:ident, $f) =>
  `(∇ λ $x => $f)
| `(∇ $x:ident : $type:term, $f) =>
  `(∇ λ $x : $type => $f)
| `(∇ $x:ident := $val:term, $f) =>
  `((∇ λ $x => $f) $val)
| `(∇ ($b:diffBinder), $f) =>
  `(∇ $b, $f)

instance (f : X → Y) [HasAdjDiff f] (x : X) : IsLin (∂† f x) := sorry

----------------------------------------------------------------------

@[simp ↓, diff]
theorem id.arg_x.adjDiff_simp
  : ∂† (λ x : X => x) = λ x dx => dx := by symdiff; simp[adjointDifferential]; done

@[simp ↓, diff]
theorem const.arg_x.adjDiff_simp
  : ∂† (λ (x : X) (i : ι) => x) = λ x f => ∑ i, f i := by simp[adjointDifferential]; done

@[simp ↓, diff]
theorem const.arg_y.adjDiff_simp (x : X)
  : ∂† (λ (y : Y) => x) = (λ y dy' => (0 : Y)) := by simp[adjointDifferential]; done

@[simp ↓ low-4, diff low-4]
theorem swap.arg_y.adjDiff_simp
  (f : ι → X → Z) [inst : ∀ i, HasAdjDiffT (f i)]
  : ∂† (λ x y => f y x) = (λ x dx' => ∑ i, (∂† (f i) x) (dx' i)) :=
by
  have := λ i => (inst i).1
  have := λ i => (inst i).2

  simp[adjointDifferential]; done

@[simp ↓ low-3, diff low-3]
theorem subst.arg_x.adjDiff_simp
  (f : X → Y → Z) [instf : HasAdjDiffNT 2 f]
  (g : X → Y) [instg : HasAdjDiffT g]
  : ∂† (λ x => f x (g x))
    =
    λ x dz =>
      let (y,dg') := ℛ g x
      -- let (dx,dy) := ∂† (uncurryN 2 f) (x,y) dz
      -- dx + dg' dy
      (∂† (λ x' => f x' y)) x dz
      +
      dg' (∂† (f x) y dz)
    :=
by
  have := instg.1
  have := instg.2
  have := instf.1
  -- these follow from instf.proof.2
  have : ∀ x y, HasAdjointT (λ dx => ∂ f x dx y) := sorry_proof
  have : ∀ x y, HasAdjointT (λ dy => ∂ (f x) y dy) := sorry_proof

  unfold adjointDifferential -- reverseDifferential, tangentMap, -comp.arg_x.parm1.adj_simp]
  sorry -- symdiff
    --sorry_proof
  -- simp (config := {singlePass := true})
  -- done

@[simp ↓ low-2, diff low-2, simp_guard g (λ x => x)]
theorem subst.arg_x.parm1.adjDiff_simp
  (a : α)
  (f : X → Y → α → Z) [HasAdjDiffNT 2 λ x y => f x y a]
  (g : X → Y) [instg : HasAdjDiffT g]
  : ∂† (λ x => f x (g x) a)
    =
    λ x dz =>
      let (y,dg') := ℛ g x
      -- let (dx,dy) := ∂† (uncurryN 2 (λ x y => f x y a)) (x,y) dz
      -- dx + dg' dy
      (∂† (λ x' => f x' y a)) x dz
      +
      dg' (∂† (λ y' => f x y' a) y dz)
    :=
by
  rw[subst.arg_x.adjDiff_simp (λ x y => f x y a) g]
  done

@[simp ↓ low-2, diff low-2, simp_guard g (λ x => x)]
theorem subst.arg_x.parm2.adjDiff_simp
  (a : α) (b : β)
  (f : X → Y → α → β → Z) [HasAdjDiffNT 2 λ x y => f x y a b]
  (g : X → Y) [instg : HasAdjDiffT g]
  : ∂† (λ x => f x (g x) a b)
    =
    λ x dz =>
      let (y,dg') := ℛ g x
      (∂† (λ x' => f x' y a b)) x dz
      +
      dg' (∂† (λ y' => f x y' a b) y dz)
    :=
by
  apply subst.arg_x.adjDiff_simp (λ x y => f x y a b) g
  done

@[simp ↓ low-2, diff low-2, simp_guard g (λ x => x)]
theorem subst.arg_x.parm3.adjDiff_simp
  (a : α) (b : β) (c : γ)
  (f : X → Y → α → β → γ → Z) [HasAdjDiffNT 2 λ x y => f x y a b c]
  (g : X → Y) [instg : HasAdjDiffT g]
  : ∂† (λ x => f x (g x) a b c)
    =
    λ x dz =>
      let (y,dg') := ℛ g x
      (∂† (λ x' => f x' y a b c)) x dz
      +
      dg' (∂† (λ y' => f x y' a b c) y dz)
    :=
by
  apply subst.arg_x.adjDiff_simp (λ x y => f x y a b c) g
  done

@[simp ↓ low-1, diff low-1, simp_guard g (λ x => x)]
theorem comp.arg_x.adjDiff_simp
  (f : Y → Z) [instf : HasAdjDiffT f]
  (g : X → Y) [instg : HasAdjDiffT g]
  : ∂† (λ x => f (g x))
    =
    λ x dz =>
      let (y,dg') := ℛ g x
      dg' ((∂† f y) dz)
  := by simp; done

@[simp ↓ low-2, diff low-2, simp_guard g₁ Prod.fst, g₂ Prod.snd]
theorem diag.arg_x.adjDiff_simp
  (f : Y₁ → Y₂ → Z) [HasAdjDiffNT 2 f]
  (g₁ : X → Y₁) [hg : HasAdjDiffT g₁]
  (g₂ : X → Y₂) [HasAdjDiffT g₂]
  : ∂† (λ x => f (g₁ x) (g₂ x))
    =
    λ x dz =>
      let (y₁,dg₁') := ℛ g₁ x
      let (y₂,dg₂') := ℛ g₂ x
      dg₁' ((∂† λ y₁' => f y₁' y₂) y₁ dz)
      +
      dg₂' ((∂† λ y₂' => f y₁ y₂') y₂ dz)
    :=
by
  rw[subst.arg_x.adjDiff_simp]
  simp only [hold,reverseDifferential]
  funext x dz
  rw[comp.arg_x.adjDiff_simp (λ y₁ => f y₁ (g₂ x))]
  simp only [reverseDifferential]
  done

@[simp ↓ low, diff low]
theorem eval.arg_f.adjDiff_simp
  (i : ι)
  : ∂† (λ (f : ι → X) => f i)
    =
    (λ f df' j => ([[i = j]] • df' : X))
:= sorry

@[simp ↓ low-1, diff low-1]
theorem eval.arg_x.parm1.adjDiff_simp
  (f : X → ι → Z) [HasAdjDiff f]
  : ∂† (λ x => f x i)
    =
    (λ x dx' => (∂† f x) (λ j => ([[i = j]] • dx' : Z)))
:=
by
  rw [comp.arg_x.adjDiff_simp (λ (x : ι → Z) => x i) f]
  simp[reverseDifferential]

--------------------------------------------------------
-- These theorems are problematic when used with simp --

@[simp ↓ low-1, diff low-1]
theorem comp.arg_x.parm1.adjDiff_simp
  (a : α)
  (f : Y → α → Z) [HasAdjDiff λ y => f y a]
  (g : X → Y) [HasAdjDiff g]
  :
    ∂† (λ x => f (g x) a)
    =
    λ x dz =>
      let (y,dg') := ℛ g x
      dg' ((∂† (hold λ y => f y a)) y dz)
:= by
  rw[subst.arg_x.parm1.adjDiff_simp]
  simp[-subst.arg_x.parm1.adjDiff_simp,hold]
  done

@[simp ↓ low-1, diff low-1]
theorem comp.arg_x.parm2.adjDiff_simp
  (a : α) (b : β)
  (f : Y → α → β → Z) [HasAdjDiff λ y => f y a b]
  (g : X → Y) [HasAdjDiff g]
  :
    ∂† (λ x => f (g x) a b)
    =
    λ x dz =>
      let (y,dg') := ℛ g x
      dg' ((∂† (hold λ y => f y a b)) y dz)
:= by
  rw[subst.arg_x.parm2.adjDiff_simp]
  simp[-subst.arg_x.parm2.adjDiff_simp,hold]
  done

@[simp ↓ low-1, diff low-1]
theorem comp.arg_x.parm3.adjDiff_simp
  (a : α) (b : β) (c : γ)
  (f : Y → α → β → γ → Z) [HasAdjDiff λ y => f y a b c]
  (g : X → Y) [HasAdjDiff g]
  :
    ∂† (λ x => f (g x) a b c)
    =
    λ x dx' =>
      let (y,dg') := ℛ g x
      dg' ((∂† (hold λ y => f y a b c)) y dx')
:= by
  rw[subst.arg_x.parm3.adjDiff_simp]
  simp[-subst.arg_x.parm3.adjDiff_simp,hold]
  done

-- TODO: fix this!!!
example (a : α) (f : Y₁ → Y₂ → α → Z) [IsSmoothT λ y₁ y₂ => f y₁ y₂ a]
  (g₁ : X → Y₁) [hg : IsSmoothT g₁] : IsSmoothT (λ x y => f (g₁ x) y a) := by (try infer_instance); admit

@[simp ↓ low-1, diff low-1] -- try to avoid using this theorem
theorem diag.arg_x.parm1.adjDiff_simp
  (a : α)
  (f : Y₁ → Y₂ → α → Z) [HasAdjDiffNT 2 λ y₁ y₂ => f y₁ y₂ a]
  (g₁ : X → Y₁) [HasAdjDiffT g₁]
  (g₂ : X → Y₂) [HasAdjDiffT g₂]
  : ∂† (λ x => f (g₁ x) (g₂ x) a)
    =
    λ x dz =>
      let (y₁,dg₁') := ℛ g₁ x
      let (y₂,dg₂') := ℛ g₂ x
      dg₁' ((∂† (hold λ y₁' => f y₁' y₂ a)) y₁ dz)
      +
      dg₂' ((∂† (hold λ y₂' => f y₁ y₂' a)) y₂ dz)
:= by
  (apply diag.arg_x.adjDiff_simp (λ y₁ y₂ => f y₁ y₂ a) g₁ g₂)

@[simp ↓ low-1, diff low-1] -- try to avoid using this theorem
theorem diag.arg_x.parm2.adjDiff_simp
  (a : α) (b : β)
  (f : Y₁ → Y₂ → α → β → Z) [HasAdjDiffNT 2 λ y₁ y₂ => f y₁ y₂ a b]
  (g₁ : X → Y₁) [HasAdjDiffT g₁]
  (g₂ : X → Y₂) [HasAdjDiffT g₂]
  : ∂† (λ x => f (g₁ x) (g₂ x) a b)
    =
    λ x dz =>
      let (y₁,dg₁') := ℛ g₁ x
      let (y₂,dg₂') := ℛ g₂ x
      dg₁' ((∂† (hold λ y₁' => f y₁' y₂ a b)) y₁ dz)
      +
      dg₂' ((∂† (hold λ y₂' => f y₁ y₂' a b)) y₂ dz)
:= by
  (apply diag.arg_x.adjDiff_simp (λ y₁ y₂ => f y₁ y₂ a b) g₁ g₂)
  done

@[simp ↓ low-1, diff low-1] -- try to avoid using this theorem
theorem diag.arg_x.parm3.adjDiff_simp
  (a : α) (b : β) (c : γ)
  (f : Y₁ → Y₂ → α → β → γ → Z) [HasAdjDiffNT 2 λ y₁ y₂ => f y₁ y₂ a b c]
  (g₁ : X → Y₁) [HasAdjDiffT g₁]
  (g₂ : X → Y₂) [HasAdjDiffT g₂]
  : ∂† (λ x => f (g₁ x) (g₂ x) a b c)
    =
    λ x dz =>
      let (y₁,dg₁') := ℛ g₁ x
      let (y₂,dg₂') := ℛ g₂ x
      dg₁' ((∂† (hold λ y₁' => f y₁' y₂ a b c)) y₁ dz)
      +
      dg₂' ((∂† (hold λ y₂' => f y₁ y₂' a b c)) y₂ dz)
:= by
  (apply diag.arg_x.adjDiff_simp (λ y₁ y₂ => f y₁ y₂ a b c) g₁ g₂)
  done

----------------------------------------------------------------------

@[simp ↓, diff]
theorem id.arg_x.revDiff_simp
  : ℛ (λ x : X => x) = λ x => (x, λ x => x) := by simp[reverseDifferential]; done

@[simp ↓, diff]
theorem const.arg_x.revDiff_simp
  : ℛ (λ (x : X) (i : ι) => x)
    =
    λ x => ((λ i => x), (λ f => ∑ i, f i))
  := by simp[reverseDifferential]; done

@[simp ↓, diff]
theorem const.arg_y.revDiff_simp (x : X)
  : ℛ (λ (y : Y) => x)
    =
    λ y =>
      (x, (λ dy' => 0))
  := by simp[reverseDifferential]; done

@[simp ↓ low-4, diff low-4]
theorem swap.arg_y.revDiff_simp
  (f : ι → X → Z) [inst : ∀ i, HasAdjDiffT (f i)]
  : ∂† (λ x y => f y x) = (λ x dx' => ∑ i, (∂† (f i) x) (dx' i)) :=
by
  have := λ i => (inst i).1
  have := λ i => (inst i).2

  simp[adjointDifferential]; done

@[simp ↓ low-3, diff low-3, simp_guard g (λ x => x)]
theorem subst.arg_x.revDiff_simp
  (f : X → Y → Z) [instf : HasAdjDiffNT 2 f]
  (g : X → Y) [instg : HasAdjDiffT g]
  : ℛ (λ x => f x (g x))
    =
    λ x =>
      let (y,dg') := ℛ g x
      let (z,df') := ℛ (uncurryN 2 f) (x,y)
      (z, λ dz' =>
           let (dx₁,dy) := df' dz'
           dx₁ + dg' dy)

    :=
by
  have := instg.1
  have := instg.2
  have := instf.1

  funext x;
  unfold reverseDifferential
  rw[subst.arg_x.adjDiff_simp]

  simp only [uncurryN, Prod.Uncurry.uncurry]
  simp only [hold, reverseDifferential]
  conv => (rhs; rw[diag.arg_x.adjDiff_simp])
  simp only [reverseDifferential,
             Prod.fst.arg_xy.adjDiff_simp,
             Prod.snd.arg_xy.adjDiff_simp,
             prod_add_elemwise,
             add_zero, zero_add]
  done

@[simp ↓ low-2, diff low-2, simp_guard g (λ x => x)]
theorem subst.arg_x.parm1.revDiff_simp
  (a : α)
  (f : X → Y → α → Z) [HasAdjDiffNT 2 λ x y => f x y a]
  (g : X → Y) [instg : HasAdjDiffT g]
  : ℛ (λ x => f x (g x) a)
    =
    λ x =>
      let (y,dg') := ℛ g x
      let (z,df') := ℛ (uncurryN 2 (λ x y => f x y a)) (x,y)
      (z, λ dz' =>
           let (dx₁,dy) := df' dz'
           dx₁ + dg' dy)
    :=
by
  apply subst.arg_x.revDiff_simp (λ x y => f x y a) g
  done

@[simp ↓ low-2, diff low-2, simp_guard g (λ x => x)]
theorem subst.arg_x.parm2.revDiff_simp
  (a : α) (b : β)
  (f : X → Y → α → β → Z) [HasAdjDiffNT 2 λ x y => f x y a b]
  (g : X → Y) [instg : HasAdjDiffT g]
  : ℛ (λ x => f x (g x) a b)
    =
    λ x =>
      let (y,dg') := ℛ g x
      let (z,df') := ℛ (uncurryN 2 (λ x y => f x y a b)) (x,y)
      (z, λ dz' =>
           let (dx₁,dy) := df' dz'
           dx₁ + dg' dy)
    :=
by
  apply subst.arg_x.revDiff_simp (λ x y => f x y a b) g
  done

@[simp ↓ low-2, diff low-2, simp_guard g (λ x => x)]
theorem subst.arg_x.parm3.revDiff_simp
  (a : α) (b : β) (c : γ)
  (f : X → Y → α → β → γ → Z) [HasAdjDiffNT 2 λ x y => f x y a b c]
  (g : X → Y) [instg : HasAdjDiffT g]
  : ℛ (λ x => f x (g x) a b c)
    =
    λ x =>
      let (y,dg') := ℛ g x
      let (z,df') := ℛ (uncurryN 2 (λ x y => f x y a b c)) (x,y)
      (z, λ dz' => let (dx₁,dy) := df' dz'; dx₁ + dg' dy)
    :=
by
  apply subst.arg_x.revDiff_simp (λ x y => f x y a b c) g
  done

-- @[simp ↓ low-10, diff low-10]
theorem uncurryN2.arg_x.diff_simp
  (f : X → Y → Z) [HasAdjDiffNT 2 f]
  : ∂† (uncurryN 2 f)
    =
    λ (x,y) dz =>
      (∂† (λ x' => f x' y) x dz, ∂† (λ y' => f x y') y dz)
  := sorry_proof

@[simp ↓ low-1, diff low-1]
theorem comp.arg_x.revDiff_simp
  (f : Y → Z) [instf : HasAdjDiffT f]
  (g : X → Y) [instg : HasAdjDiffT g]
  : ℛ (λ x => f (g x))
    =
    λ x =>
      let (y,dg') := ℛ g x
      let (z,df') := ℛ f y
      (z, λ dz => dg' (df' dz)) :=
by
  unfold reverseDifferential
  simp only [comp.arg_x.adjDiff_simp]
  simp only [reverseDifferential]
  done

@[simp ↓ low-2, diff low-2]
theorem diag.arg_x.revDiff_simp
  (f : Y₁ → Y₂ → Z) [HasAdjDiffNT 2 f]
  (g₁ : X → Y₁) [hg : HasAdjDiffT g₁]
  (g₂ : X → Y₂) [HasAdjDiffT g₂]
  : ℛ (λ x => f (g₁ x) (g₂ x))
    =
    λ x =>
      let (y₁,dg₁') := ℛ g₁ x
      let (y₂,dg₂') := ℛ g₂ x
      let (z, df') := ℛ (uncurryN 2 f) (y₁,y₂)
      (z, λ dz => let (dy₁,dy₂) := df' dz; dg₁' dy₁ + dg₂' dy₂)
      -- dg₁' ((∂† λ y₁ => f y₁ y₂) y₁ (h₁ ▸ h₂ ▸ dx'))
      -- +
      -- dg₂' ((∂† λ y₂ => f y₁ y₂) y₂ (h₂ ▸ h₁ ▸ dx'))
    :=
by
  unfold reverseDifferential
  funext x
  simp only [uncurryN, Prod.Uncurry.uncurry]
  conv => lhs; enter [2,dz]; rw [diag.arg_x.adjDiff_simp]
  conv => rhs; enter [2,dz]; rw [diag.arg_x.adjDiff_simp]
  simp only [reverseDifferential,
             Prod.fst.arg_xy.adjDiff_simp,
             Prod.snd.arg_xy.adjDiff_simp,
             prod_add_elemwise,
             add_zero, zero_add]
  done

@[simp ↓ low, diff low]
theorem eval.arg_f.revDiff_simp
  (i : ι)
  : ℛ (λ (f : ι → X) => f i)
    =
    λ f => (f i, (λ dx j => ([[i=j]] • dx : X)))
  :=
by
  simp[reverseDifferential,adjointDifferential]; done

@[simp ↓ low-1, diff low-1]
theorem eval.arg_x.parm1.revDiff_simp
  (f : X → ι → Z) [HasAdjDiff f] (i : ι)
  : ℛ (λ x => f x i)
    =
    λ x =>
      let (fx, df') := ℛ f x
      (fx i,
      λ dz => df' (λ j => ([[i=j]] • dz)))
  :=
by
  rw [comp.arg_x.revDiff_simp (λ (x : ι → Z) => x i) f]
  simp[reverseDifferential,adjointDifferential]

-- @[simp ↓]
-- theorem subst.arg_x.revDiff_simp'''
--   (f : X → Y → Z) [IsSmooth f]
--   [instfx : ∀ y, HasAdjDiff λ x => f x y]
--   [instfy : ∀ x, HasAdjDiff (f x)]
--   (g : Y → X) [instg : HasAdjDiff g]
--   : ∂† (λ y => f (g y) y)
--     =
--     λ y dy' =>
--       (∂† (λ y' => f (g y) y')) y dy'
--       +
--       (∂† g y) (∂† (λ x => f x y) (g y) dy')
--     :=
-- by
--   sorry