"""Provides scalar multipliers based on the Non Adjacent Form (NAF) recoding."""

from copy import copy
from typing import Optional, List, MutableMapping
from public import public

from pyecsca.ec.mult.base import (
    ScalarMultiplier,
    ScalarMultiplicationAction,
    ProcessingDirection,
    AccumulationOrder,
    PrecomputationAction,
    AccumulatorMultiplier,
    PrecompMultiplier,
)
from pyecsca.ec.formula import (
    AdditionFormula,
    DoublingFormula,
    ScalingFormula,
    NegationFormula,
)
from pyecsca.ec.params import DomainParameters
from pyecsca.ec.point import Point
from pyecsca.ec.scalar import naf, wnaf


@public
class BinaryNAFMultiplier(AccumulatorMultiplier, PrecompMultiplier, ScalarMultiplier):
    """
    Binary NAF (Non Adjacent Form) multiplier.

    [GECC]_ Algorithm 3.31.

    :param always: Whether the addition is always performed.
    :param direction: Whether it is LTR or RTL.
    :param accumulation_order: The order of accumulation of points.
    :param short_circuit: Whether the use of formulas will be guarded by short-circuit on inputs
                      of the point at infinity.
    :param complete: Whether it starts processing at full order-bit-length.
    """

    requires = {AdditionFormula, DoublingFormula, NegationFormula}
    optionals = {ScalingFormula}
    always: bool
    """Whether the double and add always method is used."""
    direction: ProcessingDirection
    """Whether it is LTR or RTL."""
    complete: bool
    """Whether it starts processing at full order-bit-length."""
    _point_neg: Point

    def __init__(
        self,
        add: AdditionFormula,
        dbl: DoublingFormula,
        neg: NegationFormula,
        scl: Optional[ScalingFormula] = None,
        always: bool = False,
        direction: ProcessingDirection = ProcessingDirection.LTR,
        accumulation_order: AccumulationOrder = AccumulationOrder.PeqPR,
        complete: bool = True,
        short_circuit: bool = True,
    ):
        super().__init__(
            short_circuit=short_circuit,
            accumulation_order=accumulation_order,
            add=add,
            dbl=dbl,
            neg=neg,
            scl=scl,
        )
        self.always = always
        self.direction = direction
        self.complete = complete

    def __hash__(self):
        return hash(
            (
                BinaryNAFMultiplier,
                super().__hash__(),
                self.always,
                self.direction,
                self.accumulation_order,
                self.complete,
            )
        )

    def __eq__(self, other):
        if not isinstance(other, BinaryNAFMultiplier):
            return False
        return (
            self.formulas == other.formulas
            and self.always == other.always
            and self.short_circuit == other.short_circuit
            and self.direction == other.direction
            and self.accumulation_order == other.accumulation_order
            and self.complete == other.complete
        )

    def __repr__(self):
        return f"{self.__class__.__name__}({', '.join(map(str, self.formulas.values()))}, short_circuit={self.short_circuit}, direction={self.direction.name}, accumulation_order={self.accumulation_order.name}, always={self.always}, complete={self.complete})"

    def init(self, params: DomainParameters, point: Point, bits: Optional[int] = None):
        with PrecomputationAction(params, point) as action:
            super().init(params, point, bits)
            self._point_neg = self._neg(point)
            action.exit({-1: self._point_neg})

    def _ltr(self, scalar_naf: List[int]) -> Point:
        if self.complete:
            q = copy(self._params.curve.neutral)
            while (
                len(scalar_naf) < self._bits + 1
            ):  # Pad with zeros, naf(k) can have up to one more entry han bit-length of k ([GECC]_ Theorem 3.29)
                scalar_naf.insert(0, 0)
        else:
            while scalar_naf[0] == 0:
                scalar_naf.pop(0)
            val = scalar_naf.pop(0)
            if val == 1:
                q = copy(self._point)
            elif val == -1:
                q = copy(self._point_neg)
            else:
                raise ValueError("Should not happen.")
        for val in scalar_naf:
            q = self._dbl(q)
            orig = q
            if val == 1:
                q = self._accumulate(q, self._point)
                if self.always:
                    self._accumulate(orig, self._point_neg)
            elif val == -1:
                # TODO: Whether this negation is precomputed can be a parameter
                q = self._accumulate(q, self._point_neg)
                if self.always:
                    self._accumulate(orig, self._point)
        return q

    def _rtl(self, scalar_naf: List[int]) -> Point:
        q = self._point
        r = copy(self._params.curve.neutral)
        if self.complete:
            while (
                len(scalar_naf) < self._bits + 1
            ):  # Pad with zeros, naf(k) can have up to one more entry than bit-length of k ([GECC]_ Theorem 3.29)
                scalar_naf.insert(0, 0)
        else:
            # Nothing to do here. We could skip the zeros and assign r to the first non-zero value.
            pass
        for val in reversed(scalar_naf):
            orig = r
            if val == 1:
                r = self._accumulate(r, q)
                if self.always:
                    neg = self._neg(q)
                    self._accumulate(orig, neg)
            elif val == -1:
                neg = self._neg(q)
                r = self._accumulate(r, neg)
                if self.always:
                    self._accumulate(orig, q)
            q = self._dbl(q)
        return r

    def multiply(self, scalar: int) -> Point:
        if not self._initialized:
            raise ValueError("ScalarMultiplier not initialized.")
        with ScalarMultiplicationAction(self._point, self._params, scalar) as action:
            if scalar == 0:
                return action.exit(copy(self._params.curve.neutral))
            scalar_naf = naf(scalar)
            if self.direction is ProcessingDirection.LTR:
                q = self._ltr(scalar_naf)
            elif self.direction is ProcessingDirection.RTL:
                q = self._rtl(scalar_naf)
            if "scl" in self.formulas:
                q = self._scl(q)
            return action.exit(q)


@public
class WindowNAFMultiplier(AccumulatorMultiplier, PrecompMultiplier, ScalarMultiplier):
    """
    Window NAF (Non Adjacent Form) multiplier, left-to-right.

    [GECC]_ Algorithm 3.36.

    A right-to-left variant does not make much sense.

    :param short_circuit: Whether the use of formulas will be guarded by short-circuit on inputs
                          of the point at infinity.
    :param width: The width of the window.
    :param accumulation_order: The order of accumulation of points.
    :param precompute_negation: Whether to precompute the negation of the precomputed points as well.
                                It is computed on the fly otherwise.
    :param complete: Whether it starts processing at full order-bit-length.
    """

    requires = {AdditionFormula, DoublingFormula, NegationFormula}
    optionals = {ScalingFormula}
    _points: MutableMapping[int, Point]
    _points_neg: MutableMapping[int, Point]
    precompute_negation: bool = False
    """Whether to precompute the negation of the precomputed points as well."""
    width: int
    """The width of the window."""
    complete: bool
    """Whether it starts processing at full order-bit-length."""

    def __init__(
        self,
        add: AdditionFormula,
        dbl: DoublingFormula,
        neg: NegationFormula,
        width: int,
        scl: Optional[ScalingFormula] = None,
        accumulation_order: AccumulationOrder = AccumulationOrder.PeqPR,
        precompute_negation: bool = False,
        complete: bool = True,
        short_circuit: bool = True,
    ):
        super().__init__(
            short_circuit=short_circuit,
            accumulation_order=accumulation_order,
            add=add,
            dbl=dbl,
            neg=neg,
            scl=scl,
        )
        self.width = width
        self.precompute_negation = precompute_negation
        self.complete = complete

    def __hash__(self):
        return hash(
            (
                WindowNAFMultiplier,
                super().__hash__(),
                self.width,
                self.precompute_negation,
                self.accumulation_order,
                self.complete,
            )
        )

    def __eq__(self, other):
        if not isinstance(other, WindowNAFMultiplier):
            return False
        return (
            self.formulas == other.formulas
            and self.short_circuit == other.short_circuit
            and self.width == other.width
            and self.precompute_negation == other.precompute_negation
            and self.accumulation_order == other.accumulation_order
            and self.complete == other.complete
        )

    def __repr__(self):
        return f"{self.__class__.__name__}({', '.join(map(str, self.formulas.values()))}, short_circuit={self.short_circuit}, width={self.width}, precompute_negation={self.precompute_negation}, accumulation_order={self.accumulation_order.name}, complete={self.complete})"

    def init(self, params: DomainParameters, point: Point, bits: Optional[int] = None):
        with PrecomputationAction(params, point) as action:
            super().init(params, point, bits)
            self._points = {}
            self._points_neg = {}
            current_point = point
            double_point = self._dbl(point)
            for i in range(0, 2 ** (self.width - 2)):
                self._points[2 * i + 1] = current_point
                if self.precompute_negation:
                    self._points_neg[2 * i + 1] = self._neg(current_point)
                current_point = self._add(current_point, double_point)
            action.exit({**self._points, **self._points_neg})

    def multiply(self, scalar: int) -> Point:
        if not self._initialized:
            raise ValueError("ScalarMultiplier not initialized.")
        with ScalarMultiplicationAction(self._point, self._params, scalar) as action:
            if scalar == 0:
                return action.exit(copy(self._params.curve.neutral))
            scalar_naf = wnaf(scalar, self.width)
            if self.complete:
                q = copy(self._params.curve.neutral)
                while (
                    len(scalar_naf) < self._bits + 1
                ):  # Pad with zeros, wnaf(k) can have up to one more entry han bit-length of k ([GECC]_ Theorem 3.33)
                    scalar_naf.insert(0, 0)
            else:
                while scalar_naf[0] == 0:
                    scalar_naf.pop(0)
                val = scalar_naf.pop(0)
                q = copy(self._points[val])
            for val in scalar_naf:
                q = self._dbl(q)
                if val > 0:
                    q = self._accumulate(q, self._points[val])
                elif val < 0:
                    if self.precompute_negation:
                        neg = self._points_neg[-val]
                    else:
                        neg = self._neg(self._points[-val])
                    q = self._accumulate(q, neg)
            if "scl" in self.formulas:
                q = self._scl(q)
            return action.exit(q)