local EffectTemplate = import('/lua/EffectTemplates.lua')
local RandomFloat = import('/lua/utilities.lua').GetRandomFloat
local oldUnit = Unit
Unit = Class(oldUnit) {

#################################################################
#		TELEPORT FUNCTIONS
#################################################################

    MoveThroughRollOff = function(self)
        local spin, x, y, z = self:CalculateMoveOffPoint()
        local units = {self}
        self.RollOffCommand = IssueMove(units, Vector(x, y, z))
        local rally = self.destination:GetRallyPoint()
        rx = rally[1]
        ry = rally[2]
        rz = rally[3]
        self.MoveCommand = IssueMove( units, Vector(rx, ry, rz))
    end,
    
    DirectMoveToRally = function(self)
        local units = {self}
        local rally = self.destination:GetRallyPoint()
        rx = rally[1]
        ry = rally[2]
        rz = rally[3]
        self.MoveCommand = IssueMove( units, Vector(rx, ry, rz))
    end,

	CalculateMoveOffPoint = function(self)
        local bp = self.destination:GetBlueprint().Physics.RollOffPoints
        local px, py, pz = unpack(self.destination:GetPosition())
        if not bp then return 0, px, py, pz end
        local vectorObj = self.destination:GetRallyPoint()
        local bpKey = 1
        local distance, lowest = nil
        for k, v in bp do
            distance = VDist2(vectorObj[1], vectorObj[3], v.X + px, v.Z + pz)
            if not lowest or distance < lowest then
                bpKey = k
                lowest = distance
            end
        end
        local fx, fy, fz, spin
        local bpP = bp[bpKey]
	    local unitBP = self:GetBlueprint().Display.ForcedBuildSpin
		if unitBP then
	    	spin = unitBP
		else
			spin = bpP.UnitSpin
		end
		fx = bpP.X + px
		fy = bpP.Y + py
		fz = bpP.Z + pz
		return spin, fx, fy, fz
    end,
	
    OnTeleportUnit = function(self, teleporter, location, orientation)
    	if self ~= teleporter then
	        if self.TeleportDrain then
	            RemoveEconomyEvent( self, self.TeleportDrain)
	            self.TeleportDrain = nil
	        end
	        if self.TeleportThread then
	            KillThread(self.TeleporThread)
	            self.TeleportThread = nil
	        end
			local aiBrain = self:GetAIBrain()
			local gate = aiBrain:GetUnitsAroundPoint(categories.GATE, location, 10, 'ALLY')
			for k,destination in gate do
				if destination ~= teleporter then
					loc = destination:GetPosition()
					local units = aiBrain:GetUnitsAroundPoint(categories.MOBILE, teleporter:GetPosition(), 2, 'ALLY')
					#LOG(destination.Sender)
					if table.find(units, self) and destination.Sender == false and teleporter.Reciever == false and destination.FeedbackPulse == false and teleporter.FeedbackPulse == false then
				    	self.TeleportThread = self:ForkThread(self.StartTeleportThread, teleporter, destination, loc, orientation)
				    end
				end
			end
		end
		if self == teleporter then
			#LOG('Original OnTeleportUnit')
			oldUnit.OnTeleportUnit(self, teleporter, location, orientation)
		end
    end,

    OnFailedTeleport = function(self)
    	self:RemoveFromTable()
    	#LOG('removing from table')
    	#LOG(self.teleporter.unitcount, ' || ', self.destination.unitcount)
    	if not self:IsDead() then
	    	#LOG('OnfailedTeleport')
	        if self.TeleportDrain then
	            RemoveEconomyEvent( self, self.TeleportDrain)
	            self.TeleportDrain = nil
	        end
	        if self.TeleportThread then
	            KillThread(self.TeleporThread)
	            self.TeleportThread = nil
	        end
	        self:SetWorkProgress(0.0)
	        self:SetImmobile(false)
	        self:SetBusy(false)
	        ChangeState(self, self.IdleState)
	        self:ForkThread(NilThread, self)
	        #LOG(self.UnitBeingTeleported)
	    end
    end,

    UpdateTeleportProgress = function(self, progress)
        #LOG(' UpdatingTeleportProgress ')
        self:SetWorkProgress(progress)
    end,

    StartTeleportThread = function(self, teleporter, destination, loc, orientation)
    	self.teleporter = teleporter
		self.destination = destination
		destination.Reciever = true
		destination.Sender = false
		teleporter.Sender = true
		teleporter.Reciever = false
		teleporter.destination = destination
		ChangeState(self, self.BeingTeleportedState)
        local bp = self:GetBlueprint().Economy
        local energyCost, time
        if bp then
            energyCost = bp.BuildCostEnergy * 0.3
            time = energyCost * 0.00025
        end

        #LOG('*UNIT DEBUG: TELEPORTING, energyCost = ', repr(energyCost), ' time = ', repr(time))
        self.TeleportDrain = CreateEconomyEvent(self, energyCost, 0, time, self.UpdateTeleportProgress)

        WaitFor( self.TeleportDrain  )
        self.TeleportDrain = nil

        WaitSeconds( 0.1 )

        if self.TeleportDrain then
            RemoveEconomyEvent(self, self.TeleportDrain )
            self.TeleportDrain = nil
        end

        self:SetWorkProgress(0.0)
        Warp(self, loc, orientation)
        self:PlayTeleportXInEffects()

        WaitSeconds( 0.1 )
        self.TeleportThread = nil
    end,
    
	AddToTable = function(self, teleporter) 
		#LOG('AddToTable')
		self.teleporter.unitcount = self.teleporter.unitcount + 1
		self.destination.unitcount = self.destination.unitcount + 1
		table.insert(self.destination.UnitsTeleported, self)
	end, 

	RemoveFromTable = function(self)
		#LOG('RemoveFromTable')
		if self.teleporter.unitcount == 0 or self.destination.unitcount == 0 then 
		WARN('Unit count is below 0')
		end
		self.destination.unitcount = self.destination.unitcount - 1
		self.teleporter.unitcount = self.teleporter.unitcount - 1
		for k,unit in self.destination.UnitsTeleported do
			if unit == self then
				table.remove(self.destination.UnitsTeleported,k)
			end
		end	
	end, 
	
    OnDamage = function(self, instigator, amount, vector, damageType)
        if self.CanTakeDamage then
            self:DoOnDamagedCallbacks(instigator)
            self:DoTakeDamage(instigator, amount, vector, damageType)
            if instigator ~= self then
            	local hp = self:GetHealth()
            	local maxhp = self:GetMaxHealth()
            	local damage = ( (hp/maxhp) * (hp/maxhp) *10000 ) / 7.7
            	#LOG(damage)
	            if self.Reciever == true and amount > damage then
	            	if table.getn(self.UnitsTeleported) == 0 then
						self:ForkThread(PlayDisappearFlashes, self)
						for k, gate in self.teleporters do
							gate:ForkThread(PlayDisappearFlashes, gate)
						end
					else
						for k, unit in self.UnitsTeleported do
							unit:OnKilled()
						end
					end
				end
				if self.Sender == true and amount > damage then
					if table.getn(self.destination.UnitsTeleported) == 0 then
						self.destination:ForkThread(PlayDisappearFlashes, self.destination)
						for k, gate in self.destination.teleporters do
							gate:ForkThread(PlayDisappearFlashes, gate)
						end
					else
						for k, unit in self.destination.UnitsTeleported do
							unit:OnKilled()
						end
					end
				end
				if self.UnitBeingTeleported and amount > damage then
					self:OnKilled()
				end
			end
        end
    end,
    
	OnKilled = function(self, instigator, type, overkillRatio)
		#LOG('OnKilled')
		if self.Sender == true then
			#LOG('sender died')
			self.destination:ForkThread(PlayDisappearFlashes, self.destination)
			for k, unit in self.destination.UnitsTeleported do
				unit:Destroy()
			end
			for k, gate in self.destination.teleporters do
				gate:ForkThread(PlayDisappearFlashes, gate)
			end
		end
		if self.Reciever == true then
			#LOG('reciever died')
			self:ForkThread(PlayDisappearFlashes, self)
			for k, gate in self.teleporters do
				gate:ForkThread(PlayDisappearFlashes, gate)
			end
			for k, unit in self.UnitsTeleported do
				unit:Destroy()
			end
		end
        if self.UnitBeingTeleported then
        	#LOG('self died, and is the unit')
	        self.destination:ForkThread(PlayDisappearFlashes, self.destination)
	        for k, unit in self.destination.UnitsTeleported do
	        	unit:Destroy()
	        	unit.UnitBeingTeleported = nil
	        end
	        for k, gate in self.destination.teleporters do
				gate:ForkThread(PlayDisappearFlashes, gate)
			end
		end
		oldUnit.OnKilled(self, instigator, type, overkillRatio)
	end,
    

#################################################################
#		TELEPORT EFFECTS
#################################################################

    
	PlayRecieveEffects = function(self)
		local army = self:GetArmy()
		self.RecieveEffectsBag = {}
		for k, v in EffectTemplate.GenericTeleportCharge01 do
			local fx = CreateEmitterAtEntity(self,army, v):OffsetEmitter(0,1.8,0)
			self.Trash:Add(fx)
            table.insert( self.RecieveEffectsBag, fx)
		end
		self:SetConsumptionActive(true)
		self:SetConsumptionPerSecondEnergy(2000)
	end,
	
	PlaySendEffects = function(self)
		local army = self:GetArmy()
		self.SendEffectsBag = {}
		for k, v in EffectTemplate.GenericTeleportCharge01 do
			local fx = CreateEmitterAtEntity(self,army, v):OffsetEmitter(0,1.8,0)
			self.Trash:Add(fx)
            table.insert( self.SendEffectsBag, fx)
		end
		self:SetConsumptionActive(true)
		self:SetConsumptionPerSecondEnergy(4000)
	end,
	
	PlayStreamFlashes = function(self)
		local army = self:GetArmy()
		while true and self.Teleporting == true and not self:IsDead() do
			local size1 = math.random(2,4)
			local size2 = math.random(3,5)
			local size3 = math.random(4,8)
			CreateLightParticleIntel( self, -1, army, size1, size1, 'flare_lens_add_02', 'ramp_blue_13' )
			WaitTicks(math.random(3,5))
			CreateLightParticleIntel( self, -1, army, size2, size2, 'flare_lens_add_02', 'ramp_blue_13' )
			WaitTicks(math.random(3,5))
			CreateLightParticleIntel( self, -1, army, size2, size2, 'flare_lens_add_02', 'ramp_blue_13' )
			WaitTicks(math.random(6,15))
			CreateLightParticleIntel( self, -1, army, size1, size1, 'flare_lens_add_02', 'ramp_blue_13' )
			WaitTicks(math.random(10,15))
			CreateLightParticleIntel( self, -1, army, size3, size3, 'flare_lens_add_02', 'ramp_blue_13' )
			WaitTicks(math.random(6,15))
			CreateLightParticleIntel( self, -1, army, size3, size2, 'flare_lens_add_02', 'ramp_blue_13' )
			WaitTicks(math.random(70,120))
		end
	end,
	
	PlayCloseFlashes = function(self)
		if not self:IsDead() then
			local army = self:GetArmy()
			local size1 = math.random(4,6)
			CreateLightParticleIntel( self, -1, army, size1, size1, 'flare_lens_add_02', 'ramp_blue_13' )
			local size2 = math.random(6,9)
			CreateLightParticleIntel( self, -1, army, size2, size2, 'flare_lens_add_02', 'ramp_blue_13' )
			WaitTicks(math.random(3,5))
			local size3 = math.random(9,12)
			CreateLightParticleIntel( self, -1, army, size3, size3, 'flare_lens_add_02', 'ramp_blue_13' )
		end
	end,
	
	CleanupRecieveEffectsBag = function( self )
        if self.RecieveEffectsBag then
            for keys,values in self.RecieveEffectsBag do
                values:Destroy()
            end
            self.RecieveEffectsBag = {}
            if not self:IsDead() then
            	self:SetConsumptionPerSecondEnergy(0)
            end
        end
    end,
    
	CleanupSendEffectsBag = function( self )
        if self.SendEffectsBag then
            for keys,values in self.SendEffectsBag do
                values:Destroy()
            end
            self.SendEffectsBag = {}
            if not self:IsDead() then
           		self:SetConsumptionPerSecondEnergy(0)
           	end
			for k, gate in self.destination.teleporters do
				if gate == self then
					table.remove(self.destination.teleporters,k)
				end
			end
        end
    end,
       
	PlayTeleportXInEffects = function(self)
		local army = self:GetArmy()
		local pos = self:GetPosition()
		self:PlayUnitSound('CommanderArrival')
       	for k, v in EffectTemplate.CSGTestEffect2 do
           	CreateEmitterOnEntity( self, army, v )
        end
	    CreateLightParticleIntel( self, -1, army, 7, 15, 'glow_02', 'ramp_blue_13' )
	        
	    WaitSeconds(0.3)
	        
	    CreateLightParticleIntel( self, -1, army, 15, 10, 'flare_lens_add_02', 'ramp_blue_13' )
		for k, v in EffectTemplate.CommanderQuantumGateInEnergy do
            CreateEmitterOnEntity( self, army, v )
        end
        
		WaitSeconds(0.5)
		
	    CreateLightParticleIntel( self, -1, army, 7, 15, 'glow_03', 'ramp_blue_13' )
	    
		WaitSeconds(0.5)
	
	    CreateLightParticleIntel( self, -1, army, 17, 12, 'flare_lens_add_02', 'ramp_blue_13' )
    end,
 
#################################################################
#		TELEPORT STATES
#################################################################
  
	BeingTeleportedState = State {

		Main = function(self,destination)
			#LOG('BeingTeleported')
			self:AddToTable()
        	self:SetImmobile(true)
			self:SetBusy(true)
        	self.UnitBeingTeleported = self
			self:ForkThread(TeleportingEffectsThread, self,destination)
	end,
	},
	
	TeleportSendState = State {
	
		Main = function(self)
			self.Teleporting = true
			table.insert(self.destination.teleporters, self)
       		self:ForkThread(TeleportSendEffectsThread, self)
		end,
	},
	
	TeleportRecieveState = State {
	
		Main = function(self)
			self.Teleporting = true
        	self:ForkThread(TeleportRecieveEffectsThread, self)
		end,
	},
	
}

#################################################################
#		TELEPORT THREADS
################################################################# 
	NilThread = function(self)
		self.UnitBeingTeleported = nil
	end
	
	TeleportSendEffectsThread = function(self)
		self:PlaySendEffects()
		self:PlayStreamFlashes()
	end
	
	TeleportRecieveEffectsThread = function(self)
		self:PlayRecieveEffects()
		self:PlayStreamFlashes()
	end
	
	PlayDisappearFlashes = function(self)
		local army = self:GetArmy()
		local pos = self:GetPosition()
		local size1 = math.random(4,7)
		local size2 = math.random(7,12)
		local size3 = math.random(12,16)
		CreateLightParticleIntel( self, -1, army, size1, size1, 'flare_lens_add_02', 'ramp_blue_13' )
		WaitTicks(math.random(3,5))
		for k, v in EffectTemplate.CSGTestEffect2 do
           	CreateEmitterOnEntity( self, army, v )
       	end
       	for k, v in EffectTemplate.CommanderQuantumGateInEnergy do
    		CreateEmitterOnEntity( self, army, v )
       	end
       	WaitTicks(math.random(6,15)) 
        CreateLightParticleIntel( self, -1, army, size2, size2, 'flare_lens_add_02', 'ramp_blue_13' )
		WaitTicks(math.random(6,15))
		for k, v in EffectTemplate.CommanderTeleport01 do
           	CreateEmitterOnEntity( self, army, v )
        end
        DamageArea(self, pos, 10, 2500, 'Force', true, true)	
		CreateLightParticleIntel( self, -1, army, size1, size1, 'flare_lens_add_02', 'ramp_blue_13' )
		WaitTicks(math.random(10,15))
		DamageArea(self, pos, 10, 1750, 'Force', true, true)
		CreateLightParticleIntel( self, -1, army, size3, size3, 'flare_lens_add_02', 'ramp_blue_13' )
		CreateLightParticleIntel( self, -1, army, size2, size2, 'flare_lens_add_02', 'ramp_blue_13' )
		WaitTicks(math.random(3,5))
		DamageArea(self, pos, 10, 1000, 'Force', true, true)
		WaitTicks(math.random(6,15))
		CreateLightParticleIntel( self, -1, army, size3, size2, 'flare_lens_add_02', 'ramp_blue_13' )
		self.FeedbackPulse = true
		WaitSeconds(20)
		self.FeedbackPulse = false
    end
	
	TeleportingEffectsThread = function(self)
		WaitFor(self.TeleportDrain)
 		self:HideBone(0, true)
		self:PlayTeleportXInEffects()
		WaitSeconds(0.5)
 		self:ShowBone(0, true)
		if EntityCategoryContains(categories.COMMAND * categories.AEON, self) then
			if self:HasEnhancement(self, 'ResourceAllocation') == true or self:HasEnhancement(self, 'ChronoDampener') == true or self:HasEnhancement(self, 'ResourceAllocationAdvanced') == true or self:HasEnhancement(self, 'Shield') == true or self:HasEnhancement(self, 'ShieldHeavy') == true or self:HasEnhancement(self, ' Teleporter') == true then
			else
				self:HideBone('Back_Upgrade', true)
			end
		end
		if EntityCategoryContains(categories.BUILTBYQUANTUMGATE * categories.AEON, self) then
			if self:HasEnhancement(self, 'ResourceAllocation') == true or self:HasEnhancement(self, 'Shield') == true or self:HasEnhancement(self, 'SystemIntegrityCompensator') == true or self:HasEnhancement(self, 'Teleporter') == true then
			else
				self:HideBone('Turbine', true)
			end
		end
		if EntityCategoryContains(categories.COMMAND * categories.UEF, self) then
			if self:HasEnhancement(self, 'ResourceAllocation') == true or self:HasEnhancement(self, 'HeavyAntiMatterCannon') == true  then
			else
				self:HideBone('Right_Upgrade', true)
			end
			if self:HasEnhancement(self, 'AdvancedEngineering') == true or self:HasEnhancement(self, 'DamageStablization') == true  then
			else
		       	self:HideBone('Left_Upgrade', true)
		       end
		    if self:HasEnhancement(self, 'LeftPod') == true or self:HasEnhancement(self, 'Shield') == true or self:HasEnhancement(self, 'ShieldGeneratorField') == true or self:HasEnhancement(self, 'TacticalMissile') == true or self:HasEnhancement(self, 'TacticalMissileNuke') == true or self:HasEnhancement(self, 'Teleporter') == true  then
			else
		       	self:HideBone('Back_Upgrade_B01', true)
		    end
		end
		if EntityCategoryContains(categories.BUILTBYQUANTUMGATE * categories.UEF, self) then
			if self:HasEnhancement(self, 'ResourceAllocation') == true or self:HasEnhancement(self, 'SensorRangeEnhancer') == true or self:HasEnhancement(self, 'AdvancedCoolingUpgrade') == true or self:HasEnhancement(self, 'HighExplosiveOrdnance') == true then
			else
		        self:HideBone('SAM', true)
		     end
		     if self:HasEnhancement(self, 'RadarJammer') == true or self:HasEnhancement(self, 'LeftPod') == true or self:HasEnhancement(self, 'LeftPod') == true or self:HasEnhancement(self, 'LeftPod') == true or self:HasEnhancement(self, 'LeftPod') == true then
		     else
		        self:HideBone('Jetpack', true)
		     end
		end
		if EntityCategoryContains(categories.COMMAND * categories.CYBRAN, self) then
	       	if self:HasEnhancement(self, 'MicrowaveLaserGenerator') == true or self:HasEnhancement(self, 'NaniteTorpedoTube') == true  then
	       	else
		      	self:HideBone('Back_Upgrade', true)
		    end
		    if self:HasEnhancement(self, 'ResourceAllocation') == true or self:HasEnhancement(self, 'CoolingUpgrade') == true or self:HasEnhancement(self, 'AdvancedEngineering') == true  then
	        else
	        	self:HideBone('Right_Upgrade', true)
	        end
		end
		if EntityCategoryContains(categories.BUILTBYQUANTUMGATE * categories.CYBRAN, self) then
			if self:HasEnhancement(self, 'NaniteMissileSystem') == true then
			else
	           	self:HideBone('AA_Gun', true)
	        end
-- 	        if self:HasEnhancement(self, 'MicrowaveLaserGenerator') == true then
-- 			else
        		self:HideBone('Power_Pack', true)
--       	end
--         	if self:HasEnhancement(self, 'MicrowaveLaserGenerator') == true then
-- 			else
       			self:HideBone('Rez_Protocol', true)
--         	end
--         	if self:HasEnhancement(self, 'MicrowaveLaserGenerator') == true then
-- 			else
        		self:HideBone('Torpedo', true)
--         	end
--         	if self:HasEnhancement(self, 'MicrowaveLaserGenerator') == true then
-- 			else
        		self:HideBone('Turbine', true)
--        	end
        end
        self:SetImmobile(false)
		self:SetBusy(false)
		IssueClearCommands({self})
		ChangeState(self, self.IdleState)
		#LOG('moving unit')
		IssueMoveOffFactory({self}, self.destination:GetRallyPoint())
		self.teleporter = nil
		self.destination = nil
	end
	