from abc import ABC, abstractmethod
from typing import List

from .partition import Partition


class PartitionPlan(ABC):
    @abstractmethod
    def groups(self) -> List[List[str]]:
        pass

    def weigh(self, groups: List[List[str]]) -> List[float]:
        total_rows = sum(
            [t.estimated_row_count for db in self.meta.databases() for t in db.tables]
        )
        weights = dict(
            [
                (f"{db.name}.{t.name}", t.estimated_row_count / total_rows)
                for db in self.meta.databases()
                for t in db.tables
            ]
        )
        return [sum([weights[t] for t in g]) for g in groups]

    def partitions(self, partition_count: int) -> List[Partition]:
        nodes = [Partition(self.meta, []) for _ in range(partition_count)]
        node_weights = [0 for _ in range(partition_count)]
        groups = self.groups()
        group_weights = list(zip(groups, self.weigh(groups)))
        group_weights.sort(key=lambda gw: -gw[1])
        for g, w in group_weights:
            min_idx = 0
            for i in range(len(nodes)):
                if node_weights[i] < node_weights[min_idx]:
                    min_idx = i
            nodes[min_idx].add_tables(g)
            node_weights[min_idx] += w
        return nodes