from .chart import Chart
import numpy as np
import random 
import matplotlib.pyplot as plt

class BubbleChart(Chart):
    def __init__(self, dataframe, kwargs):
        """
        Constructs all the necessary attributes for the BubbleChart object

        Parameters:
            dataframe (pandas.Dataframe): The dataframe
        """
        Chart.__init__(self, dataframe, kwargs)

    def promote_to_candidate(self):

        item_col, categorical_col = self._set_item_and_categorical()
        is_promote = self._is_var_exist(item_col, 1) and self._is_var_exist(self._numerical_column, 1)

        return is_promote

    def plot(self):
        """
        Generate visualization
        """
        if self.promote_to_candidate():
            self.draw()
        else:
            pass

    def _check_requirements(self):
        """
        Check the requirements for generating BubbleChart visualization

        Returns:
            (string) numerical_label: label column with float data type
            (string) label_name: label column with string data type
        """
        numerical_label = None
        label_name = None
        if self._is_var_exist(self._numerical_column, 1):
            numerical_label = self._numerical_column[0]
            if len(self._label_column) > 0:
                label_name = self._label_column[-1]
        
        return numerical_label, label_name

    def draw(self):
        """
        Generate table visualization
        """
        numerical_label, label_name = self._check_requirements()

        if numerical_label is not None and label_name is not None:
            self.figsize = set_figsize(self.kwargs.get('figsize'))
            bubble_chart = DrawBubbleChart(area=self.dataframe[numerical_label], bubble_spacing=2)
            bubble_chart.draw(self.dataframe[label_name], self.figsize)


class DrawBubbleChart:
    def __init__(self, area, bubble_spacing=0):
        """
        Setup for bubble collapse.

        Parameters
        ----------
        area : array-like
            Area of the bubbles.
        bubble_spacing : float, default: 0
            Minimal spacing between bubbles after collapsing.

        Notes
        -----
        If "area" is sorted, the results might look weird.
        """
        area = np.asarray(area)
        r = np.sqrt(area / np.pi)

        self.bubble_spacing = bubble_spacing
        self.bubbles = np.ones((len(area), 4))
        self.bubbles[:, 2] = r
        self.bubbles[:, 3] = area
        self.maxstep = 2 * self.bubbles[:, 2].max() + self.bubble_spacing
        self.step_dist = self.maxstep / 2

        # calculate initial grid layout for bubbles
        length = np.ceil(np.sqrt(len(self.bubbles)))
        grid = np.arange(length) * self.maxstep
        gx, gy = np.meshgrid(grid, grid)
        self.bubbles[:, 0] = gx.flatten()[:len(self.bubbles)]
        self.bubbles[:, 1] = gy.flatten()[:len(self.bubbles)]

        self.com = self.center_of_mass()
        
        # initiate color of bubbles
        group_of_items = ['#5A69AF', '#579E65', '#F9C784', '#FC944A', '#F24C00',
                          '#00B825', '#e6194b', '#3cb44b', '#ffe119', '#4363d8', 
                          '#f58231', '#911eb4', '#46f0f0', '#f032e6', '#bcf60c',
                          '#fabebe', '#008080', '#e6beff', '#9a6324', '#fffac8',
                          '#800000', '#aaffc3', '#808000', '#ffd8b1', '#000075',
                          '#808080']
        num_to_select = len(self.bubbles)
        self.colors = [random.choice(group_of_items) for _ in range(num_to_select)]

    def center_of_mass(self):
        return np.average(
            self.bubbles[:, :2], axis=0, weights=self.bubbles[:, 3]
        )

    def center_distance(self, bubble, bubbles):
        return np.hypot(bubble[0] - bubbles[:, 0],
                        bubble[1] - bubbles[:, 1])

    def outline_distance(self, bubble, bubbles):
        center_distance = self.center_distance(bubble, bubbles)
        return center_distance - bubble[2] - \
            bubbles[:, 2] - self.bubble_spacing

    def check_collisions(self, bubble, bubbles):
        distance = self.outline_distance(bubble, bubbles)
        return len(distance[distance < 0])

    def collides_with(self, bubble, bubbles):
        distance = self.outline_distance(bubble, bubbles)
        idx_min = np.argmin(distance)
        return idx_min if type(idx_min) == np.ndarray else [idx_min]

    def collapse(self, n_iterations=50):
        """
        Move bubbles to the center of mass.

        Parameters
        ----------
        n_iterations : int, default: 50
            Number of moves to perform.
        """
        for _i in range(n_iterations):
            moves = 0
            for i in range(len(self.bubbles)):
                rest_bub = np.delete(self.bubbles, i, 0)
                # try to move directly towards the center of mass
                # direction vector from bubble to the center of mass
                dir_vec = self.com - self.bubbles[i, :2]

                # shorten direction vector to have length of 1
                dir_vec = dir_vec / np.sqrt(dir_vec.dot(dir_vec))

                # calculate new bubble position
                new_point = self.bubbles[i, :2] + dir_vec * self.step_dist
                new_bubble = np.append(new_point, self.bubbles[i, 2:4])

                # check whether new bubble collides with other bubbles
                if not self.check_collisions(new_bubble, rest_bub):
                    self.bubbles[i, :] = new_bubble
                    self.com = self.center_of_mass()
                    moves += 1
                else:
                    # try to move around a bubble that you collide with
                    # find colliding bubble
                    for colliding in self.collides_with(new_bubble, rest_bub):
                        # calculate direction vector
                        dir_vec = rest_bub[colliding, :2] - self.bubbles[i, :2]
                        dir_vec = dir_vec / np.sqrt(dir_vec.dot(dir_vec))
                        # calculate orthogonal vector
                        orth = np.array([dir_vec[1], -dir_vec[0]])
                        # test which direction to go
                        new_point1 = (self.bubbles[i, :2] + orth *
                                      self.step_dist)
                        new_point2 = (self.bubbles[i, :2] - orth *
                                      self.step_dist)
                        dist1 = self.center_distance(
                            self.com, np.array([new_point1]))
                        dist2 = self.center_distance(
                            self.com, np.array([new_point2]))
                        new_point = new_point1 if dist1 < dist2 else new_point2
                        new_bubble = np.append(new_point, self.bubbles[i, 2:4])
                        if not self.check_collisions(new_bubble, rest_bub):
                            self.bubbles[i, :] = new_bubble
                            self.com = self.center_of_mass()

            if moves / len(self.bubbles) < 0.1:
                self.step_dist = self.step_dist / 2

    def draw(self, labels, figsize_input=None):
        """
        Draw the bubble plot.

        Parameters
        ----------
        ax : matplotlib.axes.Axes
        labels : list
            Labels of the bubbles.
        colors : list
            Colors of the bubbles.
        """
        self.collapse()

        #check if param figsize exist
        if figsize_input is not None:
            fig, ax = plt.subplots(figsize=figsize_input, subplot_kw=dict(aspect="equal"))
            for i in range(len(self.bubbles)):
                circ = plt.Circle(
                    self.bubbles[i, :2], self.bubbles[i, 2], 
                    color=self.colors[i],
                    )
                ax.add_patch(circ)
                ax.text(*self.bubbles[i, :2], labels[i],
                        horizontalalignment='center', verticalalignment='center')

            ax.axis("off")
            ax.relim()
            ax.autoscale_view()
            plt.show()
        else:
            fig, ax = plt.subplots(figsize=(10,8), subplot_kw=dict(aspect="equal"))
            for i in range(len(self.bubbles)):
                circ = plt.Circle(
                    self.bubbles[i, :2], self.bubbles[i, 2], 
                    color=self.colors[i],
                    )
                ax.add_patch(circ)
                ax.text(*self.bubbles[i, :2], labels[i],
                        horizontalalignment='center', verticalalignment='center')

            ax.axis("off")
            ax.relim()
            ax.autoscale_view()
            plt.show()


def set_figsize(figsize_input):
    """
    Setter of figsize based on figsize input for matplotlib chart

    Parameters:
        (tuple) figsize_input: The figsize input

    Returns:
        (tuple) figsize: The result figsize  
    """
    figsize = None
    is_numeric_value = None

    try:
        if figsize_input is not None and len(figsize_input) == 2:
            is_numeric_value = all(isinstance(v, int) or isinstance(v, float) for v in figsize_input)
        else:
            is_numeric_value = False
    except:
        is_numeric_value = False
        
    if is_numeric_value:
        figsize = figsize_input
    else:
        figsize = None

    return figsize