Untitled

/-Start Dependencies Declairation -/
module

public import Batteries.Tactic.Lint.Basic
public import Lean.Message
public import Mathlib.Init
public import Mathlib.Tactic.Common
public import Mathlib.Algebra.Group.Defs
public import Mathlib.Order.Defs.LinearOrder
public import Mathlib.Data.List.Sort
public import Mathlib.Tactic.Ring
public import Mathlib.Data.ZMod.Basic
public import Mathlib.Data.Nat.Prime.Basic
public import Mathlib.FieldTheory.Finite.Basic
public import Mathlib.Data.Nat.Log
namespace Cslib.Lint

@[expose] public section

/-!

# TimeM: Time Complexity Monad
`TimeM T α` represents a computation that produces a value of type `α` and tracks its time cost.

`T` is usually instantiated as `ℕ` to count operations, but can be instantiated as `ℝ` to count
actual wall time, or as more complex types in order to model more general costs.

## Design Principles
1. **Pure inputs, timed outputs**: Functions take plain values and return `TimeM` results
2. **Time annotations are trusted**: The `time` field is NOT verified against actual cost.
   You must manually ensure annotations match the algorithm's complexity in your cost model.
3. **Separation of concerns**: Prove correctness properties on `.ret`, prove complexity on `.time`

## Cost Model
**Document your cost model explicitly** Decide and be consistent about:
- **What costs 1 unit?** (comparison, arithmetic operation, etc.)
- **What is free?** (variable lookup, pattern matching, etc.)
- **Recursive calls:** Do you charge for the call itself?

## Notation
- **`✓`** : A tick of time, see `tick`.
- **`⟪tm⟫`** : Extract the pure value from a `TimeM` computation (notation for `tm.ret`)

## References

See [Danielsson2008] for the discussion.
-/
namespace Cslib.Algorithms.Lean

/-- A monad for tracking time complexity of computations.
`TimeM T α` represents a computation that returns a value of type `α`
and accumulates a time cost (represented as a type `T`, typically `ℕ`). -/
@[ext]
structure TimeM (T : Type*) (α : Type*) where
  /-- The return value of the computation -/
  ret : α
  /-- The accumulated time cost of the computation -/
  time : T

namespace TimeM

/-- Lifts a pure value into a `TimeM` computation with zero time cost.

Prefer to use `pure` instead of `TimeM.pure`. -/
protected def pure [Zero T] {α} (a : α) : TimeM T α :=
  ⟨a, 0⟩

instance [Zero T] : Pure (TimeM T) where
  pure := TimeM.pure

/-- Sequentially composes two `TimeM` computations, summing their time costs.

Prefer to use the `>>=` notation. -/
protected def bind {α β} [Add T] (m : TimeM T α) (f : α → TimeM T β) : TimeM T β :=
  let r := f (m.ret)
  ⟨r.ret, m.time + r.time⟩

instance [Add T] : Bind (TimeM T) where
  bind := TimeM.bind

instance : Functor (TimeM T) where
  map f x := ⟨f x.ret, x.time⟩

instance [Add T] : Seq (TimeM T) where
  seq f x := ⟨f.ret (x ()).ret, f.time + (x ()).time⟩

instance [Add T] : SeqLeft (TimeM T) where
  seqLeft x y := ⟨x.ret, x.time + (y ()).time⟩

instance [Add T] : SeqRight (TimeM T) where
  seqRight x y := ⟨(y ()).ret, x.time + (y ()).time⟩

instance [AddZero T] : Monad (TimeM T) where
  pure := Pure.pure
  bind := Bind.bind
  map := Functor.map
  seq := Seq.seq
  seqLeft := SeqLeft.seqLeft
  seqRight := SeqRight.seqRight

@[simp, grind =] theorem ret_pure {α} [Zero T] (a : α) : (pure a : TimeM T α).ret = a := rfl
@[simp, grind =] theorem ret_bind {α β} [Add T] (m : TimeM T α) (f : α → TimeM T β) :
    (m >>= f).ret = (f m.ret).ret := rfl
@[simp, grind =] theorem ret_map {α β} (f : α → β) (x : TimeM T α) : (f <$> x).ret = f x.ret := rfl
@[simp] theorem ret_seqRight {α} (x : TimeM T α) (y : Unit → TimeM T β) [Add T] :
    (SeqRight.seqRight x y).ret = (y ()).ret := rfl
@[simp] theorem ret_seqLeft {α} [Add T] (x : TimeM T α) (y : Unit → TimeM T β) :
    (SeqLeft.seqLeft x y).ret = x.ret := rfl
@[simp] theorem ret_seq {α β} [Add T] (f : TimeM T (α → β)) (x : Unit → TimeM T α) :
    (Seq.seq f x).ret = f.ret (x ()).ret := rfl

@[simp, grind =] theorem time_bind {α β} [Add T] (m : TimeM T α) (f : α → TimeM T β) :
    (m >>= f).time = m.time + (f m.ret).time := rfl
@[simp, grind =] theorem time_pure {α} [Zero T] (a : α) : (pure a : TimeM T α).time = 0 := rfl
@[simp, grind =] theorem time_map {α β} (f : α → β) (x : TimeM T α) : (f <$> x).time = x.time := rfl
@[simp] theorem time_seqRight {α} [Add T] (x : TimeM T α) (y : Unit → TimeM T β) :
    (SeqRight.seqRight x y).time = x.time + (y ()).time := rfl
@[simp] theorem time_seqLeft {α} [Add T] (x : TimeM T α) (y : Unit → TimeM T β) :
    (SeqLeft.seqLeft x y).time = x.time + (y ()).time := rfl
@[simp] theorem time_seq {α β} [Add T] (f : TimeM T (α → β)) (x : Unit → TimeM T α) :
    (Seq.seq f x).time = f.time + (x ()).time := rfl

/-- `TimeM` is lawful so long as addition in the cost is associative and absorbs zero. -/
instance [AddMonoid T] : LawfulMonad (TimeM T) := .mk'
  (id_map := fun x => rfl)
  (pure_bind := fun _ _ => by ext <;> simp)
  (bind_assoc := fun _ _ _ => by ext <;> simp [add_assoc])
  (seqLeft_eq := fun _ _ => by ext <;> simp)
  (bind_pure_comp := fun _ _ => by ext <;> simp)

/-- Creates a `TimeM` computation with a time cost. -/
def tick (c : T) : TimeM T PUnit := ⟨.unit, c⟩

@[simp, grind =] theorem ret_tick (c : T) : (tick c).ret = () := rfl
@[simp, grind =] theorem time_tick (c : T) : (tick c).time = c := rfl

/-- `✓[c] x` adds `c` ticks, then executes `x`. -/
macro "✓[" c:term "]" body:doElem : doElem => `(doElem| do TimeM.tick $c; $body:doElem)

/-- `✓ x` is a shorthand for `✓[1] x`, which adds one tick and executes `x`. -/
macro "✓" body:doElem : doElem => `(doElem| ✓[1] $body)

/-- Notation for extracting the return value from a `TimeM` computation: `⟪tm⟫` -/
scoped notation:max "⟪" tm "⟫" => (TimeM.ret tm)

end TimeM
end Cslib.Algorithms.Lean

namespace Cslib.Algorithms.Lean.TimeM



-- def fastPower (a : Nat) (n : Nat) : Nat :=
--   if (n==0) then 1
--   else
--     let half := fastPower a (n/2)
--     if n % 2 = 0 then
--       half * half
--     else
--       a * half * half
-- termination_by n
--   decreasing_by grind



def fastPower (a : Nat) (n : Nat) : TimeM ℕ ℕ := do
  if (n==0) then pure 1
  else
    ✓
    let half ← fastPower a (n/2)
    if (n % 2 = 0) then
      return half * half
    else
      return a * half * half
termination_by n
  decreasing_by grind
#eval fastPower 3 4

theorem power_one (a : Nat) : ⟪fastPower a 1⟫ = a := by
  simp[fastPower]

theorem powers_same (a n : Nat) : ⟪fastPower a n⟫ = a ^ n := by
  fun_induction fastPower a n with
  | case1 x h => simp_all
  | case2 x h ih =>
    simp at h
    by_cases heven : x % 2 = 0
    · simp[ret_bind, ret_pure, heven, if_pos, ih]
      ring
      have hx : x / 2 * 2 = x := by omega
      simp[hx]
    · simp[ret_bind, heven, ih]
      rw [show a * a ^ (x / 2) * a ^ (x / 2) = a ^ (x / 2 + x / 2 + 1) by ring]
      congr 1; omega

theorem powers_same_golfed (a n : Nat) : ⟪fastPower a n⟫ = a ^ n := by
  fun_induction fastPower a n with
  | case1 _ _ => simp_all
  | case2 x _ ih =>
    simp[ret_bind, ih]
    split_ifs with heven
    · simp[ret_pure]
      rw [← pow_add]; congr 1; omega
    · simp[ret_pure]
      rw [show a * a ^ (x/2) * a ^ (x/2) = a ^ (x/2 + x/2 + 1) by ring]
      congr 1; omega


def timePowerRec : ℕ → ℕ
| 0 => 0
| n + 1 => timePowerRec ((n + 1) / 2) + 1
termination_by n => n
decreasing_by
  omega

theorem return_less (a n : Nat) : (fastPower a n).time = timePowerRec n := by
  fun_induction fastPower with
  | case1 x h =>
    simp_all[timePowerRec]
  | case2 x h ih =>
    simp at h
    simp [time_bind, time_tick, ih]
    split_ifs with heven
    · simp [time_pure]
      obtain ⟨m, rfl⟩ := Nat.exists_eq_succ_of_ne_zero h
      simp [timePowerRec]
      omega
    · simp [time_pure]
      obtain ⟨m, rfl⟩ := Nat.exists_eq_succ_of_ne_zero h
      simp [timePowerRec]
      omega

open Nat (clog)

/-- The key arithmetic step: `clog 2` of the *halved* argument is strictly less than
`clog 2` of the full argument. This is the discrete analogue of "log of n/2 is less
than log of n by exactly 1." -/
theorem algebra_proof (x : Nat) : clog 2 ((x + 1) / 2 + 1) < clog 2 (x + 1 + 1) := by
  -- Step 1: `(x+1)/2 + 1 = (x+3)/2` — pure Nat arithmetic, omega handles it.
  have h_arg : (x + 1) / 2 + 1 = (x + 3) / 2 := by omega
  rw [h_arg]
  -- Step 2: Unfold `clog 2 (x+1+1) = clog 2 ((x+3)/2) + 1` via the clog recurrence.
  -- `Nat.clog_of_two_le : 2 ≤ k → 2 ≤ n → clog k n = clog k ((n + k - 1) / k) + 1`
  have h_two : (2 : Nat) ≤ x + 1 + 1 := by omega
  rw [Nat.clog_of_two_le (by omega) h_two]
  -- Step 3: normalize `(x+1+1+2-1)/2 = (x+3)/2` and close with omega.
  have h_norm : (x + 1 + 1 + 2 - 1) / 2 = (x + 3) / 2 := by omega
  rw [h_norm]
  omega

theorem rec_less (n : Nat) : timePowerRec n ≤ clog 2 (n + 1) := by
  fun_induction timePowerRec with
  | case1 => rfl
  | case2 x hx =>
    simp
    have h := algebra_proof x
    omega


theorem fastexp_time (a n : Nat) : (fastPower a n).time ≤ clog 2 (n + 1) := by
  rw[return_less]
  exact rec_less n

end Cslib.Algorithms.Lean.TimeM