import json
import os
from collections import defaultdict

def analyze_simulation(sim_file):
    """Analyze a single simulation file"""
    with open(sim_file, 'r') as f:
        data = json.load(f)
    winner = data.get('winner', 'Unknown')
    turns = data.get('turns', [])
    if not turns:
        return None
    # Get initial and final state
    initial_state = turns[0]['state']
    final_state = turns[-1]['state']
    # Count units by team
    def count_units(state):
        red_units = 0
        blue_units = 0
        for obj in state['objs'].values():
            if obj.get('obj_type') == 'Unit':
                if obj.get('team') == 'Red':
                    red_units += 1
                elif obj.get('team') == 'Blue':
                    blue_units += 1
        return red_units, blue_units
    initial_red, initial_blue = count_units(initial_state)
    final_red, final_blue = count_units(final_state)
    return {
        'winner': winner,
        'turns': len(turns),
        'initial_red': initial_red,
        'initial_blue': initial_blue,
        'final_red': final_red,
        'final_blue': final_blue
    }

def analyze_round(round_dir):
    """Analyze all simulations in a round"""
    results_file = os.path.join(round_dir, 'results.json')
    if not os.path.exists(results_file):
        print(f"No results.json found in {round_dir}")
        return
    with open(results_file, 'r') as f:
        results = json.load(f)
    print(f"\n=== Round {results['round_num']} Results ===")
    print(f"Winner: {results['winner']}")
    print(f"Scores: {results['scores']}")
    # Analyze a few sample simulations
    sim_files = [f for f in os.listdir(round_dir) if f.startswith('sim_') and f.endswith('.json')]
    sim_files.sort(key=lambda x: int(x.split('_')[1].split('.')[0]))
    print(f"\nTotal simulations: {len(sim_files)}")
    # Sample first, middle, and last simulations
    samples = []
    if len(sim_files) > 0:
        samples.append(sim_files[0])
    if len(sim_files) > 1:
        samples.append(sim_files[len(sim_files)//2])
    if len(sim_files) > 2:
        samples.append(sim_files[-1])
    print("\nSample simulations:")
    for sim_file in samples:
        sim_path = os.path.join(round_dir, sim_file)
        analysis = analyze_simulation(sim_path)
        if analysis:
            print(f"\n{sim_file}:")
            print(f"  Winner: {analysis['winner']}")
            print(f"  Turns: {analysis['turns']}")
            print(f"  Initial units - Red: {analysis['initial_red']}, Blue: {analysis['initial_blue']}")
            print(f"  Final units - Red: {analysis['final_red']}, Blue: {analysis['final_blue']}")

if __name__ == '__main__':
    logs_dir = '/logs/rounds'
    # Find all round directories
    round_dirs = []
    for item in os.listdir(logs_dir):
        item_path = os.path.join(logs_dir, item)
        if os.path.isdir(item_path):
            round_dirs.append(item_path)
    round_dirs.sort()
    for round_dir in round_dirs:
        analyze_round(round_dir)

"""
Round 15 - Improved spawn clearing and more aggressive combat
- Start clearing at turn 3 (earlier than turn 5)
- ALL units avoid spawn points during clearing window
- More aggressive combat to maintain unit superiority
"""

def get_robots_for_team(state, team):
    """Get all units for a team."""
    return [obj for obj in state.objs_by_team(team)]

def get_nearby_units(unit, units, radius=5):
    """Get units within radius of given unit."""
    return [u for u in units if unit.coords.distance_to(u.coords) <= radius]

def find_weakest_adjacent_enemy(unit, enemies):
    """Find the weakest adjacent enemy to attack."""
    adjacent = [(e, unit.coords.direction_to(e.coords)) 
                for e in enemies 
                if unit.coords.distance_to(e.coords) == 1]

    if not adjacent:
        return None

    adjacent.sort(key=lambda x: x[0].health)
    return adjacent[0][1]

def get_direction_counts(unit, units):
    """Count how many units are in each direction."""
    direction_counts = {}
    for u in units:
        if u.id == unit.id:
            continue
        direction = unit.coords.direction_to(u.coords)
        direction_counts[direction] = direction_counts.get(direction, 0) + 1
    return direction_counts

def get_most_common_directions(unit, units):
    """Get directions with most units."""
    counts = get_direction_counts(unit, units)
    if not counts:
        return []

    max_count = max(counts.values())
    return [d for d, c in counts.items() if c == max_count]

def is_move_valid(state, unit, direction):
    """Check if we can move in a direction."""
    try:
        new_coords = unit.coords + direction
        obj = state.obj_by_coords(new_coords)
        return obj is None
    except:
        return False

def get_move_toward_center(state, unit):
    """Get a move that takes us toward the center of the map (away from spawn points)."""
    # Map center is around (9, 9)
    center_x, center_y = 9, 9

    # Determine which direction moves us closer to center
    dx = center_x - unit.coords.x
    dy = center_y - unit.coords.y

    # Prioritize the larger difference
    directions = []
    if abs(dx) >= abs(dy):
        if dx > 0:
            directions.append(Direction.East)
        elif dx < 0:
            directions.append(Direction.West)
        if dy > 0:
            directions.append(Direction.South)
        elif dy < 0:
            directions.append(Direction.North)
    else:
        if dy > 0:
            directions.append(Direction.South)
        elif dy < 0:
            directions.append(Direction.North)
        if dx > 0:
            directions.append(Direction.East)
        elif dx < 0:
            directions.append(Direction.West)

    # Try directions that move toward center and avoid spawn points
    for direction in directions:
        if is_move_valid(state, unit, direction):
            new_coords = unit.coords + direction
            if not new_coords.is_spawn():
                return direction

    # If can't move toward center without spawn, try any non-spawn move
    all_directions = [Direction.North, Direction.South, Direction.East, Direction.West]
    for direction in all_directions:
        if is_move_valid(state, unit, direction):
            new_coords = unit.coords + direction
            if not new_coords.is_spawn():
                return direction

    # Last resort: move anywhere
    for direction in all_directions:
        if is_move_valid(state, unit, direction):
            return direction

    return None

def robot(state, unit):
    """Main robot logic."""
    turn_in_cycle = state.turn % 10

    # CRITICAL: Clear spawn points before spawn turns
    # Spawns happen at turn 0, 10, 20, 30, etc.
    # Start clearing at turn 3 to ensure spawn points are clear
    is_near_spawn_time = (turn_in_cycle >= 3)

    # If we're on a spawn point during clearing window, MUST move off
    if is_near_spawn_time and unit.coords.is_spawn():
        move_dir = get_move_toward_center(state, unit)
        if move_dir:
            return Action.move(move_dir)
        # If completely blocked, do nothing
        return None

    # Get all units
    our_robots = get_robots_for_team(state, state.our_team)
    enemy_robots = get_robots_for_team(state, state.other_team)

    if not enemy_robots:
        return None

    # Calculate aggressive multiplier - be more aggressive overall
    # Increase aggression near spawn time to push enemies away
    spawn_multiplier = 2.5 if turn_in_cycle >= 7 else 1.5
    health_multiplier = 0.5 if unit.health <= 2 else 1.0
    aggressive_multiplier = spawn_multiplier * health_multiplier

    # Get nearby units
    nearby_allies = get_nearby_units(unit, our_robots, radius=5)
    nearby_enemies = get_nearby_units(unit, enemy_robots, radius=5)

    # Calculate local strength
    ally_strength = len(nearby_allies) + sum(a.health for a in nearby_allies) / 5.0
    enemy_strength = len(nearby_enemies) + sum(e.health for e in nearby_enemies) / 5.0

    # Find adjacent enemy
    adjacent_enemy_dir = find_weakest_adjacent_enemy(unit, nearby_enemies)

    # If outnumbered locally, try to flee
    if ally_strength * aggressive_multiplier < enemy_strength:
        enemy_directions = get_most_common_directions(unit, nearby_enemies)
        flee_directions = [d.opposite for d in enemy_directions]

        ally_directions = get_most_common_directions(unit, nearby_allies)

        # Prefer directions away from enemies and toward allies
        common_directions = [d for d in flee_directions if d in ally_directions]

        # During clearing window, avoid spawn points
        for direction in common_directions:
            if is_move_valid(state, unit, direction):
                new_coords = unit.coords + direction
                if not is_near_spawn_time or not new_coords.is_spawn():
                    return Action.move(direction)

        # Just flee away from enemies
        for direction in flee_directions:
            if is_move_valid(state, unit, direction):
                new_coords = unit.coords + direction
                if not is_near_spawn_time or not new_coords.is_spawn():
                    return Action.move(direction)

        # Attack if adjacent
        if adjacent_enemy_dir:
            return Action.attack(adjacent_enemy_dir)

        return None

    # We have local advantage, be aggressive

    # Attack if adjacent
    if adjacent_enemy_dir:
        return Action.attack(adjacent_enemy_dir)

    # Move toward nearest enemy
    if enemy_robots:
        nearest_enemy = min(enemy_robots, key=lambda e: unit.coords.distance_to(e.coords))
        move_direction = unit.coords.direction_to(nearest_enemy.coords)

        # Try to move, rotating if blocked
        for _ in range(4):
            if is_move_valid(state, unit, direction):
                new_coords = unit.coords + move_direction
                # During clearing window, avoid spawn points
                if not is_near_spawn_time or not new_coords.is_spawn():
                    return Action.move(move_direction)
            move_direction = move_direction.rotate_cw

    return None