// Functions for calculating thrust values.
@lazyglobal off.

// point gravity for TWR calculations.
local G is 0.
lock G to SHIP:BODY:MU / ((SHIP:BODY:RADIUS+SHIP:ALTITUDE)^2).

// Returns the throttle value you should use to achieve the
// target TWR. If TWR can't be achieved, returns 1.0. (full throttle)
function ThrottleToTWR {
  parameter targetTWR is 1.5.

  local m is SHIP:MASS.
  return min((targetTWR*m*G)/SHIP:AVAILABLETHRUST, 1.0).
}

// Calculates the ship's current TWR.
function TWR {
  local m is ship:mass.
  local t is THROTTLE * SHIP:AVAILABLETHRUST.
  return t/(m*G).
}

// Calculate the time required to burn a given dV at a given altitude.
// Must be called while in the same SOI as the burn itself.
// Assumes a perfectly spherical Kerbal in a vacuum.
function BurnTime {
  parameter dV, m is SHIP:MASS, s is STAGE:NUMBER.

  local f is stageThrust().           // Engine Thrust (kg * m/s²)
  local Isp is stageISP().            // Engine ISP (s)
  
  local dVs is SHIP:StageDeltaV(s):VACUUM.
  if dV > dVs {
    // TODO: not 100% this is needed vs using DeltaV:DURATION.
    // Docs suggest DeltaV:DURATION is not entirely reliable, however.
    // For now we're logging both values for comparison.
    local t is burnTimeCalc(dVs, m, Isp, f).

    // debug
    print "Computed stage burn time = " + t.
    print "KSC-estimated stage burn time = " + SHIP:StageDeltaV(s):DURATION.
    // end debug

    local parts is list().
    for part in parts {
      if part:DECOUPLEDIN = s - 1 {
	set m to m - part:MASS.
      }
    }

    return t + BurnTime(dV - SHIP:STAGEDELTAV(s):VACUUM, m, s - 1).
  }

  // debug
  local t is burnTimeCalc(dV, m, Isp, f).
  print "Burn time in last utilized stage = " + t.
  // end debug

  return t.
}

// Convenience function to wrap the actual calculation for burn time.
function burnTimeCalc {
  parameter dV, m, Isp, f.

  if f = 0 {
    print "WARNING: Tried to calculate burn time with a thrust of 0. Returning 0. Your calculations are probably wrong.".
    return 0.
  }

  // TODO: this formula differs from the following:
  // t = ((M0 - Mf) * Isp * G) / f
  // which is suggested at https://www.reddit.com/r/Kos/comments/lev9pw/getting_burntime_from_next_stage/gmig0hl/?context=8&depth=9
  // are they equivalent? Is one better than the other? This one doesn't require
  // knowing final mass, which is nice.
  local g0 is CONSTANT:G0.
  return g0 * m * Isp * (1 - CONSTANT():E^(-dV/(g0*Isp))) / f.
}

// Calculate the ISP for a given stage.
// Defaults to current stage. Assumes your ship is designed so that
// engines are discarded immediately when they flame out.
function stageISP {
  parameter s is STAGE:NUMBER.

  local en is list().
  list ENGINES in en.

  local ispSum is 0.
  local eCount is 0.
  for e in en {
    if e:STAGE = s or e:STAGE > s and e:DECOUPLEDIN < s {
      set ispSum to ispSum + e:ISP.
      set eCount to eCount + 1.
    }
  }
  if eCount = 0 { return 0. }
  
  return ispSum / eCount.
}

// Calculates the total thrust for the given stage, in kN.
// Defaults to current stage. Assumes your ship is designed so that
// engines are discarded immediately when they flame out.
function stageThrust {
  parameter s is STAGE:NUMBER.

  local en is list().
  list ENGINES in en.

  local sum is 0.
  for e in en {
    if e:STAGE = s or e:STAGE > s and e:DECOUPLEDIN < s {
      set sum to sum + e:POSSIBLETHRUST.
    }
  }

  return sum.
}