module Component.PointwiseIdentity

import Component.Pointwise
import Model.Parameter
import Data.Bytes
import Model.Word64
import Std.Fin
import Std.List
import Std.Natural
import Std.Word

-- Versioned prefix encoding of typed expression trees. Tags encode structure;
-- every number and literal occupies one canonical little-endian word. Input
-- count, assignment count and destination order are part of the identity.
-- This is semantic input to checkpoint hashing, not an executable packager.
def pointwiseIdentityWord =
  (lambda unrestricted value : Nat . (dataBytesWord64LE (modelWord64FromNaturalTruncated value)))

def pointwiseUnaryIdentity =
  (lambda unrestricted operation : (family PointwiseUnary) .
    (eliminate
      PointwiseUnary
      (lambda unrestricted current : (family PointwiseUnary) . Bytes)
      operation
      (branch PointwiseNegate . b"n")
      (branch PointwiseSquare . b"s")
      (branch PointwiseRoundBF16 . b"b")
      (branch PointwiseReciprocal . b"r")
      (branch PointwiseReciprocalSquareRoot . b"i")
      (branch PointwiseSquareRoot . b"q")
      (branch PointwiseTanh . b"t")
      (branch PointwiseExponentialBase2 . b"e")))

def pointwiseBinaryIdentity =
  (lambda unrestricted operation : (family PointwiseBinary) .
    (eliminate
      PointwiseBinary
      (lambda unrestricted current : (family PointwiseBinary) . Bytes)
      operation
      (branch PointwiseAdd . b"+")
      (branch PointwiseMultiply . b"*")
      (branch PointwiseSubtract . b"-")
      (branch PointwisePositiveGate . b">")))

def pointwiseExpressionIdentity =
  (lambda unrestricted slots : Nat .
    (lambda unrestricted expression : (family PointwiseExpression StdU64 slots) .
      (pointwiseFold
        StdU64
        Bytes
        slots
        (lambda unrestricted operation : (family PointwiseUnary) .
          (lambda unrestricted child : Bytes .
            (bytes-append b"u" (bytes-append (pointwiseUnaryIdentity operation) child))))
        (lambda unrestricted operation : (family PointwiseBinary) .
          (lambda unrestricted left : Bytes .
            (lambda unrestricted right : Bytes .
              (bytes-append
                b"b"
                (bytes-append (pointwiseBinaryIdentity operation) (bytes-append left right))))))
        (lambda unrestricted bits : StdU64 . (bytes-append b"c" (dataBytesWord64LE bits)))
        expression
        (lambda unrestricted index : (family StdFin slots) .
          (bytes-append b"i" (pointwiseIdentityWord (stdFinToNatural slots index)))))))

def pointwiseProgramIdentity =
  (lambda unrestricted slots : Nat .
    (lambda unrestricted program : (family StdList (family PointwiseAssignment StdU64 slots)) .
      (bytes-append
        b"alpha-pointwise-v1;"
        (bytes-append
          (pointwiseIdentityWord slots)
          (bytes-append
            (pointwiseIdentityWord
              (stdListLength (family PointwiseAssignment StdU64 slots) program))
            (stdListFold
              (family PointwiseAssignment StdU64 slots)
              Bytes
              (lambda unrestricted assignment : (family PointwiseAssignment StdU64 slots) .
                (lambda unrestricted rest : Bytes .
                  (eliminate
                    PointwiseAssignment
                    (lambda unrestricted current : (family PointwiseAssignment StdU64 slots) .
                      Bytes)
                    assignment
                    (branch
                      PointwiseAssign
                      destination
                      expression
                      .
                      (bytes-append
                        (pointwiseIdentityWord (stdFinToNatural slots destination))
                        (bytes-append (pointwiseExpressionIdentity slots expression) rest))))))
              b""
              program))))))