fib.lean

structure Z (modulus : Nat) : Type where
  val : Fin modulus
deriving Repr

instance : Add (Z modulus) where
  add x y := { val := x.val + y.val }

instance : Mul (Z modulus) where
  mul x y := { val := x.val * y.val }

instance : ToString (Z modulus) where
  toString x := s!"{x.val}"

def Z.ofNat (n : Nat) : Z (pred_modulus + 1) where
  val := Fin.ofNat n

instance : OfNat (Z (pred_modulus + 1)) n where
  ofNat := .ofNat n

instance : Neg (Z modulus) where
  neg x := match modulus with
    | 0     => nomatch x
    | _ + 1 => x * Z.ofNat (modulus - 1)

-- Ext ≈ Z[x]/(x^2 - x - 1)
structure Ext (modulus : Nat) where
  a0 : Z modulus
  a1 : Z modulus

instance : Add (Ext modulus) where
  add x y := { a0 := x.a0 + y.a0, a1 := x.a1 + y.a1 }

instance : Mul (Ext modulus) where
  mul x y :=
    let a0 := x.a0 * y.a0 + x.a1 * y.a1
    let a1 := x.a0 * y.a1 + x.a1 * y.a0 + x.a1 * y.a1
    { a0, a1 }

def phi : Ext (pred_modulus + 1) where
  a0 := 0
  a1 := 1

class Monoid (a : Type u) extends Mul a where
  mempty : a

instance : Monoid (Ext (pred_modulus + 1)) where
  mempty := { a0 := 1, a1 := 0 }

theorem thing (n : Nat) : n / 2 + 1 < n + 2 := by
  have : n / 2 <= n := Nat.div_le_self n 2
  repeat apply Nat.succ_le_succ
  assumption

def exp [Monoid a] : a -> Nat -> a
| _, 0 => Monoid.mempty
| x, 1 => x
| x, n + 2 =>
  let multiplier := if n % 2 = 1 then x else Monoid.mempty
  multiplier * exp (x * x) (n / 2 + 1)
decreasing_by apply thing

instance : Pow (Ext (pred_modulus + 1)) Nat where
  pow := exp

def fib (n : Nat) : Z (pred_modulus + 1) := (phi ^ n).a1

def readNat : IO Nat := do
  let line <- (<- IO.getStdin).getLine
  pure line.trim.toNat!

def main : IO Unit := do
  IO.print "Modulus: "
  let modulus <- readNat

  match h:modulus with
  | 0 => IO.println "bad modulus"
  | m + 1 => do
    IO.print "Index: "
    let n <- readNat

    let Fn : Z modulus := by
      rw [h]
      exact fib n
    IO.println s!"F_{n} ≡ {Fn}  (mod {modulus})"