module LocalCorrelation.System

import Component.Pointwise
import Learning.IndependentState
import Learning.ScaledProductUpdate
import Model.Parameter
import Std.Fin
import Std.Float
import Std.List
import Std.Word

-- Two coordinates of one associative state. Each coordinate learns from an
-- immutable local observation pair. Evaluation sums the complete saved state;
-- dropping either owner changes the result and prevents valid publication.
def correlationState =
  (stdFinFirst 3)

def correlationRate =
  (stdFinNextOf 3 (stdFinFirst 2))

def correlationLeft =
  (stdFinNextOf 3 (stdFinNextOf 2 (stdFinFirst 1)))

def correlationRight =
  (stdFinNextOf 3 (stdFinNextOf 2 (stdFinNextOf 1 (stdFinFirst 0))))

def correlationAssignments =
  (scaledProductAssignments
    StdU64
    4
    correlationState
    correlationRate
    correlationLeft
    correlationRight
    (refl Nat 1))

def correlationRow =
  (lambda unrestricted state : F64 .
    (lambda unrestricted rate : F64 .
      (lambda unrestricted left : F64 .
        (lambda unrestricted right : F64 .
          (constructor
            StdList
            StdListCons
            StdU64
            (stdF64ToWord64 state)
            (constructor
              StdList
              StdListCons
              StdU64
              (stdF64ToWord64 rate)
              (constructor
                StdList
                StdListCons
                StdU64
                (stdF64ToWord64 left)
                (constructor
                  StdList
                  StdListCons
                  StdU64
                  (stdF64ToWord64 right)
                  (constructor StdList StdListEmpty StdU64)))))))))

def correlationSystemWith =
  (lambda unrestricted rate : F64 .
    (constructor
      IndependentStateSystem
      IndependentStateSystemOf
      4
      1
      correlationAssignments
      (constructor
        StdList
        StdListCons
        (family StdList StdU64)
        (correlationRow 1.0 rate 2.0 3.0)
        (constructor
          StdList
          StdListCons
          (family StdList StdU64)
          (correlationRow 2.0 rate 4.0 1.0)
          (constructor StdList StdListEmpty (family StdList StdU64))))))

def correlationSystem =
  (correlationSystemWith 0.5)

def correlationEvaluation =
  (constructor
    StdList
    StdListCons
    (family PointwiseAssignment StdU64 2)
    (constructor
      PointwiseAssignment
      PointwiseAssign
      StdU64
      2
      (stdFinFirst 1)
      (pointwiseBinary
        StdU64
        2
        (constructor PointwiseBinary PointwiseAdd)
        (constructor PointwiseExpression PointwiseInput StdU64 2 (stdFinFirst 1))
        (constructor PointwiseExpression PointwiseInput StdU64 2 (stdFinNextOf 1 (stdFinFirst 0)))))
    (constructor StdList StdListEmpty (family PointwiseAssignment StdU64 2)))