rainbowadn/rainbowadn/data/collection/trees/binary/activebinarytree.py

from typing import Generic, TypeVar

from rainbowadn.data.collection.trees.binary.binarytree import BinaryTree
from rainbowadn.data.collection.trees.comparison.comparator import (Comparator, Comparison, Left, Replace, Right)
from rainbowadn.data.collection.trees.comparison.keyedcomparator import KeyedComparator
from rainbowadn.data.collection.trees.binary.querybinarytree import QueryBinaryTree
from rainbowadn.data.collection.collection_interface.activecollectioninterface import ActiveCollectionInterface
from rainbowadn.data.collection.collection_interface.querycollectioninterface import QueryCollectionInterface
from rainbowadn.data.collection.keymetadata import KeyMetadata
from rainbowadn.data.collection.keymetadataquerycollection import KeyMetadataQueryCollection
from rainbowadn.hashing.hashpoint import HashPoint
from rainbowadn.hashing.hashresolver import HashResolver
from rainbowadn.hashing.nullability.notnull import NotNull
from rainbowadn.hashing.nullability.null import Null
from rainbowadn.hashing.nullability.nullable import Nullable
from rainbowadn.hashing.nullability.nullablereference import NullableReference

__all__ = ('ActiveBinaryTree',)

ActiveKeyType = TypeVar('ActiveKeyType')
MetaDataType = TypeVar('MetaDataType')


class ActiveBinaryTree(
    ActiveCollectionInterface[BinaryTree[KeyMetadata[ActiveKeyType, MetaDataType]], ActiveKeyType],
    Generic[ActiveKeyType, MetaDataType]
):
    def __init__(
            self,
            comparator: Comparator[ActiveKeyType],
            reference: NullableReference[BinaryTree[KeyMetadata[ActiveKeyType, MetaDataType]]]
    ):
        assert isinstance(comparator, Comparator)
        assert isinstance(reference, NullableReference)
        super().__init__(reference)
        self.comparator = comparator
        self.resolver = comparator.resolver
        assert isinstance(self.resolver, HashResolver)

    def empty_metadata(self) -> HashPoint[MetaDataType]:
        raise NotImplementedError

    def metadata(
            self,
            treel: NullableReference[BinaryTree[KeyMetadata[ActiveKeyType, MetaDataType]]],
            treer: NullableReference[BinaryTree[KeyMetadata[ActiveKeyType, MetaDataType]]],
            key: HashPoint[ActiveKeyType]
    ) -> HashPoint[MetaDataType]:
        raise NotImplementedError

    def _binary_tree(
            self,
            treel: NullableReference[BinaryTree[KeyMetadata[ActiveKeyType, MetaDataType]]],
            treer: NullableReference[BinaryTree[KeyMetadata[ActiveKeyType, MetaDataType]]],
            key: HashPoint[ActiveKeyType]
    ) -> BinaryTree[KeyMetadata[ActiveKeyType, MetaDataType]]:
        assert isinstance(treel, NullableReference)
        assert isinstance(treer, NullableReference)
        assert isinstance(key, HashPoint)
        return BinaryTree(
            treel,
            treer,
            HashPoint.of(KeyMetadata(key, self.metadata(treel, treer, key)))
        )

    def create(
            self,
            reference: NullableReference[BinaryTree[KeyMetadata[ActiveKeyType, MetaDataType]]]
    ) -> 'ActiveBinaryTree[ActiveKeyType, MetaDataType]':
        assert isinstance(reference, NullableReference)
        return type(self)(
            self.comparator,
            reference
        )

    def active_tree(
            self,
            treel: NullableReference[BinaryTree[KeyMetadata[ActiveKeyType, MetaDataType]]],
            treer: NullableReference[BinaryTree[KeyMetadata[ActiveKeyType, MetaDataType]]],
            key: HashPoint[ActiveKeyType]
    ) -> 'ActiveBinaryTree[ActiveKeyType, MetaDataType]':
        assert isinstance(treel, NullableReference)
        assert isinstance(treer, NullableReference)
        assert isinstance(key, HashPoint)
        return self.create(
            NullableReference.off(
                self._binary_tree(
                    self.create(treel).balance().reference,
                    self.create(treer).balance().reference,
                    key
                )
            )
        ).balance()

    def add(self, key: HashPoint[ActiveKeyType]) -> 'ActiveBinaryTree[ActiveKeyType, MetaDataType]':
        assert isinstance(key, HashPoint)
        reference: Nullable[
            HashPoint[BinaryTree[KeyMetadata[ActiveKeyType, MetaDataType]]]
        ] = self.reference.reference
        assert isinstance(reference, Nullable)
        if isinstance(reference, Null):
            return self.active_tree(self.reference, self.reference, key)
        elif isinstance(reference, NotNull):
            tree: BinaryTree[KeyMetadata[ActiveKeyType, MetaDataType]] = self.resolver.resolve(reference.value)
            assert isinstance(tree, BinaryTree)
            original: HashPoint[ActiveKeyType] = self.resolver.resolve(tree.key).key
            assert isinstance(original, HashPoint)
            comparison: Comparison = self.comparator.compare(original, key)
            assert isinstance(comparison, Comparison)
            if isinstance(comparison, Replace):
                return self.active_tree(tree.treel, tree.treer, key)
            elif isinstance(comparison, Left):
                return self.active_tree(self.create(tree.treel).add(key).reference, tree.treer, original)
            elif isinstance(comparison, Right):
                return self.active_tree(tree.treel, self.create(tree.treer).add(key).reference, original)
            else:
                raise TypeError
        else:
            raise TypeError

    def balance(self) -> 'ActiveBinaryTree[ActiveKeyType, MetaDataType]':
        raise NotImplementedError

    def __str__(self):
        reference: Nullable[
            HashPoint[BinaryTree[KeyMetadata[ActiveKeyType, MetaDataType]]]
        ] = self.reference.reference
        assert isinstance(reference, Nullable)
        if isinstance(reference, Null):
            return '_'
        elif isinstance(reference, NotNull):
            tree: BinaryTree[KeyMetadata[ActiveKeyType, MetaDataType]] = self.resolver.resolve(reference.value)
            assert isinstance(tree, BinaryTree)
            key: KeyMetadata[ActiveKeyType, MetaDataType] = self.resolver.resolve(tree.key)
            assert isinstance(key, KeyMetadata)
            return f'( {self.create(tree.treel)}' \
                   f' {self.resolver.resolve(key.key)}' \
                   f' {self.resolver.resolve(key.metadata)}' \
                   f' {self.create(tree.treer)} )'
        else:
            raise TypeError

    def query_tree(self) -> QueryCollectionInterface[ActiveKeyType]:
        return KeyMetadataQueryCollection(
            self.empty_metadata(),
            QueryBinaryTree(KeyedComparator(self.resolver, self.comparator), self.reference),
            self.resolver
        )