import inspect
import math
import re
import textwrap

from collections.abc   import Iterable
from functools         import wraps
from operator          import itemgetter
from typing            import Callable, Literal, Optional, overload, TypeGuard
from typing_extensions import Self

from frplib.exceptions import OperationError, StatisticError
from frplib.numeric    import as_real, as_numeric
from frplib.protocols  import Projection, Transformable
from frplib.utils      import ensure_tuple, identity, is_interactive, is_tuple, scalarize
from frplib.vec_tuples import VecTuple, as_scalar, as_scalar_strict, as_vec_tuple, is_vec_tuple

# ATTN: conversion with as_real etc in truediv, pow to prevent accidental float conversion
# This could be mitigated by eliminating ints from as_numeric*, but we'll see how this
# goes.


#
# Helpers
#


def compose2(after: 'Statistic', before: 'Statistic') -> 'Statistic':
    if after.dim == 0 or before.codim is None or after.dim == before.codim:
        def composed(*x):
            return after(before(*x))
        return Statistic(composed, dim=before.dim, codim=after.codim, name=f'{after.name}({before.name}(__))')
    raise OperationError(f'Statistics {after.name} and {before.name} are not compatible for composition.')


#
# Decorator/Wrapper to make functions auto-uncurry
#

def tuple_safe(fn: Callable, arity : Optional[int] = None) -> Callable:
    """Returns a function that can accept a single tuple or multiple individual arguments.

    Ensures that the returned function has an `arity` attribute set
    to the supplied or computed arity.
    """
    if arity is None:
        arity = len([param for param in inspect.signature(fn).parameters.values()
                     if param.kind <= inspect.Parameter.POSITIONAL_OR_KEYWORD])
    if arity == 0:
        @wraps(fn)
        def f(*x):
            if len(x) == 1 and is_tuple(x[0]):
                return as_vec_tuple(fn(x[0]))
            return as_vec_tuple(fn(x))
        setattr(f, 'arity', arity)
        return f
    elif arity == 1:
        @wraps(fn)
        def g(x):
            if is_tuple(x) and len(x) == 1:
                return as_vec_tuple(fn(x[0]))
            return as_vec_tuple(fn(x))
        setattr(g, 'arity', arity)
        return g

    @wraps(fn)
    def h(*x):
        select = itemgetter(*range(arity))
        if len(x) == 1 and is_tuple(x[0]):
            return as_vec_tuple(fn(*select(x[0])))
        return as_vec_tuple(fn(*select(x)))
    setattr(h, 'arity', arity)
    return h


#
# The Statistics Interface
#

# TODO: Need to handle dimensions of return type.
#       In particular, something like Sum returns a scalar, which needs to be wrapped in a tuple at the end ?!?
#       Also want to allow operations on tuples to be component wise and implement @ as a dot product
#       So,   Id - (1, 2, 3) is a tuple and Id @ (1, 2, 3) is a scalar wrapped
#       Add: Scalar function or wrap scalar constants in operation methods below

# Idea here is to be able to do things like
# dice | Sum > 5 and d6 ** 4 ^ Sum / 4  and d6 ** 4 ^ Mean
# or even dice ^ Sum > 5

# dice | fork(Sum, Count)               # produces tuples x :-> (sum(x), len(x))
# dice | chain(Id - 7, Abs, Id <= 3)    # should reduce to lambda x: |x - 7| <= 3
#
# Built-in statistics: Sum, Mean, Max, Min, Product, Count, Dot(c1,c2,c3,...) [dot product, repeat last coefficient ad infinitum; this is a function that produces a Statistic],
#                      Id, Proj(1,...) [func that produces a statistic, equiv to @ but allows Proj(2) - Proj(1)], Permutation(i1,i2,i3,...),
#                      Square, Cube, SumOfSquares, Sqrt, Pow(exp), Exp, Ln, Abs, Sin, Cos, Tan

# NOTE: This means that we want to accept 0-1 valued functions in conditioning predicates
# as well as boolean functions. Or maybe it's OK and just works, but probably want to convert
# boolean to 0-1 in .map().  <<<---- so events become easy

# Note: codim=0 can be used to mean the same dim as input 

## ATTN: Also implement things like __is__ and __in__ so we can do X ^ (__ in {0, 1, 2})

## ATTN: call of a statistic can check for FRP or Kind  (need root class) and call their transform method on self   Sum(X)
##   THis is great as it closely models our notation
## We still have the ^ notation for piping and can add @ for compose ("after")  so  k ^ Double @ Scalar == Double(Scalar(k)) == Double(Scalar(k))

class Statistic:
    "A function that operates on the output of an FRP."
    def __init__(self: Self,
                 fn: Callable | 'Statistic',         # Either a Statistic or a function to be turned into one
                 dim: Optional[int] = None,          # Number of arguments the function takes; 0 means tuple expected
                 codim: Optional[int] = None,        # Dimension of the codomain; None means don't know
                 name: Optional[str] = None,         # A user-facing name for the statistic
                 description: Optional[str] = None   # A description used as a __doc__ string for the Statistic
                ) -> None:
        "Note: set arity to 0 to advise that it takes the whole tuple/vector; None means try to find out"

        if isinstance(fn, Statistic):
            self.fn: Callable = fn.fn
            self.arity: int = dim if dim is not None else fn.arity
            self.codim: Optional[int] = codim if codim is not None else fn.codim
            self._name = name or fn.name
            self.__doc__: str = self.__describe__(description or fn.description or '')
            return

        f = tuple_safe(fn, dim)
        self.fn = f
        self.arity = getattr(f, 'arity')
        self.codim = codim
        self._name = name or fn.__name__ or ''
        self.__doc__ = self.__describe__(description or fn.__doc__ or '')

    def __describe__ (self, description: str, returns: Optional[str] = None) -> str:
        def splitPascal(pascal: str) -> str:
            return re.sub(r'([a-z])([A-Z])', r'\1 \2', pascal)

        my_name = splitPascal(self.__class__.__name__)
        an = 'An' if re.match(r'[AEIOU]', my_name) else 'A'
        me = f'{an} {my_name} \'{self.name}\''
        that = '' if description else ' that '
        descriptor = ' that ' + (description + '. It ' if description else '')

        scalar = ''
        if not returns:
            if self.codim == 1:
                scalar = 'returns a scalar'
            elif self.codim is not None:
                scalar = f'returns a {self.codim}-tuple'
        else:
            scalar = returns

        arity = ''
        if self.arity == 0:
            arity = 'expects a tuple'
        elif self.arity > 0:
            arity = f'expects {self.arity} argument{"s" if self.arity > 1 else ""} (or a tuple of that length)'

        conj = ' and ' if scalar and arity else that if scalar else ''
        structure = f'{arity}{conj}{scalar}.'

        return f'{me}{descriptor}{structure}'
        
    def __str__(self) -> str:
        return self.__doc__

    def __repr__(self) -> str:
        if is_interactive():  # Needed?
            return str(self)
        # ATTN! This looks like a bug
        return super().__repr__()

    @property
    def name(self) -> str:
        return self._name

    @property
    def dim(self) -> int:
        "Returns the dimension of the statistic, with 0 meaning it accepts an arbitrary tuple."
        return self.arity

    @property
    def description(self) -> str:
        return self.__doc__

    def __call__(self, *args):
        # It is important that Statistics are not Transformable!
        if len(args) == 1:
            if isinstance(args[0], Transformable):
                return args[0].transform(self)
            if isinstance(args[0], Statistic):
                return compose2(self, args[0])
        return self.fn(*args)

    # Comparisons (macros would be nice here)

    def __eq__(self, other):
        if isinstance(other, Statistic):
            def a_eq_b(*x):
                return self(*x) == other(*x)
            label = f'{other.name}(__)'
        elif callable(other):
            f = tuple_safe(other)
            def a_eq_b(*x):
                return self(*x) == f(*x)
            label = str(other)
        else:
            def a_eq_b(*x):
                return self(*x) == other
            label = str(other)

        # Break inheritance rules here, but it makes sense!
        return Condition(a_eq_b, dim=0, name=f'{self.name}(__) == {label}')

    def __ne__(self, other):
        if isinstance(other, Statistic):
            def a_ne_b(*x):
                return self(*x) != other(*x)
            label = f'{other.name}(__)'
        elif callable(other):
            f = tuple_safe(other)
            def a_ne_b(*x):
                return self(*x) != f(*x)
            label = str(other)
        else:
            def a_ne_b(*x):
                return self(*x) != other
            label = str(other)

        # Break inheritance rules here, but it makes sense!
        return Condition(a_ne_b, dim=0, name=f'{self.name}(__) != {label}')

    ## ATTN:FIX labels for methods below, so e.g., ForEach(2*__+1) prints out nicely

    def __le__(self, other):
        if isinstance(other, Statistic):
            def a_le_b(*x):
                return self(*x) <= other(*x)
            label = f'{other.name}(__)'
        elif callable(other):
            f = tuple_safe(other)
            def a_le_b(*x):
                return self(*x) <= f(*x)
            label = str(other)
        else:
            def a_le_b(*x):
                return self(*x) <= other
            label = str(other)

        # Break inheritance rules here, but it makes sense!
        return Condition(a_le_b, dim=0, name=f'{self.name}(__) <= {label}')

    def __lt__(self, other):
        if isinstance(other, Statistic):
            def a_lt_b(*x):
                return self(*x) < other(*x)
            label = f'{other.name}(__)'
        elif callable(other):
            f = tuple_safe(other)
            def a_lt_b(*x):
                return self(*x) < f(*x)
            label = str(other)
        else:
            def a_lt_b(*x):
                return self(*x) < other
            label = str(other)

        # Break inheritance rules here, but it makes sense!
        return Condition(a_lt_b, dim=0, name=f'{self.name}(__) < {label}')

    def __ge__(self, other):
        if isinstance(other, Statistic):
            def a_ge_b(*x):
                return self(*x) >= other(*x)
            label = f'{other.name}(__)'
        elif callable(other):
            f = tuple_safe(other)
            def a_ge_b(*x):
                return self(*x) >= f(*x)
            label = str(other)
        else:
            def a_ge_b(*x):
                return self(*x) >= other
            label = str(other)

        # Break inheritance rules here, but it makes sense!
        return Condition(a_ge_b, dim=0, name=f'{self.name}(__) >= {label}')

    def __gt__(self, other):
        if isinstance(other, Statistic):
            def a_gt_b(*x):
                return self(*x) > other(*x)
            label = f'{other.name}(__)'
        elif callable(other):
            f = tuple_safe(other)
            def a_gt_b(*x):
                return self(*x) > f(*x)
            label = str(other)
        else:
            def a_gt_b(*x):
                return self(*x) > other
            label = str(other)

        # Break inheritance rules here, but it makes sense!
        return Condition(a_gt_b, dim=0, name=f'{self.name}(__) > {label}')

    # Numeric Operations (still would like macros)

    def __add__(self, other):
        if isinstance(other, Statistic):
            def a_plus_b(*x):
                return self(*x) + other(*x)
            label = f'{other.name}(__)'
        elif callable(other):
            f = tuple_safe(other)
            def a_plus_b(*x):
                return self(*x) + f(*x)
            label = str(other)
        else:
            def a_plus_b(*x):
                return self(*x) + other
            label = str(other)

        return Statistic(a_plus_b, dim=0, name=f'{self.name}(__)')

    def __radd__(self, other):
        if callable(other):   # other cannot be a Statistic in __r*__
            f = tuple_safe(other)
            def a_plus_b(*x):
                return self(*x) + f(*x)
            label = str(other)
        else:
            def a_plus_b(*x):
                return self(*x) + other
            label = str(other)

        return Statistic(a_plus_b, dim=0, name=f'{label} + {self.name}(__)')

    def __sub__(self, other):
        if isinstance(other, Statistic):
            def a_minus_b(*x):
                return self(*x) - other(*x)
            label = f'{other.name}(__)'
        elif callable(other):
            f = tuple_safe(other)
            def a_minus_b(*x):
                return self(*x) - f(*x)
            label = str(other)
        else:
            def a_minus_b(*x):
                return self(*x) - other
            label = str(other)

        return Statistic(a_minus_b, dim=0, name=f'{self.name}(__) - {label}')

    def __rsub__(self, other):
        if callable(other):   # other cannot be a Statistic in __r*__
            f = tuple_safe(other)
            def a_minus_b(*x):
                return f(*x) - self(*x)
        else:
            def a_minus_b(*x):
                return other - self(*x)

        return Statistic(a_minus_b, dim=0, name=f'{str(other)} - {self.name}(__)')

    def __mul__(self, other):
        if isinstance(other, Statistic):
            def a_times_b(*x):
                return self(*x) * other(*x)
            label = f'{other.name}(__)'
        elif callable(other):
            f = tuple_safe(other)
            def a_times_b(*x):
                return self(*x) * f(*x)
            label = str(other)
        else:
            def a_times_b(*x):
                return self(*x) * as_real(as_scalar_strict(other))  # ATTN!
            label = str(other)

        return Statistic(a_times_b, dim=0, name=f'{self.name}(__) * {label}')

    def __rmul__(self, other):
        if callable(other):   # other cannot be a Statistic in __r*__
            f = tuple_safe(other)
            def a_times_b(*x):
                return self(*x) * f(*x)
        else:
            def a_times_b(*x):
                return self(*x) * as_real(as_scalar_strict(other))

        return Statistic(a_times_b, dim=0, name=f'{str(other)} * {self.name}(__)')

    def __truediv__(self, other):
        if isinstance(other, Statistic):
            def a_div_b(*x):
                return self(*x) / other(*x)
            label = f'{other.name}(__)'
        elif callable(other):
            f = tuple_safe(other)
            def a_div_b(*x):
                return self(*x) / f(*x)
            label = str(other)
        else:
            def a_div_b(*x):
                return self(*x) / as_real(as_scalar_strict(other))
            label = str(other)

        return Statistic(a_div_b, dim=0, name=f'{self.name}(__) / {label}')

    def __rtruediv__(self, other):
        if callable(other):   # other cannot be a Statistic in __r*__
            f = tuple_safe(other)
            def a_div_b(*x):
                return f(*x) / self(*x)
        else:
            def a_div_b(*x):
                return other / as_real(as_scalar_strict(self(*x)))

        return Statistic(a_div_b, dim=0, name=f'{str(other)} / {self.name}(__)')

    def __floordiv__(self, other):
        if isinstance(other, Statistic):
            def a_div_b(*x):
                return self(*x) // other(*x)
            label = f'{other.name}(__)'
        elif callable(other):
            f = tuple_safe(other)
            def a_div_b(*x):
                return self(*x) // f(*x)
            label = str(other)
        else:
            def a_div_b(*x):
                return self(*x) // other
            label = str(other)

        return Statistic(a_div_b, dim=0, name=f'{self.name}(__) // {label}')

    def __rfloordiv__(self, other):
        if callable(other):   # other cannot be a Statistic in __r*__
            f = tuple_safe(other)
            def a_div_b(*x):
                return f(*x) // self(*x)
        else:
            def a_div_b(*x):
                return other // self(*x)

        return Statistic(a_div_b, dim=0, name=f'{str(other)} // {self.name}(__)')

    def __mod__(self, other):
        if isinstance(other, Statistic):
            def a_mod_b(*x):
                return self(*x) % other(*x)
            label = f'{other.name}(__)'
        elif callable(other):
            f = tuple_safe(other)
            def a_mod_b(*x):
                return self(*x) % f(*x)
            label = str(other)
        elif self.codim == 1:
            def a_mod_b(*x):
                try:
                    return scalarize(self(*x)) % other
                except Exception as e:
                    raise OperationError(f'Could not compute {self.name} % {other}: {str(e)}')
            label = str(other)
        else:
            def a_mod_b(*x):
                val = self(*x)
                if len(val) != 1:
                    raise OperationError(f'Statistic {self.name} is not a scalar but % requires it; try using Scalar explicitly.')
                try:
                    return scalarize(self(*x)) % other
                except Exception as e:
                    raise OperationError(f'Could not compute {self.name} % {other}: {str(e)}')
            label = str(other)
        return Statistic(a_mod_b, dim=0, name=f'{self.name}(__) % {label}')

    def __rmod__(self, other):
        if callable(other):   # other cannot be a Statistic in __r*__
            f = tuple_safe(other)
            def a_mod_b(*x):
                return f(*x) % self(*x)
        else:
            def a_mod_b(*x):
                return other % self(*x)

        return Statistic(a_mod_b, dim=0, name=f'{str(other)} % {self.name}(__)')

    def __pow__(self, other):
        if isinstance(other, Statistic):
            def a_pow_b(*x):
                return self(*x) ** other(*x)
            label = f'{other.name}(__)'
        elif callable(other):
            f = tuple_safe(other)
            def a_pow_b(*x):
                return self(*x) ** f(*x)
            label = str(other)
        else:
            def a_pow_b(*x):
                return self(*x) ** as_numeric(other)
            label = str(other)

        return Statistic(a_pow_b, dim=0, name=f'{self.name}(__) ** {label}')

    def __rpow__(self, other):
        if callable(other):   # other cannot be a Statistic in __r*__
            f = tuple_safe(other)
            def a_pow_b(*x):
                return f(*x) ** self(*x)
        else:
            def a_pow_b(*x):
                return as_numeric(other) ** self(*x)

        return Statistic(a_pow_b, dim=0, name=f'{str(other)} ** {self.name}(__)')

    def __and__(self, other):
        if isinstance(other, Statistic):
            def a_and_b(*x):
                return self(*x) and other(*x)
            label = f'{self.name}(__) and {other.name}'
        elif callable(other):
            f = tuple_safe(other)
            def a_and_b(*x):
                return self(*x) and f(*x)
            label = f'{self.name}(__) and {str(other)}'
        else:
            def a_and_b(*x):
                return self(*x) and other
            label = f'{self.name}(__) and {str(other)}'

        return Statistic(a_and_b, dim=0, name=label)

    def __or__(self, other):
        if isinstance(other, Statistic):
            def a_or_b(*x):
                return self(*x) or other(*x)
            label = f'{self.name}(__) or {other.name}(__)'
        elif callable(other):
            f = tuple_safe(other)
            def a_or_b(*x):
                return self(*x) or f(*x)
            label = f'{self.name}(__) or {str(other)}'
        else:
            def a_or_b(*x):
                return self(*x) or other
            label = f'{self.name} and {str(other)}'

        return Statistic(a_or_b, dim=0, name=label)


def is_statistic(x) -> TypeGuard[Statistic]:
    return isinstance(x, Statistic)

class MonoidalStatistic(Statistic):
    def __init__(self,
                 fn: Callable | 'Statistic',        # Either a Statistic or a function to be turned into one
                 unit,                              # The unit of the monoid
                 dim: Optional[int] = None,       # Number of arguments the function takes; 0 means tuple expected
                 codim: Optional[int] = None,       # Dimension of the codomain; None means don't know
                 name: Optional[str] = None,        # A user-facing name for the statistic
                 description: Optional[str] = None  # A description used as a __doc__ string for the Statistic
                ) -> None:
        super().__init__(fn, dim, codim, name, description)
        self.unit = unit

    def __call__(self, *args):
        if len(args) == 0:
            return self.unit
        return super().__call__(*args)

class ProjectionStatistic(Statistic, Projection):
    def __init__(self,
                 # ATTN: Don't need this here, just adapt project
                 fn: Callable | 'Statistic',         # Either a Statistic or a function to be turned into one ATTN: No need for fn here!
                 onto: Iterable[int] | slice | Self, # 1-indexed projection indices
                 name: Optional[str] = None          # A user-facing name for the statistic
                ) -> None:
        codim = None
        dim = 0
        if isinstance(onto, ProjectionStatistic):
            indices: Iterable[int] | slice | 'ProjectionStatistic' = onto.subspace
            codim = onto.codim
            label = onto.name.replace('project[', '').replace(']', '')

        if isinstance(onto, Iterable):
            indices = list(onto)
            codim = len(indices)
            dim = max(indices)
            label = ", ".join(map(str, indices))
            if any([index == 0 for index in indices]):  # Negative from the end OK
                raise StatisticError('Projection indices are 1-indexed and must be non-zero')
        elif isinstance(onto, slice):
            indices = onto
            has_step = indices.step is None
            label = f'{indices.start or ""}:{indices.stop or ""}{":" if has_step else ""}{indices.step if has_step else ""}'
            # ATTN: Already converted in project; need to merge this
            #if indices.start == 0 or indices.stop == 0:
            #    raise StatisticError('Projection indices are 1-indexed and must be non-zero')
        
        description = textwrap.wrap(f'''A statistic that projects any value of dimension >= {dim or 1}
                                        to extract the {codim} components with indices {label}.''')
        # ATTN: Just pass project here, don't take an fn arg!
        super().__init__(fn, 0, codim, name, '\n'.join(description))
        self._components = indices

    @property
    def subspace(self):
        return self._components

    #ATTN: Make project() below a method here
    #ATTN?? Add minimum_dim property that specifies minimum compatible dimension; e.g., Project[3] -> 3, Project[2:-1] -> 2, Project[1,3,5] -> 5

def _ibool(x) -> Literal[0, 1]:
    return 1 if bool(x) else 0

class Condition(Statistic):
    """A condition is a statistic that returns a boolean value.

    Boolean values here are represented in the output with
    0 for false and 1 for true, though the input callable
    can return any 
    """ 
    def __init__(self,
                 predicate: Callable | 'Statistic', # Either a Statistic or a function to be turned into one
                 dim: Optional[int] = None,         # Number of arguments the function takes; 0 means tuple expected
                 name: Optional[str] = None,        # A user-facing name for the statistic
                 description: Optional[str] = None  # A description used as a __doc__ string for the Statistic
                ) -> None:
        super().__init__(predicate, dim, 1, name, description)
        self.__doc__ = self.__describe__(description or predicate.__doc__ or '', 'returns a 0-1 (boolean) value')

    def __call__(self, *args) -> Literal[0, 1] | Statistic:
        if len(args) == 1 and isinstance(args[0], Transformable):
            return args[0].transform(self)
        if isinstance(args[0], Statistic):
            return Condition(compose2(self, args[0]))
        result = super().__call__(*args)
        return _ibool(as_scalar(result))
        # if is_vec_tuple(result):
        #     return result.map(_ibool)
        # return as_vec_tuple(result).map(_ibool)


#
# Statistic decorator for easily creating a statistic out of a function
#

def statistic(
        maybe_fn : Optional[Callable] = None, # If supplied, return Statistic, else a decorator
        *,
        name: Optional[str] = None,         # A user-facing name for the statistic
        dim: Optional[int] = None,          # Number of arguments the function takes; 0 means tuple expected
        codim: Optional[int] = None,        # Dimension of the codomain; None means don't know
        description: Optional[str] = None,  # A description used as a __doc__ string for the Statistic
        monoidal = None                     # If not None, the unit for a Monoidal Statistic
) -> Statistic | Callable[[Callable], Statistic]:
    """
    """
    if maybe_fn and monoidal is None:
        return Statistic(maybe_fn, dim, codim, name, description)
    elif maybe_fn:
        return MonoidalStatistic(maybe_fn, monoidal, dim, codim, name, description)

    if monoidal is None:
        def decorator(fn: Callable) -> Statistic:     # Function to be converted to a statistic
            return Statistic(fn, dim, codim, name, description)
    else:
        def decorator(fn: Callable) -> Statistic:     # Function to be converted to a statistic
            return MonoidalStatistic(fn, monoidal, dim, codim, name, description)
    return decorator

def scalar_statistic(
        name: Optional[str] = None,         # A user-facing name for the statistic
        dim: Optional[int] = None,          # Number of arguments the function takes; 0 means tuple expected
        description: Optional[str] = None,  # A description used as a __doc__ string for the Statistic
        monoidal = None                     # If not None, the unit of a Monoidal Statistic
):
    def decorator(fn: Callable) -> Statistic:     # Function to be converted to a statistic
        return Statistic(fn, dim, 1, name, description)
    return decorator

def condition(
        maybe_predicate : Optional[Callable] = None, # If supplied, return Condition, else a decorator
        *,
        name: Optional[str] = None,         # A user-facing name for the statistic
        dim: Optional[int] = None,          # Number of arguments the function takes; 0 means tuple expected
        codim: Optional[int] = None,        # Dimension of the codomain; None means don't know
        description: Optional[str] = None,  # A description used as a __doc__ string for the Statistic
        monoidal = None                     # If not None, the unit for a Monoidal Statistic
) -> Condition | Callable[[Callable], Condition]:
    """
    """
    if maybe_predicate:
        return Condition(maybe_predicate, dim, name, description)

    def decorator(predicate: Callable) -> Condition:     # Function to be converted to a statistic
        return Condition(predicate, dim, name, description)
    return decorator


#
# Statistics Combinators
#

def fork(*statistics: Statistic) -> Statistic:
    arities = {stat.arity for stat in statistics}
    def forked(*x):
        return tuple([stat(*x) for stat in statistics])
    names = [stat.name for stat in statistics]
    return Statistic(forked, dim=next(iter(arities)) if len(arities) == 1 else 0,
                     name=f'fork({", ".join(names)})',
                     description=f'returns a tuple of the results of ({", ".join(names)})')

def chain(*statistics: Statistic) -> Statistic:
    "Compose statistics in pipeline order: (f ; g)(x) = g(f(x)), read 'f then g'."
    # ATTN: check arities compatible etc
    def chained(*x):
        state = x
        for stat in statistics:
            state = stat(*state)
        return state
    arity = statistics[0].arity if len(statistics) > 0 else None
    names = ", ".join([stat.name for stat in statistics])
    return Statistic(chained, arity, name=f'chain({names})')

def compose(*statistics: Statistic) -> Statistic:
    "Compose statistics in mathematical order: (f o g)(x) = f(g(x)), read 'f after g'."
    # ATTN: check dims and codims etc
    rev_statistics = list(statistics)
    rev_statistics.reverse()
    def composed(*x):
        state = x
        for stat in rev_statistics:
            state = stat(*state)
        return state
    arity = rev_statistics[0].arity if len(statistics) > 0 else None
    names = ", ".join([stat.name for stat in statistics])
    return Statistic(composed, arity, name=f'compose({names})')


#
# Special Numerical Values
#

infinity = math.inf  # ATTN: if needed, convert to appropriate value component type


#
# Commonly Used Statistics
#

Id = Statistic(identity, dim=0, name='identity', description='returns the value given as is')
Scalar = Statistic(lambda x: x[0] if is_tuple(x) else x, dim=1, name='scalar', description='represents a scalar value')
__ = Id # ATTN  Make this act like both Id and Scalar
_x_ = Scalar

Sum = MonoidalStatistic(sum, unit=0, dim=0, codim=1, name='sum', description='returns the sum of all the components of the given value')
Count = MonoidalStatistic(len,unit=0, dim=0, codim=1, name='count', description='returns the number of components in the given value')
Max = MonoidalStatistic(max, unit=-math.inf, dim=0, codim=1, name='max', description='returns the maximum of all components of the given value')
Min = MonoidalStatistic(min, unit=math.inf, dim=0, codim=1, name='min', description='returns the minimum of all components of the given value')
Mean = Statistic(lambda x: sum(x)/len(x), dim=0, codim=1, name='mean', description='returns the arithmetic mean of all components of the given value')
Abs = Statistic(abs, dim=1, codim=1, name='abs', description='returns the absolute value of the given number')

# Diff == takes first order diffs of the tuple
# Diffs(k) takes k-order diffs

def Constantly(x) -> Statistic:
    xtuple = ensure_tuple(x)
    return Statistic(lambda _: xtuple, dim=0, codim=len(xtuple), name=f'The constant {xtuple}')

Cos = Statistic(math.cos, dim=1, codim=1, name='cos', description='returns the cosine of a scalar argument')
Sin = Statistic(math.sin, dim=1, codim=1, name='sin', description='returns the sine of a scalar argument')
Exp = Statistic(math.exp, dim=1, codim=1, name='exp', description='returns the exponential of a scalar argument')
Log = Statistic(math.log, dim=1, codim=1, name='log', description='returns the natural logarithm of a positive scalar argument')
Log2 = Statistic(math.log2, dim=1, codim=1, name='log', description='returns the logarithm base 2 of a positive scalar argument')
Log10 = Statistic(math.log2, dim=1, codim=1, name='log', description='returns the logarithm base 10 of a positive scalar argument')


#
# Combinators
#

def ForEach( s: Statistic ) -> Statistic:
    def foreach(*x):
        if len(x) > 0 and is_tuple(x[0]):
            x = x[0]
        return tuple([s(xi) for xi in x])
    return Statistic(foreach, dim=0, name=f'applies {s.name} to every component of input value')


##
## ATTN: Need better handling of dim, arity, codim.  arity=0, etc. is fine while dim is the minimum dimension allowed
## codim=0 means returns a tuple whose equals its input length; codim=n means returns an n-tuple, codim=None is unknown
##
def Fork( stat: Statistic, *more_stats: Statistic ) -> Statistic:
    # Arities must all be the same
    if any([s.arity != stat.arity for s in more_stats]):  #ATTN!: Statistics need to distinguish arity from minimum dim, need two properties
        raise ValueError('Fork must be called on statistics with the same dimension')
    codim = 0
    if stat.codim is not None and stat.codim > 0 and all([s.codim is not None and s.codim > 0 for s in more_stats]):
        codim = stat.codim + sum(s.codim for s in more_stats)  # type: ignore
    def forked(*x):
        returns = list(stat(x))
        for s in more_stats:
            returns.extend(s(x))
        return returns
    return Statistic(forked, dim=stat.arity, codim=codim, name=f'fork({stat.name}, {", ".join([s.name for s in more_stats])})')

## ATTN: Add 
# Permute    ATTN: fix up but keeping it simple for now
def Permute(p: Iterable[int]):
    # assert p contains all unique values from 1..n
    perm = list(map(lambda k: k - 1, p))
    n = len(perm)
    @statistic(name=f'Permutation', dim=0)   # ATTN
    def permute(value):
        return VecTuple(value[perm[i]] for i in range(n))
    return permute
        


def IfThenElse( cond: Statistic
              , t: Statistic | tuple | float | int
              , f: Statistic | tuple | float | int
              ) -> Statistic:
    if not is_statistic(t):
        t = Constantly(t)
    if not is_statistic(f):
        f = Constantly(f)
    # ATTN: arity conditions need to be more nuanced
    #if cond.arity != t.arity or t.arity != f.arity or t.codim != f.codim:
    if t.codim is not None and f.codim is not None and t.codim != f.codim:
        raise ValueError(f'True and False statistics for IfElse must have matching dim and codim')  # ATTN:Replace with custom error for catching at repl
    def ifelse(*x):
        if cond(*x):
            return t(*x)
        else:
            return f(*x)
    return Statistic(ifelse, dim=cond.arity, codim=t.codim, name=f'returns {t.name} if {cond.name} is true else returns {f.name}')
        
def Not(s: Statistic) -> Condition:
    #ATTN: require s.codim == 1
    return Condition(lambda *x: 1 - s(*x), dim=s.arity, name=f'not({s.name})', description=f'returns the logical not of {s.name}')

# def And(s1: Statistic, s2: Statistic) -> Statistic:
#     #ATTN: require si.codim == 1 and s1.arity == s2.arity
#     return Statistic(lambda *x: s1(*x) and s2(*x), dim=s1.arity, codim=1, name=f'Logical And of {s1.name} and {s2.name}')
#  
# def Or(s1: Statistic, s2: Statistic) -> Statistic:
#     #ATTN: require si.codim == 1 and s1.arity == s2.arity
#     return Statistic(lambda *x: s1(*x) or s2(*x), dim=s1.arity, codim=1, name=f'Logical Or of {s1.name} and {s2.name}')

def And(*stats: Statistic) -> Condition:
    #ATTN: require si.codim == 1 and s1.arity == s2.arity
    def and_of(*x):
        val = True
        for s in stats:
            val = val and s(*x)
            if not val:
                break
        return val
    labels = ["'" + s.name + "'" for s in stats]
    return Condition(and_of, dim=max(s.arity for s in stats), name=f'({" and ".join(labels)})', description=f'returns the logical and of {", ".join(labels)}')

def Or(*stats: Statistic) -> Condition:
    #ATTN: require si.codim == 1 and s1.arity == s2.arity
    def and_of(*x):
        val = False
        for s in stats:
            val = val or s(*x)
            if val:
                break
        return val
    labels = ["'" + s.name + "'" for s in stats]
    return Condition(and_of, dim=max(s.arity for s in stats), name=f'({" or ".join(labels)})', description=f'returns the logical or of {", ".join(labels)}')

top = Condition(lambda _x: True, name='top', description='returns true for any value')

bottom = Condition(lambda _x: False, name='bottom', description='returns false for any value')


## ATTN: These should really be methods of ProjectionStatistic
## There should be no need for a callable argment in that constructor.
@overload
def project(*__indices: int) -> ProjectionStatistic:
    ...

@overload
def project(__index_tuple: Iterable[int]) -> ProjectionStatistic:
    ...

def project(*indices_or_tuple) -> ProjectionStatistic:
    """Creates a projection statistic that extracts the specified components.

       Positional variadic arguments:
         *indices_or_tuple -- a tuple of integer indices starting from 1 or a single int tuple
    """
    if len(indices_or_tuple) == 0:  # ATTN:Error here instead?
        return ProjectionStatistic(lambda _: (), (), name='Null projection')

    # ATTN:Support slice objects here
    # In that sense, it would be good if the projection statistic could also get
    # the dimension of the input tuple, then we could use Proj[2:-1] to mean
    # all but the first and Proj[1:-2] for all but the last regardless of
    # dimension.

    if isinstance(indices_or_tuple[0], slice):
        def dec_or_none(x: int | None) -> int | None:
            if x is not None and x > 0:
                return x - 1
            return x
        zindexed = indices_or_tuple[0]
        indices: slice | Iterable = slice(dec_or_none(zindexed.start),
                                          dec_or_none(zindexed.stop),
                                          zindexed.step)
        def get_indices(xs):
            return as_vec_tuple(xs[indices])
        label = str(indices)
    else:
        if isinstance(indices_or_tuple[0], Iterable):
            indices = indices_or_tuple[0]
        else:
            indices = indices_or_tuple

        def get_indices(xs):
            getter = itemgetter(*[x - 1 if x > 0 else x for x in indices if x != 0])
            return as_vec_tuple(getter(xs))
        label = ", ".join(map(str, indices))
    return ProjectionStatistic(
        get_indices,
        indices,
        name=f'project[{label}]')


class ProjectionFactory:
    @overload
    def __call__(self, *__indices: int) -> ProjectionStatistic:
        ...
     
    @overload
    def __call__(self, __index_tuple: Iterable[int]) -> ProjectionStatistic:
        ...
     
    @overload
    def __call__(self, __index_slice: slice) -> ProjectionStatistic:
        ...
     
    def __call__(self, *indices_or_tuple) -> ProjectionStatistic:
        return project(*indices_or_tuple)

    @overload
    def __getitem__(self, *__indices: int) -> ProjectionStatistic:
        ...
     
    @overload
    def __getitem__(self, __index_tuple: Iterable[int]) -> ProjectionStatistic:
        ...

    @overload
    def __getitem__(self, __index_slice: slice) -> ProjectionStatistic:
        ...
     
    def __getitem__(self, *indices_or_tuple) -> ProjectionStatistic:
        return project(*indices_or_tuple)

Proj = ProjectionFactory()

