--- todo make it parse the input unmodified

fmod OBJECT is
  sort Object .
  ops ore clay obsidian geode : -> Object [ctor] .
endfm

view Object from TRIV to OBJECT is sort Elt to Object . endv


fmod INV is
  pr NAT + OBJECT .

  sorts Item Inv .
  subsort Item < Inv .
  op nil : -> Inv [ctor] .
  op __ : Nat Object -> Item [ctor prec 10] .
  op _and_ : Inv Inv -> Inv
    [ctor comm assoc id: nil prec 11] .

  vars A B : Object .
  vars I J : Nat .

  eq 0 A = nil .
  eq (I A) and (J A) = (I + J) A .

  vars X Y : Inv .

  op _-_ : Inv Inv -> [Inv] [prec 20] .
  eq  Y         - nil       = Y .
  ceq I A and X - J A and Y = sd(I, J) A and (X - Y) if J <= I .

  op count : Object Inv -> Nat .
  eq count(A, I A and X) = I .
  eq count(A, X)         = 0 [owise] .
endfm


fmod BP-ITEM is
  pr OBJECT + INV .

  sort BpItem .
  op item : Object Inv -> BpItem [ctor] .

  var A : Object . var Xs : Inv .

  op (Each _ robot costs _ .) : Object Inv -> BpItem [prec 20] .
  eq (Each A robot costs Xs .) = item(A, Xs) .
endfm

view BpItem from TRIV to BP-ITEM is sort Elt to BpItem . endv


fmod BP is
  pr SET{BpItem} * (sort Set{BpItem} to Items, op _,_ to __ [prec 21]) .

  sort Bp .
  op bp : Nat Items -> Bp [ctor] .

  var I : Nat .
  var Rs : Items .

  op (Blueprint _ : _) : Nat Items -> Bp [prec 30] .
  eq (Blueprint I : Rs) = bp(I, Rs) .

  op index : Bp -> Nat .
  eq index(bp(I, Rs)) = I .

  op items : Bp -> Items .
  eq items(bp(I, Rs)) = Rs .
endfm

view Bp from TRIV to BP is sort Elt to Bp . endv


fmod MAYBE{T :: TRIV} is
  sort Maybe{T} .
  subsort T$Elt < Maybe{T} .
  op nothing : -> Maybe{T} [ctor] .
endfm

fmod BAG{T :: TRIV} is
  sort Bag{T} .
  subsort T$Elt < Bag{T} .

  op nil : -> Bag{T} [ctor] .
  op __  : Bag{T} Bag{T} -> Bag{T} [ctor assoc comm id: nil] .
endfm

fmod STATE is
  pr NAT .
  pr LIST{Bp} * (sort List{Bp} to Bps,
                 op __ : Bps Bps -> Bps to __ [prec 31]) .
  pr MAYBE{Object} * (sort Maybe{Object} to Making) .
  pr BAG{Object} * (sort Bag{Object} to Robots) .

  sort State .
  op running : Bp Making Robots Inv -> State [ctor] .
  subsort Nat < State .

  var Bp : Bp .
  var Rs : Robots .
  var X  : Object .
  var Y  : Making .
  var Xs : Inv .

  op new : Bp -> State .
  eq new(Bp) = running(Bp, nothing, ore, nil).

  op count : Object State -> Nat .
  eq count(X, running(Bp, Y, Rs, Xs)) = count(X, Xs) .
endfm

view State from TRIV to STATE is sort Elt to State . endv


mod RULES is
  pr STATE .

  var  Bp : Bp .
  var  Is : Items .
  var  Rs : Robots .
  var  X  : Object .
  var  Y  : Making .
  vars Xs Ys : Inv .

  op collect : Robots -> Inv .
  eq collect(nil) = nil .
  eq collect(X Rs) = 1 X and collect(Rs) .

  op qual : Bps Inv -> Nat .
  eq qual(Bp, Xs) = index(Bp) * count(geode, Xs) .

  rl [collect] : running(Bp, Y, Rs, Xs) =>
                 running(Bp, Y, Rs, Xs and collect(Rs)) .

  crl [startBuild] : running(Bp, nothing, Rs, Xs) =>
                     running(Bp, X,       Rs, Xs - Ys)
    if item(X, Ys) Is := items(Bp)
    /\ Xs - Ys :: Inv .

  rl [finishBuild] : running(Bp, X,       Rs,   Ys) =>
                     running(Bp, nothing, X Rs, Ys) .

  rl [done] : running(Bp, Y, Rs, Xs) => qual(Bp, Xs) .
endm


smod RUN is
  pr LIST{State} .
  pr RULES .

  var N  : Nat .
  var X  : Object .
  var S  : State .
  var Ss : List{State} .

  --- this won't work because e.g. in example 1 it will just keep making a
  --- million clay guys forever and never get enough ore for an obsidian guy
  strat buildBest : @ State .
  sd buildBest := startBuild[X <- geode]
          or-else startBuild[X <- obsidian]
          or-else startBuild[X <- clay]
          or-else startBuild[X <- ore] .

  strat step1 : @ State .
  sd step1 := try(buildBest) ; collect ; try(finishBuild) .

  strat stepAll : @ List{State} .
  sd stepAll := try(matchrew S Ss by S using step1, Ss using stepAll) .

  strat stepsOnly : Nat @ List{State} .
  sd stepsOnly(0)    := idle .
  sd stepsOnly(s(N)) := stepAll ; stepsOnly(N) .

  strat steps : Nat @ List{State} .
  sd steps(N) := stepsOnly(N) ; done ! .
endsm