module Runtime.NativePointwisePlan

import Component.Pointwise
import Runtime.NativeFloatPlan
import Runtime.NativeLaunchRecipeLoad
import Runtime.NativeMemoryPlan
import Runtime.NativePhysicalProgram
import Data.Bytes
import Model.Word64
import Std.Fin
import Std.List
import Std.Foundation
import Std.Natural
import Std.Word

family NativePointwiseFragment : Type 0
constructor NativePointwiseFragmentOf
field unrestricted nativePointwiseFragmentCommands : (pi unrestricted tail : (family NativePhysicalCommands) . (family NativePhysicalCommands))
field unrestricted nativePointwiseFragmentResult : Nat
field unrestricted nativePointwiseFragmentEnd : Nat

end-family

-- This is an explicitly selected native host realization. The state occupies
-- a contiguous word region followed by derived expression scratch. No GPU
-- operation can choose these commands as a fallback. Arithmetic is finite,
-- separately rounded binary64; refusal prevents publication of a partial step.
def nativePointwiseWordBytes : Nat =
  (bytes-length (dataBytesWord64LE modelWord64Zero))

def nativePointwiseAt =
  (lambda unrestricted base : Nat .
    (lambda unrestricted index : Nat .
      (naturalAdd base (naturalMultiply nativePointwiseWordBytes index))))

def nativePointwiseExtent =
  (lambda unrestricted slots : Nat .
    (lambda unrestricted program : (family StdList (family PointwiseAssignment StdU64 slots)) .
      (lambda unrestricted base : Nat .
        (nativePointwiseAt
          base
          (naturalAdd slots (pointwiseProgramTemporaryValues StdU64 slots program))))))

def nativePointwiseCode =
  (lambda unrestricted fragment : (family NativePointwiseFragment) .
    (eliminate
      NativePointwiseFragment
      (lambda unrestricted current : (family NativePointwiseFragment) .
        (pi unrestricted tail : (family NativePhysicalCommands) . (family NativePhysicalCommands)))
      fragment
      (branch NativePointwiseFragmentOf code result end . code)))

def nativePointwiseResult =
  (lambda unrestricted fragment : (family NativePointwiseFragment) .
    (eliminate
      NativePointwiseFragment
      (lambda unrestricted current : (family NativePointwiseFragment) . Nat)
      fragment
      (branch NativePointwiseFragmentOf code result end . result)))

def nativePointwiseEnd =
  (lambda unrestricted fragment : (family NativePointwiseFragment) .
    (eliminate
      NativePointwiseFragment
      (lambda unrestricted current : (family NativePointwiseFragment) . Nat)
      fragment
      (branch NativePointwiseFragmentOf code result end . end)))

def nativePointwiseUnarySupported =
  (lambda unrestricted operation : (family PointwiseUnary) .
    (eliminate
      PointwiseUnary
      (lambda unrestricted current : (family PointwiseUnary) . Nat)
      operation
      (branch PointwiseNegate . 1)
      (branch PointwiseSquare . 1)
      (branch PointwiseRoundBF16 . 0)
      (branch PointwiseReciprocal . 1)
      (branch PointwiseReciprocalSquareRoot . 0)
      (branch PointwiseSquareRoot . 1)
      (branch PointwiseTanh . 0)
      (branch PointwiseExponentialBase2 . 0)))

def nativePointwiseBinarySupported =
  (lambda unrestricted operation : (family PointwiseBinary) .
    (eliminate
      PointwiseBinary
      (lambda unrestricted current : (family PointwiseBinary) . Nat)
      operation
      (branch PointwiseAdd . 1)
      (branch PointwiseMultiply . 1)
      (branch PointwiseSubtract . 1)
      (branch PointwisePositiveGate . 0)))

def nativePointwiseSupported =
  (lambda unrestricted slots : Nat .
    (lambda unrestricted program : (family StdList (family PointwiseAssignment StdU64 slots)) .
      (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
                  (pointwiseFold
                    StdU64
                    Nat
                    slots
                    (lambda unrestricted operation : (family PointwiseUnary) .
                      (lambda unrestricted child : Nat .
                        (naturalAnd (nativePointwiseUnarySupported operation) child)))
                    (lambda unrestricted operation : (family PointwiseBinary) .
                      (lambda unrestricted left : Nat .
                        (lambda unrestricted right : Nat .
                          (naturalAnd
                            (nativePointwiseBinarySupported operation)
                            (naturalAnd left right)))))
                    (lambda unrestricted word : StdU64 . 1)
                    expression
                    (lambda unrestricted index : (family StdFin slots) . 1)))))))
        1
        program)))

def nativePointwiseAdmitted =
  (lambda unrestricted slots : Nat .
    (lambda unrestricted program : (family StdList (family PointwiseAssignment StdU64 slots)) .
      (lambda unrestricted base : Nat .
        (lambda unrestricted available : Nat .
          (naturalAnd
            (nativePointwiseSupported slots program)
            (naturalAnd
              (naturalIsZero (naturalModuloUnchecked base nativePointwiseWordBytes))
              (naturalAnd
                (naturalLessOrEqual (nativePointwiseExtent slots program base) available)
                (naturalBelowTwoPower 64 (nativePointwiseExtent slots program base)))))))))

-- Adding the sign bit modulo 2^64 flips it without rounding and preserves both
-- zero signs. Subtracting from +0 would not implement negation of +0 correctly.
def nativePointwiseNegate =
  (lambda unrestricted scratch : Nat .
    (lambda unrestricted destination : Nat .
      (lambda unrestricted source : Nat .
        (lambda unrestricted tail : (family NativePhysicalCommands) .
          (nativeFloat64RequireFinite
            (nativeMemoryLoad source)
            scratch
            b"pointwise-nonfinite"
            (nativeLaunchRecipeCommand
              (constructor
                NativePhysicalOperation
                NativePhysicalAddWord64
                (nativeMemoryState destination)
                (nativeMemoryLoad source)
                (nativeLaunchRecipeImmediate (naturalPowerOfTwo 63)))
              b"pointwise-negate"
              tail))))))

-- binary64 1.0, the dividend of a reciprocal
def nativePointwiseOneBits : Nat =
  0x3ff0000000000000

-- Binary64 under the native finite profile (Runtime.NativeFloatPlan):
-- negation flips the sign bit; square, reciprocal (1 / x) and square root are
-- single correctly rounded operations; any nonfinite operand or result
-- refuses. A reciprocal square root would be two roundings here, unlike the
-- device profiles' single approximate operation, so it is not admitted.
def nativePointwiseUnary =
  (lambda unrestricted scratch : Nat .
    (lambda unrestricted operation : (family PointwiseUnary) .
      (lambda unrestricted destination : Nat .
        (lambda unrestricted source : Nat .
          (lambda unrestricted tail : (family NativePhysicalCommands) .
            (eliminate
              PointwiseUnary
              (lambda unrestricted current : (family PointwiseUnary) .
                (family NativePhysicalCommands))
              operation
              (branch PointwiseNegate . (nativePointwiseNegate scratch destination source tail))
              (branch
                PointwiseSquare
                .
                (nativeFiniteFloat64Binary
                  b"pointwise-nonfinite"
                  scratch
                  (constructor NativePhysicalFloat64Operation NativePhysicalFloat64Multiply)
                  destination
                  source
                  source
                  tail))
              (branch PointwiseRoundBF16 . tail)
              (branch
                PointwiseReciprocal
                .
                (nativeMemoryStoreWord64
                  b"pointwise-one"
                  destination
                  (nativeLaunchRecipeImmediate nativePointwiseOneBits)
                  (nativeFiniteFloat64Binary
                    b"pointwise-nonfinite"
                    scratch
                    (constructor NativePhysicalFloat64Operation NativePhysicalFloat64Divide)
                    destination
                    destination
                    source
                    tail)))
              (branch PointwiseReciprocalSquareRoot . tail)
              (branch
                PointwiseSquareRoot
                .
                (nativeFiniteFloat64Binary
                  b"pointwise-nonfinite"
                  scratch
                  (constructor NativePhysicalFloat64Operation NativePhysicalFloat64SquareRoot)
                  destination
                  source
                  source
                  tail))
              (branch PointwiseTanh . tail)
              (branch PointwiseExponentialBase2 . tail)))))))

def nativePointwiseBinary =
  (lambda unrestricted scratch : Nat .
    (lambda unrestricted operation : (family PointwiseBinary) .
      (lambda unrestricted destination : Nat .
        (lambda unrestricted left : Nat .
          (lambda unrestricted right : Nat .
            (lambda unrestricted tail : (family NativePhysicalCommands) .
              (eliminate
                PointwiseBinary
                (lambda unrestricted current : (family PointwiseBinary) .
                  (family NativePhysicalCommands))
                operation
                (branch
                  PointwiseAdd
                  .
                  (nativeFiniteFloat64Binary
                    b"pointwise-nonfinite"
                    scratch
                    (constructor NativePhysicalFloat64Operation NativePhysicalFloat64Add)
                    destination
                    left
                    right
                    tail))
                (branch
                  PointwiseMultiply
                  .
                  (nativeFiniteFloat64Binary
                    b"pointwise-nonfinite"
                    scratch
                    (constructor NativePhysicalFloat64Operation NativePhysicalFloat64Multiply)
                    destination
                    left
                    right
                    tail))
                (branch
                  PointwiseSubtract
                  .
                  (nativeFiniteFloat64Binary
                    b"pointwise-nonfinite"
                    scratch
                    (constructor NativePhysicalFloat64Operation NativePhysicalFloat64Subtract)
                    destination
                    left
                    right
                    tail))
                (branch PointwisePositiveGate . tail))))))))

-- Unsupported branches above are excluded by nativePointwiseProgram's erased
-- admission proof. Inputs name existing slots; every other node materializes a
-- fresh temporary. Explicit assignments then make their result visible to the
-- next expression, using the same dependency order as the SM86 realization.
def nativePointwiseExpression =
  (lambda unrestricted slots : Nat .
    (lambda unrestricted base : Nat .
      (lambda unrestricted scratch : Nat .
        (lambda unrestricted expression : (family PointwiseExpression StdU64 slots) .
          (pointwiseFold
            StdU64
            (pi unrestricted next : Nat . (family NativePointwiseFragment))
            slots
            (lambda unrestricted operation : (family PointwiseUnary) .
              (lambda unrestricted operand : (pi unrestricted next : Nat . (family NativePointwiseFragment)) .
                (lambda unrestricted next : Nat .
                  (let unrestricted child =
                    (operand next)
                    in
                    (constructor
                      NativePointwiseFragment
                      NativePointwiseFragmentOf
                      (lambda unrestricted tail : (family NativePhysicalCommands) .
                        (nativePointwiseCode
                          child
                          (nativePointwiseUnary
                            scratch
                            operation
                            (nativePointwiseEnd child)
                            (nativePointwiseResult child)
                            tail)))
                      (nativePointwiseEnd child)
                      (nativePointwiseAt (nativePointwiseEnd child) 1))))))
            (lambda unrestricted operation : (family PointwiseBinary) .
              (lambda unrestricted left : (pi unrestricted next : Nat . (family NativePointwiseFragment)) .
                (lambda unrestricted right : (pi unrestricted next : Nat . (family NativePointwiseFragment)) .
                  (lambda unrestricted next : Nat .
                    (let unrestricted l =
                      (left next)
                      in
                      (let unrestricted r =
                        (right (nativePointwiseEnd l))
                        in
                        (constructor
                          NativePointwiseFragment
                          NativePointwiseFragmentOf
                          (lambda unrestricted tail : (family NativePhysicalCommands) .
                            (nativePointwiseCode
                              l
                              (nativePointwiseCode
                                r
                                (nativePointwiseBinary
                                  scratch
                                  operation
                                  (nativePointwiseEnd r)
                                  (nativePointwiseResult l)
                                  (nativePointwiseResult r)
                                  tail))))
                          (nativePointwiseEnd r)
                          (nativePointwiseAt (nativePointwiseEnd r) 1))))))))
            (lambda unrestricted bits : StdU64 .
              (lambda unrestricted next : Nat .
                (constructor
                  NativePointwiseFragment
                  NativePointwiseFragmentOf
                  (nativeMemoryStoreWord64
                    b"pointwise-literal"
                    next
                    (nativeLaunchRecipeImmediate (modelWord64Natural bits)))
                  next
                  (nativePointwiseAt next 1))))
            expression
            (lambda unrestricted index : (family StdFin slots) .
              (lambda unrestricted next : Nat .
                (constructor
                  NativePointwiseFragment
                  NativePointwiseFragmentOf
                  (lambda unrestricted tail : (family NativePhysicalCommands) . tail)
                  (nativePointwiseAt base (stdFinToNatural slots index))
                  next))))))))

def nativePointwiseCommands =
  (lambda unrestricted slots : Nat .
    (lambda unrestricted program : (family StdList (family PointwiseAssignment StdU64 slots)) .
      (lambda unrestricted base : Nat .
        (lambda unrestricted available : Nat .
          (lambda unrestricted scratch : Nat .
            (lambda unrestricted tail : (family NativePhysicalCommands) .
              (stdListFold
                (family PointwiseAssignment StdU64 slots)
                (family NativePhysicalCommands)
                (lambda unrestricted assignment : (family PointwiseAssignment StdU64 slots) .
                  (lambda unrestricted rest : (family NativePhysicalCommands) .
                    (eliminate
                      PointwiseAssignment
                      (lambda unrestricted current : (family PointwiseAssignment StdU64 slots) .
                        (family NativePhysicalCommands))
                      assignment
                      (branch
                        PointwiseAssign
                        destination
                        expression
                        .
                        (let unrestricted fragment =
                          (nativePointwiseExpression
                            slots
                            base
                            scratch
                            expression
                            (nativePointwiseAt base slots))
                          in
                          (nativePointwiseCode
                            fragment
                            (nativeFloat64RequireFinite
                              (nativeMemoryLoad (nativePointwiseResult fragment))
                              scratch
                              b"pointwise-nonfinite"
                              (nativeMemoryStoreWord64
                                b"pointwise-assignment"
                                (nativePointwiseAt base (stdFinToNatural slots destination))
                                (nativeMemoryLoad (nativePointwiseResult fragment))
                                rest))))))))
                tail
                program)))))))

-- The raw command fold above is internal to these two admission entries.
def nativePointwiseProgram =
  (lambda unrestricted slots : Nat .
    (lambda unrestricted program : (family StdList (family PointwiseAssignment StdU64 slots)) .
      (lambda unrestricted base : Nat .
        (lambda unrestricted available : Nat .
          (lambda unrestricted scratch : Nat .
            (lambda erased admitted : (equal Nat (nativePointwiseAdmitted slots program base available) 1) .
              (nativePointwiseCommands slots program base available scratch)))))))

def planNativePointwiseProgram =
  (lambda unrestricted slots : Nat .
    (lambda unrestricted program : (family StdList (family PointwiseAssignment StdU64 slots)) .
      (lambda unrestricted base : Nat .
        (lambda unrestricted available : Nat .
          (lambda unrestricted scratch : Nat .
            (lambda unrestricted tail : (family NativePhysicalCommands) .
              (nat-eliminate
                (lambda unrestricted valid : Nat .
                  (family StdResult Bytes (family NativePhysicalCommands)))
                (constructor
                  StdResult
                  StdFailure
                  Bytes
                  (family NativePhysicalCommands)
                  b"pointwise-program-admission")
                (lambda unrestricted previous : Nat .
                  (lambda unrestricted ignored : (family StdResult Bytes (family NativePhysicalCommands)) .
                    (constructor
                      StdResult
                      StdSuccess
                      Bytes
                      (family NativePhysicalCommands)
                      (nativePointwiseCommands slots program base available scratch tail))))
                (nativePointwiseAdmitted slots program base available))))))))