/Main.lean at e3480cee780b380e9923ab7433cec4d504cb0132

import Lean.Data.Json
import Lens
import Raylean

open Lens
open Raylean Types

namespace Raylean

-- Helpful for debugging
instance : ToString Vector3 := ⟨fun v ↦ s!"({v.1}, {v.2}, {v.3})"⟩

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
  pos : Nat3
  size : Nat3
  entrance : Nat3
  exit : Nat3
  variant : BuildingVariant
  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

/-- Vehicle is synonymous with person in this game -/
structure Vehicle 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⟩

def Point.empty : Point :=
  ⟨.replicate 24 .none, .replicate 24 .none⟩

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: routing table for each building, reverse grid
structure State where
  rng : Nat
  speed : Nat
  day : Nat
  time : Nat
  origin : Nat3
  grid : Std.HashMap Nat3 Point
  buildings : Array Building
  dists : Array (Std.HashMap Nat3 Nat)
  occupied : Std.HashSet Nat3
  vehicles : Array Vehicle
deriving Lean.ToJson, Lean.FromJson

makeLenses State
-- TODO open that namespace?

/-- 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 each destination using BFS -/
def mkDist : StateM State Unit := do
  let mut dists : Array (Std.HashMap Nat3 Nat) := #[]
  let s ← get
  for building in s.buildings do
    -- TODO refactor this into its own func
    let mut q : Std.Queue Nat3 := .enqueue building.entrance .empty
    let mut dist : Std.HashMap Nat3 Nat := .ofList [(building.entrance, 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 : s.grid.contains u then
        for hi : i in List.range 24 do
          let v := appd u i
          match (s.grid[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
            if hs : s.grid.contains v then
              match (s.grid[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 ()
    dists := dists.push dist
  -- TODO update dists in state

-- TODO
def tick : StateM State Unit := do
  -- Randomize list of vehicles using Fisher-Yates
  -- For each vehicle, if has dest then look at dist table and iterate through all possible moves
  -- 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)
  modify <| over State.Lens.time (· + 1)

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

def render (s : State) : IO Unit := do
  for b in s.buildings do
    let p := b.pos.toVector3
    let s := b.size.toVector3
    let o := p - s.origin.toVector3 + s / 2.0
    drawCubeV o s b.color
    drawCubeWiresV o s .black

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
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] =>
    modify <| set State.Lens.speed <| String.toNat! speed
  | "b" :: variant :: dims =>
    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
      modify <| over State.Lens.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 := 0
    speed := 1
    day := 0
    time := 0
    origin := ⟨200, 10, 200⟩
    grid := .ofList []
    dists := #[]
    buildings := #[]
    occupied := .ofList []
    vehicles := #[]
  }