import math
import random
import logging
import operator
from itertools import groupby
from collections import defaultdict

from unstdlib import to_float, to_str
import sqlalchemy as sa

from .utils import find_subclasses, get_exif

log = logging.getLogger(__name__)

class SubclassRegistry(object):
    def __init__(self, base):
        self.base = base
        self._by_name = dict(
            (cls.__name__, cls) for cls in find_subclasses(base)
        )

    def load(self, name):
        return self._by_name[name]()

#
# Stratification
#

class StratifierBase(object):
    required_attrs = []


class pdf(StratifierBase):
    def classify_page(self, page):
        page_count = page.attrs.get("exif/page-count", 1)
        weight = 1 / math.log10(10 + page_count / 2)
        return ("pdf", weight)


Stratifiers = SubclassRegistry(StratifierBase)

#
# Filtering
#

class FilterBase(object):
    pass


Filters = SubclassRegistry(FilterBase)

#
# Dimensions
#

class DimensionBase(object):
    prefix = None

    def __init__(self):
        self.name = self.prefix.rstrip("/")

    def reindex(self, env, clean=False):
        db = env.db
        if clean:
            self.clean(env)
        last = (
            db.session
                .query(db.PDF)
                .join(db.PDFAttr)
                .filter(self.get_filter(db))
                .order_by(db.PDF.id.desc())
                .first()
        )

        for pdf in db.PDF.get_unique(last_id=(last and last.id)):
            log.info("rebuilding %s index for %s", self.prefix, pdf.path)
            for page_num, key, value in self.get_values(pdf):
                pdf.set_attribute(
                    page_num,
                    key, 
                    to_float(value, value),
                )
        db.session.flush()

    def get_filter(self, db):
        return sa.and_(
            db.PDFAttr.name >= self.prefix,
            db.PDFAttr.name < self.prefix + "\x7f"
        )

    def clean(self, env):
        db = env.db
        return (
            db.session
                .query(db.PDFAttr)
                .filter(self.get_filter(db))
                .delete()
        )


class PDFDimension(DimensionBase):
    __abstract__ = True


class PDFPageDimension(DimensionBase):
    __abstract__ = True


class exif(PDFDimension):
    prefix = "exif/"

    def get_values(self, pdf):
        for (key, val) in get_exif(pdf.get_absolute_path()).items():
            yield (None, self.prefix + key, val)
            if key == "page-count":
                for pg_num in range(int(to_float(val))):
                    yield (pg_num + 1, "page", pg_num + 1)


class DimensionRegistry(object):
    def __init__(self):
        self.all_dimensions = [
            cls() for cls in find_subclasses(DimensionBase)
        ]
        self._by_prefix = dict(
            (dim.prefix, dim)
            for dim in self.all_dimensions
        )
        self._by_prefix.update(
            (dim.name, dim)
            for dim in self.all_dimensions
        )

    def get_dimension(self, name):
        return self._by_prefix[name]

    def reindex(self, env, dimensions, clean=False):
        for dim_name in dimensions:
            dim = self.get_dimension(dim_name)
            dim.reindex(env, clean=clean)


Dimensions = DimensionRegistry()


class Generator(object):
    def __init__(self, env, seed, filters, stratifiers, num):
        self.env = env

        random.seed(seed)
        self.seed = seed

        self.filters = filters
        self.stratifiers = stratifiers
        self.num = num

    @classmethod
    def from_args(cls, env, seed, filter_names, stratifier_names, num):
        return cls(
            env,
            seed,
            filter_names,
            [Stratifiers.load(s) for s in stratifier_names],
            num,
        )

    def get_test_case(self):
        db = self.env.db

        load_filter = {}
        for filter in self.filters:
            name, _, val = filter.partition("=")
            if name == "repo":
                load_filter["repo"] = val.split(",")
            else:
                raise ValueError("Unknown filter: %s" %(filter, ))
        pages = list(db.PDFPage.load_all(load_filter))

        strata = defaultdict(list)
        for page in pages:
            groups = []
            weights = []
            for stratifier in self.stratifiers:
                group, weight = stratifier.classify_page(page)
                weight = min(weight, 1.0)
                groups.append(group)
                weights.append(weight)
            weight = reduce(operator.mul, weights, 1)
            group = tuple(groups)
            strata[group].append((-weight * random.random(), page))

        for page_list in strata.values():
            page_list.sort()

        for strata_name, pages in strata.items():
            if self.num:
                pages = pages[:self.num / len(strata)]
            pages.sort(key=lambda x: (x[1].pdf_sha1, x[1].page_num))
            for _, pdf_pages in groupby(pages, key=lambda x: x[1].pdf_sha1):
                pdf = None
                page_nums = []
                for _, page in pdf_pages:
                    if pdf is None:
                        pdf = page.get_pdf()
                    page_nums.append(page.page_num)
                print to_str("%s:%s:%s # strata: %s" %(
                    pdf.repo_name,
                    pdf.path,
                    ",".join(map(str, page_nums)),
                    strata_name,
                ))
