module Bob.TrainingRun

import Bob.Model
import Bob.Learner
import Checkpoint.Envelope
import Data.Float32Bits
import Model.Config
import Model.Parameter
import Model.Word64
import Std.Natural
import Std.Physical
import Std.Word

-- The physical transport is raw contiguous FP32 P/M/V.  The optimizer step
-- remains an external receipt, matching the currently qualified native path.
family BobCheckpointTransportSchema : Type 0
constructor BobCheckpointTransportSchemaValue
field unrestricted bobCheckpointSchemaIdentity : Bytes
field unrestricted bobCheckpointParameterCount : Nat
field unrestricted bobCheckpointPlaneCount : Nat
field unrestricted bobCheckpointBytesPerScalar : Nat
field unrestricted bobCheckpointPayloadBytes : Nat
field unrestricted bobCheckpointOrderingIdentity : Bytes
field unrestricted bobCheckpointInBandStepMetadata : Nat
field unrestricted bobCheckpointExternalStepReceiptRequired : Nat

end-family

family BobTrainingRunSemantics : Type 0
constructor BobTrainingRunSemanticsValue
field unrestricted bobRunUpdatesPerInvocation : Nat
field unrestricted bobRunDeterministicInitializationIdentity : Bytes
field unrestricted bobRunFreshCheckpointInputBytes : ByteCount
field unrestricted bobRunResumeCheckpointRequired : Nat
field unrestricted bobRunFreshProcessResumeRequired : Nat
field unrestricted bobRunNativeOnlyRequired : Nat
field unrestricted bobRunAllowedHostFallbacks : Nat

end-family

def bobCheckpointPlaneCount : Nat =
  3

def bobCheckpointBytesPerScalar : Nat =
  4

-- Derived from the model's table extent, optimizer planes and representation
-- width: 256 * 256 * 3 * 4 = 786432 bytes.
def bobCheckpointPayloadByteCount : Nat =
  (nat-multiply
    bobParameterCount
    (nat-multiply bobCheckpointPlaneCount bobCheckpointBytesPerScalar))

def bobCheckpointPayloadBytes : Nat =
  bobCheckpointPayloadByteCount

def bobCheckpointSchema : (family BobCheckpointTransportSchema) =
  (record
    BobCheckpointTransportSchema
    (bobCheckpointSchemaIdentity = b"bob-raw-pmv-f32-v1")
    (bobCheckpointParameterCount = bobParameterCount)
    (bobCheckpointPlaneCount = bobCheckpointPlaneCount)
    (bobCheckpointBytesPerScalar = bobCheckpointBytesPerScalar)
    (bobCheckpointPayloadBytes = bobCheckpointPayloadBytes)
    (bobCheckpointOrderingIdentity = b"parameter-order:p-then-m-then-v")
    (bobCheckpointInBandStepMetadata = 1)
    (bobCheckpointExternalStepReceiptRequired = 0))

def bobUpdatesPerInvocation : Nat = 10

def bobTrainingRunSemantics : (family BobTrainingRunSemantics) =
  (record
    BobTrainingRunSemantics
    (bobRunUpdatesPerInvocation = bobUpdatesPerInvocation)
    (bobRunDeterministicInitializationIdentity = b"bob-seed-v1")
    (bobRunFreshCheckpointInputBytes = (stdByteCount 0))
    (bobRunResumeCheckpointRequired = 1)
    (bobRunFreshProcessResumeRequired = 1)
    (bobRunNativeOnlyRequired = 1)
    (bobRunAllowedHostFallbacks = 0))

-- The input record one invocation reads (systems/bob/tools/build_fixture.py
-- writes it): the magic and version (8 bytes each), the corpus's 256
-- (byte, target) pairs, each byte's target as a 32-bit word, and each
-- prompt's last byte as a 32-bit word.
def bobInputHeaderBytes : Nat = 16
def bobInputWordBytes : Nat = 4
def bobPromptCount : Nat = 3
def bobInputTargetsOffset : Nat =
  (naturalAdd bobInputHeaderBytes (naturalMultiply 2 bobVocabularySize))
def bobInputPromptsOffset : Nat =
  (naturalAdd bobInputTargetsOffset (naturalMultiply bobInputWordBytes bobVocabularySize))
def bobInputRecordBytes : Nat =
  (naturalAdd bobInputPromptsOffset (naturalMultiply bobInputWordBytes bobPromptCount))

-- What an invocation returns per prompt: its token and its validity, a
-- 32-bit word each
def bobPredictionBytes : Nat = 8

-- The checkpoint's envelope (Checkpoint.Envelope): the state's schema, the
-- learner (objective, optimizer, its binary32 hyperparameters' bytes), the
-- data (the input record: the model's input, the vocabulary, the prompts,
-- the extent), the initialization's seed, one chunk per plane, the updates
-- an invocation makes.  A checkpoint of any other is refused by name.
def bobIdentityPart =
  (lambda unrestricted name : Bytes .
    (lambda unrestricted value : Nat .
      (bytes-append b"|" (bytes-append name (bytes-append b"=" (naturalDecimalBytesWithin 24 value))))))

def bobCheckpointContract : (family CheckpointEnvelopeContract) =
  (constructor CheckpointEnvelopeContract CheckpointEnvelopeContractValue
    (bytes-append b"bob-raw-pmv-f32-v1|parameter-order:p-then-m-then-v" (bobIdentityPart b"payload" bobCheckpointPayloadBytes))
    (bytes-append bobObjectiveIdentity
      (bytes-append b"|" (bytes-append bobLearnerIdentity
        (bytes-append b"|" (bytes-append (dataFloat32EncodeLE bobLearningRate) (bytes-append (dataFloat32EncodeLE bobBeta1)
          (bytes-append (dataFloat32EncodeLE bobBeta2) (bytes-append (dataFloat32EncodeLE bobEpsilon) (dataFloat32EncodeLE bobWeightDecay)))))))))
    (bytes-append b"raw-byte-sequence"
      (bytes-append (bobIdentityPart b"vocabulary" bobVocabularySize)
        (bytes-append (bobIdentityPart b"prompts" bobPromptCount) (bobIdentityPart b"record" bobInputRecordBytes))))
    b"bob-seed-v1"
    bobCheckpointPayloadBytes
    (naturalMultiply bobCheckpointBytesPerScalar bobParameterCount)
    bobUpdatesPerInvocation)