# Copyright (C) 2021 Arm Limited or its affiliates. All rights reserved.
#
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the License); you may
# not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an AS IS BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Description:
# Tensor allocator based on a hill-climb search
import random
from typing import List
from typing import Set

from . import numeric_util
from .live_range import LiveRange


class LiveRangeInfo:
    def __init__(self, id: int, start_time: int, end_time: int, size: int, min_alignment: int):
        # Index of this live range
        self.id = id
        # Start time (input to the allocator algorithm)
        self.start_time = start_time
        # End time, inclusive (input to the allocator algorithm)
        self.end_time = end_time
        # Size in bytes (input to the allocator algorithm)
        self.size = size
        # Allocated address (the main output from the allocator algorithm)
        self.address: int = 0
        # End address, exclusive
        self.end_address: int = 0
        # id of predecessor live range (predecessor's end address == this lr's address)
        self.predecessor: int = 0
        # Turn at which the live range was allocated
        self.turn: int = 0
        # Max value of size_at_time (only used in the heuristic allocation)
        self.urgency = 0
        self.neighbours: List["LiveRangeInfo"] = []
        self.min_alignment = min_alignment

    def overlaps(self, addr2: int, size2: int) -> int:
        return self.address < addr2 + size2 and addr2 < self.end_address

    def is_neighbour(self, lr: "LiveRangeInfo") -> bool:
        return self.start_time <= lr.end_time and lr.start_time <= self.end_time

    def __str__(self):
        return "<LiveRangeInfo: id={}, start_time={}, end_time={}, size={}, address={}>".format(
            self.id, self.start_time, self.end_time, self.size, self.address
        )

    def __lt__(self, other) -> bool:
        if self.urgency != other.urgency:
            return self.urgency > other.urgency
        duration1 = self.end_time - self.start_time
        duration2 = other.end_time - other.start_time
        if duration1 != duration2:
            return duration1 > duration2

        if self.start_time != other.start_time:
            return self.start_time < other.start_time

        if self.size != other.size:
            return self.size > other.size

        return self.id < other.id


class HillClimbAllocator:
    """
    Implements tensor allocator using a hill climb search.

    The basic algorithm is:

    Use a heuristic allocator to find an initial allocation
    while allocation is not optimal and iterations < MAX_ITERATIONS:
        find the "bottleneck": the live range with highest end address
        find all live ranges that affected the allocation of the bottleneck
        swap the order of any two affecting live ranges
        reallocate tensors using the reordered live ranges
        if the new allocation is better: keep it, else set allocation to previous allocation
    """

    # Maximum number of iterations of the algorithm (can be override from the command line)
    MAX_ITERATIONS = 99999
    NOT_ALLOCATED = -1
    # Used for live ranges allocated at address 0
    NO_PREDECESSOR = -1
    # Special handling if best solution has not improved during this many iterations
    MAX_ITERATIONS_STUCK = 50
    # Minimum number of iterations since the last improvement (unless an optimal solution is found)
    MIN_ITERATIONS_IMPROVE = 500

    def __init__(self, live_ranges: List[LiveRange], max_iterations: int, memory_limit: int):
        # Contains the live ranges
        self.lrs: List[LiveRangeInfo] = [
            LiveRangeInfo(id, lr.start_time, lr.end_time, lr.size, lr.get_alignment())
            for id, lr in enumerate(live_ranges)
        ]
        self.lrs_at_time: List[List[LiveRangeInfo]] = []
        # The available size (input to algorithm).
        self.available_size: int = 0
        # The algorithm stops once the target size has been achieved
        self.target_size: int = 0
        # The highest end address of the best found allocation
        self.best_size: int = 1 << 63
        # For each live range: max value of size_at_time (only used in the heuristic allocation)
        self.lr_urgency = len(self.lrs) * [0]
        nr_time_slots = 1 + max(lr.end_time for lr in self.lrs)
        # Contains active live ranges at each timestamp
        self.lrs_at_time = [[] for i in range(nr_time_slots)]
        for lr in self.lrs:
            for t in range(lr.start_time, lr.end_time + 1):
                self.lrs_at_time[t].append(lr)
        # At each timestamp: accumulated size of active live ranges
        size_at_time = [sum(lr.size for lr in self.lrs_at_time[t]) for t in range(nr_time_slots)]
        # The minimum possible size, assuming all live ranges can be perfectly allocated
        self.min_required_size: int = max(size_at_time)
        # Set the maximum number of iterations the search will iterate to find a solution
        if max_iterations is None:
            self.max_iterations = self.MAX_ITERATIONS
        else:
            self.max_iterations = max_iterations
        # Defines a memory limit that the allocation must meet
        self.memory_limit = memory_limit
        if self.memory_limit < self.min_required_size:
            print(
                f"Warning: Memory limit = {self.memory_limit} is less than the minimum possible allocation size ="
                f" {self.min_required_size}"
            )
        # Calculate all neighbours + the urgency of each live range
        for lr in self.lrs:
            lr.urgency = 0
            lr.neighbours = []
            neighbours = set()
            for t in range(lr.start_time, lr.end_time + 1):
                lr.urgency = max(size_at_time[t], lr.urgency)
                for lr2 in self.lrs_at_time[t]:
                    if lr2 not in neighbours and lr != lr2:
                        neighbours.add(lr2)
                        lr.neighbours.append(lr2)

    def allocate_lr(self, lr: LiveRangeInfo):
        """
        Allocates the given live range at the smallest possible address
        """
        address = 0
        predecessor = HillClimbAllocator.NO_PREDECESSOR
        fits = False
        while not fits:
            fits = True
            # Find neighbours that overlap with address
            for lr2 in lr.neighbours:
                if lr2.address == HillClimbAllocator.NOT_ALLOCATED or lr2.end_address <= address:
                    continue
                if lr2.overlaps(address, lr.size):
                    # Overlap found increase address
                    fits = False
                    address = numeric_util.round_up(lr2.end_address, lr.min_alignment)
                    predecessor = lr2.id
        lr.address = address
        lr.end_address = address + lr.size
        lr.predecessor = predecessor

    def allocate_indices(self, indices: List[int]):
        """
        Allocates the live ranges in the order indicated by the indices;
        allocates each live range at the lowest possible address.
        """
        for lr in self.lrs:
            lr.address = HillClimbAllocator.NOT_ALLOCATED
        size = 0
        for turn, index in enumerate(indices):
            lr = self.lrs[index]
            self.allocate_lr(lr)
            lr.turn = turn
            size = max(size, lr.end_address)
            if size > self.best_size:
                # This allocation is worse than the best known allocation;
                # no need to continue
                break
        return size

    def add_predecessor_turns(self, turn_set: Set[int], turn_list: List[int], lr: LiveRangeInfo):
        """
        Adds the given live range + predecessors to the turns vector.
        Note: the turn_set is for quick detection of duplicates,
        the turn_list is to get reproduceable results
        """
        if lr.turn not in turn_set:
            turn_set.add(lr.turn)
            turn_list.append(lr.turn)
        id = lr.id
        while self.lrs[id].predecessor != HillClimbAllocator.NO_PREDECESSOR:
            id = self.lrs[id].predecessor
            turn = self.lrs[id].turn
            if turn not in turn_set:
                turn_set.add(turn)
                turn_list.append(turn)

    def attempt_bottleneck_fix(self, indices: List[int], iterations_stuck):
        """
        Finds the "bottleneck", the live range with highest end address, and reorders the indices
        such that a next allocation might lower the memory usage.

                                  ---------
                                  |       |
                                  |   D   |
                                  |       |
         ----------------------------------
         |           B                 |
         -------------------------------
         | |
         |A|                      ---
         | |                      |C|
         | |                      | |
         ---------------------------------------

        In the above example, the allocation order was [A, B, C, D] and D is the resulting bottle-neck.
        The live ranges that affected the allocation of D are the direct neighbours of D (i.e. B and C),
        and all direct and indirect predecessors of D and its neighbours
        (i.e. A, which is the predecessor of B, and indirect predecessor of D).

        By permuting the order in which the affecting live ranges are allocated, the bottleneck might
        be lowered. In the above example, almost any permutation would lower the bottleneck.
        """
        # Find the bottleneck
        max_lr = self.lrs[0]
        for lr in self.lrs[1:]:
            if lr.end_address > max_lr.end_address:
                max_lr = lr

        # Find all live ranges that affected the placement of the bottleneck live range.
        # This consists of two types of live ranges:
        # - direct neighbours of the bottleneck live range
        # - direct and indirect predecessors of these neighbours + bottleneck
        # The turns at which these live ranges were allocated are put in the turns set.
        turn_set: Set[int] = set()
        turn_list: List[int] = list()
        self.add_predecessor_turns(turn_set, turn_list, max_lr)
        for lr2 in max_lr.neighbours:
            self.add_predecessor_turns(turn_set, turn_list, lr2)

        # Non-direct neighbours that interfere with the allocation of the bottleneck are the
        # immediate cause for gaps in the allocation, and are selected with higher probability.
        non_nb_turn_list = []
        for turn in turn_list:
            lr = self.lrs[indices[turn]]
            if not max_lr.is_neighbour(lr):
                non_nb_turn_list.append(turn)
        assert turn_list
        # Pick from non-neighbour list with 30% probability
        # (magic number based on tuning)
        if random.randint(0, 100) < 30 and non_nb_turn_list:
            # Pick a live range from the "non-neighbour list"
            ix1 = non_nb_turn_list[random.randint(0, len(non_nb_turn_list) - 1)]
        else:
            # Pick any affecting live range.
            ix1 = turn_list[random.randint(0, len(turn_list) - 1)]

        ix2 = turn_list[random.randint(0, len(turn_list) - 2)]
        if ix1 == ix2:
            ix2 = turn_list[-1]
        # Swap indices
        indices[ix1], indices[ix2] = indices[ix2], indices[ix1]
        if iterations_stuck > HillClimbAllocator.MAX_ITERATIONS_STUCK:
            # The best allocation has not improved for a while, maybe improvement is not possible
            # by single-swapping indices; add more neighbour live ranges and swap up to two more indices.
            # Adding more neighbours can sometimes resolve the situation where the current bottleneck
            # is resolved, but always results in a higher bottleneck at a nearby live range.
            # Magic number is based on tuning
            for turn in non_nb_turn_list:
                for lr in self.lrs[indices[turn]].neighbours:
                    if lr.turn not in turn_set:
                        turn_set.add(lr.turn)
                        turn_list.append(lr.turn)
            ix1 = turn_list[random.randint(0, len(turn_list) - 1)]
            ix2 = turn_list[random.randint(0, len(turn_list) - 1)]
            indices[ix1], indices[ix2] = indices[ix2], indices[ix1]

    def search(self, indices: List[int]):
        """
        Search for a solution, using the given indices as initial solution.
        """
        best_indices = indices[:]
        last_improvement_iteration = 0
        i = 0
        while (self.best_size > self.memory_limit and i < self.max_iterations) or (
            i - last_improvement_iteration < self.MIN_ITERATIONS_IMPROVE
        ):
            # Reorder the indices
            self.attempt_bottleneck_fix(indices, i - last_improvement_iteration)
            # Allocate the reordered indices and check if it gave an improvement
            new_size = self.allocate_indices(indices)
            if new_size <= self.best_size:
                # The new allocation produced a new best result remember it
                if new_size < self.best_size:
                    last_improvement_iteration = i
                self.best_size = new_size
                best_indices = indices[:]
                self.allocated_addresses = [lr.address for lr in self.lrs]
                if self.best_size <= self.min_required_size:
                    # Target reached; stop
                    return
            else:
                # The new allocation produced worse result undo the change
                indices = best_indices[:]
            i += 1

    def allocate(self) -> List[int]:
        """
        Runs the allocation algorithm. Finishes when an optimal solution has been
        found or when maximum iterations have been run.
        The allocated addresses are placed in the output vector, in the same
        order as the input vector.

        Implementation note: the algorithm produces reproduceable results by using
        a well-defined random number generator with well-defined default seed,
        and using a fixed number of iterations.
        """
        random.seed(1)
        # Sort indices on priority. Note: self.lrs must be left unchanged
        indices = [lr.id for lr in sorted(self.lrs)]
        # Allocate the initial solution
        self.best_size = self.allocate_indices(indices)
        self.allocated_addresses = [lr.address for lr in self.lrs]
        if self.best_size > self.min_required_size:
            # Try to improve the heuristic allocation
            self.search(indices)
        # else the heuristic allocation returned an optimal solution; no search needed
        return self.allocated_addresses


def allocate_live_ranges(lrs: List[LiveRange], max_iterations: int, memory_limit: int) -> List[int]:
    """
    Allocates live ranges using a search based allocator.
    Returns the list of allocated addresses (one for each live range)
    """
    if not lrs:
        return []
    allocator = HillClimbAllocator(lrs, max_iterations, memory_limit)
    return allocator.allocate()