from .partition import Partition
from cached_property import cached_property
from itertools import chain


class PartitionClass:
    """ Partition Classes are used to generate partitions in a systematic way. For instance,
    a PartitionClass instnace can generate all partitions of an integer n, filter them by durfee square size,
    or crank, or etc... """

    def __init__(self, n=0):
        self.n = n

    def _generate_partitions(self, n = None):
        """ This function is about as fast as it gets for generating integer partitions
        without taking advantage of multithreading. """
        
        # Need to distinguish case where n is 0.
        if n is None:
            n = self.n
        
        a = [0 for i in range(n + 1)]
        k = 1
        y = n - 1
        while k != 0:
            x = a[k - 1] + 1
            k -= 1
            while 2 * x <= y:
                a[k] = x
                y -= x
                k += 1
            j = k + 1
            while x <= y:
                a[k] = x
                a[j] = y
                yield Partition(a[: k + 2])
                x += 1
                y -= 1
            a[k] = x + y
            y = x + y - 1
            yield Partition(a[: k + 1])

    @cached_property
    def partitions(self):
        return self._generate_partitions()

    def filter_by_attribute(self, attr, val=True):
        assert hasattr(Partition(), attr), f"The Partition class does not have attribute: {attr}."

        return (p for p in self._generate_partitions() if getattr(p, attr) == val)

    def filter_by_attribute_lists(self, attr, vals):
        assert hasattr(Partition(), attr), f"The Partition class does not have attribute: {attr}."

        return (p for p in self._generate_partitions() if getattr(p, attr) in vals)
    
    def grid_group(self, n = None):
        if n is None:
            n = self.n
            
        assert n <= self.n, "The provided n must be less than or equal to the n of the partition class."
        
        for generator in chain(self._generate_partitions(n = i) for i in range(1, n + 1)):
            yield from generator
   

class OblakClass(PartitionClass):
    """ A specific type of PartitionClass. All partitions generated by this class share a common result
    after observing the Oblak process."""

    def __init__(self, oblak=Partition()):
        super().__init__(n=oblak.sum_of_parts)
        self.oblak = oblak

    @cached_property
    def partitions(self):
        return (p for p in self._generate_partitions() if p.oblak == self.oblak)
