from collections import defaultdict

import numpy


class BanDBWrapper:
    def __init__(self, ban_db):
        self.ban_db = ban_db

    def get_all_durations(self):
        return [ban.calc_duration() for ban in self.ban_db.list()]

    def get_all_countries(self):
        return [ban.lookup_country_code() for ban in self.ban_db.list()]

    def get_durations_by_country(self):
        return [(ban.lookup_country_code(), ban.calc_duration()) for ban in self.ban_db.list()]

    def calc_average_duration_by_country(self):
        ban_durations_by_country = defaultdict(list)

        for country, ban_duration in self.get_durations_by_country():
            ban_durations_by_country[country].append(ban_duration)

        return {country: numpy.mean(ban_durations) for country, ban_durations in ban_durations_by_country.items()}