import logging
import time

import ngtpy

from annime.interface_ann import ANNInterface


class NgtANN(ANNInterface):
    """
    An implementation of ANNInterface for the NGT (Neighborhood Graph and Tree) library.
    """

    def __init__(self, dim, distance_type="L2"):
        super().__init__()
        self.index = ngtpy.Index(path=".", dimension=dim, distance_type=distance_type)
        self.dim = dim
        self.distance_type = distance_type
        self.built = False
        self.logger = logging.getLogger(__name__)

    def build_index(self, data_points, **kwargs):
        """
        Build the NGT index from the provided data points.

        Args:
            data_points (ndarray): A list of data points to index.
            **kwargs: Arbitrary keyword arguments for index configuration.
        """
        for i, vector in enumerate(data_points):
            self.index.insert(vector)
        self.index.build_index()
        self.built = True
        self.logger.info("Index built with %d data points.", len(data_points))

    def query(self, query_point, k=5, **kwargs):
        """
        Query the NGT index for the k nearest neighbors of the provided point.

        Args:
            query_point (ndarray): The query point.
            k (int): The number of nearest neighbors to return.

        Returns:
            list: The k nearest neighbors.
        """
        if not self.built:
            raise Exception("Index must be built before querying.")
        return self.index.search(query_point, size=k, with_distance=True)

    def save_index(self, filepath):
        """
        Save the built NGT index to a file.

        Args:
            filepath (str): The path to the file where the index is to be saved.
        """
        if not self.built:
            raise Exception("Index must be built before it can be saved.")
        self.index.save(filepath)
        self.logger.info("Index saved to %s.", filepath)

    def load_index(self, filepath):
        """
        Load the NGT index from a file.

        Args:
            filepath (str): The path to the file from which the index is to be loaded.
        """
        self.index = ngtpy.Index(filepath)
        self.built = True
        self.logger.info("Index loaded from %s.", filepath)

    def set_distance_metric(self, metric):
        """
        Set the distance metric for the NGT index.

        Args:
            metric (str): The distance metric to use.
        """
        self.distance_type = metric
        self.index = ngtpy.Index(path=".", dimension=self.dim, distance_type=metric)
        self.logger.info("Distance metric set to %s.", metric)

    def set_index_parameters(self, **params):
        """
        Set parameters for the NGT index.

        Args:
            **params: Arbitrary keyword arguments specific to the index configuration.
        """
        for param, value in params.items():
            setattr(self.index.property, param, value)
        self.logger.info("Index parameters set: %s", params)

    def add_items(self, data_points, ids=None):
        """
        Add items to the NGT index, optionally with specific ids.

        Args:
            data_points (ndarray): A numpy list of data points to add to the index.
            ids (list): Optional list of ids corresponding to each data point.
        """
        if ids and len(data_points) != len(ids):
            raise ValueError("Length of data_points and ids must match.")
        for i, vector in enumerate(data_points):
            idx = ids[i] if ids else i
            self.index.insert(vector, object_id=idx)
        self.logger.info("Added %d items to the index.", len(data_points))

    def delete_item(self, item_id):
        """
        Delete an item from the NGT index by id.

        Args:
            item_id (int): The id of the item to be deleted.
        """
        self.index.remove(item_id)
        self.logger.info("Item with id %d deleted from the index.", item_id)

    def clear_index(self):
        """
        Clear all items from the NGT index.
        """
        self.index.clear()
        self.built = False
        self.logger.info("Index cleared.")

    def get_item_vector(self, item_id):
        """
        Retrieve the vector of an item by id from the NGT index.

        Args:
            item_id (int): The id of the item.

        Returns:
            list: The vector of the item.
        """
        return self.index.get_object(item_id)

    def optimize_index(self):
        """
        Optimize the NGT index for better performance during queries.
        """
        self.index.build_index()
        self.logger.info("Index optimized.")

    def get_index_size(self):
        """
        Return the current size of the NGT index in terms of the number of items.

        Returns:
            int: The number of items in the index.
        """
        return self.index.size()

    def get_index_memory_usage(self):
        """
        Return the amount of memory used by the NGT index.

        Returns:
            int: The memory usage of the index.
        """
        return self.index.memory()

    def batch_query(self, query_points, k=5, include_distances=False):
        """
        Perform a batch query for multiple points, returning their k nearest neighbors.

        Args:
            query_points (list): A list of query points.
            k (int): The number of nearest neighbors to return.
            include_distances (bool): Whether to include distances in the results.

        Returns:
            list: A list of results for each query point.
        """
        results = []
        for query_point in query_points:
            result = self.query(query_point, k=k)
            if include_distances:
                results.append(result)
            else:
                results.append([x[0] for x in result])
        return results

    def parallel_query(self, query_points, k=5, num_threads=4):
        """
        Perform multiple queries in parallel, using a specified number of threads.

        Args:
            query_points (list): A list of query points.
            k (int): The number of nearest neighbors to return.
            num_threads (int): The number of threads to use.

        Returns:
            list: A list of results for each query point.
        """
        from multiprocessing import Pool
        with Pool(processes=num_threads) as pool:
            results = pool.starmap(self.query, [(query_point, k) for query_point in query_points])
        return results

    def benchmark_performance(self, queries, k=5, rounds=10):
        """
        Benchmark the query performance of the index with a set of queries repeated over several rounds.

        Args:
            queries (list): A list of query points.
            k (int): The number of nearest neighbors to return.
            rounds (int): The number of rounds to repeat the benchmark.

        Returns:
            str: The average query time.
        """
        start_time = time.time()
        for _ in range(rounds):
            for query in queries:
                self.query(query, k=k)
        end_time = time.time()
        duration = (end_time - start_time) / (len(queries) * rounds)
        self.logger.info("Benchmark completed: Avg query time = %.5f seconds", duration)
        return f"Average query time: {duration:.5f} seconds"

    def export_to_dot(self, filepath):
        """
        Export the structure of the NGT index to a DOT file for visualization.

        Args:
            filepath (str): The path to the file where the DOT representation is to be saved.

        Raises:
            NotImplementedError: NGT does not support exporting to DOT format.
        """
        raise NotImplementedError("NGT does not support exporting to DOT format.")

    def enable_logging(self, level='INFO'):
        """
        Enable detailed logging of operations within the NGT index.

        Args:
            level (str): The logging level (e.g., 'INFO', 'DEBUG').
        """
        logging.basicConfig(level=getattr(logging, level.upper()))
        self.logger.setLevel(level.upper())
        self.logger.info("Logging enabled at level: %s", level)

    def rebuild_index(self, **kwargs):
        """
        Explicitly rebuilds the entire NGT index according to the current configuration and data points.

        Args:
            **kwargs: Arbitrary keyword arguments for index configuration.
        """
        self.index.build_index()
        self.logger.info("Index rebuilt.")

    def refresh_index(self):
        """
        Refreshes the NGT index by optimizing internal structures without full rebuilding.
        """
        self.index.build_index()
        self.logger.info("Index refreshed.")

    def serialize_index(self, output_format='binary'):
        """
        Serialize the NGT index into a specified format (e.g., binary) to enable easy transmission or storage.

        Args:
            output_format (str): The format for serialization (default is 'binary').

        Returns:
            bytes: The serialized index data.

        Raises:
            ValueError: If the output format is not supported.
        """
        if output_format != 'binary':
            raise ValueError("NGT currently supports only binary serialization format.")
        temp_path = "temp_ngt_index"
        self.save_index(temp_path)
        with open(temp_path, 'rb') as file:
            serialized_data = file.read()
        return serialized_data

    def deserialize_index(self, data, input_format='binary'):
        """
        Deserialize the NGT index from a given format, restoring it to an operational state.

        Args:
            data (bytes): The serialized index data.
            input_format (str): The format of the serialized data (default is 'binary').

        Raises:
            ValueError: If the input format is not supported.
        """
        if input_format != 'binary':
            raise ValueError("NGT currently supports only binary deserialization format.")
        temp_path = "temp_load_ngt_index"
        with open(temp_path, 'wb') as file:
            file.write(data)
        self.load_index(temp_path)

    def query_radius(self, query_point, radius, sort_results=True):
        """
        Query all points within a specified distance (radius) from the query point.

        Args:
            query_point (list): The query point.
            radius (float): The radius within which to search.
            sort_results (bool): Whether to sort the results by distance (default is True).

        Returns:
            list: A list of points within the specified radius.
        """
        if not self.built:
            raise Exception("Index must be built before querying.")
        result = self.index.search(query_point, size=10000, radius=radius)
        if sort_results:
            result = sorted(result, key=lambda x: x[1])
        return result

    def nearest_centroid(self, centroids, k=1):
        """
        For each centroid provided, find the nearest k data points.

        Args:
            centroids (list): A list of centroids.
            k (int): The number of nearest neighbors to return.

        Returns:
            list: A list of results for each centroid.
        """
        results = []
        for centroid in centroids:
            neighbors = self.index.search(centroid, size=k, with_distance=True)
            results.append(neighbors)
        return results

    def incremental_update(self, new_data_points, removal_ids=None):
        """
        Update the index incrementally with new data points and optionally remove some existing points by IDs.

        Args:
            new_data_points (ndarray): A list of new data points to add to the index.
            removal_ids (list): A list of IDs of points to remove from the index (default is None).
        """
        if removal_ids:
            for item_id in removal_ids:
                self.delete_item(item_id)
        self.add_items(new_data_points)
        self.logger.info("Incremental update performed.")

    def backup_index(self, backup_location):
        """
        Create a backup of the current index state to a specified location.

        Args:
            backup_location (str): The path to the backup location.
        """
        if not self.built:
            raise Exception("Index must be built before it can be backed up.")
        self.save_index(backup_location)
        self.logger.info("Index backed up to %s.", backup_location)

    def restore_index_from_backup(self, backup_location):
        """
        Restore the index state from a backup located at the specified location.

        Args:
            backup_location (str): The path to the backup location.
        """
        self.load_index(backup_location)
        self.built = True
        self.logger.info("Index restored from %s.", backup_location)

    def apply_filter(self, filter_function):
        """
        Apply a custom filter function to all data points in the index, possibly modifying or flagging them based on
        user-defined criteria.

        Args:
            filter_function (function): A function to apply to each data point.

        Returns:
            dict: A dictionary of filtered data points.
        """
        filtered = {i: vec for i, vec in enumerate(self.index.get_objects()) if filter_function(vec)}
        return filtered

    def get_statistics(self):
        """
        Gather and return statistical data about the NGT index, such as point distribution, space utilization, etc.

        Returns:
            dict: A dictionary of statistical data.
        """
        stats = {
            'num_items': self.index.size(),
            'index_built': self.built,
            'distance_type': self.distance_type
        }
        self.logger.info("Statistics retrieved: %s", stats)
        return stats

    def register_callback(self, event, callback_function):
        """
        Register a callback function to be called on specific events.

        Args:
            event (str): The event to register the callback for.
            callback_function (function): The function to call when the event occurs.
        """
        pass

    def unregister_callback(self, event):
        """
        Unregister a previously registered callback for a specific event.

        Args:
            event (str): The event to unregister the callback for.
        """
        pass

    def list_registered_callbacks(self):
        """
        List all currently registered callbacks and the events they are associated with.

        Returns:
            list: A list of registered callbacks.
        """
        return []

    def perform_maintenance(self):
        """
        Perform routine maintenance on the NGT index to ensure optimal performance and stability.
        """
        self.index.build_index()
        self.logger.info("Maintenance performed: index verified.")

    def export_statistics(self, format='csv'):
        """
        Export collected statistical data in a specified format for analysis and reporting purposes.

        Args:
            format (str): The format for exporting statistics (default is 'csv').

        Returns:
            str: The exported statistics.

        Raises:
            ValueError: If the format is not supported.
        """
        stats = self.get_statistics()
        if format == 'csv':
            csv_data = "\n".join([f"{key},{value}" for key, value in stats.items()])
            return csv_data
        else:
            raise ValueError("Unsupported format")

    def adjust_algorithm_parameters(self, **params):
        """
        Dynamically adjust the algorithmic parameters of the NGT algorithm, facilitating on-the-fly optimization
        based on operational feedback.

        Args:
            **params: Arbitrary keyword arguments for adjusting algorithm parameters.
        """
        for param, value in params.items():
            setattr(self.index.property, param, value)
        self.logger.info("Algorithm parameters adjusted: %s", params)

    def query_with_constraints(self, query_point, constraints, k=5):
        """
        Perform a query for nearest neighbors that meet certain user-defined constraints.

        Args:
            query_point (list): The query point.
            constraints (function): A function to apply constraints to the results.
            k (int): The number of nearest neighbors to return.

        Returns:
            list: The constrained nearest neighbors.
        """
        all_results = self.query(query_point, k=k * 10)
        filtered_results = [res for res in all_results if constraints(res)]
        return filtered_results[:k]
