module Learning.IndependentState

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

-- A separable transition with a persistent prefix and immutable local inputs.
-- Each owner's complete input vector is declared in logical owner order. The
-- same typed assignment semantics serve gradient and non-gradient equations.
-- A dependency on a remote statistic cannot be hidden in this representation:
-- it requires a different composition with an explicit communication edge.
family IndependentStateSystem : Type 0
parameter erased independentSlots : Nat
constructor IndependentStateSystemOf
field unrestricted independentPersistent : Nat
field unrestricted independentProgram : (family StdList (family PointwiseAssignment StdU64 independentSlots))
field unrestricted independentInitial : (family StdList (family StdList StdU64))

end-family

def independentPersistentCount =
  (lambda erased slots : Nat .
    (lambda unrestricted system : (family IndependentStateSystem slots) .
      (eliminate
        IndependentStateSystem
        (lambda unrestricted current : (family IndependentStateSystem slots) . Nat)
        system
        (branch IndependentStateSystemOf persistent program initial . persistent))))

def independentAssignments =
  (lambda erased slots : Nat .
    (lambda unrestricted system : (family IndependentStateSystem slots) .
      (eliminate
        IndependentStateSystem
        (lambda unrestricted current : (family IndependentStateSystem slots) .
          (family StdList (family PointwiseAssignment StdU64 slots)))
        system
        (branch IndependentStateSystemOf persistent program initial . program))))

def independentInitialValues =
  (lambda erased slots : Nat .
    (lambda unrestricted system : (family IndependentStateSystem slots) .
      (eliminate
        IndependentStateSystem
        (lambda unrestricted current : (family IndependentStateSystem slots) .
          (family StdList (family StdList StdU64)))
        system
        (branch IndependentStateSystemOf persistent program initial . initial))))

def independentStateAdmitted =
  (lambda unrestricted slots : Nat .
    (lambda unrestricted system : (family IndependentStateSystem slots) .
      (eliminate
        IndependentStateSystem
        (lambda unrestricted current : (family IndependentStateSystem slots) . Nat)
        system
        (branch
          IndependentStateSystemOf
          persistent
          program
          initial
          .
          (naturalAnd
            (naturalBelowTwoPower 64 slots)
            (naturalAnd
              (naturalNonzero persistent)
              (naturalAnd
                (naturalLessOrEqual persistent slots)
                (naturalAnd
                  (naturalNonzero (stdListLength (family StdList StdU64) initial))
                  (naturalAnd
                    (stdListFold
                      (family StdList StdU64)
                      Nat
                      (lambda unrestricted row : (family StdList StdU64) .
                        (lambda unrestricted rest : Nat .
                          (naturalAnd rest (naturalEqual (stdListLength StdU64 row) slots))))
                      1
                      initial)
                    (stdListFold
                      (family PointwiseAssignment StdU64 slots)
                      Nat
                      (lambda unrestricted assignment : (family PointwiseAssignment StdU64 slots) .
                        (lambda unrestricted rest : Nat .
                          (eliminate
                            PointwiseAssignment
                            (lambda unrestricted current : (family PointwiseAssignment StdU64 slots) .
                              Nat)
                            assignment
                            (branch
                              PointwiseAssign
                              destination
                              expression
                              .
                              (naturalAnd
                                rest
                                (naturalLess (stdFinToNatural slots destination) persistent))))))
                      1
                      program))))))))))

def independentStateIdentity =
  (lambda unrestricted slots : Nat .
    (lambda unrestricted system : (family IndependentStateSystem slots) .
      (eliminate
        IndependentStateSystem
        (lambda unrestricted current : (family IndependentStateSystem slots) . Bytes)
        system
        (branch
          IndependentStateSystemOf
          persistent
          program
          initial
          .
          (bytes-append
            b"alpha-independent-state-v1;"
            (bytes-append
              (pointwiseIdentityWord persistent)
              (bytes-append
                (pointwiseProgramIdentity slots program)
                (bytes-append
                  (pointwiseIdentityWord (stdListLength (family StdList StdU64) initial))
                  (stdListFold
                    (family StdList StdU64)
                    Bytes
                    (lambda unrestricted row : (family StdList StdU64) .
                      (lambda unrestricted rest : Bytes .
                        (bytes-append
                          (stdListFold
                            StdU64
                            Bytes
                            (lambda unrestricted word : StdU64 .
                              (lambda unrestricted tail : Bytes .
                                (bytes-append (dataBytesWord64LE word) tail)))
                            b""
                            row)
                          rest)))
                    b""
                    initial)))))))))