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

local G is 9.81.
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.
function BurnTime {
  parameter dV, a, m is SHIP:MASS, s is STAGE:NUMBER.

  local Gb is SHIP:BODY:MU / ((SHIP:BODY:RADIUS+a)^2).
  local f is stageThrust().           // Engine Thrust (kg * m/s²)
  local Isp is stageISP().            // Engine ISP (s)
  
  // debug
  print "Calculating burn time.".
  print "dV = " + dV.
  print "f = " + f.
  print "m = " + m.
  print "e = " + e.
  print "Isp = " + Isp.
  print "Gb = " + Gb.
  // end debug

  if dV > SHIP:STAGEDELTAV(s) {
    local t is burnTimeCalc(SHIP:STAGEDELTAV(s), m, Gb, Isp, f).

    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), a, m, s - 1).
  }

  return burnTimeCalc(dV, m, Gb, Isp, f).
}

// Convenience function to wrap the actual calculation for burn time.
function burnTimeCalc {
  parameter dV, m, Gb, 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.
  return Gb * m * Isp * (1 - CONSTANT():E^(-dV/(Gb*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.
}