"""Typed local rules shared by numerical evaluation and all presentations."""
from dataclasses import dataclass, replace

import numpy as np

from math_expression import Expr, adjoint, literal, node, operation_expression, ref


@dataclass(frozen=True)
class ScalarLocalRule:
    formula: Expr
    operand: Expr
    derivative: Expr
    mapped_derivative: Expr
    application: Expr
    bindings: tuple[tuple[Expr, Expr], ...] = ()
    notes: tuple[str, ...] = ()


@dataclass(frozen=True)
class TensorBackwardRule:
    formula: Expr
    operand: Expr
    application: Expr
    notes: tuple[str, ...] = ()


def scalar_rule(operation, occurrence, derivative, slot_derivative=None, slots=()):
    repeated = sum(q.uid == operation.operands[occurrence].uid for q in operation.operands) > 1
    args = tuple(ref(q) for q in operation.operands)
    bindings = tuple(zip(slots, args)) if repeated else ()
    return ScalarLocalRule(
        operation_expression(operation, slots if repeated else args),
        slots[occurrence] if repeated else args[occurrence],
        slot_derivative if repeated else derivative,
        derivative, node('mul', adjoint(operation.quantity), derivative), bindings)


def tensor_rule(operation, occurrence):
    """Construct one local pullback. Its expression is also its implementation."""
    from scalar_rules import operand_slots
    slots = operand_slots(operation)
    args = slots
    shapes = [q.shape for q in operation.operands]
    child = operation.operands[occurrence]
    g, y = adjoint(operation.quantity), ref(operation.quantity)
    x = args[occurrence]
    kind, function = operation.kind, operation.function
    notes = ()
    factor = None
    product = 'elementwise' if operation.quantity.shape else 'scale'
    if kind == 'matmul':
        a, b = args
        ranks = tuple(map(len, shapes))
        if ranks == (1, 1):
            application = node('scale', g, args[1 - occurrence])
        elif ranks == (2, 1):
            application = node('outer', g, b) if occurrence == 0 else node('matmul', node('transpose', a), g)
        elif ranks == (1, 2):
            application = node('matmul', b, g) if occurrence == 0 else node('outer', a, g)
        elif ranks == (2, 2):
            application = node('matmul', g, node('transpose', b)) if occurrence == 0 else node('matmul', node('transpose', a), g)
        else:
            raise ValueError(f'Unsupported matrix product ranks: {ranks}')
    elif function == 'softmax':
        application = node('elementwise', y, node('sub', g, node('sum', node('elementwise', g, y))))
    else:
        if kind == 'neg':
            factor = literal(-1)
        elif kind == 'add':
            factor = literal(1)
        elif kind == 'sub':
            factor = literal(1 if occurrence == 0 else -1)
        elif kind in {'mul', 'dot'}:
            factor = args[1 - occurrence]
        elif kind == 'div':
            factor = node('div', literal(1), args[1]) if occurrence == 0 else node('div', node('neg', args[0]), node('pow', args[1], literal(2)))
        elif kind == 'pow':
            p = operation.exponent
            factor = (node('zeros_like', x) if child.shape else literal(0)) if p == 0 else node('mul', literal(p), node('pow', x, literal(p - 1)))
        elif function == 'pick':
            factor = node('one_hot', x, literal(operation.index))
            notes = ('Indices are zero-based; negative indices count from the end.',)
        elif function in {'sum', 'mean'}:
            factor = node('ones_like', x) if child.shape else literal(1)
            if function == 'mean' and child.shape:
                factor = node('div', factor, node('size', x))
                notes = (f'N = number of entries = {int(np.prod(child.shape))}',)
        elif function in {'max', 'min'}:
            mask = node('equal', x, y)
            factor = node('div', mask, node('sum', mask)) if child.shape else literal(1)
            notes = ('Equal-share subgradient at ties.',)
        elif function == 'sigmoid':
            factor = node('elementwise' if child.shape else 'mul', y, node('sub', literal(1), y))
        elif function == 'relu':
            factor = node('positive', x)
            notes = ('At zero, the chosen derivative is 0.',)
        elif function == 'tanh':
            factor = node('sub', literal(1), node('pow', y, literal(2)))
        elif function == 'exp':
            factor = y
        elif function == 'log':
            factor = node('div', literal(1), x)
        if factor is None:
            raise ValueError(f'Missing backward rule for {kind}/{function}')
        # A zero exponent is constant even if the incoming adjoint is nonfinite.
        application = factor if kind == 'pow' and operation.exponent == 0 else node(product, g, factor)
        if not child.shape and operation.quantity.shape:
            application = node('sum', application)
            notes += ('Sum over broadcast entries.',)
    bindings = {slot.key: ref(q) for slot, q in zip(slots, operation.operands)}
    application = application.substitute(bindings)
    if not operation.quantity.shape and not any(shapes) and factor is not None:
        # Scalar operations inside a tensor graph use the same scalar contract.
        rule = scalar_rule(operation, occurrence, factor.substitute(bindings), factor, slots)
        return replace(rule, application=application, notes=notes)
    return TensorBackwardRule(operation.formula, ref(child), application, notes)