-- VCVio/OracleComp/QueryTracking/Structures.lean

import VCVio.OracleComp.SimSemantics.SimulateQ

open OracleSpec OracleComp

namespace OracleSpec

variable {ι : Type u} {spec : OracleSpec ι}

namespace QueryCache

variable [spec.DecidableEq] [DecidableEq ι] (cache : QueryCache spec)

end QueryCache

namespace QueryCount

end QueryCount

namespace QueryLog

section getQ

variable [DecidableEq ι]

end getQ

section countQ

variable [DecidableEq ι]

end countQ

section prod

variable {ι₁ ι₂ : Type*} {spec₁ : OracleSpec ι₁} {spec₂ : OracleSpec ι₂}

end prod

end QueryLog

def QuerySeed (spec : OracleSpec ι) : Type _ :=
  (i : ι) → List (spec.range i)

namespace QuerySeed

section addValues

variable [DecidableEq ι] {i : ι} (seed : QuerySeed spec) (us : List (spec.range i))

end addValues

section takeAtIndex

variable [DecidableEq ι] (seed : QuerySeed spec) (i : ι) (n : ℕ)

-- Repository definitions used:
content: def OracleSpec (ι : Type u) : Type (max u (v + 1)) :=
  (i : ι) → Type v × Type v

content: @[inline, reducible]
protected def range (spec : OracleSpec ι) (i : ι) : Type w := (spec i).2

content: protected class DecidableEq (spec : OracleSpec ι) where
  domain_decidableEq' (i : ι) : DecidableEq (spec.domain i)
  range_decidableEq' (i : ι) : DecidableEq (spec.range i)

content: @[inline, reducible]
protected def domain (spec : OracleSpec ι) (i : ι) : Type v := (spec i).1

content: @[elab_as_elim]
protected def induction {C : OracleComp spec α → Prop}
    (oa : OracleComp spec α) (pure : (a : α) → C (pure a))
    (query_bind : (i : ι) → (t : spec.domain i) →
      (oa : spec.range i → OracleComp spec α) → (∀ u, C (oa u)) → C (query i t >>= oa))
    (failure : C failure) : C oa :=
  FreeMonad.inductionOn (Option.rec failure pure) (λ (query i t) ↦ query_bind i t) oa

content: inductive OracleQuery {ι : Type u} (spec : OracleSpec.{u,v} ι) : Type v → Type (max u v)
  | query (i : ι) (t : spec.domain i) : OracleQuery spec (spec.range i)

content: def OracleComp {ι : Type u} (spec : OracleSpec.{u,v} ι) :
    Type w → Type (max u (v + 1) w) :=
  OptionT (FreeMonad (OracleQuery.{u,v} spec))

content: inductive FreeMonad (f : Type u → Type v) (α : Type w) : Type (max (u + 1) v w)
  | protected pure (x : α) : FreeMonad f α
  | roll {β : Type u} (x : f β) (r : β → FreeMonad f α) : FreeMonad f α

content: @[elab_as_elim]
protected def inductionOn {C : FreeMonad f α → Prop}
    (pure : ∀ x, C (pure x))
    (roll : ∀ {β} (x : f β), (r : β → FreeMonad f α) →
      (∀ y, C (r y)) → C (x >>= r)) :
    (oa : FreeMonad f α) → C oa
  | FreeMonad.pure x => pure x
  | FreeMonad.roll x r => roll x _ (λ u ↦
      FreeMonad.inductionOn pure roll (r u))

/-- Add a list of values to the query seed.-/
def addValues [DecidableEq ι] {i : ι}
    (us : List (spec.range i)) (seed : QuerySeed spec) : QuerySeed spec :=
  Function.update seed i (seed i ++ us)

/-- Take only the first `n` values of the seed at index `i`. -/
def takeAtIndex [DecidableEq ι] (seed : QuerySeed spec) (i : ι) (n : ℕ) : QuerySeed spec :=
  Function.update seed i ((seed i).take n)

-- 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: def QuerySeed (spec : OracleSpec ι) : Type _ :=
  (i : ι) → List (spec.range i)

-- Local lemmas used:

-- Target theorem and proof:
lemma eq_takeAtIndex_length_iff (seed seed' : QuerySeed spec) (i : ι) :
    seed = seed'.takeAtIndex i (seed i).length ↔
      seed' = seed.addValues ((seed' i).drop (seed i).length)
:= by
  refine ⟨λ h ↦ QuerySeed.ext _ _ (λ j ↦ ?_), λ h ↦ ?_⟩
  · by_cases hj : j = i
    · induction hj
      rw [h]
      suffices (seed j).length ≤ (seed' j).length
      by simp [this]
      simpa using congr_arg List.length (congr_fun h j)
    · rw [h]
      simp [hj]
  · rw [h]
    simp