module PartitionedRegression.Coordinates

import Learning.ScaledGradientUpdate
import Checkpoint.RegionManifest
import Representation.Partition
import Std.List
import Std.Natural

-- One prediction depends on both parameter owners. No participant can train
-- its coordinate from its own prediction alone: its adjoint requires the
-- residual of the sum. Device arithmetic is a separate admitted realization.
def coordinatePrediction =
  (lambda erased value : Type 0 .
    (lambda unrestricted add : (pi unrestricted left : value . (pi unrestricted right : value . value)) .
      (lambda unrestricted multiply : (pi unrestricted left : value . (pi unrestricted right : value . value)) .
        (lambda unrestricted leftWeight : value .
          (lambda unrestricted rightWeight : value .
            (lambda unrestricted leftInput : value .
              (lambda unrestricted rightInput : value .
                (add (multiply leftWeight leftInput) (multiply rightWeight rightInput)))))))))

-- Persistent vectors have the same logical coordinate domain. Their physical
-- addresses and owners are supplied by placement; the epoch/update/cursor are
-- the common capture prefix, not duplicated parameter fields.
def coordinateCount : Nat =
  2

def coordinateWordBytes : Nat =
  8

def coordinateVectorBytes =
  (naturalMultiply coordinateCount coordinateWordBytes)

def coordinateRegions =
  (constructor
    StdList
    StdListCons
    (family CheckpointLogicalRegion)
    (constructor
      CheckpointLogicalRegion
      CheckpointLogicalRegionOf
      b"parameters"
      b"f64-le"
      coordinateVectorBytes)
    (constructor
      StdList
      StdListCons
      (family CheckpointLogicalRegion)
      (constructor
        CheckpointLogicalRegion
        CheckpointLogicalRegionOf
        b"scaled-gradient-accumulators"
        b"f64-le"
        coordinateVectorBytes)
      (constructor StdList StdListEmpty (family CheckpointLogicalRegion))))

def coordinateInterval =
  (lambda unrestricted coordinate : Nat .
    (constructor
      LogicalInterval
      LogicalIntervalOf
      (naturalMultiply coordinate coordinateWordBytes)
      (naturalMultiply (succ coordinate) coordinateWordBytes)))

def coordinateIntervals =
  (stdListMap
    Nat
    (family LogicalInterval)
    coordinateInterval
    (stdListNaturalsBelow coordinateCount))

def coordinatePartitions =
  (stdListMap
    Nat
    (family CheckpointRegionPartition)
    (lambda unrestricted ordinal : Nat .
      (constructor
        CheckpointRegionPartition
        CheckpointRegionPartitionOf
        ordinal
        coordinateIntervals))
    (stdListNaturalsBelow (stdListLength (family CheckpointLogicalRegion) coordinateRegions)))