module Representation.Partition

import Std.Foundation
import Std.List
import Std.Natural

-- Coordinates name logical items, not addresses, ranks, or tensor layouts.
-- A caller may use a stable ordinal enumeration of keys or population IDs.
-- Half-open intervals allow empty assignments without inventing padding work.
family LogicalInterval : Type 0
constructor LogicalIntervalOf
field unrestricted logicalIntervalBegin : Nat
field unrestricted logicalIntervalEnd : Nat

end-family

family PartitionFailure : Type 0
constructor PartitionNoParticipants
constructor PartitionParticipantOutside
constructor PartitionExtentOutsideWord

end-family

def intervalBegin =
  (lambda unrestricted interval : (family LogicalInterval) .
    (eliminate
      LogicalInterval
      (lambda unrestricted current : (family LogicalInterval) . Nat)
      interval
      (branch LogicalIntervalOf begin end . begin)))

def intervalEnd =
  (lambda unrestricted interval : (family LogicalInterval) .
    (eliminate
      LogicalInterval
      (lambda unrestricted current : (family LogicalInterval) . Nat)
      interval
      (branch LogicalIntervalOf begin end . end)))

def intervalExtent =
  (lambda unrestricted interval : (family LogicalInterval) .
    (naturalSaturatingSubtract (intervalEnd interval) (intervalBegin interval)))

-- A disjoint intersection is an empty interval at the larger begin. Preserve
-- begin <= end so downstream consumers can reuse the ordinary interval laws.
def intervalIntersection =
  (lambda unrestricted left : (family LogicalInterval) .
    (lambda unrestricted right : (family LogicalInterval) .
      (let unrestricted begin =
        (naturalSelect
          (naturalLess (intervalBegin left) (intervalBegin right))
          (intervalBegin right)
          (intervalBegin left))
        in
        (let unrestricted end =
          (naturalMinimum (intervalEnd left) (intervalEnd right))
          in
          (constructor
            LogicalInterval
            LogicalIntervalOf
            begin
            (naturalSelect (naturalLess end begin) begin end))))))

def intervalWithin =
  (lambda unrestricted extent : Nat .
    (lambda unrestricted interval : (family LogicalInterval) .
      (naturalAnd
        (naturalLessOrEqual (intervalBegin interval) (intervalEnd interval))
        (naturalLessOrEqual (intervalEnd interval) extent))))

-- Validate ordered coverage in one pass over fragments, not over coordinates.
-- Checking only summed lengths would accept an overlap paired with a hole.
-- Requiring the next begin to equal the previous end detects both, including
-- at the boundaries. Unordered input must be canonicalized by its owner first.
def intervalPartitionValid :
  (pi unrestricted extent : Nat .
    (pi unrestricted intervals : (family StdList (family LogicalInterval)) . Nat)) =
  (lambda unrestricted extent : Nat .
    (lambda unrestricted intervals : (family StdList (family LogicalInterval)) .
      (app
        (eliminate
          StdList
          (lambda unrestricted current : (family StdList (family LogicalInterval)) .
            (pi unrestricted cursor : Nat . Nat))
          intervals
          (branch StdListEmpty . (lambda unrestricted cursor : Nat . (naturalEqual cursor extent)))
          (branch
            StdListCons
            head
            tail
            induction
            .
            (lambda unrestricted cursor : Nat .
              (naturalAnd
                (naturalEqual cursor (intervalBegin head))
                (naturalAnd (intervalWithin extent head) (induction (intervalEnd head)))))))
        0)))

-- Project one role from a regular partition in constant arithmetic work.
-- The quotient/remainder law keeps every endpoint <= extent; intermediate
-- multiplication cannot overflow for an admitted role and U64 extent.
-- No list of 200,000 ranks and no recursion on an integer magnitude is built.
def balancedIntervalUnchecked =
  (lambda unrestricted extent : Nat .
    (lambda unrestricted participants : Nat .
      (lambda unrestricted role : Nat .
        (app
          (lambda unrestricted quotient : Nat .
            (app
              (lambda unrestricted remainder : Nat .
                (app
                  (lambda unrestricted begin : Nat .
                    (constructor
                      LogicalInterval
                      LogicalIntervalOf
                      begin
                      (naturalAdd begin (naturalAdd quotient (naturalLess role remainder)))))
                  (naturalAdd (naturalMultiply quotient role) (naturalMinimum role remainder))))
              (naturalModuloUnchecked extent participants)))
          (naturalDivideUnchecked extent participants)))))

def balancedInterval =
  (lambda unrestricted extent : Nat .
    (lambda unrestricted participants : Nat .
      (lambda unrestricted role : Nat .
        (nat-eliminate
          (lambda unrestricted flag : Nat .
            (family StdResult (family PartitionFailure) (family LogicalInterval)))
          (constructor
            StdResult
            StdFailure
            (family PartitionFailure)
            (family LogicalInterval)
            (constructor PartitionFailure PartitionExtentOutsideWord))
          (lambda unrestricted predecessor : Nat .
            (lambda unrestricted induction : (family StdResult (family PartitionFailure) (family LogicalInterval)) .
              (nat-eliminate
                (lambda unrestricted flag : Nat .
                  (family StdResult (family PartitionFailure) (family LogicalInterval)))
                (constructor
                  StdResult
                  StdFailure
                  (family PartitionFailure)
                  (family LogicalInterval)
                  (constructor PartitionFailure PartitionNoParticipants))
                (lambda unrestricted previous : Nat .
                  (lambda unrestricted rest : (family StdResult (family PartitionFailure) (family LogicalInterval)) .
                    (nat-eliminate
                      (lambda unrestricted flag : Nat .
                        (family StdResult (family PartitionFailure) (family LogicalInterval)))
                      (constructor
                        StdResult
                        StdFailure
                        (family PartitionFailure)
                        (family LogicalInterval)
                        (constructor PartitionFailure PartitionParticipantOutside))
                      (lambda unrestricted before : Nat .
                        (lambda unrestricted ignored : (family StdResult (family PartitionFailure) (family LogicalInterval)) .
                          (constructor
                            StdResult
                            StdSuccess
                            (family PartitionFailure)
                            (family LogicalInterval)
                            (balancedIntervalUnchecked extent participants role))))
                      (naturalLess role participants))))
                (naturalNonzero participants))))
          (naturalLessOrEqual extent 18446744073709551615)))))