Lee Ambrose Park

# Transform 1: Join soldier records with equipment hand receipts
# Produces one row per soldier-equipment pairing for readiness tracking
@transform_df(
    Output("/SPO-T/datasets/cult/soldier_equipment_joined"),
    soldiers=Input("/SPO-T/datasets/cult/soldiers_cleaned"),
    equipment=Input("/SPO-T/datasets/cult/equipment_status"),
)
def compute(soldiers, equipment):
    return soldiers.join(
        equipment,
        soldiers.dodid == equipment.hand_receipt_holder,
        "left"
    ).select(
        "dodid", "rank", "last_name", "unit",
        "equipment_lin", "equipment_nsn", "status", "location"
    )

# Transform 2: Clean and validate incoming GCSS-Army equipment data
# Standardizes LIN/NSN formats, flags missing serials, deduplicates
@transform_df(
    Output("/SPO-T/datasets/cult/equipment_status"),
    raw_equip=Input("/SPO-T/datasets/cult/gcss_equipment_raw"),
)
def compute(raw_equip):
    from pyspark.sql import functions as F

    cleaned = raw_equip.filter(
        F.col("lin").isNotNull() & F.col("nsn").isNotNull()
    ).withColumn(
        "lin", F.upper(F.trim(F.col("lin")))
    ).withColumn(
        "nsn", F.regexp_replace(F.col("nsn"), "[^0-9]", "")
    ).withColumn(
        "has_serial", F.col("serial_number").isNotNull()
    )

    # Deduplicate: keep latest record per LIN + serial
    return cleaned.dropDuplicates(["lin", "serial_number"])

# Transform 3: Aggregate readiness metrics per unit
# Window function ranks units by equipment fill rate
@transform_df(
    Output("/SPO-T/datasets/cult/unit_readiness_rollup"),
    joined=Input("/SPO-T/datasets/cult/soldier_equipment_joined"),
)
def compute(joined):
    from pyspark.sql import functions as F
    from pyspark.sql.window import Window

    unit_stats = joined.groupBy("unit").agg(
        F.count("dodid").alias("total_soldiers"),
        F.countDistinct("equipment_lin").alias("unique_lins"),
        F.avg(F.when(F.col("status") == "FMC", 1).otherwise(0))
         .alias("fmc_rate"),
    )

    # Rank units by FMC (Fully Mission Capable) rate
    w = Window.orderBy(F.desc("fmc_rate"))
    return unit_stats.withColumn(
        "readiness_rank", F.rank().over(w)
    )