import random
import json
import os
import sys
import time
from tqdm import tqdm
from multiprocessing import Pool

GRID_SIZE = 10
SHIPS = {"Carrier": 5, "Battleship": 4, "Cruiser": 3, "Submarine": 3, "Destroyer": 2}
POPULATION_SIZE = 200
GENERATIONS = 50
BASE_MUTATION_RATE = 0.1
GAMES_PER_INDIVIDUAL = 20
LEARNING_FILE = "battleship_learning.json"

MUTATION_MULTIPLIERS = {
    "adjacent_hit_weight": 0.5,
    "ship_size_weight": 1.5,
    "parity_weight": 1.5,
    "hunt_mode_threshold": 1.0,
    "prob_grid_weight": 1.0,
    "avoid_edges_weight": 1.0,
}


class Board:
    def __init__(self):
        self.grid = [[None for _ in range(GRID_SIZE)] for _ in range(GRID_SIZE)]
        self.ships_left = set(SHIPS.keys())
        self.place_all_ships()

    def place_all_ships(self):
        for ship, size in SHIPS.items():
            while True:
                row = random.randint(0, GRID_SIZE - 1)
                col = random.randint(0, GRID_SIZE - 1)
                horizontal = random.choice([True, False])
                if self.is_valid_placement(row, col, size, horizontal):
                    self.place_ship(row, col, size, horizontal, ship)
                    break

    def is_valid_placement(self, row, col, size, horizontal):
        if horizontal:
            if col + size > GRID_SIZE:
                return False
            return all(self.grid[row][c] is None for c in range(col, col + size))
        else:
            if row + size > GRID_SIZE:
                return False
            return all(self.grid[r][col] is None for r in range(row, row + size))

    def place_ship(self, row, col, size, horizontal, ship):
        if horizontal:
            for c in range(col, col + size):
                self.grid[row][c] = ship
        else:
            for r in range(row, row + size):
                self.grid[r][col] = ship

    def is_sunk(self, ship, hits):
        return all(
            (r, c) in hits
            for r in range(GRID_SIZE)
            for c in range(GRID_SIZE)
            if self.grid[r][c] == ship
        )

    def update_sunk_ships(self, hits):
        for ship in list(self.ships_left):
            if self.is_sunk(ship, hits):
                self.ships_left.remove(ship)

    def get_remaining_ships(self):
        return [size for ship, size in SHIPS.items() if ship in self.ships_left]


class BattleshipStrategy:
    ADJACENT_HIT_WEIGHT_RANGE = (1, 20)
    SHIP_SIZE_WEIGHT_RANGE = (0, 10)
    PARITY_WEIGHT_RANGE = (0, 10)
    HUNT_MODE_THRESHOLD_RANGE = (0.1, 1)
    PROB_GRID_WEIGHT_RANGE = (0, 1)
    AVOID_EDGES_WEIGHT_RANGE = (0, 1)

    def __init__(self):
        self.adjacent_hit_weight = random.uniform(*self.ADJACENT_HIT_WEIGHT_RANGE)
        self.ship_size_weight = random.uniform(*self.SHIP_SIZE_WEIGHT_RANGE)
        self.parity_weight = random.uniform(*self.PARITY_WEIGHT_RANGE)
        self.hunt_mode_threshold = random.uniform(*self.HUNT_MODE_THRESHOLD_RANGE)
        self.prob_grid_weight = random.uniform(*self.PROB_GRID_WEIGHT_RANGE)
        self.avoid_edges_weight = random.uniform(*self.AVOID_EDGES_WEIGHT_RANGE)
        self._probability_grid = None

    def serialize(self, include_grid=False):
        data = {
            "adjacent_hit_weight": self.adjacent_hit_weight,
            "ship_size_weight": self.ship_size_weight,
            "parity_weight": self.parity_weight,
            "hunt_mode_threshold": self.hunt_mode_threshold,
            "prob_grid_weight": self.prob_grid_weight,
            "avoid_edges_weight": self.avoid_edges_weight,
        }
        if include_grid:
            data["probability_grid"] = self.probability_grid
        return data

    @property
    def probability_grid(self):
        if self._probability_grid is None:
            self.reset_probability_grid()
        return self._probability_grid

    def reset_probability_grid(self):
        self._probability_grid = [
            [1 for _ in range(GRID_SIZE)] for _ in range(GRID_SIZE)
        ]

    def update_probability_grid(self, hits, misses):
        for i in range(GRID_SIZE):
            for j in range(GRID_SIZE):
                if (i, j) in hits:
                    self.probability_grid[i][j] = 0
                elif (i, j) in misses:
                    self.probability_grid[i][j] = -1
                else:
                    self.probability_grid[i][j] = 1

        for i, j in hits:
            for di, dj in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
                ni, nj = i + di, j + dj
                if (
                    0 <= ni < GRID_SIZE
                    and 0 <= nj < GRID_SIZE
                    and (ni, nj) not in hits
                    and (ni, nj) not in misses
                ):
                    self.probability_grid[ni][nj] *= self.adjacent_hit_weight

        for i in range(GRID_SIZE):
            for j in range(GRID_SIZE):
                if i == 0 or i == GRID_SIZE - 1 or j == 0 or j == GRID_SIZE - 1:
                    self.probability_grid[i][j] *= 1 - self.avoid_edges_weight

        max_prob = max(max(row) for row in self.probability_grid if max(row) > 0)
        if max_prob > 0:
            for i in range(GRID_SIZE):
                for j in range(GRID_SIZE):
                    if self.probability_grid[i][j] > 0:
                        self.probability_grid[i][j] = int(
                            (self.probability_grid[i][j] / max_prob) * 100
                        )

    def get_next_move(self, board, hits, misses):
        self.update_probability_grid(hits, misses)
        probabilities = [
            [self.probability_grid[i][j] for j in range(GRID_SIZE)]
            for i in range(GRID_SIZE)
        ]
        remaining_ships = [size for ship, size in SHIPS.items() if size > len(hits)]

        for i in range(GRID_SIZE):
            for j in range(GRID_SIZE):
                if (i, j) in hits or (i, j) in misses:
                    probabilities[i][j] = 0
                    continue

                for size in remaining_ships:
                    if self.can_fit_ship(board, i, j, size):
                        probabilities[i][j] *= self.ship_size_weight

                if (i + j) % 2 == 0:
                    probabilities[i][j] *= self.parity_weight

        max_prob = max(max(row) for row in probabilities)
        for i in range(GRID_SIZE):
            for j in range(GRID_SIZE):
                probabilities[i][j] *= self.prob_grid_weight
                if max_prob > 0:
                    probabilities[i][j] += (1 - self.prob_grid_weight) * (
                        probabilities[i][j] / max_prob
                    )

        max_prob = max(max(row) for row in probabilities)
        candidates = [
            (i, j)
            for i in range(GRID_SIZE)
            for j in range(GRID_SIZE)
            if probabilities[i][j] == max_prob
        ]

        return random.choice(candidates)

    def can_fit_ship(self, board, row, col, size):
        if col + size <= GRID_SIZE and all(
            board[row][c] == 0 for c in range(col, col + size)
        ):
            return True
        if row + size <= GRID_SIZE and all(
            board[r][col] == 0 for r in range(row, row + size)
        ):
            return True
        return False


def play_game(strategy, use_ensemble=False):
    board = Board()
    hits = set()
    misses = set()
    turns = 0

    elite_strategies = load_elite_strategies() if use_ensemble else None

    if strategy is not None:
        strategy.reset_probability_grid()

    while board.ships_left and turns < 100:
        if use_ensemble:
            i, j = get_ensemble_move(elite_strategies, board, hits, misses)
        else:
            i, j = strategy.get_next_move(board.grid, hits, misses)
            strategy.update_probability_grid(hits, misses)

        turns += 1
        if board.grid[i][j] is not None:
            hits.add((i, j))
            ship = board.grid[i][j]
            board.update_sunk_ships(hits)
        else:
            misses.add((i, j))

    return turns


def crossover(parent1, parent2):
    child = BattleshipStrategy()
    child.adjacent_hit_weight = random.choice(
        [parent1.adjacent_hit_weight, parent2.adjacent_hit_weight]
    )
    child.ship_size_weight = random.choice(
        [parent1.ship_size_weight, parent2.ship_size_weight]
    )
    child.parity_weight = random.choice([parent1.parity_weight, parent2.parity_weight])
    child.hunt_mode_threshold = random.choice(
        [parent1.hunt_mode_threshold, parent2.hunt_mode_threshold]
    )
    child.prob_grid_weight = random.choice(
        [parent1.prob_grid_weight, parent2.prob_grid_weight]
    )
    child.avoid_edges_weight = random.choice(
        [parent1.avoid_edges_weight, parent2.avoid_edges_weight]
    )
    return child


def mutate(individual):
    if (
        random.random()
        < BASE_MUTATION_RATE * MUTATION_MULTIPLIERS["adjacent_hit_weight"]
    ):
        individual.adjacent_hit_weight = random.uniform(
            *BattleshipStrategy.ADJACENT_HIT_WEIGHT_RANGE
        )
    if random.random() < BASE_MUTATION_RATE * MUTATION_MULTIPLIERS["ship_size_weight"]:
        individual.ship_size_weight = random.uniform(
            *BattleshipStrategy.SHIP_SIZE_WEIGHT_RANGE
        )
    if random.random() < BASE_MUTATION_RATE * MUTATION_MULTIPLIERS["parity_weight"]:
        individual.parity_weight = random.uniform(
            *BattleshipStrategy.PARITY_WEIGHT_RANGE
        )
    if (
        random.random()
        < BASE_MUTATION_RATE * MUTATION_MULTIPLIERS["hunt_mode_threshold"]
    ):
        individual.hunt_mode_threshold = random.uniform(
            *BattleshipStrategy.HUNT_MODE_THRESHOLD_RANGE
        )
    if random.random() < BASE_MUTATION_RATE * MUTATION_MULTIPLIERS["prob_grid_weight"]:
        individual.prob_grid_weight = random.uniform(
            *BattleshipStrategy.PROB_GRID_WEIGHT_RANGE
        )
    if (
        random.random()
        < BASE_MUTATION_RATE * MUTATION_MULTIPLIERS["avoid_edges_weight"]
    ):
        individual.avoid_edges_weight = random.uniform(
            *BattleshipStrategy.AVOID_EDGES_WEIGHT_RANGE
        )
    return individual


def evaluate_strategy(strategy):
    total_turns = 0
    for _ in range(GAMES_PER_INDIVIDUAL):
        total_turns += play_game(strategy)
    return total_turns


def genetic_algorithm(continue_training=False):
    if continue_training and os.path.exists(LEARNING_FILE):
        with open(LEARNING_FILE, "r") as f:
            data = json.load(f)
            population = [BattleshipStrategy() for _ in range(POPULATION_SIZE)]
            for i, strategy_data in enumerate(data):
                population[i].__dict__.update(strategy_data)
        tqdm.write("Continuing training from existing population.")
    else:
        if os.path.exists(LEARNING_FILE):
            timestamp = time.strftime("%Y%m%d-%H%M%S")
            os.rename(LEARNING_FILE, f"{LEARNING_FILE}.{timestamp}")
            tqdm.write(f"Moved existing model to {LEARNING_FILE}.{timestamp}")
        population = [BattleshipStrategy() for _ in range(POPULATION_SIZE)]
        tqdm.write("Starting with a new population.")

    best_fitness = float("inf")
    generations_without_improvement = 0

    progress_bar = tqdm(
        total=GENERATIONS * POPULATION_SIZE, file=sys.stderr, desc="Overall Progress"
    )

    for generation in range(GENERATIONS):
        with Pool() as pool:
            fitness_scores = list(pool.imap(evaluate_strategy, population))
            progress_bar.update(POPULATION_SIZE)

        population = [
            x
            for _, x in sorted(
                zip(fitness_scores, population), key=lambda pair: pair[0]
            )
        ]

        current_best_fitness = min(fitness_scores)
        avg_fitness = sum(fitness_scores) / len(fitness_scores)

        if current_best_fitness < best_fitness:
            best_fitness = current_best_fitness
            generations_without_improvement = 0
        else:
            generations_without_improvement += 1

        tqdm.write(
            f"Generation {generation + 1}: Best fitness = {current_best_fitness}, Avg fitness = {avg_fitness:.2f}, Best overall = {best_fitness}"
        )

        elitism_count = POPULATION_SIZE // 20
        new_population = population[:elitism_count]

        while len(new_population) < POPULATION_SIZE:
            tournament_size = 5
            parent1 = min(
                random.sample(population, tournament_size),
                key=lambda x: fitness_scores[population.index(x)],
            )
            parent2 = min(
                random.sample(population, tournament_size),
                key=lambda x: fitness_scores[population.index(x)],
            )

            child = crossover(parent1, parent2)
            new_population.append(mutate(child))

        population = new_population

        elite_strategies = population[:elitism_count]
        elite_data = [
            strategy.serialize(include_grid=True) for strategy in elite_strategies
        ]

        elite_file = "elites.json"
        with open(elite_file, "w") as f:
            json.dump(elite_data, f, indent=4)

        try:
            with open(LEARNING_FILE, "w") as f:
                json.dump(
                    [strategy.serialize() for strategy in population], f, indent=4
                )
        except Exception as e:
            tqdm.write(f"Error saving generation {generation + 1}: {e}")

        if generations_without_improvement >= 20:
            tqdm.write("No improvement for 20 generations. Stopping early.")
            break

    progress_bar.close()

    return population[0]


def load_best_strategy():
    if os.path.exists(LEARNING_FILE):
        with open(LEARNING_FILE, "r") as f:
            data = json.load(f)
            best_strategy = BattleshipStrategy()
            best_strategy.__dict__.update(data[0])
        return best_strategy
    else:
        return BattleshipStrategy()


def load_elite_strategies(num_elites=10):
    if os.path.exists(LEARNING_FILE):
        with open(LEARNING_FILE, "r") as f:
            data = json.load(f)
            elite_strategies = []
            for strategy_data in data[:num_elites]:
                strategy = BattleshipStrategy()
                strategy.__dict__.update(strategy_data)
                elite_strategies.append(strategy)
        return elite_strategies
    else:
        return [BattleshipStrategy()]


def get_ensemble_move(elite_strategies, board, hits, misses):
    probabilities = [[0 for _ in range(GRID_SIZE)] for _ in range(GRID_SIZE)]
    remaining_ships = board.get_remaining_ships()

    for strategy in elite_strategies:
        for i in range(GRID_SIZE):
            for j in range(GRID_SIZE):
                if (i, j) in hits or (i, j) in misses:
                    continue

                prob = 1

                for di, dj in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
                    if (
                        0 <= i + di < GRID_SIZE
                        and 0 <= j + dj < GRID_SIZE
                        and (i + di, j + dj) in hits
                    ):
                        prob *= strategy.adjacent_hit_weight

                for size in remaining_ships:
                    if strategy.can_fit_ship(board.grid, i, j, size):
                        prob *= strategy.ship_size_weight

                if (i + j) % 2 == 0:
                    prob *= strategy.parity_weight

                if i == 0 or i == GRID_SIZE - 1 or j == 0 or j == GRID_SIZE - 1:
                    prob *= 1 - strategy.avoid_edges_weight

                probabilities[i][j] += prob * strategy.prob_grid_weight

    max_prob = max(max(row) for row in probabilities)
    candidates = [
        (i, j)
        for i in range(GRID_SIZE)
        for j in range(GRID_SIZE)
        if probabilities[i][j] == max_prob
    ]

    return random.choice(candidates)


def main():
    if len(sys.argv) < 2:
        print(
            "Usage: python battleship_genetic2.py [train|train_continue|sample|sample_multi|sample_ensemble]"
        )
        sys.exit(1)

    mode = sys.argv[1]

    if mode == "train":
        best_strategy = genetic_algorithm(continue_training=False)
        print("Training completed. Best strategy saved.")
    elif mode == "train_continue":
        best_strategy = genetic_algorithm(continue_training=True)
        print("Continued training completed. Best strategy saved.")
    elif mode == "sample":
        best_strategy = load_best_strategy()
        turns = play_game(best_strategy)
        print(turns)
    elif mode == "sample_multi":
        best_strategy = load_best_strategy()
        num_games = 100
        total_turns = sum(play_game(best_strategy) for _ in range(num_games))
        avg_turns = total_turns / num_games
        print(f"Average turns over {num_games} games: {avg_turns:.2f}")
    elif mode == "sample_ensemble":
        turns = play_game(None, use_ensemble=True)
        print(turns)
    else:
        print(
            "Invalid mode. Use 'train', 'train_continue', 'sample', 'sample_multi', or 'sample_ensemble'."
        )
        sys.exit(1)


if __name__ == "__main__":
    main()
