#!/usr/bin/env python3

"""
The core IR types for datalog.
"""

from collections import namedtuple
from typing import Sequence, Tuple, Union


class Constant(namedtuple("Constant", ["value"])):
  """Representation of a constant for interpreter dispatching."""


class LVar(namedtuple("LVar", ["name"])):
  """Representation of an LVar for interpreter dispatching."""


class Rule(namedtuple("Rule", ["pattern", "clauses"])):
  """Representation of an Rule for the interpreter."""

  def __new__(cls, pattern, clauses):
    return super(cls, Rule).__new__(cls, pattern, tuple(clauses))

  @property
  def used_vars(self):
    return {e for e in self.pattern if isinstance(e, LVar)}

  @property
  def bound_vars(self):
    return {e
            for c in self.clauses
            for e in c
            if isinstance(e, LVar)}

  @property
  def free_vars(self):
    return {v for v in self.used_vars if v not in self.bound_vars}


# Logical type for 'Tuple' as we're gonna use it for the rest of the thing.
LTuple = Tuple[Union[Constant, LVar]]
CTuple = Tuple[Constant]


class Dataset(object):
  """A set of tuples and rules which can be queried."""

  def __init__(self,
               tuples: Sequence[CTuple],
               rules: Sequence[Rule]):
    self.__tuples = tuples
    self.__rules = rules

  def tuples(self) -> Sequence[CTuple]:
    for t in self.__tuples:
      yield t

  def rules(self) -> Sequence[Rule]:
    for r in self.__rules:
      yield r

  def merge(self, other: "Dataset") -> "Dataset":
    """Merge two datasets together, returning a new one."""

    return type(self)(list({*self.tuples(), *other.tuples()}),
                      [*self.rules(), *other.rules()])


class CachedDataset(Dataset):
  """An extension of the dataset which features a cache of rule produced tuples.

  Note that this cache is lost when merging datasets - which ensures correctness.
  """

  # Inherits tuples, rules, merge

  def __init__(self, tuples, rules):
    super(__class__, self).__init__(tuples, rules)
    # The cache is a mapping from a Rule to tuples produced by it.
    self.__cache = {}

  def scan_cache(self, rule_tuple):
    if rule_tuple in self.__cache:
      return iter(self.__cache.get(rule_tuple))

  def cache_tuple(self, rule_tuple, tuple: CTuple):
    coll = self.__cache.get(rule_tuple, list())
    self.__cache[rule_tuple] = coll
    if tuple not in coll:
      coll.append(tuple)

class TableIndexedDataset(CachedDataset):
  """An extension of the Dataset type which features both a cache and an index by table & length.

  The index allows more efficient scans by maintaining 'table' style partitions.
  It does not support user-defined indexing schemes.

  Note that index building is delayed until an index is scanned.
  """

  # From Dataset:
  #   tuples, rules, merge
  # from CachedDataset:
  #   cache_tuple, scan_cache

  @staticmethod
  def __key(t: LTuple) -> str:
   assert isinstance(t[0], Constant)
   return f"{t[0].value}_{len(t)}"

  def __init__(self, tuples, rules):
    super(__class__, self).__init__(tuples, rules)
    self.__index = {}

  def __build_indices(self):
    if not self.__index:
      for t in self.tuples():
        key = self.__key(t)
        # FIXME: Walrus operator???
        coll = self.__index[key] = self.__index.get(key, list())
        coll.append(t)

  def scan_index(self, t: LTuple) -> Sequence[CTuple]:
    self.__build_indices()
    for t in self.__index.get(self.__key(t), []):
      yield t


class PartlyIndexedDataset(TableIndexedDataset):
  """An extension of the Dataset type which features both a cache and and a full index by table,
  length, tuple index and value.

  The index allows extremely efficient scans when elements of the tuple are known.

  """

  # From Dataset:
  #   tuples, rules, merge
  # from CachedDataset:
  #   cache_tuple, scan_cache
  # from IndexedDataset / TableIndexedDataset:
  #   scan_index

  @staticmethod
  def __key(t: LTuple, i: int) -> str:
   assert isinstance(t[0], Constant)
   return (f"{t[0].value}_{len(t)}_{i}", t[i])

  def __init__(self, tuples, rules, index_prefix=999):
    super(__class__, self).__init__(tuples, rules)
    self.__index_prefix = index_prefix
    self.__index = {}

  def __build_indices(self):
    if not self.__index:
      self.__index = index = {}
      # Index by single value
      for t in self.tuples():
        for e, i in zip(t, range(self.__index_prefix)):
          key = self.__key(t, i)
          # FIXME: Walrus operator???
          coll = index[key] = index.get(key, list())
          coll.append(t)

  def scan_index(self, t: LTuple) -> Sequence[CTuple]:
    self.__build_indices()

    default_key = self.__key(t, 0)
    column_indices = []
    for e, i in zip(t, range(self.__index_prefix)):
      if isinstance(e, Constant):
        _k = self.__key(t, i)
        v = self.__index.get(_k)
        if v:
          column_indices.append((_k, v))
        else:
          # If there's no such index, then there's no such tuple. Abort.
          return iter([])

    if column_indices:
      sorted(column_indices, key=lambda x: len(x[1]))
      key, l = column_indices[-1]
    else:
      key = default_key
      l = self.__index[key] = self.__index.get(key, list())

    return iter(l)
