-- SSA/Projects/FullyHomomorphicEncryption/Basic.lean

import Mathlib.RingTheory.Polynomial.Quotient

import Mathlib.RingTheory.Ideal.Defs

import Mathlib.RingTheory.Ideal.Basic

import Mathlib.Data.ZMod.Defs

import Mathlib.Data.ZMod.Basic

import Mathlib.Algebra.MonoidAlgebra.Basic

import Mathlib.Algebra.Polynomial.RingDivision

import Mathlib.Data.Finset.Sort

import Mathlib.Data.List.ToFinsupp

import Mathlib.Data.List.Basic

import Mathlib.Tactic.Cases

import LeanMLIR.Framework

open Polynomial -- for R[X] notation

section CommRing

variable (q t : Nat) [Fact (q > 1)] (n : Nat)

noncomputable def f : (ZMod q)[X] := X^(2^n) + 1

abbrev R := (ZMod q)[X] ⧸ (Ideal.span {f q n})

abbrev R.fromPoly {q n : Nat} : (ZMod q)[X] →+* R q n := Ideal.Quotient.mk (Ideal.span {f q n})

end CommRing

section Representative

variable (q t n : Nat)

private noncomputable def R.representative' :
    R q n → (ZMod q)[X] := Function.surjInv (R.surjective_fromPoly q n)

noncomputable def R.representative :
    R q n → (ZMod q)[X] := fun x => R.representative' q n x %ₘ (f q n)

variable [Fact (q > 1)]

end Representative

section Coeff

noncomputable def R.coeff {q n} (a : R q n) (i : Nat) : ZMod q :=
  Polynomial.coeff a.representative i

noncomputable def R.monomial {q n : Nat} (c : ZMod q) (i : Nat): R q n :=
  R.fromPoly (Polynomial.monomial i c)

noncomputable def R.fromTensor {q n} (coeffs : List Int) : R q n :=
  coeffs.zipIdx.foldl (init := 0) fun res (c, i) =>
    res + R.monomial ↑c i

end Coeff

section FinnSupp

variable {q n : Nat}

noncomputable def R.fromTensorFinsupp (q : Nat) (coeffs : List Int) : (ZMod q)[X] :=
  Polynomial.ofFinsupp (List.toFinsupp (coeffs.map Int.cast))

end FinnSupp

section Tensor

variable {q n : Nat} [Fact (q > 1)]

-- Repository definitions used:
content: ...

-- Repository lemmas used:
content: theorem List.getElem_append_left{α : Type u_1} {l₁ l₂ : List α} {i : Nat} (hn : i < l₁.length) :\n(l₁ ++ l₂)[i] = l₁[i]

-- Local definitions used:
content: noncomputable def f : (ZMod q)[X] := X^(2^n) + 1

content: abbrev R := (ZMod q)[X] ⧸ (Ideal.span {f q n})

content: abbrev R.fromPoly {q n : Nat} : (ZMod q)[X] →+* R q n := Ideal.Quotient.mk (Ideal.span {f q n})

content: private noncomputable def R.representative' :
    R q n → (ZMod q)[X] := Function.surjInv (R.surjective_fromPoly q n)

content: noncomputable def R.representative :
    R q n → (ZMod q)[X] := fun x => R.representative' q n x %ₘ (f q n)

content: noncomputable def R.coeff {q n} (a : R q n) (i : Nat) : ZMod q :=
  Polynomial.coeff a.representative i

content: noncomputable def R.monomial {q n : Nat} (c : ZMod q) (i : Nat): R q n :=
  R.fromPoly (Polynomial.monomial i c)

content: noncomputable def R.fromTensor {q n} (coeffs : List Int) : R q n :=
  coeffs.zipIdx.foldl (init := 0) fun res (c, i) =>
    res + R.monomial ↑c i

content: noncomputable def R.fromTensorFinsupp (q : Nat) (coeffs : List Int) : (ZMod q)[X] :=
  Polynomial.ofFinsupp (List.toFinsupp (coeffs.map Int.cast))

-- Local lemmas used:
content: theorem f_deg_eq : (f q n).degree = 2^n

content: theorem f_monic : Monic (f q n)

content: theorem R.fromPoly_rep'_eq_ideal :
    forall a : (ZMod q)[X],
      ∃ i ∈ Ideal.span {f q n}, (R.fromPoly (n:=n) a).representative' = a + i

content: theorem R.representative_fromPoly :
    forall a : (ZMod q)[X], (R.fromPoly (n:=n) a).representative = a %ₘ (f q n)

content: @[simp]
theorem R.representative_fromPoly_eq (x : (ZMod q)[X]) (DEGREE: x.degree < (f q n).degree) :
   R.representative q n (R.fromPoly (n:=n) x) = x

content: theorem Polynomial.degree_toFinsupp [Semiring M] [DecidableEq M]
  (xs : List M) :
  degree { toFinsupp := List.toFinsupp (l := xs) } ≤ List.length xs

content: theorem R.fromTensorFinsupp_degree (q : Nat) (coeffs : List Int):
  (R.fromTensorFinsupp q coeffs).degree ≤ coeffs.length

content: theorem R.fromTensorFinsupp_coeffs (coeffs : List Int) :
  Polynomial.coeff (fromTensorFinsupp q coeffs) i = ↑(List.getD coeffs i 0)

content: theorem R.fromTensorFinsupp_concat_monomial (c : Int) (cs : List Int) :
    (R.fromTensorFinsupp q (cs ++ [c])) =
      (R.fromTensorFinsupp q cs) +
        (Polynomial.monomial cs.length (Int.cast c : (ZMod q)))

content: theorem R.fromTensor_eq_fromTensorFinsupp_fromPoly {coeffs : List Int} :
    R.fromTensor (q := q) (n := n) coeffs =
  R.fromPoly (q := q) (n := n) (R.fromTensorFinsupp q coeffs)

-- Target theorem and proof:
theorem R.coeff_fromTensor (tensor : List Int)
    (htensorlen : tensor.length < 2^n) :
    (R.fromTensor (q := q) (n := n) tensor).coeff i = (tensor.getD i 0)
:= by
  rw [fromTensor_eq_fromTensorFinsupp_fromPoly]
  have hfromTensorFinsuppDegree := fromTensorFinsupp_degree q tensor
  rw [coeff, representative_fromPoly_eq]
  apply fromTensorFinsupp_coeffs
  case DEGREE =>
    generalize htensor_degree : degree (fromTensorFinsupp q tensor) = tensor_degree
    rw [f_deg_eq]
    cases tensor_degree
    case bot =>
      rw [WithBot.lt_def]
      simp
      exists 2^n
    norm_cast
    case coe tensor_degree =>
      /- I hate this coercion stuff -/
      apply WithBot.coe_strictMono
      norm_cast
      have htrans : tensor_degree  ≤ List.length tensor := by
        rw [htensor_degree] at hfromTensorFinsuppDegree
        rw [← WithBot.coe_le_coe]
        assumption
      apply Nat.lt_of_le_of_lt htrans htensorlen