from typing import Iterable, List, Mapping
from collections import Counter
import logging
from tqdm import tqdm
from tqdm.contrib.concurrent import process_map

import ray
from ray.actor import ActorHandle

from dbca.base import Compound
from dbca.sample_set import SampleSet
from dbca.storage import SampleStore

@ray.remote
class FullSampleSetBuilder(SampleSet):
    """
    Represents the full pool of Samples from which the train and test splits will be created.
    """
    def __init__(self, sample_store: SampleStore, pid: int = 0):
        """[summary]

        Args:
            sample_store (SampleStore): Store containing sample objects.
            pid (int): ID of the Ray process.
        """
        #super(FullSampleSetBuilder, self).__init__()
        super().__init__()
        self._sample_store = sample_store  
        sample_ids = sample_store.sample_ids

        # compute all atom weights and record all compounds.
        for sample_id in sample_ids:
            self._add_sample(sample_id)     

        print(f"Initialized loader process {pid}")

  
    def get_compounds_by_uids(self, c_uids: List[str]) -> List[Compound]:
        """ 
        Retrieve Compound objects keyed by compound unique ids.
        """
        return [self._sample_store.get_compound_by_uid(c_uid) for c_uid in c_uids]

    def get_local_compounds(self):
        return list(self.local_compounds_by_type.keys())
        
    def get_atom_weights(self):
        return self.atom_weights

    def get_compound_weights(self):
        return self.atom_weights_by_sample, self.compound_weights, \
            self.compound_weights_by_sample, self.compound_weights_by_type

    def _add_sample(self, sample_id: str):
        """
        Add sample with id `sample_id` from storage to set and update atom weights
        based on sample atoms. Compound weights not computed here, only after all samples are loaded.
        """

        # update atom weight records based on sample atoms
        sample_atom_weights = Counter(self._sample_store.get_sample_atoms(sample_id))
        self.atom_weights_by_sample[sample_id] = sample_atom_weights
        self.atom_weights += sample_atom_weights

        # update compound records - weight not computed here, only after all samples are loaded
        sample_c_refs_dict = self._sample_store.get_sample_compounds_key_sample(sample_id)
        self.local_compounds_by_samp[sample_id] = sample_c_refs_dict
        by_compound_types = self._sample_store.get_sample_compounds_key_type(sample_id)
        
        updated_c_types = list(sample_c_refs_dict.keys())
        
        for compound_type in updated_c_types:
            self.local_compounds_by_type[compound_type].update(by_compound_types[compound_type])

    def weight_compound(self, compound_type: str, pba: ActorHandle):
        """
        Compute compound weight for the given compound type.

        Parameters
        ----------
        str: 
            compound type to load
        ActorHandle:
            progress bar actor for tracking progress

        """
        self.compound_weights[compound_type] = self.calc_compound_weight_in_sample_set(compound_type)
        pba.update.remote(1)

    
    def calc_compound_weight_in_sample_set(self, compound_type: str) -> float:        
        """
        
        Calculate the weight of compound `compound_type` in sample set, by summing over weights of compound in
        all samples in which it is present.
        """
        weight = 0
        samples_with_compound = self.local_compounds_by_type.get(compound_type)

        for sample_id in samples_with_compound:
            sample_weight = self.calc_compound_weight_in_sample(compound_type, sample_id)
            self.compound_weights_by_type[compound_type][sample_id] = sample_weight
            self.compound_weights_by_sample[sample_id][compound_type] = sample_weight
            weight += sample_weight
            
        return weight
    
            
    def get_c_uids_by_type(self, compound_type: str) -> List[str]:
        """
        Get compound uids for all compounds of type `compound_type`. The c_uid keys the actual
        Compound object.
        """
        c_uids = []
        for sample_c_uids in self.local_compounds_by_type[compound_type].values():
            c_uids += sample_c_uids
        return c_uids

    def calc_max_occur_supergraph_prob(self, c_uid: str) -> float:
        """ 
        For a given compound occurrence `c_uid`, find the supergraph co-occuring
        with it the most often across all samples, and return the co-occurrence probability.
        The higher this probability, the less interesting the occurrence.
        """
        compound = self._sample_store.get_compound_by_uid(c_uid)
        # get all occurences for compound of this type
        occs_c_uids = self.get_c_uids_by_type(str(compound))
        occs_compounds = self.get_compounds_by_uids(occs_c_uids)
        max_co_occur_prob = 0.0

        # to find maximal co-occurring super-graph, we can just consider all "super-edges" w.r.t
        # the compound.
        for edge in compound.super_edges():
            co_occur_prob = sum([c.is_super_edge(edge) for c in occs_compounds]) / len(occs_compounds)
            if co_occur_prob > max_co_occur_prob:
                max_co_occur_prob = co_occur_prob
        return max_co_occur_prob
            
    
    def calc_compound_weight_in_sample(self, compound_type: str, sample_id: str) -> float:
        """
        Calculate the weight of compound of type `compound_type` in `sample_id`, taking weight of the occurrence
        with maximal weight if there exist multiple occurrences.

        The weight of the occurrence is the complement of the maximum empirical probability from 
        `calc_max_occur_supergraph_prob`.

        Args:
            compound_type (str): Compound type to calcualte weight for.
            sample_id (str): id of sample to consider.

        Returns:
            float: Compound weight in sample.
        """
        sample_compound_occs = self.local_compounds_by_samp[sample_id][compound_type]
        max_weight = 0.0
        for c_uid in sample_compound_occs:
            # taking complement of maximum co-occurrence probability
            weight = (1 - self.calc_max_occur_supergraph_prob(c_uid))
            max_weight = weight if weight > max_weight else max_weight
        
        return max_weight
    
     
     
    
    

    
