monotonicity

A city-building and transport simulation game written in Lean

  1. 1
  2. 2
  3. 3
  4. 4
  5. 5
  6. 6
  7. 7
  8. 8
  9. 9
  10. 10
  11. 11
  12. 12
  13. 13
  14. 14
  15. 15
  16. 16
  17. 17
  18. 18
  19. 19
  20. 20
  21. 21
  22. 22
  23. 23
  24. 24
  25. 25
  26. 26
  27. 27
  28. 28
  29. 29
  30. 30
  31. 31
  32. 32
  33. 33
  34. 34
  35. 35
  36. 36
  37. 37
  38. 38
  39. 39
  40. 40
  41. 41
  42. 42
  43. 43
  44. 44
  45. 45
  46. 46
  47. 47
  48. 48
  49. 49
  50. 50
  51. 51
  52. 52
  53. 53
  54. 54
  55. 55
  56. 56
  57. 57
  58. 58
  59. 59
  60. 60
  61. 61
  62. 62
  63. 63
  64. 64
  65. 65
  66. 66
  67. 67
  68. 68
  69. 69
  70. 70
  71. 71
  72. 72
  73. 73
  74. 74
  75. 75
  76. 76
  77. 77
  78. 78
  79. 79
  80. 80
  81. 81
  82. 82
  83. 83
  84. 84
  85. 85
  86. 86
  87. 87
  88. 88
  89. 89
  90. 90
  91. 91
  92. 92
  93. 93
  94. 94
  95. 95
  96. 96
  97. 97
  98. 98
  99. 99
  100. 100
  101. 101
  102. 102
  103. 103
  104. 104
  105. 105
  106. 106
  107. 107
  108. 108
  109. 109
  110. 110
  111. 111
  112. 112
  113. 113
  114. 114
  115. 115
  116. 116
  117. 117
  118. 118
  119. 119
  120. 120
  121. 121
  122. 122
  123. 123
  124. 124
  125. 125
  126. 126
  127. 127
  128. 128
  129. 129
  130. 130
  131. 131
  132. 132
  133. 133
  134. 134
  135. 135
  136. 136
  137. 137
  138. 138
  139. 139
  140. 140
  141. 141
  142. 142
  143. 143
  144. 144
  145. 145
  146. 146
  147. 147
  148. 148
  149. 149
  150. 150
  151. 151
  152. 152
  153. 153
  154. 154
  155. 155
  156. 156
  157. 157
  158. 158
  159. 159
  160. 160
  161. 161
  162. 162
  163. 163
  164. 164
  165. 165
  166. 166
  167. 167
  168. 168
  169. 169
  170. 170
  171. 171
  172. 172
  173. 173
  174. 174
  175. 175
  176. 176
  177. 177
  178. 178
  179. 179
  180. 180
  181. 181
  182. 182
  183. 183
  184. 184
  185. 185
  186. 186
  187. 187
  188. 188
  189. 189
  190. 190
  191. 191
  192. 192
  193. 193
  194. 194
  195. 195
  196. 196
  197. 197
  198. 198
  199. 199
  200. 200
  201. 201
  202. 202
  203. 203
  204. 204
  205. 205
  206. 206
  207. 207
  208. 208
  209. 209
  210. 210
  211. 211
  212. 212
  213. 213
  214. 214
  215. 215
  216. 216
  217. 217
  218. 218
  219. 219
  220. 220
  221. 221
  222. 222
  223. 223
  224. 224
  225. 225
  226. 226
  227. 227
  228. 228
  229. 229
  230. 230
  231. 231
  232. 232
  233. 233
  234. 234
  235. 235
  236. 236
  237. 237
  238. 238
  239. 239
  240. 240
  241. 241
  242. 242
  243. 243
  244. 244
  245. 245
  246. 246
  247. 247
  248. 248
  249. 249
  250. 250
  251. 251
  252. 252
  253. 253
  254. 254
  255. 255
  256. 256
  257. 257
  258. 258
  259. 259
  260. 260
  261. 261
  262. 262
  263. 263
  264. 264
  265. 265
  266. 266
  267. 267
  268. 268
  269. 269
  270. 270
  271. 271
  272. 272
  273. 273
  274. 274
  275. 275
  276. 276
  277. 277
  278. 278
  279. 279
  280. 280
  281. 281
  282. 282
  283. 283
  284. 284
  285. 285
  286. 286
  287. 287
  288. 288
  289. 289
  290. 290
  291. 291
  292. 292
  293. 293
  294. 294
  295. 295
  296. 296
  297. 297
  298. 298
  299. 299
  300. 300
  301. 301
  302. 302
  303. 303
  304. 304
  305. 305
  306. 306
  307. 307
  308. 308
  309. 309
  310. 310
  311. 311
  312. 312
  313. 313
  314. 314
  315. 315
  316. 316
  317. 317
  318. 318
  319. 319
  320. 320
  321. 321
  322. 322
  323. 323
  324. 324
  325. 325
  326. 326
  327. 327
  328. 328
  329. 329
  330. 330
  331. 331
  332. 332
  333. 333
  334. 334
  335. 335
  336. 336
  337. 337
  338. 338
  339. 339
  340. 340
  341. 341
  342. 342
  343. 343
  344. 344
  345. 345
  346. 346
  347. 347
  348. 348
  349. 349
  350. 350
  351. 351
  352. 352
  353. 353
  354. 354
  355. 355
  356. 356
  357. 357
  358. 358
  359. 359
  360. 360
  361. 361
  362. 362
  363. 363
  364. 364
  365. 365
  366. 366
  367. 367
  368. 368
  369. 369
  370. 370
  371. 371
  372. 372
  373. 373
  374. 374
  375. 375
  376. 376
  377. 377
  378. 378
  379. 379
  380. 380
  381. 381
  382. 382
  383. 383
  384. 384
  385. 385
  386. 386
  387. 387
  388. 388
  389. 389
  390. 390
  391. 391
  392. 392
  393. 393
  394. 394
  395. 395
  396. 396
  397. 397
  398. 398
  399. 399
  400. 400
  401. 401
  402. 402
  403. 403
  404. 404
  405. 405
  406. 406
  407. 407
  408. 408
  409. 409
  410. 410
  411. 411
  412. 412
  413. 413
  414. 414
  415. 415
  416. 416
  417. 417
  418. 418
  419. 419
  420. 420
  421. 421
  422. 422
  423. 423
  424. 424
  425. 425
  426. 426
  427. 427
  428. 428
  429. 429
  430. 430
  431. 431
  432. 432
  433. 433
  434. 434
  435. 435
  436. 436
  437. 437
  438. 438
  439. 439
  440. 440
  441. 441
  442. 442
  443. 443
  444. 444
  445. 445
  446. 446
  447. 447
  448. 448
  449. 449
  450. 450
  451. 451
  452. 452
  453. 453
  454. 454
  455. 455
  456. 456
  457. 457
  458. 458
  459. 459
  460. 460
  461. 461
  462. 462
  463. 463
  464. 464
  465. 465
  466. 466
  467. 467
  468. 468
  469. 469
  470. 470
  471. 471
  472. 472
  473. 473
  474. 474
  475. 475
  476. 476
  477. 477
  478. 478
  479. 479
  480. 480
  481. 481
  482. 482
  483. 483
  484. 484
  485. 485
  486. 486
  487. 487
  488. 488
  489. 489
  490. 490
  491. 491
  492. 492
  493. 493
  494. 494
  495. 495
  496. 496
  497. 497
  498. 498
  499. 499
  500. 500
  501. 501
  502. 502
  503. 503
  504. 504
  505. 505
  506. 506
  507. 507
  508. 508
  509. 509
  510. 510
  511. 511
  512. 512
  513. 513
  514. 514
  515. 515
  516. 516
  517. 517
  518. 518
  519. 519
  520. 520
  521. 521
  522. 522
  523. 523
  524. 524
  525. 525
  526. 526
  527. 527
  528. 528
  529. 529
  530. 530
  531. 531
  532. 532
  533. 533
  534. 534
  535. 535
  536. 536
  537. 537
  538. 538
  539. 539
  540. 540
  541. 541
  542. 542
  543. 543
  544. 544
  545. 545
  546. 546
  547. 547
  548. 548
  549. 549
  550. 550
  551. 551
  552. 552
  553. 553
  554. 554
  555. 555
  556. 556
  557. 557
  558. 558
  559. 559
  560. 560
  561. 561
  562. 562
  563. 563
  564. 564
  565. 565
  566. 566
  567. 567
  568. 568
  569. 569
  570. 570
  571. 571
  572. 572
  573. 573
  574. 574
  575. 575
  576. 576
  577. 577
  578. 578
  579. 579
  580. 580
  581. 581
  582. 582
  583. 583
  584. 584
  585. 585
  586. 586
  587. 587
  588. 588
  589. 589
  590. 590
  591. 591
  592. 592
  593. 593
  594. 594
  595. 595
  596. 596
  597. 597
  598. 598
  599. 599
  600. 600
  601. 601
  602. 602
  603. 603
  604. 604
  605. 605
  606. 606
  607. 607
  608. 608
  609. 609
  610. 610
  611. 611
  612. 612
  613. 613
  614. 614
  615. 615
  616. 616
  617. 617
  618. 618
  619. 619
  620. 620
  621. 621
  622. 622
  623. 623
  624. 624
  625. 625
  626. 626
  627. 627
  628. 628
  629. 629
  630. 630
  631. 631
  632. 632
  633. 633
  634. 634
  635. 635
  636. 636
  637. 637
  638. 638
  639. 639
  640. 640
  641. 641
  642. 642
  643. 643
  644. 644
  645. 645
  646. 646
  647. 647
  648. 648
  649. 649
  650. 650
  651. 651
  652. 652
  653. 653
  654. 654
  655. 655
  656. 656
  657. 657
  658. 658
  659. 659
  660. 660
  661. 661
  662. 662
  663. 663
  664. 664
  665. 665
  666. 666
  667. 667
  668. 668
  669. 669
  670. 670
  671. 671
  672. 672
  673. 673
  674. 674
  675. 675
  676. 676
  677. 677
  678. 678
  679. 679
  680. 680
  681. 681
  682. 682
  683. 683
  684. 684
  685. 685
  686. 686
  687. 687
  688. 688
  689. 689
  690. 690
  691. 691
  692. 692
  693. 693
  694. 694
  695. 695
  696. 696
  697. 697
  698. 698
  699. 699
  700. 700
  701. 701
  702. 702
  703. 703
  704. 704
  705. 705
  706. 706
  707. 707
  708. 708
  709. 709
  710. 710
  711. 711
  712. 712
  713. 713
  714. 714
  715. 715
  716. 716
  717. 717
  718. 718
  719. 719
  720. 720
  721. 721
  722. 722
  723. 723
  724. 724
  725. 725
  726. 726
  727. 727
  728. 728
  729. 729
  730. 730
  731. 731
  732. 732
  733. 733
  734. 734
  735. 735
  736. 736
  737. 737
  738. 738
  739. 739
  740. 740
  741. 741
  742. 742
  743. 743
  744. 744
  745. 745
  746. 746
  747. 747
  748. 748
  749. 749
  750. 750
  751. 751
  752. 752
  753. 753
  754. 754
  755. 755
  756. 756
  757. 757
  758. 758
  759. 759
  760. 760
  761. 761
  762. 762
  763. 763
  764. 764
  765. 765
  766. 766
  767. 767
  768. 768
  769. 769
  770. 770
  771. 771
  772. 772
  773. 773
  774. 774
  775. 775
  776. 776
  777. 777
  778. 778
  779. 779
  780. 780
  781. 781
  782. 782
  783. 783
  784. 784
  785. 785
  786. 786
  787. 787
  788. 788
  789. 789
  790. 790
  791. 791
  792. 792
  793. 793
  794. 794
  795. 795
  796. 796
  797. 797
  798. 798
  799. 799
  800. 800
  801. 801
  802. 802
  803. 803
  804. 804
  805. 805
  806. 806
  807. 807
  808. 808
  809. 809
  810. 810
  811. 811
  812. 812
  813. 813
  814. 814
  815. 815
  816. 816
  817. 817
  818. 818
  819. 819
  820. 820
  821. 821
  822. 822
  823. 823
  824. 824
  825. 825
  826. 826
  827. 827
  828. 828
  829. 829
  830. 830
  831. 831
  832. 832
  833. 833
  834. 834
  835. 835
  836. 836
  837. 837
  838. 838
  839. 839
  840. 840
  841. 841
  842. 842
  843. 843
  844. 844
  845. 845
  846. 846
  847. 847
  848. 848
  849. 849
  850. 850
  851. 851
  852. 852
  853. 853
  854. 854
  855. 855
  856. 856
  857. 857
  858. 858
  859. 859
  860. 860
  861. 861
  862. 862
  863. 863
  864. 864
  865. 865
  866. 866
  867. 867
  868. 868
  869. 869
  870. 870
  871. 871
  872. 872
  873. 873
  874. 874
  875. 875
  876. 876
  877. 877
  878. 878
  879. 879
  880. 880
  881. 881
  882. 882
  883. 883
  884. 884
  885. 885
  886. 886
  887. 887
  888. 888
  889. 889
  890. 890
  891. 891
  892. 892
  893. 893
  894. 894
  895. 895
  896. 896
  897. 897
  898. 898
  899. 899
  900. 900
  901. 901
  902. 902
  903. 903
  904. 904
  905. 905
  906. 906
  907. 907
  908. 908
  909. 909
  910. 910
  911. 911
  912. 912
  913. 913
  914. 914
  915. 915
  916. 916
  917. 917
  918. 918
  919. 919
  920. 920
  921. 921
  922. 922
  923. 923
  924. 924
  925. 925
  926. 926
  927. 927
  928. 928
  929. 929
  930. 930
  931. 931
  932. 932
  933. 933
  934. 934
  935. 935
  936. 936
  937. 937
  938. 938
  939. 939
  940. 940
  941. 941
  942. 942
  943. 943
  944. 944
  945. 945
  946. 946
  947. 947
  948. 948
  949. 949
  950. 950
  951. 951
  952. 952
  953. 953
  954. 954
  955. 955
  956. 956
  957. 957
  958. 958
  959. 959
  960. 960
  961. 961
  962. 962
  963. 963
  964. 964
  965. 965
  966. 966
  967. 967
  968. 968
  969. 969
  970. 970
  971. 971
  972. 972
  973. 973
  974. 974
  975. 975
  976. 976
  977. 977
  978. 978
  979. 979
  980. 980
  981. 981
  982. 982
  983. 983
  984. 984
  985. 985
  986. 986
  987. 987
  988. 988
  989. 989
  990. 990
  991. 991
  992. 992
  993. 993
  994. 994
  995. 995
  996. 996
  997. 997
  998. 998
  999. 999
  1000. 1000
  1001. 1001
  1002. 1002
  1003. 1003
  1004. 1004
  1005. 1005
  1006. 1006
  1007. 1007
  1008. 1008
  1009. 1009
  1010. 1010
  1011. 1011
  1012. 1012
  1013. 1013
  1014. 1014
  1015. 1015
  1016. 1016
  1017. 1017
  1018. 1018
  1019. 1019
  1020. 1020
  1021. 1021
  1022. 1022
  1023. 1023
  1024. 1024
  1025. 1025
  1026. 1026
  1027. 1027
  1028. 1028
  1029. 1029
  1030. 1030
  1031. 1031
  1032. 1032
  1033. 1033
  1034. 1034
  1035. 1035
  1036. 1036
  1037. 1037
  1038. 1038
  1039. 1039
  1040. 1040
  1041. 1041
  1042. 1042
  1043. 1043
  1044. 1044
  1045. 1045
  1046. 1046
  1047. 1047
  1048. 1048
  1049. 1049
  1050. 1050
  1051. 1051
  1052. 1052
  1053. 1053
  1054. 1054
  1055. 1055
  1056. 1056
  1057. 1057
  1058. 1058
  1059. 1059
  1060. 1060
  1061. 1061
  1062. 1062
  1063. 1063
  1064. 1064
  1065. 1065
  1066. 1066
  1067. 1067
  1068. 1068
  1069. 1069
  1070. 1070
  1071. 1071
  1072. 1072
  1073. 1073
  1074. 1074
  1075. 1075
  1076. 1076
  1077. 1077
  1078. 1078
  1079. 1079
  1080. 1080
  1081. 1081
  1082. 1082
  1083. 1083
  1084. 1084
  1085. 1085
  1086. 1086
  1087. 1087
  1088. 1088
  1089. 1089
  1090. 1090
  1091. 1091
  1092. 1092
  1093. 1093
  1094. 1094
  1095. 1095
  1096. 1096
  1097. 1097
  1098. 1098
  1099. 1099
  1100. 1100
  1101. 1101
  1102. 1102
  1103. 1103
  1104. 1104
  1105. 1105
  1106. 1106
  1107. 1107
  1108. 1108
  1109. 1109
  1110. 1110
  1111. 1111
  1112. 1112
  1113. 1113
  1114. 1114
  1115. 1115
  1116. 1116
  1117. 1117
  1118. 1118
  1119. 1119
  1120. 1120
  1121. 1121
  1122. 1122
  1123. 1123
  1124. 1124
  1125. 1125
  1126. 1126
  1127. 1127
  1128. 1128
  1129. 1129
  1130. 1130
  1131. 1131
  1132. 1132
  1133. 1133
  1134. 1134
  1135. 1135
  1136. 1136
  1137. 1137
  1138. 1138
  1139. 1139
  1140. 1140
  1141. 1141
  1142. 1142
  1143. 1143
  1144. 1144
  1145. 1145
  1146. 1146
  1147. 1147
  1148. 1148
  1149. 1149
  1150. 1150
  1151. 1151
  1152. 1152
  1153. 1153
  1154. 1154
  1155. 1155
  1156. 1156
  1157. 1157
  1158. 1158
  1159. 1159
  1160. 1160
  1161. 1161
  1162. 1162
  1163. 1163
  1164. 1164
  1165. 1165
  1166. 1166
  1167. 1167
  1168. 1168
  1169. 1169
  1170. 1170
  1171. 1171
  1172. 1172
  1173. 1173
  1174. 1174
  1175. 1175
  1176. 1176
  1177. 1177
  1178. 1178
  1179. 1179
  1180. 1180
  1181. 1181
  1182. 1182
  1183. 1183
  1184. 1184
  1185. 1185
  1186. 1186
  1187. 1187
  1188. 1188
  1189. 1189
  1190. 1190
  1191. 1191
  1192. 1192
  1193. 1193
  1194. 1194
  1195. 1195
  1196. 1196
  1197. 1197
  1198. 1198
  1199. 1199
  1200. 1200
  1201. 1201
  1202. 1202
  1203. 1203
  1204. 1204
  1205. 1205
  1206. 1206
  1207. 1207
  1208. 1208
import Lean.Data.Json
import Raylean

-- This fixes the "unexpected error when elaborating 'let'" bug
set_option backward.do.legacy false

open Raylean Types

/-
## Stuff missing from Raylean
-/

namespace Raylean

-- https://github.com/raysan5/raylib/blob/aaacda6e147031f2af0cfb6c1fd7e64d761ddb1f/src/raylib.h#L567
def Flags.window_resizable : UInt64 := 0x00000004
def Flags.window_highdpi : UInt64 := 0x00002000

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

-- Arithmetic for `Vector3`s
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 Float Vector3 Vector3 where
  hMul c a := a.x * c, a.y * c, a.z * c

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

def Types.Vector3.dist (a b : Vector3) :=
  Float.sqrt <| (a.x - b.x) ^ 2 + (a.y - b.y) ^ 2 + (a.z - b.z) ^ 2

/-- 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 Float.pi := 4 * Float.atan 1

def Float.hypot (x y : Float) := Float.sqrt <| x ^ 2 + y ^ 2

def Vector.modify (V : Vector α n) (i : Nat) (f : α  α) : Vector α n := by
  have : (V.toArray.modify i f).size = n := by grind
  exact this  (V.toArray.modify i f |>.toVector)

/-
## Basic utilities
-/

/-- Scale factor for grid to Raylib coordinates
The purpose of this is ~~to give me headaches~~ because the Raylib third-person camera has a fixed speed -/
def scaleN := 10
def scale := scaleN.toFloat

def fps := 60

-- Each cell of the grid corresponds to 5m, so peeps move at 15m/s (≈33 mph)
def ticksPerSecond := 3

def initialScreenWidth := 960
def initialScreenHeight := 640

/-- Keyboard sensitivity -/
def sensitivity := 5

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

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

instance : Add Nat3 where
  add a b := a.x + b.x, a.y + b.y, a.z + b.z

instance : Sub Nat3 where
  sub a b := a.x - b.x, a.y - b.y, a.z - b.z

instance : HMul Nat Nat3 Nat3 where
  hMul c a := a.x * c, a.y * c, a.z * c

instance : HMul Nat3 Nat Nat3 where
  hMul a c := a.x * c, a.y * c, a.z * c

instance : HDiv Nat3 Nat Nat3 where
  hDiv a c := a.x / c, a.y / c, a.z / c

def Nat3.swizzle (a : Nat3) : Nat3 := a.z, a.y, a.x

/-
## MonotonoCity types
-/

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

instance : ToString BuildingVariant where
  toString
    | .house => "house"
    | .apartment => "apartment"
    | .office => "office"
    | .shop => "shop"
    | .factory => "factory"

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
  spots : Nat
  isDeleted : Bool
deriving Inhabited, Lean.ToJson, Lean.FromJson

namespace Building

def capacity (size : Nat3) : BuildingVariant  Nat
  | .house => 1
  | .apartment => size.x * size.y * size.z / 2
  | .office => 2 * size.x * size.y * size.z
  | .shop => size.x * size.y * size.z / 4
  | .factory => size.x * size.y * size.z / 32

def isResidential (b : Building) :=
  match b.variant with
  | .house | .apartment => true
  | _ => false

def 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

def cost (b : Building) :=
  -- Make taller buildings quadratically more expensive
  100 * b.size.x * b.size.y * b.size.y * b.size.z *
    match b.variant with
    | .house => 1
    | .apartment => 2
    | .office => 2
    | .shop => 5
    | .factory => 10

def checkCollide (pos size pos' size' : Nat3) :=
  -- 1D collision
  let f (a b c d : Nat) := !(b  c || d  a)
  -- Do it for all 3 dimensions
  (f pos.x (pos.x + size.x) pos'.x (pos'.x + size'.x) && f pos.y (pos.y + size.y) pos'.y (pos'.y + size'.y) && f pos.z (pos.z + size.z) pos'.z (pos'.z + size'.z))

#guard checkCollide 10, 10, 10 5, 5, 5 12, 8, 14 5, 5, 5

#guard !checkCollide 10, 10, 10 5, 5, 5 0, 8, 14 5, 5, 5

end Building

structure Peep where
  home : Nat
  work : Nat
  pos : Nat3
  dest : Nat
  dir : Option Nat
  brokenDown : Bool
  isCommercial : Bool
deriving BEq, Lean.ToJson, Lean.FromJson

/-- Grid helper functions -/
def dx (i : Nat) := Int.ofNat (i / 9 % 3) - 1

def dy (i : Nat) := Int.ofNat (i % 3) - 1

def dz (i : Nat) := Int.ofNat (i / 3 % 3) - 1

def di (x y z : Int) := 9 * x + y + 3 * z + 13 |>.toNat

/-- These functions are inverses kinda -/
example : di (dx i) (dy i) (dz i) = i % 27 := by
  simp [di, dx, dy, dz]
  grind

-- `i` and `26 - i` should be in opposite directions
#guard (List.range 27 |>.map fun i  dx i == -(dx (26 - i)) && dy i == -(dy (26 - i)) && dz i == -(dz (26 - i))).and

def Nat3.appdk (p : Nat3) (i k : Nat) : Nat3 :=
  p.x + k * dx i |>.toNat, p.y + k * dy i |>.toNat, p.z + k * dz i |>.toNat

def Nat3.appd (p : Nat3) (i : Nat) : Nat3 :=
  appdk p i 1

structure TrafficLight where
  redLen : Nat
  greenLen : Nat
  shift : Nat
deriving BEq, Lean.ToJson, Lean.FromJson

instance : ToString TrafficLight where
  toString t := s!"Red: {t.redLen}, Green: {t.greenLen}, Shift: {t.shift}"

def TrafficLight.isRed (t : TrafficLight) (ticks : Nat) : Bool :=
  (ticks / ticksPerSecond - t.shift) % (t.redLen + t.greenLen) < t.redLen

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

def Road.cost (r : Road) (dir y : Nat) :=
  (if y == 0 && dy dir == 0 then 1 else 10) *
    match r with
    | none => 0
    | low => 2
    | high => 10

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}"

inductive Control
  | none
  | yield
  | trafficLight (t : TrafficLight)
deriving Inhabited, BEq, Lean.ToJson, Lean.FromJson

instance : ToString Control where
  toString
    | .none => ""
    | .yield => "Yield"
    | .trafficLight t => toString t

def Control.isRed (ticks : Nat)
  | trafficLight t => t.isRed ticks
  | _ => false

structure Point where
  name : String
  control : Control
  e : Vector Road 27
deriving Inhabited, Lean.ToJson, Lean.FromJson

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

structure State where
  rng : StdGen
  ticks : Nat
  speed : Nat
  money : Nat
  origin : Nat3
  grid : Std.HashMap Nat3 Point
  buildings : Array Building
  unfull : Vector (Array Nat) 2
  shops : Array Nat
  dists : Array (Std.HashMap Nat3 Nat)
  peeps : Array Peep
  occupied : Std.HashSet Nat3
deriving Lean.ToJson, Lean.FromJson

/-
## doTick and friends
-/

/-- Lift vanilla `StateM` into `StateT` (state monad wrapped around something else) -/
instance [Monad m] : MonadLift (StateM σ) (StateT σ m) where
  monadLift x := modifyGet <| StateT.run x

/-- Lift exception throwing into the IO monad -/
instance : MonadLift (Except String) IO where
  monadLift x := .ofExcept x

/-- For lifting `StateT State (Except String)` to `StateT State IO` -/
instance [MonadLift m n] [Monad n] : MonadLift (StateT σ m) (StateT σ n) where
  monadLift x s := monadLift (x s)

/-- 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 })

/-- Macro for easily setting a specific field of the state -/
macro "setf" field:ident val:term : term =>
  let lval := .node .none `Lean.Parser.Term.structInstLVal #[field.raw, Lean.mkNullNode]
  `(modify fun s  { s with $lval := $val })
-- For some reason `notation "setf" field val => modifyf field (fun _ ↦ val)` doesn't work

/-- 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]!
    for hi : i in List.range 27 do
      let v := u.appd (26 - i)
      if hg : g.contains v then
        match g[v].e[i]'(by grind) with
        | .low =>
          if !dist.contains v then
            dist := dist.insert v (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`
          let v' := v.appd (26 - i)
          if hg : g.contains v' then
            match g[v'].e[i]'(by grind) with
            | .high =>
              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
  setf dists #[]
  for building in ( get).buildings do
    if !building.isDeleted then
      modifyf dists (·.push <| mkDist ( get).grid building.entrance)
    else
      modifyf dists (·.push <| .ofList [])

/-- 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
  setf rng 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

abbrev homeToWorkProb := 1000
abbrev toShopProb := 500
abbrev factoryToShopProb := 25
abbrev workToHomeProb := 1000
abbrev breakdownProb := 5000
abbrev shopToHomeProb := 100
abbrev repairProb := 100

/-- Run one iteration of the game -/
def doTick : StateM State Unit := do
  let mut peeps := #[]
  -- Occupied after this tick
  let mut occupied : Std.HashSet Nat3 := .ofList []
  -- Occupied after half this tick (all coordinates multiplied by 2)
  let mut occupiedMid : Std.HashSet Nat3 := .ofList []
  let origPeeps :=
    let g := ( get).grid
    ( ( get).peeps.shuffle).partition (fun p 
      if h : g.contains p.pos then g[p.pos].control != .yield else true)
  for p in origPeeps.1 ++ origPeeps.2 do
    let u := p.pos
    let s :=  get
    if u == s.buildings[p.dest]!.entrance then
      if p.dir.isSome then
        modifyf money (· + (if p.isCommercial then 1000 else 100))
      let mut p' := { p with dir := none }
      if p.dest == p.home then
        -- At home, go to work or shops
        if ( rand homeToWorkProb) == 0 then
          p' := { p' with dest := p.work, isCommercial := false }
        else if h : ( rand toShopProb) == 0 && !s.shops.isEmpty then
          let shopIdx  rand s.shops.size
          p' := { p' with dest := shopIdx, isCommercial := false }
      else if p.dest == p.work then
        -- At work, go to shops if factory worker else go home
        if h : s.buildings[p.work]!.variant == .factory && !s.shops.isEmpty && ( rand factoryToShopProb) == 0 then
          let shopIdx  rand s.shops.size
          p' := { p' with dest := shopIdx, isCommercial := true }
        else if ( rand workToHomeProb) == 0 then
          p' := { p with dest := p.home, isCommercial := false }
      else
        -- At shop, go to work if factory else go home
        if s.buildings[p.work]!.variant == .factory then
          p' := { p with dest := p.work, isCommercial := false }
        else if ( rand shopToHomeProb) == 0 then
          p' := { p with dest := p.home, isCommercial := false }
      -- Start the journey!
      if p'.dest != p.dest then
        p' := { p' with pos := s.buildings[p.dest]!.exit }
      peeps := peeps.push p'
    else if (if h : s.grid.contains u then s.grid[u].control.isRed s.ticks else false) then
      -- Traffic light
      occupied := occupied.insert u
      peeps := peeps.push ({ p with dir := none })
    else if p.brokenDown then
      occupied := occupied.insert u
      peeps := peeps.push ({ p with dir := none, brokenDown := ( rand repairProb) != 0 })
    else
      let dist := s.dists[p.dest]!
      let mut neighbors := #[]
      if hs : s.grid.contains u then
        for hi : i in List.range 27 do
          let v := u.appd i
          if hd : dist.contains v then
            match s.grid[u].e[i]'(by grind) with
            | .none =>
              pure ()
            | .low =>
              neighbors := neighbors.push (dist[v], i)
            | .high =>
              neighbors := neighbors.push (dist[v], i)
              if hv : s.grid.contains v then
                let v' := v.appd i
                if hd : dist.contains v' then
                  match s.grid[v].e[i]'(by grind) with
                  | .high =>
                    neighbors := neighbors.push (dist[v'], i + 27)
                  | _ =>
                    pure ()
      neighbors := neighbors.qsort (fun a b  a.1 < b.1 || (a.1 == b.1 && a.2 < b.2))
      let mut p' := p
      for (d, i) in neighbors do
        -- if d > neighbors[0]!.1 + 100 || d > 3 * neighbors[0]!.1 then
          -- Don't move if this makes us take a really long detour
          -- TODO: better heuristic
          -- break
        let v := u.appd i
        if i < 27 then
          if s.occupied.contains v || occupied.contains v || occupiedMid.contains (u + v) then
            continue
          occupied := occupied.insert v
          occupiedMid := occupiedMid.insert (u + v)
          p' := { p with pos := v, dir := some i }
        else
          let v':= u.appdk i 2
          if s.occupied.contains v || occupied.contains v || occupiedMid.contains (u + v) || s.occupied.contains v' || occupied.contains v' || occupiedMid.contains (v + v') then
            continue
          occupied := occupied.insert v
          occupiedMid := occupiedMid.insert (u + v)
          occupied := occupied.insert v'
          occupiedMid := occupiedMid.insert (v + v')
          p' := { p with pos := v', dir := some i }
        break
      if p' == p then
        occupied := occupied.insert u
        p' := { p with pos := u, dir := none }
      peeps := peeps.push { p' with brokenDown := ( rand breakdownProb) == 0 }
  setf peeps peeps
  setf occupied occupied
  modifyf ticks (· + 1)

/-
## Commands
-/

/-- 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

def endPointsToCorners (start stop : Nat3) : Nat3 × Nat3 :=
  (min start.x stop.x, min start.y stop.y, min start.z stop.z, max start.x stop.x, max start.y stop.y, max start.z stop.z)

/-- Delete roads in a region -/
def delete (start stop : Nat3) : StateM State Unit := do
  let (start, stop) := endPointsToCorners start stop
  -- Delete roads
  for x in [start.x:stop.x + 1] do
    for y in [start.y:stop.y + 1] do
      for z in [start.z:stop.z + 1] do
        modifyf grid (·.erase x, y, z)
  -- Delete buildings if intersects center
  modifyf buildings (·.map fun b 
    if Building.checkCollide (b.pos + b.size / 2) 1, 1, 1 start (stop - start) then { b with isDeleted := true} else b)
  let buildings := ( get).buildings
  -- Note that we don't add back spots, this is because ~~I'm lazy~~ to penalize deleting buildings
  modifyf peeps (·.filterMap fun p 
    if !buildings[p.home]!.isDeleted && !buildings[p.work]!.isDeleted then
      some (if buildings[p.dest]!.isDeleted then { p with dest := p.home }else p)
    else
      none)
  modifyf shops (·.filter fun idx  !buildings[idx]!.isDeleted)
  modifyf unfull (fun uf  #v[uf[0].filter fun idx  !buildings[idx]!.isDeleted, uf[1].filter fun idx  !buildings[idx]!.isDeleted])
  mkDists

/-- Spend some money! -/
def spend (cost : Nat) : StateT State (Except String) Unit := do
  if cost > ( get).money then
    throw s!"Need ${cost}, only have ${(← get).money}"
  modifyf money (· - cost)

/-- Add a building to the state -/
def addBuilding (variant : BuildingVariant) (pos size entrance exit : Nat3) : StateT State (Except String) Unit := do
  if entrance == exit then
    throw "Building entrance cannot be in same position as exit"
  let notOnSide (p : Nat3) :=
    p.x != pos.x && p.x != pos.x + size.x && p.z != pos.z && p.z != pos.z + size.z
  if notOnSide entrance || notOnSide exit then
    throw "Building entrance or exit not on side of building"
  for x in [pos.x:pos.x + size.x + 1] do
    for y in [pos.y:pos.y + size.y + 1] do
      for z in [pos.z:pos.z + size.z + 1] do
        if ( get).grid.contains x, y, z then
          throw "Building collides with existing road"
  for b' in ( get).buildings do
    if Building.checkCollide pos size b'.pos b'.size then
      throw "Building collides with existing building"
    if Building.checkCollide (entrance - 1, 0, 1) 2, 1, 2 b'.pos b'.size then
      throw "Building entrance collides with existing building"
    if Building.checkCollide (exit - 1, 0, 1) 2, 1, 2 b'.pos b'.size then
      throw "Building exit collides with existing building"
  let b : Building := {
    variant
    pos
    size
    entrance
    exit
    spots := Building.capacity size variant
    isDeleted := false
  }
  spend b.cost
  -- Randomly iterate through empty spots in buildings of the opposite kind and create new peeps
  let mut spots := b.spots
  while spots > 0 do
    let choices := (( get).unfull[b.isResidential.toNat]'(by grind [Bool.toNat_lt]))
    if h : !choices.isEmpty then
      let choiceIdx  rand choices.size
      let idx := choices[choiceIdx]
      if b.isResidential then
        modifyf peeps (·.push {
          home := ( get).buildings.size
          work := idx
          pos := b.exit
          dest := idx
          dir := none
          brokenDown := false
          isCommercial := false
        })
      else
        let newIdx := ( get).buildings.size
        modifyf peeps (·.push {
          home := idx
          work := newIdx
          pos := ( get).buildings[idx]!.exit
          dest := newIdx
          dir := none
          brokenDown := false
          isCommercial := false
        })
      spots := spots - 1
      modifyf buildings (·.modify idx (fun b  { b with spots := b.spots - 1 }))
      if ( get).buildings[idx]!.spots == 0 then
        modifyf unfull (·.modify b.isResidential.toNat (·.eraseIdx! choiceIdx))
    else
      break
  let newIdx := ( get).buildings.size
  modifyf buildings (·.push { b with spots })
  if spots > 0 then
    modifyf unfull (·.modify (!b.isResidential).toNat (·.push newIdx))
  if b.variant == .shop then
    modifyf shops (·.push newIdx)
  modifyf dists (·.push <| mkDist ( get).grid b.entrance)

def max3 (a b c : Nat) := max a (max b c)

def diff (a b : Nat) := if a < b then b - a else a - b

/-- Snap `a` to the nearest small `b` multiple -/
def snap (a : Int) (b : Nat) :=
  if a  -b / 2 then
    -1
  else if a  b / 2 then
     0
  else 1

/-- Get direction and length between two endpoints -/
def endpointsToRoad (start stop : Nat3) :=
  let length := max3 (diff start.x stop.x) (diff start.y stop.y) (diff start.z stop.z)
  let dir := di (snap (Int.ofNat stop.x - start.x) length) (snap (Int.ofNat stop.y - start.y) length) (snap (Int.ofNat stop.z - start.z) length)
  (length, dir)

-- Sanity check
#guard endpointsToRoad 5, 5, 5 1, 5, 2 == (4, di (-1) 0 (-1))

/-- Add roads to the state -/
def addRoad (start stop : Nat3) (isHigh : Bool) : StateT State (Except String) Unit := do
  let (length, dir) := endpointsToRoad start stop
  if dir == di 0 0 0 || dir == di 0 1 0 || dir == di 0 (-1) 0 then
    throw "Road cannot go straight up or down"
  let road : Road := if isHigh then .high else .low
  spend <| length * road.cost dir start.y
  -- TODO: Better heuristic here
  let mut name := ""
  let stop := start.appdk dir length
  for u in [start, stop, { start with x := start.x + 1 }, { start with z := start.z + 1 }, { start with x := start.x - 1 }, { start with z := start.z - 1 }] do
    if ( get).grid.contains u then
      -- Yeah this is not ideal but Lean doesn't know the two `← get`s are the same
      name := ( get).grid[u]!.name
  if name == "" then
    if isHigh then
      name := s!"Highway {(← rand 998) + 1}"
    else
      have : 0 < street_names.size := by native_decide
      name := s!"{street_names[← rand street_names.size]} Street"
  for i in List.range (length + 1) do
    let v := start.appdk dir i
    modifyf grid fun g  Id.run do
      let mut g := g
      if !g.contains v then
        g := g.insert v name, .none, .replicate 27 .none
      g.modify v (fun p 
        if i < length then { p with e := p.e.set! dir road } else p)
  mkDists

def addMultiRoad (start stop : Nat3) (isHigh : Bool) : StateT State (Except String) Unit := do
  if start.y != stop.y then
    throw "Multilane road must remain at same level"
  -- We can't overwrite `start` and `stop` here since we need to preserve the direction of the road
  let y := start.y
  let swizzled := diff start.x stop.x < diff start.z stop.z
  let (start', stop') :=
    if swizzled then
      let (start', stop') := endPointsToCorners start stop
      (start'.swizzle, stop'.swizzle)
    else
      endPointsToCorners start stop
  let addRoad' p p' :=
    (if swizzled then
      if start.z < stop.z then addRoad p.swizzle p'.swizzle else addRoad p'.swizzle p.swizzle
    else
      if start.x < stop.x then addRoad p p' else addRoad p' p) isHigh
  for i in [start'.z:stop'.z + 1] do
    -- Straight lines
    addRoad' start'.x, y, i stop'.x, y, i
  -- Diagonal segments
  let d := stop'.z - start'.z
  for i in [start'.x:stop'.x + d] do
    if i < start'.x + d then
      addRoad' start'.x, y, start'.z + i - start'.x i, y, start'.z
      addRoad' start'.x, y, stop'.z + start'.x - i i, y, stop'.z
    else if i > stop'.x then
      addRoad' i - d, y, start'.z stop'.x, y, start'.z + stop'.x + d - i
      addRoad' i - d, y, stop'.z stop'.x, y, stop'.z + i - stop'.x - d
    else
      addRoad' i - d, y, start'.z i, y, stop'.z
      addRoad' i - d, y, stop'.z i, y, start'.z

def addYield (pos : Nat3) : StateT State (Except String) Unit := do
  spend 10
  if !( get).grid.contains pos then
    throw s!"Could not place yield at {pos}"
  modifyf grid (·.modify pos fun p  { p with control := .yield })

def addTrafficLight (pos : Nat3) (t : TrafficLight) : StateT State (Except String) Unit := do
  spend 100
  if !( get).grid.contains pos then
    throw s!"Could not place traffic light at {pos}"
  modifyf grid (·.modify pos fun p  { p with control := .trafficLight t })

def addIntersection (pos : Nat3) : StateT State (Except String) Unit := do
  -- Main roads
  addRoad (pos + 0, 0, 1) (pos + 3, 0, 1) false
  addRoad (pos + 3, 0, 1) (pos + 0, 0, 1) false
  addRoad (pos + 1, 0, 0) (pos + 1, 0, 3) false
  addRoad (pos + 1, 0, 3) (pos + 1, 0, 0) false
  -- Left turns
  addRoad (pos + 1, 0, 1) (pos + 2, 0, 2) false
  addRoad (pos + 2, 0, 2) (pos + 1, 0, 1) false
  addRoad (pos + 1, 0, 2) (pos + 2, 0, 1) false
  addRoad (pos + 2, 0, 1) (pos + 1, 0, 2) false
  -- Traffic lights
  addTrafficLight (pos + 0, 0, 2) 10, 10, 0
  addTrafficLight (pos + 3, 0, 1) 10, 10, 0
  addTrafficLight (pos + 2, 0, 3) 10, 10, 10
  addTrafficLight (pos + 1, 0, 0) 10, 10, 10

/-- Currently unused -/
def addIntersection' (pos : Nat3) : StateT State (Except String) Unit := do
  -- Main roads
  addRoad (pos + 0, 0, 3) (pos + 5, 0, 3) false
  addRoad (pos + 5, 0, 2) (pos + 0, 0, 2) false
  addRoad (pos + 2, 0, 0) (pos + 2, 0, 5) false
  addRoad (pos + 3, 0, 5) (pos + 3, 0, 0) false
  -- Right turns
  addRoad (pos + 2, 0, 0) (pos + 0, 0, 2) false
  addRoad (pos + 0, 0, 3) (pos + 2, 0, 5) false
  addRoad (pos + 3, 0, 5) (pos + 5, 0, 3) false
  addRoad (pos + 5, 0, 2) (pos + 3, 0, 0) false
  -- Left turns
  addRoad (pos + 2, 0, 2) (pos + 3, 0, 3) false
  addRoad (pos + 3, 0, 3) (pos + 2, 0, 2) false
  addRoad (pos + 2, 0, 3) (pos + 3, 0, 2) false
  addRoad (pos + 3, 0, 2) (pos + 2, 0, 3) false
  -- Traffic lights
  addTrafficLight (pos + 1, 0, 3) 10, 10, 0
  addTrafficLight (pos + 4, 0, 2) 10, 10, 0
  addTrafficLight (pos + 2, 0, 1) 10, 10, 10
  addTrafficLight (pos + 3, 0, 4) 10, 10, 10
  -- Yields
  addYield (pos + 1, 0, 1)
  addYield (pos + 4, 0, 1)
  addYield (pos + 1, 0, 4)
  addYield (pos + 4, 0, 4)

inductive Action
  | delete
  | road (isHigh : Bool)
  | multiRoad (isHigh : Bool)
  | building (variant : BuildingVariant)
  | yield
  | trafficLight (phase : Bool)
  | intersection

instance : ToString Action where
  toString
    | .delete => "Delete"
    | .road isHigh => s!"Build {if isHigh then "highway" else "road"}"
    | .multiRoad isHigh => s!"Build multilane {if isHigh then "highway" else "road"}"
    | .building variant => s!"Build {variant}"
    | .yield => "Build yield"
    | .trafficLight phase => s!"Build traffic light (phase {if phase then 2 else 1})"
    | .intersection => "Build intersection"

def keyAction : List (Char × Action) := [
  ('0', .delete),
  ('1', .road false),
  ('2', .road true),
  ('3', .multiRoad false),
  ('4', .multiRoad true),
  ('5', .building .house),
  ('6', .building .apartment),
  ('7', .building .office),
  ('8', .building .shop),
  ('9', .building .factory),
  ('Z', .yield),
  ('X', .trafficLight false),
  ('C', .trafficLight true),
  ('V', .intersection),
]

-- TODO: Add height to existing building
def handleCmd (cmd : String) : StateT State IO Unit := do
  match cmd.split ' ' |>.toStringList with
  | "h" :: _ | "help" :: _ | "?" :: _ =>
    IO.println "HOTKEYS"
    IO.println "WASD: Move around"
    IO.println <| "\n".intercalate <| keyAction.map fun (c, a)  s!"{c}: {a}"
  | "s" :: path :: _ =>
    IO.FS.writeFile path <| Lean.toJson ( get) |>.compress
  | "l" :: path :: _  =>
    set <|  loadState path
  | "v" :: speed :: _  =>
    setf speed (sensitivity * String.toNat! speed)
  | "d" :: dims =>
    if h : dims.length = 6 then
      let dims := dims.map String.toNat!
      have : dims.length = 6 := by grind
      delete dims[0], dims[1], dims[2] dims[3], dims[4], dims[5]
  | "b" :: variant :: dims =>
    let variant := BuildingVariant.ofString? variant
    if h : dims.length = 12 && variant.isSome then
      let dims := dims.map String.toNat!
      have : dims.length = 12 := by grind
      addBuilding (variant.get (by grind)) dims[0], dims[1], dims[2] dims[3], dims[4], dims[5] dims[6], dims[7], dims[8] dims[9], dims[10], dims[11]
    else
      throw <| .userError "Failed to parse add building command"
  | "r" :: dims =>
    if h : dims.length = 7 then
      let dims := dims.map String.toNat!
      have : dims.length = 7 := by grind
      addRoad dims[0], dims[1], dims[2] dims[3], dims[4], dims[5] (dims[6] == 1)
    else
      throw <| .userError "Failed to parse add road command"
  | "m" :: dims =>
    if h : dims.length = 7 then
      let dims := dims.map String.toNat!
      have : dims.length = 7 := by grind
      addMultiRoad dims[0], dims[1], dims[2] dims[3], dims[4], dims[5] (dims[6] == 1)
    else
      throw <| .userError "Failed to parse add multilane road command"
  | "y" :: dims =>
    if h : dims.length = 3 then
      let dims := dims.map String.toNat!
      have : dims.length = 3 := by grind
      addYield dims[0], dims[1], dims[2]
    else
      throw <| .userError "Failed to parse add yield command"
  | "t" :: dims =>
    if h : dims.length = 6 then
      let dims := dims.map String.toNat!
      have : dims.length = 6 := by grind
      addTrafficLight dims[0], dims[1], dims[2] dims[3], dims[4], dims[5]
    else
      throw <| .userError "Failed to parse add traffic light command"
  | "i" :: dims =>
    if h : dims.length = 3 then
      let dims := dims.map String.toNat!
      have : dims.length = 3 := by grind
      addIntersection dims[0], dims[1], dims[2]
    else
      throw <| .userError "Failed to parse add intersection command"
  | "n" :: oldName :: newName :: _ =>
    modifyf grid (·.map fun _ p 
      if p.name == oldName then { p with name := newName } else p)
  | _ =>
    throw <| .userError "Command not found"

/-
## Rendering and UI
-/

def Nat3.toVector3 (pos : Nat3) : Vector3 :=
  pos.x.toFloat / scale, pos.y.toFloat / scale, pos.z.toFloat / scale

def Nat3.toVector3Shift (pos : Nat3) (s : State) : Vector3 :=
  pos.toVector3 - s.origin.toVector3

/-- Low `speed` means paused -/
def maxFrames (speed : Nat) :=
  if speed < sensitivity then 2 ^ 32 else fps / ticksPerSecond / (speed / sensitivity)

/-- Draw the game state -/
def render (s : State) (camera : Camera3D) (frames : Nat) : IO Unit := do
  -- Render buildings
  for b in s.buildings do
    if !b.isDeleted then
      let sizeV3 := b.size.toVector3
      let posV3 := b.pos.toVector3Shift s + sizeV3 / 2.0
      let entranceV3 := b.entrance.toVector3Shift s + 0, 0.05, 0
      let exitV3 := b.exit.toVector3Shift s + 0, 0.05, 0
      drawCubeV posV3 sizeV3 b.color
      drawCubeWiresV posV3 sizeV3 .black
      drawCubeV entranceV3 0.1, 0.1, 0.1 .green
      drawCubeV exitV3 0.1, 0.1, 0.1 .red
  -- Render building names
  endMode3D
  for b in s.buildings do
    if !b.isDeleted then
      if h : s.grid.contains b.entrance then
        let name := s.grid[b.entrance].name
        let sizeV3 := b.size.toVector3
        let posV3 := b.pos.toVector3Shift s + sizeV3 / 2.0
        let pos2D  getWorldToScreen (posV3 + 0, 0.2, 0) camera
        let address :=
          if b.entrance.x == b.pos.x || b.entrance.x == b.pos.x + b.size.x then b.entrance.z else b.entrance.x
        drawText s!"{address} {name} ({b.variant})" pos2D.x.toUInt64.toNat pos2D.y.toUInt64.toNat 10 Color.black
  beginMode3D camera
  -- Render roads
  for (pos, pt) in s.grid do
    let posV3 := pos.toVector3Shift s
    for hi : i in List.range 27 do
      let stop := pos.appd i
      let stopV3 := stop.toVector3Shift s
      match pt.e[i]'(by grind) with
      | .none =>
        pure ()
      | .low =>
        drawCylinderEx posV3 stopV3 0.01 0 32 Color.Raylean.pink
        -- Fake shadows
        if pos.y != s.origin.y || stop.y != s.origin.y then
          drawLine3D { posV3 with y := 0 } { stopV3 with y := 0 } Color.Raylean.gray
      | .high =>
        drawCylinderEx posV3 stopV3 0.01 0 32 Color.Raylean.gold
        -- Fake shadows
        if pos.y != s.origin.y || stop.y != s.origin.y then
          drawLine3D { posV3 with y := 0 } { stopV3 with y := 0 } Color.Raylean.gray
      match pt.control with
      | .yield =>
        drawCubeV (posV3 + 0, 0.075, 0) 0.02, 0.02, 0.02 .yellow
      | .trafficLight t =>
        drawCubeV (posV3 + 0, 0.075, 0) 0.02, 0.02, 0.02 (if t.isRed s.ticks then .red else .green)
      | .none =>
        pure ()
  -- Render peeps
  for peep in s.peeps do
    -- Unfortunately we can't use `Nat3.toVector3` here because they're `Int`s
    let posV3 :=
      match peep.dir with
      | none =>
        peep.pos.toVector3Shift s
      | some i =>
        let m := (maxFrames s.speed).toFloat
        let ivec : Vector3 := .ofInt (dx i) / scale, .ofInt (dy i) / scale, .ofInt (dz i) / scale
        peep.pos.toVector3Shift s - (m - frames.toFloat - 1) / m * (if i < 27 then 1 else 2) * ivec
    drawCube (posV3 + 0, 0.035, 0) 0.075 0.075 0.075 (if peep.isCommercial then Color.Raylean.beige else Color.Raylean.skyblue)
    drawCubeWires (posV3 + 0, 0.035, 0) 0.075 0.075 0.075 .black

/-- Get 3D coordinates at level `y` of 2D screen position
This assumes `up = ⟨0, 1, 0⟩ ∧ projection = .perspective` -/
def project (p : Vector2) (camera : Camera3D) (y width height : Float) : Vector3 :=
  let hor := (p.x - width / 2) / (height / 2)
  let vert := (p.y - height / 2) / (height / 2)
  let focal := 1 / .tan (camera.fovy * .pi / 360)
  -- Horizontal angle
  let hora := Float.atan2 hor focal
  -- Vertical angle
  let verta := Float.atan2 vert focal
  -- Compute current radial angle
  let r := (camera.position.x - camera.target.x).hypot (camera.position.z - camera.target.z)
  let h := camera.position.y - camera.target.y
  let θ := Float.atan2 r h
  -- Compute multiplier for radial move based on `θ - verta`
  let mul := .tan (θ - verta) * h / r
  -- Radial move
  let dradial := mul * Vector3.mk (camera.target.x - camera.position.x) 0 (camera.target.z - camera.position.z)
  -- Now compute the tangential move
  let d := camera.position.dist camera.target
  -- The term `h / (h + r * .tan verta)` is intuitively how much horizontal lines get squished or stretched
  let mul := d * .tan hora / r * h / (h + r * .tan verta)
  -- Tangential move
  let dtang := mul * Vector3.mk (camera.target.z - camera.position.z) 0 (camera.position.x - camera.target.x)
  -- Now scale to level `y`
  (dradial - dtang) * (camera.position.y - y) / (camera.position.y - camera.target.y) + camera.position.x, y, camera.position.z

/-- Inverse of `Nat3.toVector3Shift` -/
def Raylean.Types.Vector3.toNat3Shift (posV3 : Vector3) (s : State) : Nat3 :=
  let posV3 := posV3 + s.origin.toVector3
  let round (f : Float) := ((scale * f).round + 1e-9).toUInt64.toNat
  round posV3.x, round posV3.y, round posV3.z

/-- Get position of mouse at level `y` (relative to origin) -/
def getMouse3D (y : Int) (camera : Camera3D) : IO Vector3 := do
  let pos2  getMousePosition
  let windowScale  getWindowScaleDPI
  -- `.dot` is just element-wise product, not the dot product!
  return project (pos2.dot windowScale) camera (Float.ofInt (y / sensitivity) / scale) ( getScreenWidth).toFloat ( getScreenHeight).toFloat

def spawnInitBuildings : StateT State IO Unit := do
  let randPos : StateM State Nat3 := do
    return ( get).origin.x - 50 + ( rand 100), 10, ( get).origin.z - 50 + ( rand 100)
  let randPosExt : StateM State Nat3 := do
    repeat
      let x  rand 200
      let z  rand 200
      if 50 < x && x < 150 && 50 < z && z < 150 then
        continue
      return ( get).origin.x - 100 + x, 10, ( get).origin.z - 100 + z
  let sideToPos (pos size : Nat3) (side : Fin 4) (offset : Nat) :=
    pos +
      (match side with
      | 0 => offset, 0, 0
      | 1 => size.x, 0, offset
      | 2 => size.x - offset, 0, size.z
      | 3 => 0, 0, size.z - offset)
  let randBuilding (variant : BuildingVariant) (size : Nat3) (offset : Nat) (isClose : Bool) : StateT State IO Unit := do
    let pos  (if isClose then randPos else randPosExt)
    let size := if ( rand 2) == 0 then size else size.swizzle
    let side  rand 4
    try
      addBuilding variant pos size (sideToPos pos size side offset) (sideToPos pos size side (offset + 1))
    catch e =>
      IO.println e
  -- Spawn the commerical buildings first so peeps' workplaces get evenly distributed among them
  for i in [:3] do
    randBuilding .office 8, 2, 4 2 true
  for i in [:15] do
    randBuilding .shop 8, 1, 5 3 true
  for i in [:10] do
    randBuilding .factory 10, 5, 10 5 false
  for i in [:50] do
    randBuilding .house 5, 1, 3 1 false
  for i in [:10] do
    randBuilding .apartment 10, 3, 5 1 true

def endpointsToPosSize (pos pos' : Nat3) : Nat3 × Nat3 :=
  (min pos.x pos'.x, min pos.y pos'.y, min pos.z pos'.z, diff pos.x pos'.x, diff pos.y pos'.y, diff pos.z pos'.z)

def doAction (actionState : List Nat3) (curAction : Action) (mousePos : Nat3) (newHeight : Nat) : StateT State IO (List Nat3) := do
  match actionState, curAction with
  | pos :: _, .delete =>
    delete pos mousePos
  | pos :: _, .road isHigh =>
    addRoad pos mousePos isHigh
  | pos :: _, .multiRoad isHigh =>
    addMultiRoad pos mousePos isHigh
  | pos :: pos' :: entrance :: _, .building variant =>
    let (pos, size) := endpointsToPosSize pos { pos' with y := pos.y + newHeight / sensitivity }
    addBuilding variant pos size entrance mousePos
  | _, .yield =>
    addYield mousePos
  | _, .trafficLight phase =>
    addTrafficLight mousePos 10, 10, if phase then 10 else 0
  | _, .intersection =>
    addIntersection mousePos
  | _, _ =>
    return actionState ++ [mousePos]
  return []

def renderAction (actionState : List Nat3) (curAction : Action) (mousePos : Nat3) (s : State) (newHeight : Nat) : StateT State IO Unit := do
  let mousePosV3 := mousePos.toVector3Shift s
  match actionState, curAction with
  | pos :: _, .delete =>
    let posV3 := pos.toVector3Shift s
    drawCubeV ((posV3 + mousePosV3) / 2.0) (mousePosV3 - posV3) Color.Raylean.maroon
  | pos :: _, .road isHigh =>
    let (length, dir) := endpointsToRoad pos mousePos
    let stop := pos.appdk dir length
    let posV3 := pos.toVector3Shift s
    let stopV3 := stop.toVector3Shift s
    drawLine3D posV3 stopV3 (if isHigh then Color.Raylean.gold else Color.Raylean.pink)
    -- Fake shadows
    if pos.y != s.origin.y || stop.y != s.origin.y then
      drawLine3D { posV3 with y := 0 } { stopV3 with y := 0 } Color.Raylean.gray
  | pos :: _, .multiRoad isHigh =>
    let posV3 := pos.toVector3Shift s
    drawCubeV ((posV3 + mousePosV3) / 2.0) (mousePosV3 - posV3) (if isHigh then Color.Raylean.gold else Color.Raylean.pink)
  | poses, .building _ =>
    let poses := poses ++ [mousePos]
    if h : poses.length > 1 then
      let (pos, size) := endpointsToPosSize poses[0] { poses[1] with y := poses[0].y + newHeight / sensitivity }
      let posV3 := pos.toVector3Shift s
      drawCubeV (posV3 + size.toVector3 / 2.0) size.toVector3 Color.Raylean.lightgray
      drawCubeWiresV (posV3 + size.toVector3 / 2.0) size.toVector3 .black
    if h : poses.length > 2 then
      drawCubeV (poses[2].toVector3Shift s + 0, 0.05, 0) 0.1, 0.1, 0.1 .green
    if h : poses.length > 3 then
      drawCubeV (poses[3].toVector3Shift s + 0, 0.05, 0) 0.1, 0.1, 0.1 .red
  | _, .yield | _, .trafficLight _ =>
    drawCubeV (mousePosV3 + 0, 0.075, 0) 0.02, 0.02, 0.02 Color.Raylean.lime
  | _, .intersection =>
    drawCubeV (mousePosV3 + 0.15, 0, 0.15) 0.3, 0, 0.3 Color.Raylean.gray
  | _, _ =>
    pure ()

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

def padTime (n : Nat) :=
  if n < 10 then s!"0{n}" else s!"{n}"

def gameLoop : StateT State IO Unit := do
  spawnInitBuildings
  setf money 50000
  let mut camera : Camera3D := {
    position := scale, scale, scale
    target := 0, 0, 0
    up := 0, 1, 0
    fovy := 45
    projection := .perspective
  }
  let stdin  IO.getStdin
  let mut task  IO.asTask <| getInput stdin
  let mut actionState : List Nat3 := []
  let mut curAction := .road false
  let mut y : Int := 0
  let mut newHeight : Nat := sensitivity
  let mut frames := 0
  while !( windowShouldClose) do
    -- Fix `up` to prevent the Q and E keys from messing it up
    camera := {  updateCamera camera .thirdPerson with up := 0, 1, 0 }
    if  isKeyDown Key.up then
      y := y + 1
    if  isKeyDown Key.down then
      y := y - 1
    if ( isMouseButtonPressed MouseButton.left) then do
      let mousePos := ( getMouse3D y camera).toNat3Shift ( get)
      try
        actionState  doAction actionState curAction mousePos newHeight
      catch e =>
        IO.println e
    if ( isMouseButtonPressed MouseButton.right) then do
      actionState := []
    if ( isKeyDown Key.left) then do
      modifyf speed (· - 1)
    if ( isKeyDown Key.right) then do
      modifyf speed (· + 1)
    if ( isKeyDown '-'.toNat) then do
      newHeight := newHeight - 1
    if ( isKeyDown '='.toNat) then do
      newHeight := newHeight + 1
    for (key, action) in keyAction do
      if ( isKeyDown key.toNat) then do
        curAction := action
    if  IO.hasFinished task then
      let cmd  (.ofExcept task.get)
      try
        handleCmd cmd
      catch e =>
        IO.println e
      task  IO.asTask <| getInput stdin
    if frames  maxFrames ( get).speed then
      doTick
      frames := 0
    else
      frames := frames + 1
    let s  get
    renderFrame do
      clearBackground Color.white
      let mousePosV3  getMouse3D y camera
      let mousePos := mousePosV3.toNat3Shift s
      renderWithCamera camera do
        render s camera frames
        drawGrid (s.origin.x / scaleN * 2) 1
        renderAction actionState curAction mousePos s newHeight
      drawText s!"Active action: {curAction}" 10 10 20 .black
      drawText s!"Time: {s.ticks / ticksPerSecond / 60 / 60}:{padTime <| s.ticks / ticksPerSecond / 60 % 60}:{padTime <| s.ticks / ticksPerSecond % 60}" 10 40 20 .black
      drawText s!"Money: ${s.money}" 10 70 20 .black
      drawText s!"Population: {s.peeps.size}" 10 100 20 .black
      drawText s!"Speed: {s.speed / sensitivity}" 10 130 20 .black
      drawText s!"{mousePos}" 10 160 20 .black
      if hs : s.grid.contains mousePos then
        drawText s!"{s.grid[mousePos].name} {s.grid[mousePos].control}" 10 190 20 .black
      drawFPS (( getScreenWidth) - 100) 10
      -- TODO: GUI buttons?
  closeWindow

def main (args : List String) : IO Unit := do
  setConfigFlags <| Flags.window_resizable ||| Flags.window_highdpi
  initWindow initialScreenWidth initialScreenHeight "MonotoniCity"
  setTargetFPS fps
  let mut s := {
    rng := mkStdGen ( IO.rand 0 (2 ^ 32))
    ticks := 0
    speed := sensitivity
    money := 2 ^ 32 -- The initial value doesn't matter since we spawn a bunch of buildings first
    origin := 300, 10, 300
    grid := .ofList []
    buildings := #[]
    unfull := #v[#[], #[]]
    shops := #[]
    dists := #[]
    occupied := .ofList []
    peeps := #[]
  }
  if h : args.length > 0 then
    try
      s  loadState args[0]
    catch e =>
      IO.println s!"Failed to load save {args[0]} with error {e}"
  gameLoop.run' s
  -- We can't use `exit` because that won't kill the `getInput` task
  IO.Process.forceExit 0