-- authored by fynngodau 2026
-- depends on: Cslib @ 012cc2f5dddcea6262fa9b6cb16e93b6e6db3e95

import Cslib.Computability.Machines.Turing.SingleTape.Deterministic
import Cslib.Foundations.Data.RelatesInSteps

/-!

This file defines time constructability for Turing Machines, albeit the concept
is not very useful for single-tape TMs.

Time constructability, in general, is needed as prerequisite for the
hierarchy theorems (Stearns, Hartmanis, Lewis 1965) and padding lemmas (in terms
of time complexity, respectively).

The main result in this file is that there are arbitrarily large linear functions
that are time constructable (by repeatedly walking back and forth on the tape).
For this, we explicitly give a machine with runtime exactly `2n + 2`, and compose
it with itself as often as desired.

!-/

open Relation

open Cslib.Turing
open BiTape StackTape
open SingleTapeTM
variable {Symbol : Type} [Inhabited Symbol] [Fintype Symbol]


/-!
## Time Constructability / Timing Machines

This section defines Turing Machines that exactly take the time described by bundled function,
i.e. which time-construct such function.
-/

section TimeConstructable

/-- A proof of `tm` outputting `l'` on input `l` in exactly `m` steps. -/
def OutputsInTime (tm : SingleTapeTM Symbol) (l l' : List Symbol) (m : ℕ) :=
  RelatesInSteps tm.TransitionRelation (initCfg tm l) (haltCfg tm l') m

/-- A Turing machine + a time function +
a proof it terminates in exactly `time(input.length)` steps. -/
structure TimeConstructable (time : ℕ → ℕ) where
  /-- the underlying bundled SingleTapeTM -/
  tm : SingleTapeTM Symbol
  /-- proof this machine terminates in exactly `time(input.length)` steps -/
  haltsInTime (a) : ∃ b, OutputsInTime tm a b (time a.length)


/-- Restricted version of time constructability where the machine leaves the tape untouched. -/
structure TimeConstructableRestricted (time : ℕ → ℕ) where
  tm : SingleTapeTM Symbol
  haltsInTime (a) : OutputsInTime tm a a (time a.length)

def restricted_to_unrestricted {t : ℕ → ℕ} (ht : TimeConstructableRestricted (Symbol := Symbol) t)
  : TimeConstructable (Symbol := Symbol) t := {
  tm := ht.tm
  haltsInTime a := ⟨a, ht.haltsInTime a⟩
}

/-!

Copied & adapted from `Deterministic.lean` to talk about _exact_ timing
of composed machines.

-/
section compComputerLemmas'

variable (tm1 tm2 : SingleTapeTM Symbol) (cfg1 : tm1.Cfg) (cfg2 : tm2.Cfg)

lemma compComputer_q₀_eq : (compComputer tm1 tm2).q₀ = Sum.inl tm1.q₀ := rfl

/--
Convert a `Cfg` over the first input machine to a config over the composed machine.
Note it may transition to the start state of the second machine if the first machine halts.
-/
private def toCompCfg_left : (compComputer tm1 tm2).Cfg :=
  match cfg1.state with
  | some q => ⟨some (Sum.inl q), cfg1.BiTape⟩
  | none => ⟨some (Sum.inr tm2.q₀), cfg1.BiTape⟩

/-- Convert a `Cfg` over the second input machine to a config over the composed machine -/
private def toCompCfg_right : (compComputer tm1 tm2).Cfg :=
  ⟨Option.map Sum.inr cfg2.state, cfg2.BiTape⟩

/-- The initial configuration for the composed machine, with the first machine starting. -/
private def initialCfg (input : List Symbol) : (compComputer tm1 tm2).Cfg :=
  ⟨some (Sum.inl tm1.q₀), BiTape.mk₁ input⟩

/-- The intermediate configuration for the composed machine,
after the first machine halts and the second machine starts. -/
private def intermediateCfg (intermediate : List Symbol) : (compComputer tm1 tm2).Cfg :=
  ⟨some (Sum.inr tm2.q₀), BiTape.mk₁ intermediate⟩

/-- The final configuration for the composed machine, after the second machine halts. -/
private def finalCfg (output : List Symbol) : (compComputer tm1 tm2).Cfg :=
  ⟨none, BiTape.mk₁ output⟩

/-- The left converting function commutes with steps of the machines. -/
private theorem map_toCompCfg_left_step (hcfg1 : cfg1.state.isSome) :
    Option.map (toCompCfg_left tm1 tm2) (tm1.step cfg1) =
      (compComputer tm1 tm2).step (toCompCfg_left tm1 tm2 cfg1) := by
  cases cfg1 with | mk state BiTape => cases state with
    | none => grind
    | some q =>
      simp only [step, toCompCfg_left, compComputer]
      generalize hM : tm1.tr q BiTape.head = result
      obtain ⟨⟨wr, dir⟩, nextState⟩ := result
      #adaptation_note
      /-- A grind regression found moving to nightly-2026-03-31 (changes from lean#13166) -/
      cases nextState <;> (simp_all; rfl)

/-- The right converting function commutes with steps of the machines. -/
private theorem map_toCompCfg_right_step :
    Option.map (toCompCfg_right tm1 tm2) (tm2.step cfg2) =
      (compComputer tm1 tm2).step (toCompCfg_right tm1 tm2 cfg2) := by
  cases cfg2 with
  | mk state BiTape =>
    cases state with
    | none => rfl
    | some q =>
      generalize hM : tm2.tr q BiTape.head = result
      obtain ⟨⟨wr, dir⟩, nextState⟩ := result
      simp only [compComputer]
      grind [toCompCfg_right, step, compComputer]

/--
Simulation for the first phase of the composed computer.
When the first machine runs from start to halt, the composed machine
runs from start (with Sum.inl state) to Sum.inr tm2.q₀ (the start of the second phase).
This takes the same number of steps because the halt transition becomes a transition to the
second machine.
-/
private theorem comp_left_relatesInSteps (input intermediate : List Symbol) (t : ℕ)
    (htm1 :
      RelatesInSteps tm1.TransitionRelation
        (tm1.initCfg input)
        (tm1.haltCfg intermediate)
        t) :
    RelatesInSteps (compComputer tm1 tm2).TransitionRelation
      (initialCfg tm1 tm2 input)
      (intermediateCfg tm1 tm2 intermediate)
      t := by
  simp only [initialCfg, intermediateCfg, initCfg, haltCfg] at htm1 ⊢
  refine RelatesInSteps.map (toCompCfg_left tm1 tm2) ?_ htm1
  intro a b hab
  have ha : a.state.isSome := by
    simp only [TransitionRelation, step] at hab
    cases a with | mk state _ => cases state <;> simp_all
  have h1 := map_toCompCfg_left_step tm1 tm2 a ha
  rw [hab, Option.map_some] at h1
  exact h1.symm

/--
Simulation for the second phase of the composed computer.
When the second machine runs from start to halt, the composed machine
runs from Sum.inr tm2.q₀ to halt.
-/
private theorem comp_right_relatesInSteps (intermediate output : List Symbol) (t : ℕ)
    (htm2 :
      RelatesInSteps tm2.TransitionRelation
        (tm2.initCfg intermediate)
        (tm2.haltCfg output)
        t) :
    RelatesInSteps (compComputer tm1 tm2).TransitionRelation
      (intermediateCfg tm1 tm2 intermediate)
      (finalCfg tm1 tm2 output)
      t := by
  simp only [intermediateCfg, finalCfg, initCfg, haltCfg] at htm2 ⊢
  refine RelatesInSteps.map (toCompCfg_right tm1 tm2) ?_ htm2
  intro a b hab
  grind [map_toCompCfg_right_step tm1 tm2 a]

/--
Piece it together…
-/
private theorem comp_relatesInSteps (input intermediate output : List Symbol) (t u : ℕ)
    (htm1 :
      RelatesInSteps tm1.TransitionRelation
        (tm1.initCfg input)
        (tm1.haltCfg intermediate)
        t)
    (htm2 :
      RelatesInSteps tm2.TransitionRelation
        (tm2.initCfg intermediate)
        (tm2.haltCfg output)
        u) :
    RelatesInSteps (compComputer tm1 tm2).TransitionRelation
      (initialCfg tm1 tm2 input)
      (finalCfg tm1 tm2 output)
      (t + u) :=
  RelatesInSteps.trans
    (comp_left_relatesInSteps _ _ input intermediate t htm1)
    (comp_right_relatesInSteps _ _ intermediate output u htm2)


/-!

From the above lemmas, we get a restricted version of composition for free,
with the restriction to leave the tape as-is as per above.
(We could lift the restriction for the second machine, but would then not
get a restricted machine as a result, which is useful for composability.)

-/
def TimeConstructable.restricted_comp {t u : ℕ → ℕ}
    (ht : TimeConstructableRestricted (Symbol := Symbol) t)
    (hu : TimeConstructableRestricted (Symbol := Symbol) u) :
    (TimeConstructableRestricted (Symbol := Symbol) (fun n ↦ t n + u n)) where
      tm := compComputer ht.tm hu.tm
      haltsInTime (a) := comp_relatesInSteps _ _ _ _ _ _ _ (ht.haltsInTime a) (hu.haltsInTime a)

end compComputerLemmas'
section Examples

/-- Time constructability of the constant 1 function. -/
def TimeConstructable.one : TimeConstructableRestricted (Symbol := Symbol) (fun _ ↦ 1) where
  tm := idComputer
  haltsInTime a := by apply RelatesInSteps.single; rfl
def fgtm : SingleTapeTM Symbol := {
    State := Turing.Dir
    stateFintype := {
      elems := {Turing.Dir.left, Turing.Dir.right},
      complete x := by cases x <;> simp
    }
    q₀ := .right
    tr q h := match h with
    | some s => ({ symbol := s, movement := q }, some q)
    | none =>
      match q with
      | .right => ({ symbol := none, movement := some .left }, some .left)
      | .left  => ({ symbol := none, movement := some .right }, none)
  }

/-! Time constructibility of `2n + 2`. -/
def TimeConstructable.twox : TimeConstructableRestricted (Symbol := Symbol) (fun n ↦ 2 + 2 * n) where
  tm := fgtm
  haltsInTime a := by
    unfold OutputsInTime initCfg haltCfg BiTape.mk₁
    match a with
    | [] => 
        let tr1 := fgtm.TransitionRelation (fgtm.initCfg (Symbol := Symbol) []) ({state := some Turing.Dir.left, BiTape := nil})
        have z : tr1 := by simp [tr1, fgtm, TransitionRelation, initCfg, BiTape.mk₁, BiTape.nil, BiTape.write, BiTape.optionMove, BiTape.move, BiTape.moveLeft, StackTape.nil, StackTape.head, StackTape.tail, StackTape.cons]
        let tr2 := fgtm.TransitionRelation ({state := some Turing.Dir.left, BiTape := nil}) (fgtm.haltCfg (Symbol := Symbol) [])
        have z2 : tr2 := by simp [tr2, fgtm, TransitionRelation, initCfg, BiTape.mk₁, BiTape.nil, BiTape.write, BiTape.optionMove, BiTape.move, BiTape.moveRight, StackTape.nil, StackTape.head, StackTape.tail, StackTape.cons, haltCfg]
        simp
        apply RelatesInSteps.tail
        · apply RelatesInSteps.tail
          · apply RelatesInSteps.refl
          · exact z
        · exact z2
    | h :: t =>
      have prf (left : StackTape Symbol) : RelatesInSteps fgtm.TransitionRelation
          { state := some Turing.Dir.right, BiTape := { head := some h, left := left, right := mapSome t } }
          { state := some Turing.Dir.left, BiTape := { head := some h, left := left, right := mapSome t } } (1 + 2 * t.length + 1) := by
        induction t generalizing left h with
        | nil =>
          let tr1 := fgtm.TransitionRelation
            {state := some Turing.Dir.right, BiTape := { head := some h, left := left,        right := nil }}
            {state := some Turing.Dir.right, BiTape := { head := none,   left := left.cons h, right := nil}}
          have z : tr1 := by simp [tr1, fgtm, TransitionRelation, mk₁, BiTape.nil, write, optionMove, move, moveLeft, StackTape.nil, StackTape.head, StackTape.tail, cons,moveRight]
          let tr2 := fgtm.TransitionRelation
            {state := some Turing.Dir.right, BiTape := { head := none,   left := left.cons h, right := nil}}
            {state := some Turing.Dir.left,  BiTape := { head := some h, left := left,        right := nil}}
          have z2 : tr2 := by simp [tr2, fgtm, TransitionRelation, mk₁, BiTape.nil, write, optionMove, move, StackTape.nil, StackTape.head, StackTape.tail, cons, moveRight, haltCfg, moveLeft]
          apply RelatesInSteps.tail
          · apply RelatesInSteps.tail
            · apply RelatesInSteps.refl
            · exact z
          · exact z2
        | cons e a th =>
          simp [fgtm]
          let tr1 := fgtm.TransitionRelation
            { state := some Turing.Dir.right, BiTape := { head := some h, left := left,        right := mapSome (e :: a)} }
            { state := some Turing.Dir.right, BiTape := { head := some e, left := left.cons h, right := mapSome a       } }
          have z : tr1 := by simp [tr1, fgtm, TransitionRelation, mk₁, BiTape.nil, write, optionMove, move, moveLeft, StackTape.nil, StackTape.head, StackTape.tail, cons,moveRight, toList, mapSome]
          have ind := th e (left.cons h)
          let tr2 := fgtm.TransitionRelation
            { state := some Turing.Dir.left, BiTape := { head := some e, left := left.cons h, right := mapSome a       } }
            { state := some Turing.Dir.left, BiTape := { head := some h, left := left,        right := mapSome (e :: a)} }
          have z2 : tr2 := by simp [tr2, fgtm, TransitionRelation, mk₁, BiTape.nil, write, optionMove, move, StackTape.nil, StackTape.head, StackTape.tail, cons, moveRight, haltCfg, moveLeft, mapSome]
          have lo2l := RelatesInSteps.single (r := fgtm.TransitionRelation) z
          have lo3l := RelatesInSteps.trans lo2l ind
          have lo4l := RelatesInSteps.single (r := fgtm.TransitionRelation) z2
          have whol := RelatesInSteps.trans lo3l lo4l
          simp [fgtm] at whol ⊢
          have math : (1 + (1 + 2 * a.length + 1) + 1) = (1 + 2 * (a.length + 1) + 1) := by ring
          rw [math] at whol
          exact whol
      simp
      let tr1 := fgtm.TransitionRelation
        { state := some Turing.Dir.left, BiTape := { head := some h, left := nil, right := mapSome t          } }
        { state := some Turing.Dir.left, BiTape := { head := none,   left := nil, right := (mapSome t).cons h } }
      have z : tr1 := by simp [tr1, fgtm, TransitionRelation, mk₁, BiTape.nil, write, optionMove, move, moveLeft, StackTape.nil, StackTape.head, StackTape.tail, cons,moveRight, toList, mapSome]
      let tr2 := fgtm.TransitionRelation
        { state := some Turing.Dir.left, BiTape := { head := none,   left := nil, right := (mapSome t).cons h } }
        { state := none,                 BiTape := { head := some h, left := nil, right := mapSome t } }
      have z2 : tr2 := by simp [tr2, fgtm, TransitionRelation, mk₁, BiTape.nil, write, optionMove, move, StackTape.nil, StackTape.head, StackTape.tail, cons, moveRight, haltCfg, moveLeft, mapSome]
      have zstp := RelatesInSteps.single (r := fgtm.TransitionRelation) z
      have c1   := RelatesInSteps.trans (prf nil) zstp
      have t2sp := RelatesInSteps.single (r := fgtm.TransitionRelation) z2
      have c2   := RelatesInSteps.trans c1 t2sp
      simp [fgtm]
      have math : (1 + 2 * t.length + 1 + 1 + 1) = (2 + 2 * (t.length + 1)) := by ring
      rw [math] at c2
      apply c2

/--
Repeat the previous construction by repeatedly applying the restricted composition lemma
from above.
-/
def iterate_linear_constructable (z : ℕ) : TimeConstructableRestricted (Symbol := Symbol)
  (fun n ↦ (2 * (z + 1)) + (2 * (z + 1) * n))
  := match z with
| 0 => by
  simp only [zero_add, mul_one]
  exact TimeConstructable.twox
| n + 1 => by
  let prev := iterate_linear_constructable n
  have math : (fun n_1 ↦ 2 * (n + 1 + 1) + 2 * (n + 1 + 1) * n_1) =
              (fun n_1 ↦ (2 * (n + 1) + 2 * (n + 1) * n_1) + (2 + 2 * n_1)) := by ring_nf
  rw [math]
  exact TimeConstructable.restricted_comp (Symbol := Symbol) prev TimeConstructable.twox

/- this delivers you larger and larger linear functions that are time constructable
   in the given model =)  -/
#check iterate_linear_constructable 21

end Examples
end TimeConstructable