/**** 
* Plugins
****/ 
var tween = LK.import("@upit/tween.v1");

/**** 
* Classes
****/ 
var Bullet = Container.expand(function (startX, startY, targetX, targetY, damage, faction) {
	var self = Container.call(this);
	self.x = startX;
	self.y = startY;
	self.damage = damage;
	self.faction = faction;
	self.speed = 8;
	var dx = targetX - startX;
	var dy = targetY - startY;
	var distance = Math.sqrt(dx * dx + dy * dy);
	self.velocityX = dx / distance * self.speed;
	self.velocityY = dy / distance * self.speed;
	self.maxDistance = distance;
	self.traveled = 0;
	var bulletGraphics = self.attachAsset('bullet', {
		anchorX: 0.5,
		anchorY: 0.5
	});
	self.update = function () {
		self.x += self.velocityX;
		self.y += self.velocityY;
		self.traveled += self.speed;
		// Check if bullet reached target area or max distance
		if (self.traveled >= self.maxDistance) {
			self.explode();
		}
	};
	self.explode = function () {
		// Find units in explosion radius
		var explosionRadius = 30;
		var explosionPlayed = false;
		for (var i = 0; i < allUnits.length; i++) {
			var unit = allUnits[i];
			if (unit.faction !== self.faction && unit.health > 0) {
				var dx = unit.x - self.x;
				var dy = unit.y - self.y;
				var distance = Math.sqrt(dx * dx + dy * dy);
				if (distance < explosionRadius) {
					var destroyed = unit.takeDamage(self.damage);
					if (destroyed && !explosionPlayed) {
						// Play unit-specific destruction sound
						var soundName = unit.type + 'Destroy';
						LK.getSound(soundName).play();
						explosionPlayed = true;
					}
				}
			}
		}
		// Check cities
		var enemyCities = self.faction === 'red' ? blueCities : redCities;
		for (var i = 0; i < enemyCities.length; i++) {
			var city = enemyCities[i];
			if (city.health > 0) {
				var dx = city.x - self.x;
				var dy = city.y - self.y;
				var distance = Math.sqrt(dx * dx + dy * dy);
				if (distance < explosionRadius) {
					var destroyed = city.takeDamage(self.damage);
					if (destroyed && !explosionPlayed) {
						LK.getSound('cityDestroy').play();
						explosionPlayed = true;
					}
				}
			}
		}
		// Remove bullet
		for (var i = 0; i < bullets.length; i++) {
			if (bullets[i] === self) {
				bullets.splice(i, 1);
				break;
			}
		}
		self.destroy();
	};
	return self;
});
var City = Container.expand(function (faction) {
	var self = Container.call(this);
	self.faction = faction;
	self.health = 1000;
	self.maxHealth = 1000;
	self.lastProduction = 0;
	self.productionCooldown = 180; // 3 seconds at 60fps
	var cityGraphics = self.attachAsset('city', {
		anchorX: 0.5,
		anchorY: 0.5
	});
	// Color based on faction
	if (faction === 'red') {
		cityGraphics.tint = 0xff4444;
	} else {
		cityGraphics.tint = 0x4444ff;
	}
	// Health bar background
	var healthBarBg = self.attachAsset('healthBarBg', {
		anchorX: 0.5,
		anchorY: 0.5,
		y: -50
	});
	// Health bar
	var healthBar = self.attachAsset('healthBar', {
		anchorX: 0.5,
		anchorY: 0.5,
		y: -50
	});
	self.takeDamage = function (damage) {
		self.health -= damage;
		if (self.health <= 0) {
			self.health = 0;
			return true; // City destroyed
		}
		// Update health bar
		var healthPercent = self.health / self.maxHealth;
		healthBar.width = 80 * healthPercent;
		return false;
	};
	self.produceUnit = function () {
		if (LK.ticks - self.lastProduction < self.productionCooldown) return;
		// Check if we've reached the maximum unit limit per faction
		var aliveUnits = 0;
		for (var i = 0; i < allUnits.length; i++) {
			if (allUnits[i].health > 0 && allUnits[i].faction === self.faction) {
				aliveUnits++;
			}
		}
		if (aliveUnits >= 10) return; // Don't produce more units if we have 10 or more per faction
		// Choose random unit type
		var unitTypes = ['infantry', 'tank', 'artillery'];
		var unitType = unitTypes[Math.floor(Math.random() * unitTypes.length)];
		var unit = new Unit(unitType, self.faction);
		// Find valid position that doesn't overlap with other units
		var validPosition = false;
		var attempts = 0;
		var maxAttempts = 20;
		var unitRadius = 40; // Minimum distance between units
		while (!validPosition && attempts < maxAttempts) {
			var angle = Math.random() * Math.PI * 2;
			var distance = 100 + Math.random() * 100;
			unit.x = self.x + Math.cos(angle) * distance;
			unit.y = self.y + Math.sin(angle) * distance;
			// Keep within bounds
			unit.x = Math.max(50, Math.min(1998, unit.x));
			unit.y = Math.max(50, Math.min(2682, unit.y));
			// Check for collision with existing units
			validPosition = true;
			for (var i = 0; i < allUnits.length; i++) {
				var otherUnit = allUnits[i];
				if (otherUnit.health > 0) {
					var dx = unit.x - otherUnit.x;
					var dy = unit.y - otherUnit.y;
					var distance = Math.sqrt(dx * dx + dy * dy);
					if (distance < unitRadius) {
						validPosition = false;
						break;
					}
				}
			}
			attempts++;
		}
		// Only spawn if valid position found
		if (validPosition) {
			allUnits.push(unit);
			game.addChild(unit);
			self.lastProduction = LK.ticks;
		}
	};
	self.update = function () {
		if (self.health <= 0) return;
		self.produceUnit();
	};
	return self;
});
var Unit = Container.expand(function (type, faction) {
	var self = Container.call(this);
	self.type = type;
	self.faction = faction;
	self.target = null;
	self.lastShot = 0;
	self.health = 100;
	self.maxHealth = 100;
	self.moveSpeed = 2;
	self.range = 100;
	self.damage = 20;
	self.shootCooldown = 60;
	// Set unit stats based on type
	if (type === 'infantry') {
		self.moveSpeed = 2;
		self.health = 60;
		self.maxHealth = 60;
		self.damage = 15;
		self.range = 80;
		self.shootCooldown = 40;
	} else if (type === 'tank') {
		self.moveSpeed = 2.5;
		self.health = 120;
		self.maxHealth = 120;
		self.damage = 30;
		self.range = 100;
		self.shootCooldown = 60;
	} else if (type === 'artillery') {
		self.moveSpeed = 1;
		self.health = 80;
		self.maxHealth = 80;
		self.damage = 50;
		self.range = 180;
		self.shootCooldown = 90;
	}
	var assetName = type;
	if (faction === 'red') {
		assetName = 'red' + type.charAt(0).toUpperCase() + type.slice(1);
	}
	var unitGraphics = self.attachAsset(assetName, {
		anchorX: 0.5,
		anchorY: 0.5
	});
	// Color based on faction
	if (faction === 'red') {
		unitGraphics.tint = 0xff4444;
	} else {
		unitGraphics.tint = 0x4444ff;
	}
	// Health bar background
	var healthBarBg = self.attachAsset('healthBarBg', {
		anchorX: 0.5,
		anchorY: 0.5,
		y: -50
	});
	// Health bar
	var healthBar = self.attachAsset('healthBar', {
		anchorX: 0.5,
		anchorY: 0.5,
		y: -50
	});
	self.takeDamage = function (damage) {
		self.health -= damage;
		if (self.health <= 0) {
			self.health = 0;
			return true; // Unit destroyed
		}
		// Update health bar
		var healthPercent = self.health / self.maxHealth;
		healthBar.width = 80 * healthPercent;
		return false;
	};
	self.findTarget = function () {
		var closestDistance = Infinity;
		var closestTarget = null;
		// Find closest enemy unit
		for (var i = 0; i < allUnits.length; i++) {
			var unit = allUnits[i];
			if (unit.faction !== self.faction && unit.health > 0) {
				var dx = unit.x - self.x;
				var dy = unit.y - self.y;
				var distance = Math.sqrt(dx * dx + dy * dy);
				if (distance < closestDistance) {
					closestDistance = distance;
					closestTarget = unit;
				}
			}
		}
		// Check enemy cities
		var enemyCities = self.faction === 'red' ? blueCities : redCities;
		for (var i = 0; i < enemyCities.length; i++) {
			var city = enemyCities[i];
			if (city.health > 0) {
				var dx = city.x - self.x;
				var dy = city.y - self.y;
				var distance = Math.sqrt(dx * dx + dy * dy);
				if (distance < closestDistance) {
					closestDistance = distance;
					closestTarget = city;
				}
			}
		}
		return closestTarget;
	};
	self.update = function () {
		if (self.health <= 0) return;
		self.target = self.findTarget();
		if (self.target && self.target.health > 0) {
			var dx = self.target.x - self.x;
			var dy = self.target.y - self.y;
			var distance = Math.sqrt(dx * dx + dy * dy);
			if (distance > self.range) {
				// Move towards target
				var moveX = dx / distance * self.moveSpeed;
				var moveY = dy / distance * self.moveSpeed;
				// Calculate new position
				var newX = self.x + moveX;
				var newY = self.y + moveY;
				// Check for collision with other units and apply separation
				var unitRadius = 45; // Minimum distance between units
				var separationForce = {
					x: 0,
					y: 0
				};
				var nearbyUnits = 0;
				for (var i = 0; i < allUnits.length; i++) {
					var otherUnit = allUnits[i];
					if (otherUnit !== self && otherUnit.health > 0) {
						var otherDx = newX - otherUnit.x;
						var otherDy = newY - otherUnit.y;
						var otherDistance = Math.sqrt(otherDx * otherDx + otherDy * otherDy);
						if (otherDistance < unitRadius) {
							// Apply separation force
							if (otherDistance > 0) {
								separationForce.x += otherDx / otherDistance;
								separationForce.y += otherDy / otherDistance;
								nearbyUnits++;
							} else {
								// Units are on same position, push in random direction
								var randomAngle = Math.random() * Math.PI * 2;
								separationForce.x += Math.cos(randomAngle);
								separationForce.y += Math.sin(randomAngle);
								nearbyUnits++;
							}
						}
					}
				}
				// Apply separation if needed
				if (nearbyUnits > 0) {
					separationForce.x /= nearbyUnits;
					separationForce.y /= nearbyUnits;
					// Normalize and apply separation
					var sepLength = Math.sqrt(separationForce.x * separationForce.x + separationForce.y * separationForce.y);
					if (sepLength > 0) {
						separationForce.x = separationForce.x / sepLength * self.moveSpeed * 0.5;
						separationForce.y = separationForce.y / sepLength * self.moveSpeed * 0.5;
						newX += separationForce.x;
						newY += separationForce.y;
					}
				}
				// Update position
				self.x = newX;
				self.y = newY;
				// Keep units within bounds
				self.x = Math.max(40, Math.min(2008, self.x));
				self.y = Math.max(40, Math.min(2692, self.y));
			} else {
				// In range, shoot
				if (LK.ticks - self.lastShot > self.shootCooldown) {
					self.shoot();
					self.lastShot = LK.ticks;
				}
			}
		}
	};
	self.shoot = function () {
		if (!self.target || self.target.health <= 0) return;
		var bullet = new Bullet(self.x, self.y, self.target.x, self.target.y, self.damage, self.faction);
		bullets.push(bullet);
		game.addChild(bullet);
		// Play unit-specific shooting sound
		var shootSoundName = self.type + 'Shoot';
		LK.getSound(shootSoundName).play();
	};
	return self;
});

/**** 
* Initialize Game
****/ 
var game = new LK.Game({
	backgroundColor: 0x228B22
});

/**** 
* Game Code
****/ 
// Game state
var allUnits = [];
var bullets = [];
var redCities = [];
var blueCities = [];
var gameStarted = false;
var showingMenu = true;
// Add territory background image
var territoryBackground = LK.getAsset('territoryBackground', {
	anchorX: 0.5,
	anchorY: 0.5
});
territoryBackground.x = 2048 / 2;
territoryBackground.y = 2732 / 2;
game.addChild(territoryBackground);
// UI
var statusText = new Text2('Europa 1941: AI War', {
	size: 40,
	fill: 0xFFFFFF
});
statusText.anchor.set(0.5, 0);
LK.gui.top.addChild(statusText);
// Main Menu UI
var menuBackground = LK.getAsset('menuBackground', {
	anchorX: 0.5,
	anchorY: 0.5
});
menuBackground.alpha = 1.0;
// Scale to cover full screen
menuBackground.width = 2048;
menuBackground.height = 2732;
LK.gui.center.addChild(menuBackground);
var menuTitle = LK.getAsset('titleImage', {
	anchorX: 0.5,
	anchorY: 0.5
});
menuTitle.y = -100;
LK.gui.center.addChild(menuTitle);
var startButton = LK.getAsset('startButton', {
	anchorX: 0.5,
	anchorY: 0.5
});
startButton.y = 150;
LK.gui.center.addChild(startButton);
var instructionsText = new Text2('Red vs Blue AI armies\nDestroy all enemy cities to win', {
	size: 25,
	fill: 0xAAAAA
});
instructionsText.anchor.set(0.5, 0.5);
instructionsText.y = 250;
LK.gui.center.addChild(instructionsText);
// Animate start button
function animateStartButton() {
	tween(startButton, {
		scaleX: 1.2,
		scaleY: 1.2
	}, {
		duration: 800,
		yoyo: true,
		repeat: Infinity
	});
}
animateStartButton();
// Play menu music
LK.playMusic('menuMusic');
// BETA text in menu
var betaText = new Text2('BETA', {
	size: 80,
	fill: 0x0000FF
});
betaText.anchor.set(0, 0);
betaText.x = -900;
betaText.y = -1200;
LK.gui.center.addChild(betaText);
// Red city counter button
var redCityButton = new Container();
var redButtonBg = redCityButton.attachAsset('cityCounterButton', {
	anchorX: 0.5,
	anchorY: 0.5
});
redButtonBg.tint = 0xff4444;
var redCityText = new Text2('0', {
	size: 40,
	fill: 0xFFFFFF
});
redCityText.anchor.set(0, 0.5);
redCityText.x = 70;
redCityButton.addChild(redCityText);
redCityButton.x = 180;
redCityButton.y = 100;
LK.gui.topLeft.addChild(redCityButton);
// Blue city counter button
var blueCityButton = new Container();
var blueButtonBg = blueCityButton.attachAsset('cityCounterButton', {
	anchorX: 0.5,
	anchorY: 0.5
});
blueButtonBg.tint = 0x4444ff;
var blueCityText = new Text2('0', {
	size: 40,
	fill: 0xFFFFFF
});
blueCityText.anchor.set(0, 0.5);
blueCityText.x = 70;
blueCityButton.addChild(blueCityText);
blueCityButton.x = -180;
blueCityButton.y = 100;
LK.gui.topRight.addChild(blueCityButton);
// Button click handlers
redCityButton.down = function (x, y, obj) {
	console.log("Red city button clicked!");
	// Add any specific functionality for red city button here
};
blueCityButton.down = function (x, y, obj) {
	console.log("Blue city button clicked!");
	// Add any specific functionality for blue city button here
};
// Initialize cities
function initializeCities() {
	var minDistance = 120; // Minimum distance between cities
	// Red cities (left side)
	for (var i = 0; i < 5; i++) {
		var city = new City('red');
		var attempts = 0;
		var maxAttempts = 50;
		var validPosition = false;
		while (!validPosition && attempts < maxAttempts) {
			city.x = 150 + Math.random() * 500; // Random x between 150-650
			city.y = 200 + Math.random() * 2300; // Random y between 200-2500
			validPosition = true;
			// Check distance from other red cities
			for (var j = 0; j < redCities.length; j++) {
				var dx = city.x - redCities[j].x;
				var dy = city.y - redCities[j].y;
				var distance = Math.sqrt(dx * dx + dy * dy);
				if (distance < minDistance) {
					validPosition = false;
					break;
				}
			}
			attempts++;
		}
		redCities.push(city);
		game.addChild(city);
	}
	// Blue cities (right side)
	for (var i = 0; i < 5; i++) {
		var city = new City('blue');
		var attempts = 0;
		var maxAttempts = 50;
		var validPosition = false;
		while (!validPosition && attempts < maxAttempts) {
			city.x = 1400 + Math.random() * 500; // Random x between 1400-1900
			city.y = 200 + Math.random() * 2300; // Random y between 200-2500
			validPosition = true;
			// Check distance from other blue cities
			for (var j = 0; j < blueCities.length; j++) {
				var dx = city.x - blueCities[j].x;
				var dy = city.y - blueCities[j].y;
				var distance = Math.sqrt(dx * dx + dy * dy);
				if (distance < minDistance) {
					validPosition = false;
					break;
				}
			}
			attempts++;
		}
		blueCities.push(city);
		game.addChild(city);
	}
}
function updateUI() {
	var redAlive = 0;
	var blueAlive = 0;
	// Count alive cities
	for (var i = 0; i < redCities.length; i++) {
		if (redCities[i].health > 0) {
			redAlive++;
		}
	}
	for (var i = 0; i < blueCities.length; i++) {
		if (blueCities[i].health > 0) {
			blueAlive++;
		}
	}
	redCityText.setText(redAlive.toString());
	blueCityText.setText(blueAlive.toString());
	// Check win conditions
	if (redAlive === 0) {
		statusText.setText('Blue AI Wins!');
		returnToMenu();
	} else if (blueAlive === 0) {
		statusText.setText('Red AI Wins!');
		returnToMenu();
	}
}
function cleanupDeadUnits() {
	for (var i = allUnits.length - 1; i >= 0; i--) {
		if (allUnits[i].health <= 0) {
			allUnits[i].destroy();
			allUnits.splice(i, 1);
		}
	}
}
// Function to return to main menu
function returnToMenu() {
	// Reset game state
	gameStarted = false;
	showingMenu = true;
	statusText.setText('Europa 1941: AI War');
	// Clear all game objects
	for (var i = 0; i < allUnits.length; i++) {
		allUnits[i].destroy();
	}
	allUnits = [];
	for (var i = 0; i < bullets.length; i++) {
		bullets[i].destroy();
	}
	bullets = [];
	for (var i = 0; i < redCities.length; i++) {
		redCities[i].destroy();
	}
	redCities = [];
	for (var i = 0; i < blueCities.length; i++) {
		blueCities[i].destroy();
	}
	blueCities = [];
	// Show menu elements
	menuBackground.visible = true;
	menuTitle.visible = true;
	startButton.visible = true;
	instructionsText.visible = true;
	betaText.visible = true;
	// Reset UI counters
	redCityText.setText('0');
	blueCityText.setText('0');
	// Stop background music and play menu music
	LK.stopMusic();
	LK.playMusic('menuMusic');
}
// Function to start the game
function startGame() {
	// Hide menu elements
	menuBackground.visible = false;
	menuTitle.visible = false;
	startButton.visible = false;
	instructionsText.visible = false;
	betaText.visible = false;
	showingMenu = false;
	// Stop menu music and play background music
	LK.stopMusic();
	LK.playMusic('backgroundMusic');
	// Start game initialization
	LK.setTimeout(function () {
		initializeCities();
		gameStarted = true;
		statusText.setText('Battle in Progress...');
	}, 500);
}
// Touch handler for menu
game.down = function (x, y, obj) {
	if (showingMenu) {
		startGame();
	}
};
// Game update loop
game.update = function () {
	if (showingMenu || !gameStarted) return;
	// Update all units
	for (var i = 0; i < allUnits.length; i++) {
		if (allUnits[i].health > 0) {
			allUnits[i].update();
		}
	}
	// Update all cities
	for (var i = 0; i < redCities.length; i++) {
		redCities[i].update();
	}
	for (var i = 0; i < blueCities.length; i++) {
		blueCities[i].update();
	}
	// Update all bullets
	for (var i = 0; i < bullets.length; i++) {
		bullets[i].update();
	}
	// Cleanup dead units
	cleanupDeadUnits();
	// Update UI
	updateUI();
};

===================================================================
--- original.js
+++ change.js
@@ -564,12 +564,12 @@
 	blueCityText.setText(blueAlive.toString());
 	// Check win conditions
 	if (redAlive === 0) {
 		statusText.setText('Blue AI Wins!');
-		LK.showYouWin();
+		returnToMenu();
 	} else if (blueAlive === 0) {
 		statusText.setText('Red AI Wins!');
-		LK.showYouWin();
+		returnToMenu();
 	}
 }
 function cleanupDeadUnits() {
 	for (var i = allUnits.length - 1; i >= 0; i--) {
@@ -578,8 +578,44 @@
 			allUnits.splice(i, 1);
 		}
 	}
 }
+// Function to return to main menu
+function returnToMenu() {
+	// Reset game state
+	gameStarted = false;
+	showingMenu = true;
+	statusText.setText('Europa 1941: AI War');
+	// Clear all game objects
+	for (var i = 0; i < allUnits.length; i++) {
+		allUnits[i].destroy();
+	}
+	allUnits = [];
+	for (var i = 0; i < bullets.length; i++) {
+		bullets[i].destroy();
+	}
+	bullets = [];
+	for (var i = 0; i < redCities.length; i++) {
+		redCities[i].destroy();
+	}
+	redCities = [];
+	for (var i = 0; i < blueCities.length; i++) {
+		blueCities[i].destroy();
+	}
+	blueCities = [];
+	// Show menu elements
+	menuBackground.visible = true;
+	menuTitle.visible = true;
+	startButton.visible = true;
+	instructionsText.visible = true;
+	betaText.visible = true;
+	// Reset UI counters
+	redCityText.setText('0');
+	blueCityText.setText('0');
+	// Stop background music and play menu music
+	LK.stopMusic();
+	LK.playMusic('menuMusic');
+}
 // Function to start the game
 function startGame() {
 	// Hide menu elements
 	menuBackground.visible = false;