Balanced Tanks with Dry Staging - KontrolSystem2

asparagus_stage.to2

use { Vessel } from ksp::vessel
use { CONSOLE } from ksp::console
use { current_time, sleep, wait_until } from ksp::game
use { create_resource_transfer, FlowDirection } from ksp::resource
use { format } from core::str
use { max, min, clamp } from core::math

// Balance Methane, Oxidizer and Hydrogen; trigger dry stages
fn balance_up_tanks(vessel: Vessel) -> Result<Unit, string> = {
    CONSOLE.clear()
    // Capacity, Storage and Target
    let cap : float[][] = (0..vessel.staging.count).map(fn(i) -> [0.0, 0.0, 0.0])
    let sto : float[][] = (0..vessel.staging.count).map(fn(i) -> [0.0, 0.0, 0.0])
    let tgt : float[][] = (0..vessel.staging.count).map(fn(i) -> [0.0, 0.0, 0.0])

    // Calculate capacity and storage of all stages
    for(part in vessel.parts) {
        if(part.resources.list.length == 0) continue
        if (part.decouple_stage < -1) continue
        for (res in part.resources.list) {
            if (res.resource.id >= 6 && res.resource.id <= 8) {
                cap[part.decouple_stage + 1][res.resource.id - 6] += res.capacity_units
                sto[part.decouple_stage + 1][res.resource.id - 6] += res.stored_units
            }
        }
    }

    // Calculate how much fuel we can pump up
    for (fuel in 0...2) {
        for (stage in 0..vessel.staging.count) {
            tgt[stage][fuel] = sto[stage][fuel]
        }
        if (vessel.staging.count > 1) {
            for (stage_in in 0..vessel.staging.count - 1) {
                for (stage_out in stage_in + 1 .. vessel.staging.count) {                    
                    let need = cap[stage_in][fuel] - tgt[stage_in][fuel]
                    if (need > 0) {
                        let supply = tgt[stage_out][fuel]
                        if (supply > need) {
                            tgt[stage_in][fuel] += need
                            tgt[stage_out][fuel] -= need                        
                        } else {
                            tgt[stage_in][fuel] += supply
                            tgt[stage_out][fuel] -= supply                        
                        }
                    }
                }
            }
        }
    }

    // Log out all our stage values
    CONSOLE.print_line("Storage / Target / Capacity")
    for (stage in 0..vessel.staging.count) {
        CONSOLE.print_line(stage.to_string() + ":" + 
            format("{0,6:##0.00}>{1,6:##0.00}/{2,6:##0.00} {3,6:##0.00}>{4,6:##0.00}/{5,6:##0.00} {6,6:##0.00}>{7,6:##0.00}/{8,6:##0.00}", 
                [sto[stage][0], tgt[stage][0], cap[stage][0], sto[stage][1], tgt[stage][1], cap[stage][1], sto[stage][2], tgt[stage][2], cap[stage][2]]
            )
        )
    }

    if (!has_resources_remaining(vessel, vessel.staging.current)) {
        // trigger if there's empty stages
        CONSOLE.print_line("STAGE DRY")
        wait_until(fn() -> vessel.staging.ready)
        vessel.staging.next()        
        sleep(0.25)
        wait_until(fn() -> vessel.staging.ready)
    } else {
        // Transfer our fuels
        let rt = create_resource_transfer()

        for(part in vessel.parts) {
            if(part.resources.list.length == 0) continue
            for (res in part.resources.list) {
                if (res.resource.id < 6 || res.resource.id > 8) continue
                let stage = part.decouple_stage + 1
                let ridx = res.resource.id - 6
                if (cap[stage][ridx] > 0) {
                    // stage_fill allows us to average out the tanks in that stage
                    let stage_fill = tgt[stage][ridx] / cap[stage][ridx]
                    let target = res.capacity_units * stage_fill
                    let amount = target - res.stored_units
                    if (amount > 0) {
                        rt.add_resource(FlowDirection.FLOW_INBOUND, res, amount)
                    } 
                    else if (amount < 0) {
                        rt.add_resource(FlowDirection.FLOW_OUTBOUND, res, -amount)
                    }
                }
            }
        }

        // Do the actual transfer
        // One second seems sufficient - lower numbers can leave small amounts in the tanks, higher numbers reduces responsiveness
        rt.start()

        for(i in 1..20) {
            if (rt.is_running) {
                CONSOLE.print((i%10).to_string())
            }
            else
            {
                if (i>5) continue
                CONSOLE.print(".")
            }
            sleep(0.1)
        }

        rt.stop()
    }
}

pub sync fn has_resources_remaining(vessel: Vessel, stage: int) -> bool = {
    // parts with resources -> decoupling after indicated stage -> That has more than trace amounts
    // TODO: Should filter out by resource type, not everything matters, but we don't want to dump solid rockets either
    vessel.parts.filter(fn(part) -> part.resources.list.length > 0)
        .filter(fn(part) -> part.decouple_stage == stage - 1)
        .filter(fn(part) -> part.resources.list.filter(fn(res) -> res.stored_units > 0.0001).length > 0)
        .length > 0
}

pub fn main_flight(vessel: Vessel) -> Result<Unit, string> = {
    while (true) {
        balance_up_tanks(vessel)
    }
}