#! /usr/bin/env python3

"""Note: unlike the course slides, the rows and columns (ranks and
files) of the puzzle are counted as {0, ..., N - 1} here (rather than
starting from 1) because this is the more common convention in
programming.

This is not an efficient implementation, as it is meant to be easy to
understand.

The code does not follow the pseudo-code in the lecture 100% due to
the need to gather statistics.

Note that the code is not fully generic in the way that state-space
search code is. Some of the assumptions about the problem are baked
directly into the algorithms, for example the fact that we are dealing
with a minimization problem.
"""

import math
import random


def get_random_candidate(board_size):
    return [random.randrange(board_size) for _ in range(board_size)]


def get_number_of_conflicts(candidate):
    num_conflicts = 0
    size = len(candidate)
    for index1, row1 in enumerate(candidate):
        for index2 in range(index1 + 1, size):
            row2 = candidate[index2]
            if row1 == row2 or abs(row2 - row1) == index2 - index1:
                num_conflicts += 1
    return num_conflicts


def is_solution(candidate):
    return get_number_of_conflicts(candidate) == 0


def get_random_neighbor(candidate):
    size = len(candidate)
    file_moved = random.randrange(size)
    new_rank = random.randrange(size)
    neighbor = list(candidate)
    neighbor[file_moved] = new_rank
    return neighbor


def h(candidate):
    return get_number_of_conflicts(candidate)


def schedule(t):
    return 2 / math.log(t + 1)


def simulated_annealing(board_size, max_steps):
    curr = get_random_candidate(board_size)
    is_solved = False
    for t in range(1, max_steps + 1):
        # The test for solutions and their quality is different from
        # the slides because we have a pure search problem and hence
        # there is no point in comparing the quality of solutions --
        # all solutions are equally good and we can stop as soon as we
        # find one.
        if is_solution(curr):
            is_solved = True
            break
        temperature = schedule(t)
        next = get_random_neighbor(curr)
        # The sign of delta_e is reversed compared to the lecture
        # slides because we want to minimize, not maximize.
        delta_e = h(curr) - h(next)
        if delta_e >= 0 or random.random() < math.exp(delta_e / temperature):
            curr = next
    return is_solved, t


def main(board_size, num_trials, update_interval, max_steps):
    num_solved = 0
    total_steps = 0
    for trial in range(1, num_trials + 1):
        if trial % update_interval == 0:
            print(f"trial {trial} out of {num_trials}...")
        solved, num_steps = simulated_annealing(
            board_size=board_size, max_steps=max_steps)
        num_solved += int(solved)
        total_steps += num_steps
    percent_solved = 100.0 * num_solved / num_trials
    print(f"{num_solved} out of {num_trials} trials solved ({percent_solved}%)")
    print(f"average number of steps: {total_steps / num_trials}")


if __name__ == "__main__":
    main(board_size=8, num_trials=1000, update_interval=10, max_steps=10000)