FtD - Lua Controlled Plane and Missiles

--[[
Advanced aerial AI, version 2.8
Created by Madwand 8/22/2015
Use and modify this code however you like, however please credit me
if you use this AI or a derivative of it in a tournament, or you publish a blueprint
using it. Also let me know if you make any significant improvements,
I can add them back into the code.

Documentation on the options is at http://pastebin.com/GyE59Y3a
--]]

-- BASIC OPTIONS
AngleBeforeTurn = 5
AngleBeforeRoll = 10
AttackRunDistance = 800
AbortRunDistance = 300
ClosingDistance = 1000
ForceAttackTime = 15
CruiseAltitude = 200
MaxElevationAngle = 30
CruiseSpeed = 5
AttackRunSpeed = 5
EscapeSpeed = 5
RollingSpeed = 3
ClosingSpeed = 5
OrbitSpawn = true

-- AIRSHIP/HELICOPTER OPTIONS
UseAltitudeJets = false
UseSpinners = false
MinHelicopterBladeSpeed = 20
MaxHelicopterBladeSpeed = 30

-- TERRAIN AVOIDANCE OPTIONS
TerrainAvoidanceStrategy = 1.
MinAltitude = 100
MaxAltitude = 400
TerrainLookahead = {0,1,2,4}
MaxTerrainSpeed = 5

-- WATER START OPTIONS
DeployAlt = 5
ReleaseAlt = 15
EnableEnginesAlt = 10

-- COLLISION AVOIDANCE OPTIONS
CollisionDetectionRange = 300
CollisionDetectionHeight = 20
CollisionAngle = 20

-- "DOGFIGHTING" OPTIONS
MatchTargetAltitude = false
MatchAltitudeRange = 600
MatchAltitudeOffset = 50

-- ADVANCED OPTIONS
MODE = 2
AltitudeClamp = .2
PitchDamping = 90
YawDamping = 90
RollDamping = 45
DefaultRollAngle = 0
UsePropulsionBalancing = true

--Static variables. Do not change these.
DRIVELEFT = 0
DRIVERIGHT = 1
DRIVEMAIN = 8
DRIVEUP = 4
DRIVEDOWN = 5
ROLLLEFT = 2
ROLLRIGHT = 3

firstpass = true
yawing = ""
pitching = ""
rolling = ""
vertical = ""
state = "cruise"
DebugMode = false
onattackrun = false
WaterStarted = false
NextAttackTime = 0
OverTerrain = false
Roll = 0
Pitch = 0
Yaw = 0
DRoll = 0
DPitch = 0
DYaw = 0
PitchCorrection = 0
CurrentSpeed = 0

--Try to pitch (or yaw, if necessary) to a given angle of elevation
function AdjustElevationToAngle(I,DesiredElevation)
  local PitchSpeedCheck = math.abs(DesiredElevation-Pitch)/PitchDamping
  local YawSpeedCheck = math.abs(DesiredElevation-Pitch)/YawDamping
  if (Pitch > DesiredElevation+5) then
    if (math.abs(Roll)>135 and PitchSpeedCheck>DPitch) then
      I:RequestControl(MODE,DRIVEUP,1)
      pitching = "Pup"
    elseif (Roll<-45 and YawSpeedCheck>DYaw) then
      I:RequestControl(MODE,DRIVERIGHT,1)
      yawing = "right"
    elseif (Roll>45 and YawSpeedCheck>-DYaw) then
      I:RequestControl(MODE,DRIVELEFT,1)
      yawing = "left"
    elseif (math.abs(Roll)<45 and PitchSpeedCheck>-DPitch) then
      I:RequestControl(MODE,DRIVEDOWN,1)
      pitching = "Pdown"
    end
  elseif (Pitch < DesiredElevation-5) then
    if (math.abs(Roll)>135 and PitchSpeedCheck>-DPitch) then
      I:RequestControl(MODE,DRIVEDOWN,1)
      pitching = "Pdown"
    elseif (Roll<-45 and YawSpeedCheck>-DYaw) then
      I:RequestControl(MODE,DRIVELEFT,1)
      yawing = "left"
    elseif (Roll>45 and YawSpeedCheck>DYaw) then
      I:RequestControl(MODE,DRIVERIGHT,1)
      yawing = "right"
    elseif (math.abs(Roll)<45 and PitchSpeedCheck>DPitch) then
      I:RequestControl(MODE,DRIVEUP,1)
      pitching = "Pup"
    end
  end
end

--Try to climb or decend to a particular altitude
function AdjustAltitudeTo(I, DesiredAltitude)
  local Alt = I:GetConstructCenterOfMass().y

  AdjustAltitudeJets(I,DesiredAltitude)

  local Angle = math.min(math.abs(Alt-DesiredAltitude)*AltitudeClamp,MaxElevationAngle)
  if ((Alt<MinAltitude and Pitch<0) or (Alt>MaxAltitude and Pitch>0)) then
    Angle = MaxElevationAngle
  end -- pull up/down with full strength if we are below/above min/max altitudes

  if (Alt < DesiredAltitude) then
    if (Pitch < 0) then
      AdjustElevationToAngle(I,MaxElevationAngle)
    else
      AdjustElevationToAngle(I,Angle)
    end
  else
    if (Pitch > 0) then
      AdjustElevationToAngle(I,-MaxElevationAngle)
    else
      AdjustElevationToAngle(I,-Angle)
    end
  end

  AdjustHelicopterBlades(I,DesiredAltitude)
end

function AdjustAltitudeJets(I,DesiredAltitude)
  if not UseAltitudeJets then return end
  local Alt = I:GetConstructCenterOfMass().y
  if (Alt < DesiredAltitude) then
    I:RequestThrustControl(4)
    vertical = "up"
  else
    I:RequestThrustControl(5)
    vertical = "down"
  end
end

function AdjustHelicopterBlades(I,DesiredAltitude)
  if not UseSpinners then return end
  local Alt = I:GetConstructCenterOfMass().y
  if (Alt < DesiredAltitude) then
    for spinner = 0, I:GetSpinnerCount() - 1 do
      I:SetSpinnerContinuousSpeed(spinner, MaxHelicopterBladeSpeed)
      vertical = "up"
    end
  else
    for spinner = 0, I:GetSpinnerCount() - 1 do
      I:SetSpinnerContinuousSpeed(spinner, MinHelicopterBladeSpeed)
      vertical = "down"
    end
  end
end

-- Get the current stats about angles of the vehicle
function GetAngleStats(I)
  Roll = I:GetConstructRoll()
-- transform roll into roll we are familiar with from the game
  if (Roll > 180) then Roll = Roll-360 end

  Pitch = I:GetConstructPitch()
-- transform pitch into pitch we are familiar with from the game
  if (Pitch > 180) then
    Pitch = 360 - Pitch
  elseif (Pitch>0) then
    Pitch = -Pitch
  end

-- Many vehicles need to keep their nose pitched upwards relative to velocity.
-- This uses basic machine learning to calculate the upwards pitch correction.
  VPitch = math.deg(math.asin(I:GetVelocityVectorNormalized().y)) -- observed pitch from velocity
  PitchCorrection = PitchCorrection + .01*(Pitch-VPitch-PitchCorrection)
  PitchCorrection = math.min(math.max(PitchCorrection,-MaxElevationAngle),MaxElevationAngle)
  Pitch = Pitch - PitchCorrection

  local AngularVelocity = I:GetLocalAngularVelocity()
  DRoll = AngularVelocity.z
  DPitch = -AngularVelocity.x
  DYaw = AngularVelocity.y
end

-- Try to roll the craft and pull up on pitch controls to face a given azimuth (relative to the vehicles facing)
-- Will try to set roll angle at one that will maintain a given altitude
function RollTowardsAzimuth(I,Azimuth,DesiredAltitude)
  local Alt = I:GetConstructCenterOfMass().y
-- depending on our altitude, RollAngle will be set to 90 +/- 20 degrees to climb or descend as appropriate
  RollAngle = math.min(110,math.max(70,(90 + (Alt-DesiredAltitude)*.66)))

  if ((Azimuth<0 and (Azimuth>-150 or Roll<0)) or
      (Azimuth>150 and Roll<0)) then  -- roll towards the left, otherwise towards the right
    RollAngle = -RollAngle
  end

  AdjustRollToAngle(I,RollAngle)

  if (Roll >= RollAngle-10 and Roll <= RollAngle+10) then -- start pitching
    I:RequestControl(MODE,DRIVEUP,1)
    pitching = "up"
  end

  --I:LogToHud(string.format("%.2f %.2f %.2f", Azimuth, RollAngle, Roll))
end

-- Roll the vehicle to a specified roll angle
function AdjustRollToAngle(I,Angle)
  local RollSpeedCheck = math.abs(Angle-Roll)/RollDamping
  if (Roll < Angle and RollSpeedCheck>DRoll) then
    I:RequestControl(MODE,ROLLLEFT,1)
    rolling = "Rleft"
  elseif (Roll > Angle and RollSpeedCheck>-DRoll) then
    I:RequestControl(MODE,ROLLRIGHT,1)
    rolling = "Rright"
  end
end

-- Yaw the vehicle towards a given targets aziumth (relative to the vehicle's facing)
-- This function wil try to keep the nose "AngleBeforeTurn" degrees away from directly
-- pointing at the given aziumuth
function YawTowardsTarget(I, Azimuth, OffsetAngle)
  if (Azimuth > 0) then
    return YawTowardsAzimuth(I, Azimuth-OffsetAngle)
  end
  return YawTowardsAzimuth(I, Azimuth+OffsetAngle)
end

-- Yaw the vehicle towards a given aziumth (relative to the vehicle's facing)
function YawTowardsAzimuth(I,Azimuth)
  if (math.abs(Roll)<45) then
    local YawSpeedCheck = math.abs(Azimuth)/YawDamping
    if (Azimuth > 0 and YawSpeedCheck>-DYaw) then
      I:RequestControl(MODE,DRIVELEFT,1)
      yawing = "left"
    elseif (Azimuth < 0 and YawSpeedCheck>DYaw) then
      I:RequestControl(MODE,DRIVERIGHT,1)
      yawing = "right"
    end
    return true
  end
  return false
end

-- No enemies, just cruise along
function Cruise(I)
  if OrbitSpawn then
    SpawnInfo = I:GetTargetPositionInfoForPosition(0, SpawnPos.x, 0, SpawnPos.z)
    NavigateToPoint(I, false, SpawnInfo)
  else
    state = "cruise"
    AdjustRollToAngle(I, DefaultRollAngle)
    SetSpeed(I, CruiseSpeed)
    DesiredAltitude = AdjustAltitudeForTerrain(I, CruiseAltitude, false)
    AdjustAltitudeTo(I, DesiredAltitude)
    NextAttackTime = I:GetTime()+ForceAttackTime
  end
end

-- Given the current velocity vector and an azimuth, get the expected terrain height in that direction
function GetTerrainAltitude(I, VelocityV, Angle)
  local TerrainAltitude = -999
  MyPos = I:GetConstructCenterOfMass()
  for i,Lookahead in ipairs(TerrainLookahead) do
    Position = Quaternion.Euler(0,Angle,0) * VelocityV * Lookahead
    Position = MyPos + VelocityV * Lookahead
    TerrainAltitude = math.max(TerrainAltitude,I:GetTerrainAltitudeForPosition(Position))
  end
  return TerrainAltitude
end

-- Adjust a given altitude to compensate for terrain, according to "TerrainAvoidanceStrategy"
-- and other terrain-avoidance-specific parameters.
-- Last parameter forces the use of "TerrainAvoidanceStrategy" 2
-- returns new altitude that should be used
function AdjustAltitudeForTerrain(I, Altitude, StrategyOverride)
  VelocityV = I:GetVelocityVector()
  TerrainAltitude = GetTerrainAltitude(I, VelocityV, 0)

  OverTerrain = false
  if (TerrainAltitude + MinAltitude > Altitude or TerrainAltitude>0) then -- we'll need to avoid the terrain
    OverTerrain = TerrainAltitude + MinAltitude > Altitude

    if(StrategyOverride or TerrainAvoidanceStrategy == 2) then -- don't change altitude unless needed
      Altitude = math.max(Altitude, TerrainAltitude + MinAltitude)
    elseif(TerrainAvoidanceStrategy == 1) then -- just add Altitude to terrain height
      Altitude = math.max(math.min(TerrainAltitude+Altitude, MaxAltitude), TerrainAltitude + MinAltitude)
    end
  end

  --I:LogToHud(string.format("TA: %.2f Here: %.2f Alt: %.2f", TerrainAltitude, I:GetTerrainAltitudeForLocalPosition(0, 0, 0), Altitude))

  return Altitude
end

-- given information about the target, returns the angle to try to approach
-- the target at.
function CollisionAvoidance(I, Pos)
  local Alt = I:GetConstructCenterOfMass().y
  if (Pos.Range<CollisionDetectionRange and
      math.abs(Pos.Azimuth)<CollisionAngle and
      math.abs(Pos.ElevationForAltitudeComponentOnly)<CollisionDetectionHeight) then
    --I:LogToHud("Collision warning.")
    return math.max(CollisionAngle,AngleBeforeTurn)
  end
  return AngleBeforeTurn
end

-- Perform a water start check. Returns true if movement is permitted.
function WaterStartCheck(I)
  local Alt = I:GetConstructCenterOfMass().y

  if (Alt < DeployAlt) then
    I:Component_SetBoolLogicAll(0, true)
    WaterStarted = true
  elseif (WaterStarted and Alt > ReleaseAlt) then
    I:Component_SetBoolLogicAll(0, false)
    WaterStarted = false
  end

  if (not WaterStarted or Alt>=EnableEnginesAlt) then
    return true
  end
  return false
end

function SetSpeed(I, CurrentSpeed)
  if UsePropulsionBalancing then
    I:RequestThrustControl(0,CurrentSpeed/5)
  else
    I:RequestControl(MODE,DRIVEMAIN,CurrentSpeed)
  end
end

-- given the range, azimuth off the nose, and height of a target, navigates to it
function NavigateToPoint(I, Engaged, Pos)
  local DesiredAltitude = CruiseAltitude
  local MatchingAltitude = Engaged and MatchTargetAltitude and Pos.GroundDistance < MatchAltitudeRange
  if MatchingAltitude then
    DesiredAltitude = math.max(math.min(Pos.AltitudeAboveSeaLevel+MatchAltitudeOffset,MaxAltitude),MinAltitude)
  end

  if (TerrainAvoidanceStrategy>0) then
    DesiredAltitude = AdjustAltitudeForTerrain(I, DesiredAltitude, MatchingAltitude)
  end

  CurrentSpeed = EscapeSpeed
  if onattackrun then
    CurrentSpeed = RollingSpeed
    if (math.abs(Pos.Azimuth) > AngleBeforeRoll) then -- roll to turn
      RollTowardsAzimuth(I, Pos.Azimuth, DesiredAltitude)
      state = "rolling"
    else -- yaw to turn
      AdjustRollToAngle(I, DefaultRollAngle)
      if (Pos.Range > ClosingDistance) then
        if (YawTowardsAzimuth(I, Pos.Azimuth)) then
          CurrentSpeed = ClosingSpeed
        end
        state = "closing in"
      else
        local YawAngle = CollisionAvoidance(I, Pos)
        if (YawTowardsTarget(I, Pos.Azimuth, YawAngle)) then
          CurrentSpeed = AttackRunSpeed
        end
        state = "yawing in"
      end
    end

    if (Pos.Range < AbortRunDistance) then
      onattackrun = false
      NextAttackTime = I:GetTime()+ForceAttackTime
    end
  else -- not on attack run, just go forwards until we're out of range
    state = "escaping"
    AdjustRollToAngle(I, DefaultRollAngle)
    local YawAngle = CollisionAvoidance(I, Pos)
    if (YawAngle>AngleBeforeTurn and YawAngle>Pos.Azimuth) then
      YawTowardsTarget(I, Pos.Azimuth, YawAngle)
    end
    if (Pos.Range > AttackRunDistance or I:GetTime()>NextAttackTime) then
      onattackrun = true
    end
  end

  if OverTerrain then
    CurrentSpeed = math.min(CurrentSpeed,MaxTerrainSpeed)
  end
  if not Engaged then
    CurrentSpeed = math.min(CurrentSpeed,CruiseSpeed)
  end
  SetSpeed(I, CurrentSpeed)

  AdjustAltitudeTo(I, DesiredAltitude)
end

-- Calculate all movement for the vehicle
function Movement(I)
  if (firstpass) then
    firstpass = false
    MaxElevationAngle = math.max(MaxElevationAngle,1)
    SpawnPos = I:GetConstructCenterOfMass()
  end

  GetAngleStats(I)

  yawing = ""
  pitching = ""
  rolling = ""
  vertical = ""

  if (I:GetNumberOfMainframes() > 0 and I:GetNumberOfTargets(0) > 0) then
    local TargetPos = I:GetTargetPositionInfo(0,0)
    if TargetPos.Valid then
      NavigateToPoint(I, true, TargetPos)
    else
      Cruise(I)
    end
  else
    Cruise(I)
  end

  if DebugMode then
    I:LogToHud(string.format("%s %s %s %s %d", state, yawing, pitching, rolling, CurrentSpeed))
    --I:LogToHud(string.format("%.2f %.2f %.2f", Yaw, Pitch, Roll))
  end
end






mainframeNumber_=0
warningInfo_={}

-- Side Functions
function Speed(vel)
   return math.sqrt(vel.x^2 + vel.y^2 + vel.z^2)
end
function Distance(a,b)
   return math.sqrt((a.x - b.x)^2 + (a.y - b.y)^2 + (a.z - b.z)^2)
end

-- Aim Point
function Predict(P_Missile,V_Missile,P_Target,V_Target)
   D = Distance(P_Missile,P_Target)
   S = Speed(V_Target-V_Missile)
   Time =  D/S

   if(Time>0 and Time < 200) then
      P_Target = P_Target + (V_Target*Time)
   end

   return P_Target
end





-- Main update function. Everything starts here.
function Update(I)
  if WaterStartCheck(I) then
    Movement(I)
  end

    -- Loops through the tranceivers and missiles
  for transceiverNumber=0,I:GetLuaTransceiverCount() do
    for missileNumber=0,I:GetLuaControlledMissileCount(transceiverNumber) do
      -- Gets currently selected missile information
      missileInfo=I:GetLuaControlledMissileInfo(transceiverNumber, missileNumber)
      targetInfo=I:GetTargetInfo(mainframeNumber_,0)
      predictPosition=Predict(missileInfo.Position,missileInfo.Velocity,targetInfo.Position,targetInfo.Velocity)
      I:SetLuaControlledMissileAimPoint(transceiverNumber, missileNumber, predictPosition.x, predictPosition.y, predictPosition.z)
    end
  end
end