#!/usr/bin/env python3
"""
Comprehensive Balance Optimizer for Village Simulation

This script runs simulations and optimizes config values for:
- High survival rate (target: >50% at end)
- Religion diversity (no single religion >60%)
- Faction survival (all factions have living members)
- Active market (trades happening, money circulating)
- Oil industry activity (drilling and refining)

Usage:
    python tools/optimize_balance.py [--iterations 20] [--steps 1000]
    python tools/optimize_balance.py --quick-test
    python tools/optimize_balance.py --analyze-current
"""

import argparse
import json
import random
import re
import sys
from collections import defaultdict
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Optional

# Add parent directory for imports
sys.path.insert(0, str(Path(__file__).parent.parent))

from backend.config import get_config, reload_config
from backend.core.engine import GameEngine
from backend.core.logger import reset_simulation_logger
from backend.domain.action import reset_action_config_cache
from backend.domain.resources import reset_resource_cache


@dataclass
class BalanceMetrics:
    """Comprehensive metrics for simulation balance."""
    total_turns: int = 0
    initial_population: int = 0
    final_population: int = 0
    
    # Survival tracking
    deaths_by_cause: dict = field(default_factory=lambda: defaultdict(int))
    population_over_time: list = field(default_factory=list)
    
    # Religion tracking
    religion_counts: dict = field(default_factory=lambda: defaultdict(int))
    conversions: int = 0
    
    # Faction tracking
    faction_counts: dict = field(default_factory=lambda: defaultdict(int))
    wars_declared: int = 0
    peace_treaties: int = 0
    
    # Market tracking
    total_listings: int = 0
    total_trades: int = 0
    trade_volume: int = 0
    trade_value: int = 0
    trades_by_resource: dict = field(default_factory=lambda: defaultdict(int))
    
    # Action diversity
    action_counts: dict = field(default_factory=lambda: defaultdict(int))
    
    # Oil industry
    oil_drilled: int = 0
    fuel_refined: int = 0
    
    # Economy
    money_circulation: list = field(default_factory=list)
    avg_wealth: list = field(default_factory=list)
    wealth_gini: list = field(default_factory=list)
    
    @property
    def survival_rate(self) -> float:
        """Final survival rate."""
        if self.initial_population == 0:
            return 0
        return self.final_population / self.initial_population
    
    @property
    def religion_diversity(self) -> float:
        """Religion diversity score (0-1, higher = more diverse)."""
        if not self.religion_counts:
            return 0
        total = sum(self.religion_counts.values())
        if total == 0:
            return 0
        max_count = max(self.religion_counts.values())
        # Perfect diversity = 20% each (5 religions), worst = 100% one religion
        return 1.0 - (max_count / total)
    
    @property
    def dominant_religion_pct(self) -> float:
        """Percentage held by dominant religion."""
        if not self.religion_counts:
            return 0
        total = sum(self.religion_counts.values())
        if total == 0:
            return 0
        return max(self.religion_counts.values()) / total
    
    @property
    def factions_alive(self) -> int:
        """Number of factions with living members."""
        return len([f for f, c in self.faction_counts.items() if c > 0])
    
    @property
    def faction_diversity(self) -> float:
        """Faction diversity (0-1)."""
        if not self.faction_counts:
            return 0
        alive = self.factions_alive
        # We have 5 non-neutral factions
        return alive / 5.0
    
    @property
    def market_activity(self) -> float:
        """Market activity score."""
        if self.total_turns == 0:
            return 0
        trades_per_turn = self.total_trades / self.total_turns
        # Target: 0.3 trades per turn per 10 agents
        return min(1.0, trades_per_turn / 0.3)
    
    @property
    def trade_diversity(self) -> float:
        """How many different resources are being traded."""
        resources_traded = len([r for r, c in self.trades_by_resource.items() if c > 0])
        return resources_traded / 6.0  # 6 tradeable resources
    
    @property
    def oil_industry_activity(self) -> float:
        """Oil industry health score."""
        total_oil_ops = self.oil_drilled + self.fuel_refined
        # Target: 5% of actions should be oil-related
        total_actions = sum(self.action_counts.values())
        if total_actions == 0:
            return 0
        return min(1.0, (total_oil_ops / total_actions) / 0.05)
    
    @property
    def economy_health(self) -> float:
        """Overall economy health."""
        if not self.avg_wealth:
            return 0
        final_wealth = self.avg_wealth[-1]
        # Target: average wealth should stay above 50
        return min(1.0, final_wealth / 50)
    
    def score(self) -> float:
        """Calculate overall balance score (0-100)."""
        score = 0
        
        # Survival rate (0-30 points) - CRITICAL
        # Target: at least 30% survival
        survival_score = min(30, self.survival_rate * 100)
        score += survival_score
        
        # Religion diversity (0-15 points)
        # Target: no single religion > 50%
        religion_score = self.religion_diversity * 15
        score += religion_score
        
        # Faction survival (0-15 points)  
        # Target: at least 4 of 5 factions alive
        faction_score = self.faction_diversity * 15
        score += faction_score
        
        # Market activity (0-15 points)
        market_score = self.market_activity * 15
        score += market_score
        
        # Trade diversity (0-10 points)
        trade_div_score = self.trade_diversity * 10
        score += trade_div_score
        
        # Oil industry (0-10 points)
        oil_score = self.oil_industry_activity * 10
        score += oil_score
        
        # Economy health (0-5 points)
        econ_score = self.economy_health * 5
        score += econ_score
        
        return score


def run_simulation(config_overrides: dict, num_steps: int = 1000, num_agents: int = 100) -> BalanceMetrics:
    """Run a simulation with custom config and return metrics."""
    # Apply config overrides
    config_path = Path("config.json")
    with open(config_path) as f:
        config = json.load(f)
    
    # Deep merge overrides
    for section, values in config_overrides.items():
        if section in config:
            config[section].update(values)
        else:
            config[section] = values
    
    # Save temp config
    temp_config = Path("config_temp.json")
    with open(temp_config, 'w') as f:
        json.dump(config, f, indent=2)
    
    # Reload config
    reload_config(str(temp_config))
    reset_action_config_cache()
    reset_resource_cache()
    
    # Initialize engine - need to set initial_agents BEFORE reset() calls initialize()
    GameEngine._instance = None  # Reset singleton
    engine = GameEngine()
    # Note: reset() already calls world.initialize(), so we must set initial_agents first
    # Get the config and modify it before reset
    sim_config = get_config()
    engine.world.config.initial_agents = num_agents
    # Reset creates a new world and initializes it
    from backend.core.world import World, WorldConfig
    world_config = WorldConfig(initial_agents=num_agents)
    engine.reset(config=world_config)
    
    # Suppress logging
    import logging
    logging.getLogger("simulation").setLevel(logging.ERROR)
    
    metrics = BalanceMetrics()
    metrics.initial_population = num_agents
    
    # Run simulation
    for step in range(num_steps):
        if not engine.is_running:
            break
        
        turn_log = engine.next_step()
        metrics.total_turns += 1
        
        # Track population
        living = len(engine.world.get_living_agents())
        metrics.population_over_time.append(living)
        
        # Track money
        agents = engine.world.get_living_agents()
        if agents:
            total_money = sum(a.money for a in agents)
            avg_money = total_money / len(agents)
            metrics.money_circulation.append(total_money)
            metrics.avg_wealth.append(avg_money)
            
            # Gini coefficient
            moneys = sorted([a.money for a in agents])
            n = len(moneys)
            if n > 1 and total_money > 0:
                sum_of_diffs = sum(abs(m1 - m2) for m1 in moneys for m2 in moneys)
                gini = sum_of_diffs / (2 * n * total_money)
            else:
                gini = 0
            metrics.wealth_gini.append(gini)
        
        # Process actions
        for action_data in turn_log.agent_actions:
            decision = action_data.get("decision", {})
            result = action_data.get("result", {})
            action_type = decision.get("action", "unknown")
            
            metrics.action_counts[action_type] += 1
            
            # Track specific actions
            if action_type == "drill_oil" and result.get("success"):
                for res in result.get("resources_gained", []):
                    if res.get("type") == "oil":
                        metrics.oil_drilled += res.get("quantity", 0)
            
            elif action_type == "refine" and result.get("success"):
                for res in result.get("resources_gained", []):
                    if res.get("type") == "fuel":
                        metrics.fuel_refined += res.get("quantity", 0)
            
            elif action_type == "preach" and result.get("success"):
                if "converted" in result.get("message", "").lower():
                    metrics.conversions += 1
            
            elif action_type == "declare_war" and result.get("success"):
                metrics.wars_declared += 1
            
            elif action_type == "make_peace" and result.get("success"):
                metrics.peace_treaties += 1
            
            elif action_type == "trade" and result.get("success"):
                message = result.get("message", "")
                if "Listed" in message:
                    metrics.total_listings += 1
                elif "Bought" in message:
                    match = re.search(r"Bought (\d+) (\w+) for (\d+)c", message)
                    if match:
                        qty = int(match.group(1))
                        res = match.group(2)
                        value = int(match.group(3))
                        metrics.total_trades += 1
                        metrics.trade_volume += qty
                        metrics.trade_value += value
                        metrics.trades_by_resource[res] += 1
        
        # Process deaths
        for death_name in turn_log.deaths:
            for agent in engine.world.agents:
                if agent.name == death_name and agent.death_reason:
                    metrics.deaths_by_cause[agent.death_reason] += 1
                    break
    
    # Collect final stats
    living_agents = engine.world.get_living_agents()
    metrics.final_population = len(living_agents)
    
    # Count religions and factions
    for agent in living_agents:
        metrics.religion_counts[agent.religion.religion.value] += 1
        metrics.faction_counts[agent.diplomacy.faction.value] += 1
    
    # Cleanup
    engine.logger.close()
    temp_config.unlink(missing_ok=True)
    
    return metrics


def generate_balanced_config() -> dict:
    """Generate a config focused on balance."""
    return {
        "agent_stats": {
            "start_hunger": random.randint(85, 95),
            "start_thirst": random.randint(85, 95),
            "hunger_decay": random.randint(1, 2),
            "thirst_decay": random.randint(1, 3),
            "heat_decay": random.randint(1, 2),
            "faith_decay": random.randint(1, 2),
            "critical_threshold": round(random.uniform(0.12, 0.20), 2),
        },
        "resources": {
            "meat_hunger": random.randint(40, 55),
            "berries_hunger": random.randint(12, 18),
            "water_thirst": random.randint(50, 70),
            "fire_heat": random.randint(25, 40),
        },
        "actions": {
            "hunt_success": round(random.uniform(0.75, 0.90), 2),
            "drill_oil_success": round(random.uniform(0.70, 0.85), 2),
            "hunt_meat_min": random.randint(3, 4),
            "hunt_meat_max": random.randint(5, 7),
            "gather_min": random.randint(4, 5),
            "gather_max": random.randint(6, 8),
            # preach_convert_chance is in actions, not religion
            "preach_convert_chance": round(random.uniform(0.03, 0.08), 2),
        },
        "religion": {
            "conversion_resistance": round(random.uniform(0.65, 0.85), 2),
            "zealot_threshold": round(random.uniform(0.80, 0.92), 2),
            "same_religion_bonus": round(random.uniform(0.08, 0.15), 2),
            "different_religion_penalty": round(random.uniform(0.02, 0.06), 2),
        },
        "diplomacy": {
            "starting_relations": random.randint(55, 70),
            "relation_decay": random.randint(0, 1),
            "trade_relation_boost": random.randint(6, 10),
            "war_exhaustion_rate": random.randint(8, 15),
            "war_threshold": random.randint(15, 25),
        },
        "economy": {
            "buy_efficiency_threshold": round(random.uniform(0.80, 0.95), 2),
            "min_wealth_target": random.randint(25, 50),
            "max_price_markup": round(random.uniform(1.4, 1.8), 1),
        },
    }


def mutate_config(config: dict, mutation_rate: float = 0.3) -> dict:
    """Mutate a configuration."""
    new_config = json.loads(json.dumps(config))
    
    for section, values in new_config.items():
        for key, value in values.items():
            if random.random() < mutation_rate:
                if isinstance(value, int):
                    delta = max(1, abs(value) // 4)
                    new_config[section][key] = max(0, value + random.randint(-delta, delta))
                elif isinstance(value, float):
                    delta = abs(value) * 0.15
                    new_val = value + random.uniform(-delta, delta)
                    new_config[section][key] = round(max(0.01, min(0.99, new_val)), 2)
    
    return new_config


def crossover_configs(config1: dict, config2: dict) -> dict:
    """Crossover two configurations."""
    new_config = {}
    for section in set(config1.keys()) | set(config2.keys()):
        if section in config1 and section in config2:
            new_config[section] = {}
            for key in set(config1[section].keys()) | set(config2[section].keys()):
                if random.random() < 0.5 and key in config1[section]:
                    new_config[section][key] = config1[section][key]
                elif key in config2[section]:
                    new_config[section][key] = config2[section][key]
        elif section in config1:
            new_config[section] = config1[section].copy()
        else:
            new_config[section] = config2[section].copy()
    return new_config


def print_metrics(metrics: BalanceMetrics, detailed: bool = True):
    """Print metrics in a readable format."""
    print(f"\n  📊 Balance Score: {metrics.score():.1f}/100")
    print(f"  ├─ Survival: {metrics.survival_rate*100:.0f}% ({metrics.final_population}/{metrics.initial_population})")
    print(f"  ├─ Religion: {metrics.religion_diversity*100:.0f}% diversity (dominant: {metrics.dominant_religion_pct*100:.0f}%)")
    print(f"  ├─ Factions: {metrics.factions_alive}/5 alive ({metrics.faction_diversity*100:.0f}%)")
    print(f"  ├─ Market: {metrics.total_trades} trades, {metrics.total_listings} listings")
    print(f"  ├─ Trade diversity: {metrics.trade_diversity*100:.0f}%")
    print(f"  ├─ Oil industry: {metrics.oil_drilled} oil, {metrics.fuel_refined} fuel")
    print(f"  └─ Economy: avg wealth ${metrics.avg_wealth[-1]:.0f}" if metrics.avg_wealth else "  └─ Economy: N/A")
    
    if detailed:
        print(f"\n  📋 Death causes:")
        for cause, count in sorted(metrics.deaths_by_cause.items(), key=lambda x: -x[1])[:5]:
            print(f"     - {cause}: {count}")
        
        print(f"\n  🏛️ Religions:")
        for religion, count in sorted(metrics.religion_counts.items(), key=lambda x: -x[1]):
            print(f"     - {religion}: {count}")
        
        print(f"\n  ⚔️ Factions:")
        for faction, count in sorted(metrics.faction_counts.items(), key=lambda x: -x[1]):
            print(f"     - {faction}: {count}")


def optimize_balance(iterations: int = 20, steps_per_sim: int = 1000, population_size: int = 8):
    """Run genetic optimization for balance."""
    print("\n" + "=" * 70)
    print("🧬 BALANCE OPTIMIZER - Finding Optimal Configuration")
    print("=" * 70)
    print(f"  Iterations: {iterations}")
    print(f"  Steps per simulation: {steps_per_sim}")
    print(f"  Population size: {population_size}")
    print(f"  Agents per simulation: 100")
    print("=" * 70)
    
    # Create initial population
    population = []
    
    # Start with a well-balanced baseline
    baseline = {
        "agent_stats": {
            "start_hunger": 92,
            "start_thirst": 92,
            "hunger_decay": 1,
            "thirst_decay": 2,
            "heat_decay": 1,
            "faith_decay": 1,
            "critical_threshold": 0.15,
        },
        "resources": {
            "meat_hunger": 50,
            "berries_hunger": 15,
            "water_thirst": 65,
            "fire_heat": 35,
        },
        "actions": {
            "hunt_success": 0.85,
            "drill_oil_success": 0.80,
            "hunt_meat_min": 4,
            "hunt_meat_max": 6,
            "gather_min": 4,
            "gather_max": 7,
            "preach_convert_chance": 0.05,
        },
        "religion": {
            "conversion_resistance": 0.75,
            "zealot_threshold": 0.88,
            "same_religion_bonus": 0.10,
            "different_religion_penalty": 0.03,
        },
        "diplomacy": {
            "starting_relations": 65,
            "relation_decay": 0,
            "trade_relation_boost": 8,
            "war_exhaustion_rate": 12,
            "war_threshold": 18,
        },
        "economy": {
            "buy_efficiency_threshold": 0.88,
            "min_wealth_target": 35,
            "max_price_markup": 1.5,
        },
    }
    population.append(baseline)
    
    # Add survival-focused variant
    survival_focused = json.loads(json.dumps(baseline))
    survival_focused["agent_stats"]["hunger_decay"] = 1
    survival_focused["agent_stats"]["thirst_decay"] = 1
    survival_focused["resources"]["meat_hunger"] = 55
    survival_focused["resources"]["berries_hunger"] = 18
    survival_focused["resources"]["water_thirst"] = 70
    population.append(survival_focused)
    
    # Add religion-balanced variant
    religion_balanced = json.loads(json.dumps(baseline))
    religion_balanced["religion"]["conversion_resistance"] = 0.82
    religion_balanced["actions"]["preach_convert_chance"] = 0.03
    religion_balanced["religion"]["zealot_threshold"] = 0.92
    population.append(religion_balanced)
    
    # Add diplomacy-stable variant
    diplomacy_stable = json.loads(json.dumps(baseline))
    diplomacy_stable["diplomacy"]["relation_decay"] = 0
    diplomacy_stable["diplomacy"]["starting_relations"] = 70
    diplomacy_stable["diplomacy"]["war_exhaustion_rate"] = 15
    population.append(diplomacy_stable)
    
    # Fill rest with random
    while len(population) < population_size:
        population.append(generate_balanced_config())
    
    best_config = None
    best_score = 0
    best_metrics = None
    
    for gen in range(iterations):
        print(f"\n📍 Generation {gen + 1}/{iterations}")
        print("-" * 50)
        
        scored_population = []
        for i, config in enumerate(population):
            sys.stdout.write(f"\r  Evaluating config {i + 1}/{len(population)}...")
            sys.stdout.flush()
            
            metrics = run_simulation(config, steps_per_sim, num_agents=100)
            score = metrics.score()
            scored_population.append((config, metrics, score))
        
        # Sort by score
        scored_population.sort(key=lambda x: x[2], reverse=True)
        
        # Print top results
        print(f"\r  Top configs this generation:")
        for i, (config, metrics, score) in enumerate(scored_population[:3]):
            print(f"\n  #{i + 1}: Score {score:.1f}")
            print_metrics(metrics, detailed=False)
        
        # Track best overall
        if scored_population[0][2] > best_score:
            best_config = scored_population[0][0]
            best_score = scored_population[0][2]
            best_metrics = scored_population[0][1]
            print(f"\n  ⭐ New best score: {best_score:.1f}")
        
        # Create next generation
        new_population = []
        
        # Keep top 2 (elitism)
        new_population.append(scored_population[0][0])
        new_population.append(scored_population[1][0])
        
        # Crossover and mutate
        while len(new_population) < population_size:
            parent1 = random.choice(scored_population[:4])[0]
            parent2 = random.choice(scored_population[:4])[0]
            child = crossover_configs(parent1, parent2)
            child = mutate_config(child, mutation_rate=0.25)
            new_population.append(child)
        
        population = new_population
    
    print("\n" + "=" * 70)
    print("🏆 OPTIMIZATION COMPLETE")
    print("=" * 70)
    
    print(f"\n  Best Score: {best_score:.1f}/100")
    print_metrics(best_metrics, detailed=True)
    
    print("\n  📝 Best Configuration:")
    print("-" * 50)
    print(json.dumps(best_config, indent=2))
    
    # Save optimized config
    output_path = Path("config_balanced.json")
    
    with open("config.json") as f:
        full_config = json.load(f)
    
    for section, values in best_config.items():
        if section in full_config:
            full_config[section].update(values)
        else:
            full_config[section] = values
    
    with open(output_path, 'w') as f:
        json.dump(full_config, f, indent=2)
    
    print(f"\n  ✅ Saved optimized config to: {output_path}")
    print("  To apply: cp config_balanced.json config.json")
    
    return best_config, best_metrics


def analyze_current_config(steps: int = 500):
    """Analyze the current configuration."""
    print("\n" + "=" * 70)
    print("📊 ANALYZING CURRENT CONFIGURATION")
    print("=" * 70)
    
    metrics = run_simulation({}, steps, num_agents=100)
    print_metrics(metrics, detailed=True)
    
    # Provide recommendations
    print("\n" + "=" * 70)
    print("💡 RECOMMENDATIONS")
    print("=" * 70)
    
    if metrics.survival_rate < 0.3:
        print("\n  ⚠️ LOW SURVIVAL RATE")
        print("     - Reduce hunger_decay and thirst_decay")
        print("     - Increase food resource values (meat_hunger, berries_hunger)")
        print("     - Lower critical_threshold")
    
    if metrics.dominant_religion_pct > 0.6:
        print("\n  ⚠️ RELIGION DOMINANCE")
        print("     - Increase conversion_resistance (try 0.75+)")
        print("     - Reduce preach_convert_chance (try 0.05)")
        print("     - Increase zealot_threshold (try 0.88+)")
    
    if metrics.factions_alive < 4:
        print("\n  ⚠️ FACTIONS DYING OUT")
        print("     - Set relation_decay to 0 or 1")
        print("     - Increase starting_relations (try 65+)")
        print("     - Increase war_exhaustion_rate (try 10+)")
    
    if metrics.total_trades < metrics.total_turns * 0.1:
        print("\n  ⚠️ LOW MARKET ACTIVITY")
        print("     - Increase buy_efficiency_threshold (try 0.9)")
        print("     - Lower min_wealth_target")
        print("     - Reduce max_price_markup")
    
    if metrics.oil_drilled + metrics.fuel_refined < 50:
        print("\n  ⚠️ LOW OIL INDUSTRY")
        print("     - Increase drill_oil_success (try 0.80)")
        print("     - Check that factions with oil bonus survive")
    
    return metrics


def quick_test(steps: int = 500):
    """Quick test with a balanced preset."""
    print("\n🧪 Quick Test with Balanced Preset")
    print("-" * 50)
    
    test_config = {
        "agent_stats": {
            "start_hunger": 92,
            "start_thirst": 92,
            "hunger_decay": 1,
            "thirst_decay": 2,
            "heat_decay": 1,
            "faith_decay": 1,
            "critical_threshold": 0.15,
        },
        "resources": {
            "meat_hunger": 50,
            "berries_hunger": 16,
            "water_thirst": 65,
            "fire_heat": 35,
        },
        "actions": {
            "hunt_success": 0.85,
            "drill_oil_success": 0.80,
            "hunt_meat_min": 4,
            "hunt_meat_max": 6,
            "gather_min": 4,
            "gather_max": 7,
            "preach_convert_chance": 0.04,
        },
        "religion": {
            "conversion_resistance": 0.78,
            "zealot_threshold": 0.90,
            "same_religion_bonus": 0.10,
            "different_religion_penalty": 0.03,
        },
        "diplomacy": {
            "starting_relations": 65,
            "relation_decay": 0,
            "trade_relation_boost": 8,
            "war_exhaustion_rate": 12,
            "war_threshold": 18,
        },
        "economy": {
            "buy_efficiency_threshold": 0.90,
            "min_wealth_target": 30,
            "max_price_markup": 1.5,
        },
    }
    
    print("\n  Testing config:")
    print(json.dumps(test_config, indent=2))
    
    metrics = run_simulation(test_config, steps, num_agents=100)
    print_metrics(metrics, detailed=True)
    
    return metrics


def main():
    parser = argparse.ArgumentParser(description="Optimize Village Simulation balance")
    parser.add_argument("--iterations", "-i", type=int, default=15, help="Optimization iterations")
    parser.add_argument("--steps", "-s", type=int, default=800, help="Steps per simulation")
    parser.add_argument("--population", "-p", type=int, default=8, help="Population size for GA")
    parser.add_argument("--quick-test", "-q", action="store_true", help="Quick test balanced preset")
    parser.add_argument("--analyze-current", "-a", action="store_true", help="Analyze current config")
    
    args = parser.parse_args()
    
    if args.quick_test:
        quick_test(args.steps)
    elif args.analyze_current:
        analyze_current_config(args.steps)
    else:
        optimize_balance(args.iterations, args.steps, args.population)


if __name__ == "__main__":
    main()