import os.path as osp
from collections import defaultdict
import itertools as it
import logging

from wepy.reporter.dashboard import DashboardReporter

import numpy as np
import pandas as pd

class WExploreDashboardReporter(DashboardReporter):
    """A text based report of a wepy simulation with specific information
    related to the WExplore resampler.

    This includes an overview of the number of regions at each level
    of the hierarchy, a log of when and under what conditions new
    regions were defined, and a table showing the populations of
    walkers and total weights in the leaf regions.

    """

    SUGGESTED_EXTENSIONS = ("wexplore.dash.org",)

    RESAMPLER_TEMPLATE = \
"""* WExplore
Max Number of Regions: {max_n_regions}
Max Region Sizes: {max_region_sizes}
Number of Regions per level:

{regions_per_level}

** Region Hierarchy
Defined Regions with the number of child regions per parent region:
{region_hierarchy}

** Leaf Region Table
{leaf_region_table}

** WExplore Log

{wexplore_log}
* Walker Table
{walker_table}
"""


    def __init__(self,
                 step_time=None, # seconds
                 max_n_regions=None,
                 max_region_sizes=None,
                 bc_cutoff_distance=None,
                 **kwargs
                ):
        """Constructor for the WExplore dashboard reporter.

        Parameters
        ----------

        step_time : float
            The length of the time in each dynamics step.

        max_n_regions : tuple of int
            The maximum number of regions allowed at each level of the
            region hierarchy.

        max_region_sizes : tuple of int
            The distance cutoff at each level of the hierarchy used to
            determine when a new region needs to be created.

        bc_cutoff_distance : float
            The distance for which a walker will be warped.

        """

        super().__init__(step_time, bc_cutoff_distance, **kwargs)

        assert max_n_regions is not None, "number of regions per level for WExplore must be given"
        self.max_n_regions = max_n_regions

        assert max_region_sizes is not None, "region sizes for WExplore must be given"
        self.max_region_sizes = max_region_sizes

        self.n_levels = len(self.max_n_regions)



        # WExplore

        # resampler
        self.root_region = ()
        init_leaf_region = tuple([0 for i in range(self.n_levels)])
        self.region_ids = [init_leaf_region]
        self.regions_per_level = []
        self.children_per_region = {}

        # resampling
        self.walker_assignments = []
        self.walker_image_distances = []
        self.curr_region_probabilities = defaultdict(int)
        self.curr_region_counts = defaultdict(int)

        #wexplore
        self.branch_records = []


    def _update_resampler_values(self, cycle_idx, resampling_data, resampler_data):
        """Update values associated with the WExplore resampler.

        Parameters
        ----------
        cycle_idx : int
            The index of the last completed cycle.

        resampling_data : list of dict of str : value
            List of records specifying the resampling to occur at this
            cycle.

        resampler_data : list of dict of str : value
            List of records specifying the changes to the state of the
            resampler in the last cycle.

        """

        # the region assignments for walkers
        assignments = []
        # re-initialize the current weights dictionary
        self.curr_region_probabilities = defaultdict(int)
        self.curr_region_counts = defaultdict(int)
        for walker_record in resampling_data:

            assignment = tuple(walker_record['region_assignment'])
            walker_idx = walker_record['walker_idx'][0]
            assignments.append((walker_idx, assignment))

            # calculate the probabilities and counts of the regions
            # given the current distribution of walkers
            self.curr_region_probabilities[assignment] += self.walker_weights[walker_idx]
            self.curr_region_counts[assignment] += 1

        # sort them to get the walker indices in the right order
        assignments.sort()
        # then just get the assignment since it is sorted
        self.walker_assignments = [assignment for walker, assignment in assignments]

        # add to the records for region creation in WExplore
        for resampler_record in resampler_data:

            # get the values
            new_leaf_id = tuple(resampler_record['new_leaf_id'])
            branching_level = resampler_record['branching_level'][0]
            walker_image_distance = resampler_record['distance'][0]

            # add the new leaf id to the list of regions in the order they were created
            self.region_ids.append(new_leaf_id)

            # make a new record for a branching event which is:
            # (region_id, level branching occurred, distance of walker that triggered the branching)
            branch_record = (new_leaf_id,
                             branching_level,
                             walker_image_distance)

            # save it in the records
            self.branch_records.append(branch_record)

        # count the number of child regions each region has
        self.children_per_region = {}
        all_regions = self.leaf_regions_to_all_regions(self.region_ids)
        for region_id in all_regions:
            # if its a leaf region it has no children
            if len(region_id) == self.n_levels:
                self.children_per_region[region_id] = 0

            # all others we cound how many children it has
            else:
                # get all regions that have this one as a root
                children_idxs = set()
                for poss_child_id in all_regions:

                    # get the root at the level of this region for the child
                    poss_child_root = poss_child_id[0:len(region_id)]
                    # if the root is the same we keep it without
                    # counting children below the next level, but we skip the same region
                    if (poss_child_root == region_id) and (poss_child_id != region_id):

                        try:
                            child_idx = poss_child_id[len(region_id)]
                        except IndexError:
                            import ipdb; ipdb.set_trace()

                        children_idxs.add(child_idx)

                # count the children of this region
                self.children_per_region[region_id] = len(children_idxs)

        # count the number of regions at each level
        self.regions_per_level = [0 for i in range(self.n_levels)]
        for region_id, n_children in self.children_per_region.items():
            level = len(region_id)

            # skip the leaves
            if level == self.n_levels:
                continue

            self.regions_per_level[level] += n_children

    def _walker_table_string(self):
        # make the table of walkers using pandas, using the order here
        # TODO add the image distances
        walker_table_colnames = ('weight', 'assignment', 'progress') #'image_distances'
        walker_table_d = {}
        walker_table_d['weight'] = self.walker_weights
        walker_table_d['assignment'] = self.walker_assignments
        walker_table_d['progress'] = self.walker_distance_to_prot
        walker_table_df = pd.DataFrame(walker_table_d, columns=walker_table_colnames)
        walker_table_str = walker_table_df.to_string()

        return walker_table_str



    def leaf_regions_to_all_regions(self, region_ids):
        """

        Parameters
        ----------
        region_ids :


        Returns
        -------

        """
        # make a set of all the regions starting with the root region
        regions = set([self.root_region])
        for region_id in region_ids:
            for i in range(len(region_id)):
                regions.add(region_id[0:i+1])

        regions = list(regions)
        regions.sort()

        return regions

    def _resampler_string(self):

        regions = self.leaf_regions_to_all_regions(self.region_ids)
        region_children = [self.children_per_region[region] for region in regions]
        region_children_pairs = it.chain(*zip(regions, region_children))
        region_hierarchy = '\n'.join(['{}     {}' for i in range(len(regions))]).format(*region_children_pairs)

        # make a table for the regions
        region_table_colnames = ('region', 'n_walkers', 'curr_weight')
        region_table_d = {}
        region_table_d['region'] = self.region_ids
        region_table_d['n_walkers'] = [self.curr_region_counts[region] for region in self.region_ids]
        region_table_d['curr_weight'] = [self.curr_region_probabilities[region] for region in self.region_ids]
        leaf_region_table_df = pd.DataFrame(region_table_d, columns=region_table_colnames)
        leaf_region_table_df.set_index('region', drop=True)
        leaf_region_table_str = leaf_region_table_df.to_string()

        # table for aggregeated worker stats


        # log of branching events
        branching_table_colnames = ('new_leaf_id', 'branching_level', 'trigger_distance')
        branching_table_df = pd.DataFrame(self.branch_records, columns=branching_table_colnames)
        branching_table_str = branching_table_df.to_string()


        resampler_string = self.RESAMPLER_TEMPLATE.format(
            max_n_regions=self.max_n_regions,
            max_region_sizes=self.max_region_sizes,
            regions_per_level=self.regions_per_level,
            region_hierarchy=region_hierarchy,
            leaf_region_table=leaf_region_table_str,
            wexplore_log=branching_table_str,
            walker_table= self._walker_table_string()
        )

        return resampler_string
