/Main.lean at 9e7dd794c7ea92b8829fc31dc8bb21633f833e17

import Lean.Data.Json
import Raylean

open Raylean Types

namespace Raylean

-- Helpful for debugging
instance : ToString Vector3 := ⟨fun a ↦ s!"({a.x}, {a.y}, {a.z})"⟩

instance : Add Vector3 where
  add a b := ⟨a.x + b.x, a.y + b.y, a.z + b.z⟩

instance : Sub Vector3 where
  sub a b := ⟨a.x - b.x, a.y - b.y, a.z - b.z⟩

instance : HMul Vector3 Float Vector3 where
  hMul a c := ⟨a.x * c, a.y * c, a.z * c⟩

instance : HDiv Vector3 Float Vector3 where
  hDiv a c := ⟨a.x / c, a.y / c, a.z / c⟩

/-- This is in Raylib but not Raylean so let's just define it ourselves -/
def drawCubeV (pos size : Vector3) (color : Color) :=
  drawCube pos size.x size.y size.z color

/-- Same as above -/
def drawCubeWiresV (pos size : Vector3) (color : Color) :=
  drawCubeWires pos size.x size.y size.z color

end Raylean

def fps := 60
def screenWidth := 960
def screenHeight := 640

structure Nat3 where
  x : Nat
  y : Nat
  z : Nat
deriving BEq, Hashable, Lean.ToJson, Lean.FromJson

inductive BuildingVariant
  | house
  | apartment
  | office
  | shop
  | factory
deriving Lean.ToJson, Lean.FromJson

def BuildingVariant.ofString? : String → Option BuildingVariant
  | "h" | "house" => some .house
  | "a" | "apartment" => some .apartment
  | "o" | "office" => some .office
  | "s" | "shop" => some .shop
  | "f" | "factory" => some .factory
  | _ => none

structure Building where
  variant : BuildingVariant
  pos : Nat3
  size : Nat3
  entrance : Nat3
  exit : Nat3
  occupants : Nat
deriving Lean.ToJson, Lean.FromJson

def Building.capacity (b : Building) :=
  match b.variant with
  | .house => 1
  | .apartment => b.size.x * b.size.y * b.size.z / 2
  | .office => 0
  | .shop => 0
  | .factory => 0

def Building.color (b : Building) :=
  match b.variant with
  | .house => Color.Raylean.red
  | .apartment => Color.Raylean.orange
  | .office => Color.Raylean.blue
  | .shop => Color.Raylean.purple
  | .factory => Color.Raylean.green

structure Peep where
  home : Nat
  work : Nat
  pos : Nat3
  dest : Nat3
deriving Lean.ToJson, Lean.FromJson

inductive Road
  | none
  | low
  | high
deriving Lean.ToJson, Lean.FromJson

instance [Lean.ToJson α] : Lean.ToJson (Vector α n) where
  toJson := Array.toJson ∘ Vector.toArray

instance [Lean.FromJson α] : Lean.FromJson (Vector α n) where
  fromJson? j := do
    let A ← Array.fromJson? j
    if h : A.size = n then
      return h ▸ A.toVector
    else
      throw s!"expected size {n}, got {A.size}"

structure Point where
  name : String
  ein : Vector Road 24
  eout : Vector Road 24
deriving Lean.ToJson, Lean.FromJson

def dx : Vector Int 24 :=
  List.replicate 9 (-1) ++ List.replicate 6 0 ++ List.replicate 9 1 |>.toArray.toVector

def dy : Vector Int 24 :=
  List.replicate 8 [-1, 0, 1] |>.flatten.toArray.toVector

def dz : Vector Int 24 :=
  #v[-1, -1, -1, 0, 0, 0, 1, 1, 1, -1, -1, -1, 1, 1, 1, -1, -1, -1, 0, 0, 0, 1, 1, 1]

-- Roads should not go straight up or straight down
#guard (List.range 24 |>.mapFinIdx fun i _ hi ↦ dx[i] != 0 || dz[i] != 0).and

-- `i` and `23 - i` should be in opposite directions
#guard (List.range 24 |>.mapFinIdx fun i _ hi ↦ dx[i] == (-dx[23 - i]) && dy[i] == (-dy[23 - i]) && dz[i] == (-dz[23 - i])).and

def appd (p : Nat3) (i : Nat) (hi : i < 24 := by grind) : Nat3 :=
  ⟨p.x + dx[i] |>.toNat, p.y + dy[i] |>.toNat, p.z + dz[i] |>.toNat⟩

instance : Lean.ToJson StdGen where
  toJson x := Lean.toJson (x.1, x.2)

instance : Lean.FromJson StdGen where
  fromJson? j := do
    let (x : Nat × Nat) ← Lean.fromJson? j
    return ⟨x.1, x.2⟩

instance [BEq α] [Hashable α] [Lean.ToJson α] : Lean.ToJson (Std.HashSet α) where
  toJson := List.toJson ∘ Std.HashSet.toList

instance [BEq α] [Hashable α] [Lean.FromJson α] : Lean.FromJson (Std.HashSet α) where
  fromJson? j := .ofList <$> List.fromJson? j

instance [BEq α] [Hashable α] [Lean.ToJson α] [Lean.ToJson β] : Lean.ToJson (Std.HashMap α β) where
  toJson := List.toJson ∘ Std.HashMap.toList

instance [BEq α] [Hashable α] [Lean.FromJson α] [Lean.FromJson β] : Lean.FromJson (Std.HashMap α β) where
  fromJson? j := .ofList <$> List.fromJson? j

-- TODO: Money? Eh can do that later
structure State where
  rng : StdGen
  speed : Nat
  day : Nat
  time : Nat
  origin : Nat3
  grid : Std.HashMap Nat3 Point
  buildings : Array Building
  dists : Array (Std.HashMap Nat3 Nat)
  peeps : Array Peep
  occupied : Std.HashSet Nat3
deriving Lean.ToJson, Lean.FromJson

/-- Macro for easily updating a specific field of the state -/
macro "modifyf" field:ident fn:term : term =>
  let lval := ⟨.node .none `Lean.Parser.Term.structInstLVal #[field.raw, Lean.mkNullNode]⟩
  `(modify fun s ↦ { s with $lval := $fn s.$field })

/-- Generate array of street names at compile time -/
elab "get_street_names" : term => do
  return Lean.toExpr <|
    (← IO.FS.readFile "street-names.txt").split '\n' |>.toStringArray

def street_names := get_street_names

theorem queue_dequeue_isSome_if_not_isEmpty {q : Std.Queue α} (h : ¬q.isEmpty) : q.dequeue?.isSome := by
  rw [Std.Queue.dequeue?]
  by_cases q.dList = []
  · have : q.eList ≠ [] := by grind [Std.Queue.isEmpty]
    have : q.eList.reverse ≠ [] := by simp [this]
    grind
  · grind

/-- Precompute distances to `start` using BFS -/
def mkDist (g : Std.HashMap Nat3 Point) (start : Nat3) := Id.run do
  let mut q := Std.Queue.enqueue start .empty
  let mut dist := Std.HashMap.ofList [(start, 0)]
  while hq : ¬q.isEmpty do
    let uq := q.dequeue?.get (queue_dequeue_isSome_if_not_isEmpty hq)
    let u := uq.1
    q := uq.2
    let d := dist[u]!
    if hs : g.contains u then
      for hi : i in List.range 24 do
        let v := appd u i
        match (g[u]'hs).ein[i]'(by grind) with
        | .low =>
          if !dist.contains v then
            dist := dist.insert u (d + 1)
            q := q.enqueue v
        | .high =>
          if !dist.contains v then
            dist := dist.insert v (d + 1)
            q := q.enqueue v
          -- Try traveling another unit in direction `i`
          if hs : g.contains v then
            match (g[v]'hs).ein[i]'(by grind) with
            | .high =>
              let v' := appd u i
              if !dist.contains v' then
                dist := dist.insert v' (d + 1)
                q := q.enqueue v'
            | _ =>
              pure ()
        | .none =>
          pure ()
  return dist

/-- Precompute all distances -/
def mkDists : StateM State Unit := do
  modifyf dists (fun _ ↦ #[])
  for building in (← get).buildings do
    modifyf dists (·.push <| mkDist (← get).grid building.entrance)

/-- Generate a random nat in [0, n) (with a slight bias towards smaller numbers)
`randNat` is more sophisticated but doesn't bundle a bounds proof, so let's just use modulo for simplicity
-/
def rand (n : Nat) (hn : 0 < n := by grind) : StateM State (Fin n) := do
  let (ret, rng) := stdNext (← get).rng
  modifyf rng (fun _ ↦ rng)
  return ⟨ret % n, Nat.mod_lt ret hn⟩

/-- Shuffle an array using the Fisher-Yates algorithm -/
def Array.shuffle (A : Array α) : StateM State (Array α) := do
  let mut A' := A.toVector
  for hi : i in [1:A'.size] do
    let j ← rand (i + 1)
    A' := A'.swap i j
  return A'.toArray

/-- Run one iteration of the game randomly -/
def tick : StateM State Unit := do
  let peeps ← (← get).peeps.shuffle
  for peep in peeps do
    -- TODO
  -- For each vehicle, if has dest then look at dist table and iterate through all possible moves
  -- If at dest then remove
  -- Cannot do a move if occupied currently or after this tick
  -- Do that move
  -- Set the new list of vehicles with the new position and direction (for animating)
  if 12 * 60 * 60 < (← get).time then
    modifyf time (fun _ ↦ 0)
    modifyf day (· + 1)
  else
    modifyf time (· + 1)

def Nat3.toVector3 (p : Nat3) : Vector3 :=
  ⟨p.x.toFloat / 10, p.y.toFloat / 10, p.z.toFloat / 10⟩

/-- Draw the game state -/
def render (s : State) : IO Unit := do
  for b in s.buildings do
    let size := b.size.toVector3
    let pos := b.pos.toVector3 - s.origin.toVector3 + size / 2.0
    drawCubeV pos size b.color
    drawCubeWiresV pos size .black
  for (pos, pt) in s.grid do
    if pos.x % 10 == 0 && pos.z % 10 == 0 then
      -- https://www.raylib.com/examples/core/loader.html?name=core_world_screen

      drawText
    let pos' := pos.toVector3 - s.origin.toVector3
    for i in [:24] do
      match pt.eout[i] with
      | .none =>
        pure ()
      | .low =>



def getInput (stdin : IO.FS.Stream) := do
  IO.print "> "
  return (← stdin.getLine).trimAsciiEnd.toString

/-- Load game state from file -/
def loadState (path : String) : IO State := do
  let serialized ← IO.FS.readFile path
  let json ← .ofExcept <| Lean.Json.parse serialized
  .ofExcept <| Lean.fromJson? json

-- TODO: Build roads
-- TODO: rename roads
def handleCmd (cmd : String) : StateT State IO Unit := do
  match cmd.split ' ' |>.toStringList with
  | ["s", path] =>
    IO.FS.writeFile path <| Lean.toJson (← get) |>.compress
  | ["l", path] =>
    set <| ← loadState path
  | ["v", speed] =>
    modifyf speed fun _ ↦ String.toNat! speed
  | ["i"] =>
    IO.println s!"Population: {(← get).peeps.size}"
  | "b" :: variant :: dims =>
    -- TODO refactor into own function, check collisions, update dists
    let variant := BuildingVariant.ofString? variant
    if h : dims.length = 6 && variant.isSome then
      let dims := dims.map String.toNat!
      have : dims.length = 6 := by grind
      modifyf buildings (·.push ⟨⟨dims[0], dims[1], dims[2]⟩, ⟨dims[3], dims[4], dims[5]⟩, variant.get (by grind), 0⟩)
    else
      throw <| .userError "Failed to parse build command"
  | _ =>
    throw <| .userError "Command not found"

def gameLoop : StateT State IO Unit := do
  let mut camera : Camera3D := {
    position := ⟨10, 10, 10⟩
    target := ⟨0, 0, 0⟩
    up := ⟨0, 1, 0⟩
    fovy := 45
    projection := .perspective
  }
  let stdin ← IO.getStdin
  let mut task ← IO.asTask <| getInput stdin
  while !(← windowShouldClose) do
    camera ← updateCamera camera .thirdPerson
    if ← IO.hasFinished task then
      let cmd ← (.ofExcept task.get)
      try
        handleCmd cmd
      catch e =>
        IO.println e
      task ← IO.asTask <| getInput stdin
    renderFrame do
      drawFPS (screenWidth - 100) 10
      clearBackground Color.white
      renderWithCamera camera do
        let s ← get
        drawGrid (s.origin.x / 10) 1
        render s
  closeWindow

def main : IO Unit := do
  -- This constant is FLAG_WINDOW_HIGHDPI || FLAG_WINDOW_RESIZABLE
  -- https://github.com/raysan5/raylib/blob/aaacda6e147031f2af0cfb6c1fd7e64d761ddb1f/src/raylib.h#L567
  setConfigFlags 0x00002004
  initWindow screenWidth screenHeight "LeanTTD"
  setTargetFPS fps
  gameLoop.run' {
    rng := mkStdGen (← IO.rand 0 (2 ^ 32))
    speed := 1
    day := 0
    time := 0
    origin := ⟨200, 10, 200⟩
    grid := .ofList []
    dists := #[]
    buildings := #[]
    occupied := .ofList []
    peeps := #[]
  }