module Coppelius.TrainingRun

import Checkpoint.Envelope
import Coppelius.Learner
import Coppelius.Model
import Model.Config
import Model.Parameter
import Model.Word64
import Std.Natural
import Std.Physical
import Std.Word

-- Coppelius trains under its run policy (below): one process, a fixed
-- number of AdamW steps, each on the next context of a token stream.  Loss
-- is telemetry, never a stopping condition.  A zero-byte checkpoint input
-- selects fresh deterministic initialization; a full raw P/M/V image
-- selects continuation.
family CoppeliusTrainingRunSemantics : Type 0
constructor CoppeliusTrainingRunSemanticsValue
field unrestricted coppeliusRunUpdatesPerInvocation : Nat
field unrestricted coppeliusRunCheckpointInterval : Nat
field unrestricted coppeliusRunLossControlsTermination : Nat
field unrestricted coppeliusRunDeterministicInitialization : Nat
field unrestricted coppeliusRunFreshCheckpointInputBytes : ByteCount
field unrestricted coppeliusRunResumeCheckpointRequired : Nat
field unrestricted coppeliusRunCheckpointOutputRequired : Nat
field unrestricted coppeliusRunFreshProcessResumeRequired : Nat
field unrestricted coppeliusRunNativeOnlyRequired : Nat
field unrestricted coppeliusRunAllowedHostFallbacks : Nat

end-family

-- The checkpoint's payload: a contiguous FP32 P/M/V image.  The file wraps
-- it in the envelope (coppeliusCheckpointContract, below), whose header
-- carries the update count in band.
family CoppeliusCheckpointTransportSchema : Type 0
constructor CoppeliusCheckpointTransportSchemaValue
field unrestricted coppeliusCheckpointSchemaIdentity : Bytes
field unrestricted coppeliusCheckpointParameterCount : StdU64
field unrestricted coppeliusCheckpointPlaneCount : Nat
field unrestricted coppeliusCheckpointBytesPerScalar : Nat
field unrestricted coppeliusCheckpointPayloadBytes : ByteCount
field unrestricted coppeliusCheckpointOrderingIdentity : Bytes
field unrestricted coppeliusCheckpointInBandStepMetadata : Nat
field unrestricted coppeliusCheckpointExternalStepReceiptRequired : Nat

end-family

-- ---- the run policy ----
-- One process takes `coppeliusRunSteps` AdamW steps, step s on context
-- c0 + s of the token stream the request names (c0 the contexts its
-- checkpoint has consumed: 0 fresh), and publishes its checkpoint after
-- every `coppeliusRunCheckpointInterval` steps -- the last publication at
-- its end.  A checkpoint goes to the request's checkpoint path through its
-- temporary: written whole, digested, synced, then renamed over the path,
-- so the path always holds a complete checkpoint (Platform.Linux.Nvidia.
-- PlanHostRequest).  The loop is compiled: the host issues one update
-- submission of the plan per step (Coppelius.Build.NativeHost).
def coppeliusRunSteps : Nat =
  10000

def coppeliusRunCheckpointInterval : Nat =
  1000

-- publications a run makes: the interval divides the steps
def coppeliusRunPublications : Nat =
  (naturalDivideUnchecked coppeliusRunSteps coppeliusRunCheckpointInterval)

def coppeliusRunIntervalDividesSteps :
  (equal Nat (naturalMultiply coppeliusRunPublications coppeliusRunCheckpointInterval) coppeliusRunSteps) =
  (refl Nat coppeliusRunSteps)

-- the plan's update submissions: one, issued once per step
def coppeliusUpdatePieces : Nat =
  1

def coppeliusFreshCheckpointInputBytes : ByteCount =
  (stdByteCount 0)

def coppeliusTrainingRunSemantics : (family CoppeliusTrainingRunSemantics) =
  (record
    CoppeliusTrainingRunSemantics
    (coppeliusRunUpdatesPerInvocation = coppeliusRunSteps)
    (coppeliusRunCheckpointInterval = coppeliusRunCheckpointInterval)
    (coppeliusRunLossControlsTermination = 0)
    (coppeliusRunDeterministicInitialization = 1)
    (coppeliusRunFreshCheckpointInputBytes = coppeliusFreshCheckpointInputBytes)
    (coppeliusRunResumeCheckpointRequired = 1)
    (coppeliusRunCheckpointOutputRequired = 1)
    (coppeliusRunFreshProcessResumeRequired = 1)
    (coppeliusRunNativeOnlyRequired = 1)
    (coppeliusRunAllowedHostFallbacks = 0))

def coppeliusCheckpointParameterCount : StdU64 =
  57_705_472

def coppeliusCheckpointPayloadBytesNatural : Nat = 692_465_664
def coppeliusCheckpointPayloadBytes : ByteCount =
  (stdByteCount (modelWord64FromNaturalTruncated coppeliusCheckpointPayloadBytesNatural))

-- The input record a step consumes, as the kernels' parameter blocks place
-- it in the staging region: the context's 1024 input tokens, and its 1024
-- targets at 0x1000.  Nothing else: the rotary tables are computed from the
-- model (Coppelius.Build.Graph.cgRotaryTables), and a fresh model is drawn
-- from the run's seed.
def coppeliusInputSequence : Nat = 1024
def coppeliusInputTargetsOffset : Nat = 0x1000
def coppeliusInputRecordBytes : Nat =
  (naturalAdd coppeliusInputTargetsOffset (naturalMultiply 4 coppeliusInputSequence))

-- The token stream the host reads the records from: a token cache (the
-- tokenizer's output), a 16-byte header (the source's modification time
-- and size) then little-endian u32 token ids.  Context c is ids
-- [1024 c, 1024 c + 1024]: its inputs the first 1024, its targets the last
-- 1024 -- the same window one token on.
def coppeliusTokenStreamHeaderBytes : Nat = 16
def coppeliusTokenBytes : Nat = 4
def coppeliusContextStreamBytes : Nat = (naturalMultiply coppeliusTokenBytes coppeliusInputSequence)

-- the seed a fresh model is drawn from
def coppeliusRunSeed : Nat =
  20260925

-- What the kernels leave in the staging region for the host to return (the
-- R6 result layout, R9 appended the host timestamp table): the loss block
-- at 0x180000 (16 KiB, four 4 KiB windows: the last step's per-row
-- cross-entropy -- the target logit gathered into `correct` makes the
-- cross-entropy rows the loss; before that they held the log-sum-exp --
-- and its embedding gradient, then the embedding parameters and their AdamW
-- first moment after the final step), the three next-token predictions at
-- 0x1F0000 (three u64), and the result record at 0x200000 -- its first
-- 147,456 bytes are written after the forward submission, before training
-- (the marker below is put at its start first), the next 49,152 after the
-- run.  The result file is, in the order the host writes it: the first
-- 147,456 bytes of the record (the forward on the run's first context,
-- before any step); every step's per-row cross-entropy (1024 binary32
-- words, one 4 KiB window per step, appended as the step completes: the
-- update's loss copy leaves it at 0x190000); then, after the run, the
-- checkpoint input's stat (160 bytes, with gpGet, gpPut and the first
-- semaphore word at 96, 100 and 104), the loss block, 49,152 bytes of the
-- record (row 1023's logits after the final step), the semaphore region
-- (32 bytes per submission of the plan: the push segment's final release,
-- the fence the host waits for, and its pre-launch release, each with the
-- device's timestamp -- a submission issued again leaves its last issue's),
-- 64 bytes of the error notifier, the host's 160-byte status block, the
-- timestamp table (64 bytes per submission the run issues, in issue order).
-- A step's loss is the forward on a context no earlier step has trained
-- on: the loss on unseen text.
def coppeliusResultWindowBytes : Nat = 0x1000
def coppeliusResultLossRecordOffset : Nat = 0x190000
def coppeliusResultLossBlockOffset : Nat = 0x180000
def coppeliusResultLossBlockBytes : Nat = 16384
def coppeliusResultPredictionsOffset : Nat = 0x1F0000
def coppeliusResultPredictionsBytes : Nat = 24
-- the rows whose next token the invocation predicts, in the order their
-- predictions (token, validity: 8 bytes each) and logits are returned
def coppeliusPredictionRow0 : Nat = 255
def coppeliusPredictionRow1 : Nat = 511
def coppeliusPredictionRow2 : Nat = 1023
def coppeliusPredictionRows : Nat = 3
def coppeliusResultPredictionBytes : Nat = 8
def coppeliusResultRecordOffset : Nat = 0x200000
def coppeliusResultRecordForwardBytes : Nat = 147456
def coppeliusResultRecordFinalBytes : Nat = 49152
def coppeliusResultRecordMarker : Bytes = b"\xbe\xba\xfe\xca\xbe\xba\xfe\xca"
def coppeliusResultErrorNotifierBytes : Nat = 64
def coppeliusResultStatGPGetOffset : Nat = 96
def coppeliusResultStatGPPutOffset : Nat = 100
def coppeliusResultStatSemaphoreOffset : Nat = 104

def coppeliusRawCheckpointSchemaIdentity : Bytes =
  b"coppelius-raw-pmv-f32-v1"

def coppeliusParameterOrderingIdentity : Bytes =
  b"parameter-order:p-then-m-then-v"

def coppeliusCheckpointSchema : (family CoppeliusCheckpointTransportSchema) =
  (record
    CoppeliusCheckpointTransportSchema
    (coppeliusCheckpointSchemaIdentity = coppeliusRawCheckpointSchemaIdentity)
    (coppeliusCheckpointParameterCount = coppeliusCheckpointParameterCount)
    (coppeliusCheckpointPlaneCount = 3)
    (coppeliusCheckpointBytesPerScalar = 4)
    (coppeliusCheckpointPayloadBytes = coppeliusCheckpointPayloadBytes)
    (coppeliusCheckpointOrderingIdentity = coppeliusParameterOrderingIdentity)
    (coppeliusCheckpointInBandStepMetadata = 1)
    (coppeliusCheckpointExternalStepReceiptRequired = 0))

-- The checkpoint's envelope (Checkpoint.Envelope): the state's schema (the
-- raw P/M/V layout, its order, its extent), the learner (objective,
-- optimizer, hyperparameters), the data (the token stream: vocabulary,
-- context, its layout), the initialization (the generator and its seed),
-- chunks of the staging window's size, a run's steps.  A checkpoint of any
-- other is refused by name.
def coppeliusCheckpointChunkBytes : Nat = 4194304

def coppeliusIdentityPart =
  (lambda unrestricted name : Bytes .
    (lambda unrestricted value : Nat .
      (bytes-append b"|" (bytes-append name (bytes-append b"=" (naturalDecimalBytesWithin 24 value))))))

def coppeliusCheckpointContract : (family CheckpointEnvelopeContract) =
  (constructor CheckpointEnvelopeContract CheckpointEnvelopeContractValue
    (bytes-append coppeliusRawCheckpointSchemaIdentity
      (bytes-append b"|" (bytes-append coppeliusParameterOrderingIdentity (coppeliusIdentityPart b"payload" coppeliusCheckpointPayloadBytesNatural))))
    (bytes-append coppeliusObjectiveIdentity
      (bytes-append b"|halves=" (bytes-append coppeliusHalfFormatIdentity
      (bytes-append b"|" (bytes-append coppeliusLearnerIdentity
        (bytes-append (coppeliusIdentityPart b"lr" coppeliusLearningRateNumerator) (bytes-append (coppeliusIdentityPart b"/" coppeliusLearningRateDenominator)
        (bytes-append (coppeliusIdentityPart b"beta1" coppeliusBeta1Numerator) (bytes-append (coppeliusIdentityPart b"/" coppeliusBeta1Denominator)
        (bytes-append (coppeliusIdentityPart b"beta2" coppeliusBeta2Numerator) (bytes-append (coppeliusIdentityPart b"/" coppeliusBeta2Denominator)
        (bytes-append (coppeliusIdentityPart b"decay" coppeliusWeightDecayNumerator) (bytes-append (coppeliusIdentityPart b"/" coppeliusWeightDecayDenominator)
        (bytes-append (coppeliusIdentityPart b"epsilon" coppeliusEpsilonNumerator) (coppeliusIdentityPart b"/" coppeliusEpsilonDenominator)))))))))))))))
    (bytes-append b"coppelius-token-stream"
      (bytes-append (coppeliusIdentityPart b"vocabulary" (stdU32ToNatural modelCoppeliusVocab))
        (bytes-append (coppeliusIdentityPart b"context" coppeliusInputSequence)
          (bytes-append (coppeliusIdentityPart b"header" coppeliusTokenStreamHeaderBytes) (coppeliusIdentityPart b"token" coppeliusTokenBytes)))))
    (bytes-append b"coppelius-initialization:random-normal" (coppeliusIdentityPart b"seed" coppeliusRunSeed))
    coppeliusCheckpointPayloadBytesNatural
    coppeliusCheckpointChunkBytes
    coppeliusRunSteps)