-- SSA/Projects/FullyHomomorphicEncryption/Statements.lean

import SSA.Projects.FullyHomomorphicEncryption.Basic

import Batteries.Data.List.Lemmas

import Mathlib.Data.List.Basic

namespace Poly

open Polynomial in

open Polynomial in

end Poly

namespace Poly

-- Repository definitions used:
content: syntax "neg" : MLIR.Pretty.uniform_op

content: macro_rules
  | `(mlir_op| $res:mlir_op_operand = const ($x)
      $[: $outer_type]? ) => do
      let outer_type ← outer_type.getDM `(mlir_type| _)
      `(mlir_op| $res:mlir_op_operand = "const"()
          {val = $x:num : $outer_type} : ($outer_type) -> ($outer_type) )

content: macro_rules
  | `(mlir_op| $res:mlir_op_operand = li ($x)
     $[: $outer_type]?  ) => do
      let outer_type ← outer_type.getDM `(mlir_type| _ )
      `(mlir_op| $res:mlir_op_operand = "li"()
          {imm = $x:num : $outer_type } : ($outer_type) -> ($outer_type))

content: macro_rules
| `(mlir_op| $res:mlir_op_operand = $op1:MLIR.Pretty.RV.opWithImmediate $reg1 , $x : $t)  => do
    let some opName := MLIR.EDSL.Pretty.extractOpName op1.raw
      | Macro.throwUnsupported
    `(mlir_op| $res:mlir_op_operand = $opName ($reg1) {imm = $x:num : $t} : ($t) -> ($t) )

content: macro_rules
| `(mlir_op| $res:mlir_op_operand = $op1:MLIR.Pretty.RV.opWithShamt $reg1 , $x  : $t )  => do
    let some opName := MLIR.EDSL.Pretty.extractOpName op1.raw
      | Macro.throwUnsupported
    `(mlir_op| $res:mlir_op_operand = $opName ($reg1) {shamt = $x:num : $t}  : ($t) -> ($t) )

content: noncomputable def R.repLength {q n} (a : R q n) : Nat := match
  Polynomial.degree a.representative with
    | none => 0
    | some d => d + 1

content: noncomputable def R.representative :
    R q n → (ZMod q)[X] := fun x => R.representative' q n x %ₘ (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 f : (ZMod q)[X] := X^(2^n) + 1

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.monomial {q n : Nat} (c : ZMod q) (i : Nat): R q n :=
  R.fromPoly (Polynomial.monomial i c)

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

content: noncomputable def R.toTensor {q n} (a : R q n) : List Int :=
  List.range a.repLength |>.map fun i =>
        a.coeff i |>.toInt

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

content: def ZMod.toInt (x : ZMod q) : Int := ZMod.cast x

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

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

-- Repository lemmas used:
content: @[simp]
theorem R.fromPoly_kernel_eq_zero (x : (ZMod q)[X]) : R.fromPoly (n := n) (f q n * x) = 0

content: theorem R.representative_fromPoly :
    forall a : (ZMod q)[X], (R.fromPoly (n:=n) a).representative = a %ₘ (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 f_monic : Monic (f q n)

content: 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)

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

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

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

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)

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: @[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 f_deg_eq : (f q n).degree = 2^n

content: @[simp]
theorem R.fromPoly_representative [Fact (q > 1)]:
    forall a : R q n, (R.fromPoly (n:=n) (R.representative q n a)) = a

content: theorem R.repLength_leq_representative_degree_plus_1 (a : R q n) :
  a.repLength ≤ (R.representative q n a).natDegree + 1

content: theorem R.toTensor_length {q n} (a : R q n) :
    (R.toTensor a).length = a.repLength

-- Local definitions used:

-- Local lemmas used:
content: theorem eq_iff_rep_eq [Fact (q > 1)] (a b : R q n) :
    a.representative = b.representative ↔ a = b

content: theorem from_poly_zero : R.fromPoly (0 : (ZMod q)[X]) (n := n) = (0 : R q n)

content: theorem rep_zero [Fact (q > 1)]: R.representative q n 0 = 0

content: theorem R.toTensor_getD [Fact (q > 1)] (a : R q n) (i : Nat) :
    a.toTensor.getD i 0 = (a.coeff i).toInt

content: theorem R.toTensor_getD' [hqgt1 : Fact (q > 1)] (a : R q n) (i : Nat) :
    ↑(a.toTensor.getD i 0) = a.coeff i

content: theorem eq_iff_coeff_eq [hqgt1 : Fact (q > 1)] (a b : R q n) :
    a = b ↔ Polynomial.coeff a.representative = Polynomial.coeff b.representative

content: theorem toTensor_fromTensor [hqgt1 : Fact (q > 1)] (tensor : List Int) (i : Nat)
  (htensorlen : List.length tensor < 2 ^ n) :
  (R.fromTensor tensor (q:=q) (n :=n)).toTensor.getD i 0 = (tensor.getD i 0) % q

-- Target theorem and proof:
theorem fromTensor_toTensor [hqgt1 : Fact (q > 1)] (a : R q n)
    (adeg : (R.representative q n a).natDegree + 1 < 2^n) :
  R.fromTensor a.toTensor = a
:= by
  cases h : Polynomial.degree (R.representative q n a) with
    | bot =>
        have h' :=  Polynomial.degree_eq_bot.1 h
        rw [← rep_zero] at h'
        have h'' := (eq_iff_rep_eq _ _).1 h'
        simp only [R.fromTensor, R.toTensor, R.repLength]
        rw [h, h'']
        simp
    | coe deg =>
        apply (eq_iff_coeff_eq _ _).2
        have hCoeff := R.toTensor_getD' (q := q) (n := n)
        unfold R.coeff at hCoeff
        apply funext
        intro i
        rw [← hCoeff, ← hCoeff]
        rw [toTensor_fromTensor]
        rw [← ZMod.coe_intCast]
        norm_cast
        · simp only [R.toTensor_length]
          have hdeg := R.repLength_leq_representative_degree_plus_1 a
          linarith