from lmi_proteins.utils import write_to_file, isfile, get_current_datetime, create_dir, file_exists
import numpy as np
from lmi_proteins.indexes import BaseIndex
import pandas as pd
from typing import Dict, Tuple, List
import json
from os import path
import logging


class Evaluator:
    """ Runs an LMI experiment, evaluates it and dumps experiment info."""
    def __init__(
        self,
        index: BaseIndex,
        knns: Dict[str, Dict],
        queries: List,
        config: Dict[str, str],
        n_knns=None,
        job_id=None
    ):
        """
        Arguments:
            index (BaseIndex):
                Trained index - LMI, or instantiated - M-index or M-tree
            job_id (str):
                Identifier of the experiment run
            config (Dict[str, str]):
                Experiment's configuration
        """
        self.index = index

        assert len(set(knns.keys())) >= len(queries), \
            "`queries` file is larger than the `knns` file, hence not all queries can be evaluated."
        self.gt_knns = knns
        self.queries = queries
        self.range = 0

        self.stop_conditions_perc = config['experiment']['search-stop-conditions']
        self.stop_conditions = [
            int(index.data.X.shape[0]*cond) for cond in self.stop_conditions_perc
        ]
        if n_knns:
            self.n_knns = n_knns
        else:
            self.n_knns = len(knns[list(knns.keys())[0]])
        self.model = self.index.classifier
        if job_id:
            self.job_id = path.join('outputs', job_id)
        else:
            self.job_id = path.join('outputs', get_current_datetime())
        create_dir(self.job_id)

        self.LOG = logging.getLogger(__name__)
        self.LOG.setLevel(logging.INFO)

    def bucket_occupancy(self):
        bucket_info = {}
        bucket_info['n_buckets'] = len(self.index.objects_in_buckets)
        bucket_values = list(self.index.objects_in_buckets.values())
        bucket_info['mean_bucket_occupancy'] = np.array(bucket_values).mean()
        bucket_info['std_bucket_occupancy'] = np.array(bucket_values).std()
        with open(f'{self.job_id}/bucket-summary.json', 'w') as f:
            json.dump(bucket_info, f, indent=4)

    def get_buckets(self, gt_knns: Dict[str, Dict[str, float]], query: str, q_score: float) -> Tuple[List[int], List[int]]:
        """Gets bucket identifiers of query's 30 nearest neighbors within `self.index`.
        Assumes that all of the neighbors can be found within the structure.

        Parameters:
            gt_knns (Dict[str, Dict[str, float]]): Dictionary of queries' 30 nearest neighbors
            query (str): Current query's identifier (as a string)

        Returns:
            Tuple[List[int], List[int]]:
            List of the buckets, without the nan (placeholder) values if present
            and number of their occurencies within the 30 NNs
        """
        nan_placeholder = -1
        knn_keys = self.get_n_in_range(gt_knns, query, q_score)
        try:
            knn_buckets, counts = np.unique(
                np.nan_to_num(
                    self.index.data.y.loc[knn_keys][self.index.pred_labels].to_numpy(
                        dtype=np.float16,
                        na_value=np.nan
                    ),
                    nan=nan_placeholder
                ),
                axis=0,
                return_counts=True
            )
        except KeyError:
            print(f'Some of knn_keys: {knn_keys} are not in the dataset used by LMI.')
        return [tuple([int(b) for b in bucket if b != nan_placeholder]) for bucket in knn_buckets], counts

    def get_n_in_range(self, gt_knns: Dict[str, Dict[str, float]], query: str, range: float):
        return [k for (k, v) in gt_knns[query].items() if v >= range]   
    
    def run_evaluate(self, q_score = 0.9, stop_condition_index = 0):
        search_results = []
        times = []
        results_filename = f'{self.job_id}/search.csv'
        if not file_exists(results_filename):
            write_to_file(
                results_filename,
                'query,condition,range_score,time,accuracy,visited_objects,found,total'  # noqa: E231
            )
        self.LOG.info(f'Results will be saved to {self.job_id}.')
        self.range = q_score
        results_dataframe_list = []
        
        for query, query_data in self.queries:
            search_results, times, visited_objects_all = self.index.search(query_data, self.stop_conditions)
            knn_buckets, counts = self.get_buckets(self.gt_knns, str(query), q_score)
            nrange = self.get_n_in_range(self.gt_knns, query, q_score)
            
            if len(nrange) == 1:
                continue
            
            for bucket in search_results[stop_condition_index]:
                proteins_in_bucket = list(self.index.data.y[(self.index.data.y["L1_pred"] == bucket[0]) & (self.index.data.y["L2_pred"] == bucket[1])].index)
                for prot in proteins_in_bucket:
                    distance = 1
                    if prot in self.gt_knns[query]:
                        distance = 1 - self.gt_knns[query][prot]
                    results_dataframe_list.append({"query": query, "result": prot, "distance": distance})
            
            for search_result, condition, time, visited_objects in zip(
                search_results, self.stop_conditions, times, visited_objects_all
            ):
                score = 0
                for bucket, n_of_occurences in zip(knn_buckets, counts):
                    if bucket in search_result:
                        score += n_of_occurences

                if len(nrange) > 1:
                    write_to_file(
                        f'{self.job_id}/search.csv',f'{query},'  # noqa: E231
                        f'{condition},{score / len(nrange)},{time},{score / visited_objects},{visited_objects},{score},{len(nrange)}'  # noqa: E231 
                    )
        search_results = pd.DataFrame(results_dataframe_list)
        search_results.to_pickle(f'{self.job_id}/search_results.pkl')
        return search_results
        
                
            

    def generate_summary(self):
        """Creates the `summary.json` file with aggregate information about the results.
        """
        info_dict = {}
        info_dict['model'] = self.model
        info_dict['range'] = self.range
        info_dict['experiment'] = self.job_id
        info_dict['stop_conditions_perc'] = self.stop_conditions_perc

        assert isfile(f'{self.job_id}/search.csv'), \
            f'`{self.job_id}/search.csv` does not exist, cannot generate summary.'
        search_df = pd.read_csv(f'{self.job_id}/search.csv')
        condition_agg = search_df.groupby('condition')
        scores = condition_agg['range_score'].mean().values
        accuracies = condition_agg['accuracy'].mean().values
        conditions = condition_agg['condition'].mean().values
        times = condition_agg['time'].mean().values
        visited_objects = condition_agg['visited_objects'].mean().values

        info_dict['results'] = {}
        for cond, score, accuracy, time, vis_obj in zip(conditions, scores, accuracies, times, visited_objects):
            info_dict['results'][str(cond)] = {}
            info_dict['results'][str(cond)]['time'] = time
            info_dict['results'][str(cond)]['score'] = score
            info_dict['results'][str(cond)]['accuracy'] = accuracy
            info_dict['results'][str(cond)]['visited_objects'] = int(vis_obj)

        total_range_query = search_df[['query','total']].groupby('query').mean().value_counts()
        info_dict['total_range_counts'] = {}
        for index, value in total_range_query.items():
            info_dict['total_range_counts'][index[0]] = value

        with open(f'{self.job_id}/summary.json', 'w') as f:
            json.dump(info_dict, f, indent=4)
