from __future__ import annotations

from typing import Dict, Union, Mapping, Protocol, Any, Optional, List

import numpy as np

class _TensorLike(Protocol):
    """Protocol defining what methods and operations a Generic Tensor can perform."""

    def __mul__(self: _TensorLike, other: Any) -> _TensorLike:
        """Multiplies the self tensor with the other tensor and returns the result."""
        ...

    def __sub__(self: _TensorLike, other: Any) -> _TensorLike:
        """Subtracts the other tensor from the self tensor and returns the result."""
        ...

    def squeeze(self: _TensorLike, axis: Optional[Any] = None) -> _TensorLike:
        """Returns a tensor with all the dimensions of input of size 1 removed."""
        ...

_Weights = Mapping[str, _TensorLike]

class Test:
    _other_worker_shares: List[int]

    def _encode_secret(self, secret: Union[_TensorLike, np.ndarray]) -> np.ndarray:
        return np.ones(5)

    def _reconstruct_secret(self, shares: List[Union[np.ndarray, int]]) -> np.ndarray:
        return np.ones(5)

    def _encode_and_reconstruct_state_dict(
        self, secret_state_dict: Union[_Weights, Mapping[str, np.ndarray]]
    ) -> Dict[str, np.ndarray]:
        """Encodes and reconstructs state dict from own and worker shares.

        Encrypts `secret_state_dict` using `own_shares` then reconstructs it using
        `self._other_worker_shares`.
        """
        encrypted_state_dict = [
            self._encode_secret(v) for _, v in secret_state_dict.items()
        ]
        reconstructed = [
            self._reconstruct_secret([*self._other_worker_shares, value])
            for value in encrypted_state_dict
        ]
        x = list(secret_state_dict)
        z = zip(x, reconstructed)
        d = dict(z)
        return d