module Component.Linear

import Std.Foundation
import Std.Natural
import Std.List
import Representation.Schema

-- PRD 19 P3/P4: a linear map is a MECHANISM. It declares what it computes and
-- which parameters it owns; the parameter schema is DERIVED from that
-- declaration, keyed by the structural component path, so no definition
-- anywhere carries a hand-written parameter list (POC-001, LIB-003, REG-014).
-- Bias is a field of the map rather than a separate component: PRD 19 L25's
-- variants turn it on and off, and a bias without its map has no meaning.
family LinearComponent : Type 0
constructor LinearComponentOf
field unrestricted linearComponentName : Bytes
field unrestricted linearComponentInputs : Nat
field unrestricted linearComponentOutputs : Nat
field unrestricted linearComponentBias : (family StdBool)

end-family

def linearComponentNameOf =
  (lambda unrestricted component : (family LinearComponent) .
    (eliminate
      LinearComponent
      (lambda unrestricted current : (family LinearComponent) . Bytes)
      component
      (branch LinearComponentOf name inputs outputs bias . name)))

-- the weight's logical extent is inputs x outputs elements of the representation
def linearWeightStorage =
  (lambda unrestricted inputs : Nat .
    (lambda unrestricted outputs : Nat .
      (constructor
        LogicalStorage
        LogicalStorageOf
        (nat-multiply inputs outputs)
        4
        (constructor LogicalOwnership LogicalOwned)
        (constructor LogicalLifetime LogicalPersistent))))

def linearBiasStorage =
  (lambda unrestricted outputs : Nat .
    (constructor
      LogicalStorage
      LogicalStorageOf
      outputs
      4
      (constructor LogicalOwnership LogicalOwned)
      (constructor LogicalLifetime LogicalPersistent)))

def linearParameterAt =
  (lambda unrestricted path : (family StdList Bytes) .
    (lambda unrestricted shape : (family StdList Nat) .
      (lambda unrestricted storage : (family LogicalStorage) .
        (constructor
          ParameterSchema
          ParameterSchemaOf
          (constructor ParameterIdentity ParameterIdentityOf path)
          (constructor ParameterShape ParameterShapeOf shape)
          b"f32"
          storage))))

def componentPath =
  (lambda unrestricted prefix : (family StdList Bytes) .
    (lambda unrestricted name : Bytes .
      (lambda unrestricted leaf : Bytes .
        (stdListAppend
          Bytes
          prefix
          (constructor
            StdList
            StdListCons
            Bytes
            name
            (constructor StdList StdListCons Bytes leaf (constructor StdList StdListEmpty Bytes)))))))

-- The parameters a linear map owns under a path prefix: its weight, and its
-- bias exactly when the map declares one.
def linearParameters =
  (lambda unrestricted prefix : (family StdList Bytes) .
    (lambda unrestricted component : (family LinearComponent) .
      (eliminate
        LinearComponent
        (lambda unrestricted current : (family LinearComponent) .
          (family StdList (family ParameterSchema)))
        component
        (branch
          LinearComponentOf
          name
          inputs
          outputs
          bias
          .
          (constructor
            StdList
            StdListCons
            (family ParameterSchema)
            (linearParameterAt
              (componentPath prefix name b"weight")
              (constructor
                StdList
                StdListCons
                Nat
                inputs
                (constructor StdList StdListCons Nat outputs (constructor StdList StdListEmpty Nat)))
              (linearWeightStorage inputs outputs))
            (eliminate
              StdBool
              (lambda unrestricted current : (family StdBool) .
                (family StdList (family ParameterSchema)))
              bias
              (branch
                StdTrue
                .
                (constructor
                  StdList
                  StdListCons
                  (family ParameterSchema)
                  (linearParameterAt
                    (componentPath prefix name b"bias")
                    (constructor
                      StdList
                      StdListCons
                      Nat
                      outputs
                      (constructor StdList StdListEmpty Nat))
                    (linearBiasStorage outputs))
                  (constructor StdList StdListEmpty (family ParameterSchema))))
              (branch StdFalse . (constructor StdList StdListEmpty (family ParameterSchema)))))))))

def linearComponentInputsOf =
  (lambda unrestricted component : (family LinearComponent) .
    (eliminate
      LinearComponent
      (lambda unrestricted current : (family LinearComponent) . Nat)
      component
      (branch LinearComponentOf name inputs outputs bias . inputs)))

def linearComponentOutputsOf =
  (lambda unrestricted component : (family LinearComponent) .
    (eliminate
      LinearComponent
      (lambda unrestricted current : (family LinearComponent) . Nat)
      component
      (branch LinearComponentOf name inputs outputs bias . outputs)))