module Bob.Prediction

import Bob.Model
import Data.Float32Bits
import Model.Config
import Model.Word32
import Std.Float
import Std.Foundation
import Std.List
import Std.Natural
import Std.Word

-- Prediction is a total semantic operation over a checked 256-logit row.
-- Equal maxima keep the earlier byte, giving deterministic greedy decoding.
family BobPredictionSemantics : Type 0
constructor BobPredictionSemanticsValue
field unrestricted bobPredictionIdentity : Bytes
field unrestricted bobPredictionContextBytes : Nat
field unrestricted bobPredictionOutputDomain : Nat
field unrestricted bobPredictionTieBreakIdentity : Bytes

end-family

family BobArgmaxState : Type 0
constructor BobArgmaxEmpty
constructor BobArgmaxValue
field unrestricted bobArgmaxPosition : Nat
field unrestricted bobArgmaxLogit : F32

end-family

def bobPredictionSemantics : (family BobPredictionSemantics) =
  (record
    BobPredictionSemantics
    (bobPredictionIdentity = b"greedy-next-byte")
    (bobPredictionContextBytes = bobContextLength)
    (bobPredictionOutputDomain = bobVocabularySize)
    (bobPredictionTieBreakIdentity = b"lowest-byte-on-equal-logit"))

-- IEEE total-order key for finite binary32 values.  Negative encodings are
-- complemented; non-negative encodings have the sign bit set.  Unsigned word
-- comparison then follows numerical order for every admitted finite logit.
def bobF32OrderKey =
  (lambda unrestricted value : F32 .
    (app
      (lambda unrestricted bits : StdU32 .
        (nat-eliminate
          (lambda unrestricted negative : Nat . StdU32)
          (stdU32Xor bits 0x80000000)
          (lambda unrestricted predecessor : Nat .
            (lambda unrestricted induction : StdU32 . (stdU32Not bits)))
          (stdF32IsNegative value)))
      (stdF32ToWord32 value)))

def bobArgmaxStep =
  (lambda unrestricted position : Nat .
    (lambda unrestricted logit : F32 .
      (lambda unrestricted state : (family BobArgmaxState) .
        (eliminate
          BobArgmaxState
          (lambda unrestricted current : (family BobArgmaxState) . (family BobArgmaxState))
          state
          (branch BobArgmaxEmpty . (constructor BobArgmaxState BobArgmaxValue position logit))
          (branch
            BobArgmaxValue
            bestPosition
            bestLogit
            .
            (nat-eliminate
              (lambda unrestricted currentIsNotSmaller : Nat . (family BobArgmaxState))
              (constructor BobArgmaxState BobArgmaxValue bestPosition bestLogit)
              (lambda unrestricted predecessor : Nat .
                (lambda unrestricted induction : (family BobArgmaxState) .
                  (constructor BobArgmaxState BobArgmaxValue position logit)))
              (naturalIsZero (stdU32LessThan (bobF32OrderKey logit) (bobF32OrderKey bestLogit)))))))))

def bobGreedyIndex =
  (lambda unrestricted logits : (family StdList F32) .
    (stdListFold
      F32
      (pi unrestricted position : Nat . (family BobArgmaxState))
      (lambda unrestricted logit : F32 .
        (lambda unrestricted rest : (pi unrestricted position : Nat . (family BobArgmaxState)) .
          (lambda unrestricted position : Nat .
            (bobArgmaxStep position logit (rest (succ position))))))
      (lambda unrestricted position : Nat . (constructor BobArgmaxState BobArgmaxEmpty))
      logits
      0))

def bobArgmaxPositionOrZero =
  (lambda unrestricted state : (family BobArgmaxState) .
    (eliminate
      BobArgmaxState
      (lambda unrestricted current : (family BobArgmaxState) . Nat)
      state
      (branch BobArgmaxEmpty . 0)
      (branch BobArgmaxValue position logit . position)))

def bobGreedyByte =
  (lambda unrestricted logits : (family StdList F32) .
    (nat-to-byte (bobArgmaxPositionOrZero (bobGreedyIndex logits))))

def bobEmptyLogits : (family StdList F32) =
  (constructor StdList StdListEmpty F32)

-- Reference forward semantics: select the table row named by the current raw
-- byte.  Shape admission below makes the fallback unreachable for a valid
-- model value; it remains explicit so the function is total.
def bobSelectLogitsRow =
  (lambda unrestricted table : (family StdList (family StdList F32)) .
    (lambda unrestricted currentByte : Byte .
      (stdOptionValueOr
        (family StdList F32)
        bobEmptyLogits
        (stdListIndex (family StdList F32) table (byte-to-nat currentByte)))))

def bobTableRowsValid =
  (lambda unrestricted table : (family StdList (family StdList F32)) .
    (stdListFold
      (family StdList F32)
      Nat
      (lambda unrestricted row : (family StdList F32) .
        (lambda unrestricted rest : Nat .
          (naturalAnd (naturalEqual (stdListLength F32 row) bobVocabularySize) rest)))
      1
      table))

def bobTableShapeValid =
  (lambda unrestricted table : (family StdList (family StdList F32)) .
    (naturalAnd
      (naturalEqual (stdListLength (family StdList F32) table) bobVocabularySize)
      (bobTableRowsValid table)))