module Learning.ScaledProductUpdate

import Component.Pointwise
import Std.Fin
import Std.List
import Std.Natural

-- A local correlation update (often called a Hebbian update). The arithmetic
-- order is explicit: form the observation product, scale it, then add to state.
-- No derivative, objective, optimizer state or target realization is required.
def accumulateScaledProduct =
  (lambda erased scalar : Type 0 .
    (lambda unrestricted add : (pi unrestricted left : scalar . (pi unrestricted right : scalar . scalar)) .
      (lambda unrestricted multiply : (pi unrestricted left : scalar . (pi unrestricted right : scalar . scalar)) .
        (lambda unrestricted coefficient : scalar .
          (lambda unrestricted previous : scalar .
            (lambda unrestricted left : scalar .
              (lambda unrestricted right : scalar .
                (add previous (multiply coefficient (multiply left right))))))))))

def scaledProductSlotsIndependent =
  (lambda erased slots : Nat .
    (lambda unrestricted state : (family StdFin slots) .
      (lambda unrestricted coefficient : (family StdFin slots) .
        (lambda unrestricted left : (family StdFin slots) .
          (lambda unrestricted right : (family StdFin slots) .
            (naturalAnd
              (pointwiseSlotsSeparate slots state coefficient)
              (naturalAnd
                (pointwiseSlotsSeparate slots state left)
                (pointwiseSlotsSeparate slots state right))))))))

-- Read-only observations may alias each other; they must not alias the state
-- being mutated. The same independent-slot obligation used by scaled-gradient
-- compositions prevents a placement from silently changing the local rule.
def scaledProductAssignments =
  (lambda erased word : Type 0 .
    (lambda erased slots : Nat .
      (lambda unrestricted state : (family StdFin slots) .
        (lambda unrestricted coefficient : (family StdFin slots) .
          (lambda unrestricted left : (family StdFin slots) .
            (lambda unrestricted right : (family StdFin slots) .
              (lambda erased independent : (equal Nat (scaledProductSlotsIndependent slots state coefficient left right) 1) .
                (let unrestricted input =
                  (lambda unrestricted index : (family StdFin slots) .
                    (constructor PointwiseExpression PointwiseInput word slots index))
                  in
                  (constructor
                    StdList
                    StdListCons
                    (family PointwiseAssignment word slots)
                    (constructor
                      PointwiseAssignment
                      PointwiseAssign
                      word
                      slots
                      state
                      (accumulateScaledProduct
                        (family PointwiseExpression word slots)
                        (pointwiseBinary word slots (constructor PointwiseBinary PointwiseAdd))
                        (pointwiseBinary word slots (constructor PointwiseBinary PointwiseMultiply))
                        (input coefficient)
                        (input state)
                        (input left)
                        (input right)))
                    (constructor StdList StdListEmpty (family PointwiseAssignment word slots)))))))))))