import inspect
import torch
import itertools
import types
from typing import Union, Tuple, Optional
import geoopt.utils
from geoopt.manifolds.base import Manifold, ScalingInfo
import functools

__all__ = ["Scaled"]

def rescale_value(value, scaling, power):
    return value * scaling ** power if power != 0 else value

def rescale(function, scaling_info):
    if scaling_info is ScalingInfo.NotCompatible:

        def stub(self, *args, **kwargs):
            raise NotImplementedError(
                "Scaled version of '{}' is not available".format(function.__name__)

        return stub
    signature = inspect.signature(function)

    def rescaled_function(self, *args, **kwargs):
        params = signature.bind(self.base, *args, **kwargs)
        arguments = params.arguments
        for k, power in scaling_info.kwargs.items():
            arguments[k] = rescale_value(arguments[k], self.scale, power)
        params = params.__class__(signature, arguments)
        results = function(*params.args, **params.kwargs)
        if not scaling_info.results:
            # do nothing
            return results
        wrapped_results = []
        is_tuple = isinstance(results, tuple)
        results = geoopt.utils.make_tuple(results)
        for i, (res, power) in enumerate(
            itertools.zip_longest(results, scaling_info.results, fillvalue=0)
            wrapped_results.append(rescale_value(res, self.scale, power))
        if not is_tuple:
            wrapped_results = wrapped_results[0]
            wrapped_results = results.__class__(wrapped_results)
        return wrapped_results

    return rescaled_function

class Scaled(Manifold):
    Scaled manifold.

    Scales all the distances on tha manifold by a constant factor. Scaling may be learnable
    since the underlying representation is canonical.

    Here is a simple example of radius 2 Sphere

    >>> import geoopt, torch, numpy as np
    >>> sphere = geoopt.Sphere()
    >>> radius_2_sphere = Scaled(sphere, 2)
    >>> p1 = torch.tensor([-1., 0.])
    >>> p2 = torch.tensor([0., 1.])
    >>> np.testing.assert_allclose(sphere.dist(p1, p2), np.pi / 2)
    >>> np.testing.assert_allclose(radius_2_sphere.dist(p1, p2), np.pi)

    def __init__(self, manifold: Manifold, scale=1.0, learnable=False):
        self.base = manifold
        scale = torch.as_tensor(scale, dtype=torch.get_default_dtype())
        scale = scale.requires_grad_(False)
        if not learnable:
            self.register_buffer("_scale", scale)
            self.register_buffer("_log_scale", None)
            self.register_buffer("_scale", None)
            self.register_parameter("_log_scale", torch.nn.Parameter(scale.log()))
        # do not rebuild scaled functions very frequently, save them

        for method, scaling_info in self.base.__scaling__.items():
            # register rescaled functions as bound methods of this particular instance
            unbound_method = getattr(self.base, method).__func__  # unbound method
                method, types.MethodType(rescale(unbound_method, scaling_info), self)

    def scale(self) -> torch.Tensor:
        if self._scale is None:
            return self._log_scale.exp()
            return self._scale

    def log_scale(self) -> torch.Tensor:
        if self._log_scale is None:
            return self._scale.log()
            return self._log_scale

    # propagate all important stuff
    reversible = property(lambda self: self.base.reversible)
    ndim = property(lambda self: self.base.ndim)
    name = "Scaled"
    __scaling__ = property(lambda self: self.base.__scaling__)

    # Make AbstractMeta happy, to be fixed in __init__
    retr = NotImplemented
    expmap = NotImplemented

    def __getattr__(self, item):
            return super().__getattr__(item)
        except AttributeError as original:
                # propagate only public methods and attributes, ignore buffers, parameters, etc
                if isinstance(self.base, Scaled) and item in self._base_attributes:
                    return self.base.__getattr__(item)
                    return self.base.__getattribute__(item)
            except AttributeError as e:
                raise original from e

    def _base_attributes(self):
        if isinstance(self.base, Scaled):
            return self.base._base_attributes
            base_attributes = set(dir(self.base.__class__))
            base_attributes |= set(self.base.__dict__.keys())
            return base_attributes

    def __dir__(self):
        return list(set(super().__dir__()) | self._base_attributes)

    def __repr__(self):
        extra = self.base.extra_repr()
        if extra:
            return + "({})({}) manifold".format(, extra)
            return + "({}) manifold".format(

    def _check_shape(self, shape: Tuple[int], name: str) -> Tuple[bool, Optional[str]]:
        return self.base._check_shape(shape, name)

    def _check_point_on_manifold(
        self, x: torch.Tensor, *, atol=1e-5, rtol=1e-5
    ) -> Union[Tuple[bool, Optional[str]], bool]:
        return self.base._check_point_on_manifold(x, atol=atol, rtol=rtol)

    def _check_vector_on_tangent(
        self, x: torch.Tensor, u: torch.Tensor, *, atol=1e-5, rtol=1e-5
    ) -> Union[Tuple[bool, Optional[str]], bool]:
        return self.base._check_vector_on_tangent(x, u, atol=atol, rtol=rtol)

    # stuff that should remain the same but we need to override it
    def inner(
        x: torch.Tensor,
        u: torch.Tensor,
        v: torch.Tensor = None,
    ) -> torch.Tensor:
        return self.base.inner(x, u, v, keepdim=keepdim, **kwargs)

    def norm(
        self, x: torch.Tensor, u: torch.Tensor, *, keepdim=False, **kwargs
    ) -> torch.Tensor:
        return self.base.norm(x, u, keepdim=keepdim, **kwargs)

    def proju(self, x: torch.Tensor, u: torch.Tensor, **kwargs) -> torch.Tensor:
        return self.base.proju(x, u, **kwargs)

    def projx(self, x: torch.Tensor, **kwargs) -> torch.Tensor:
        return self.base.projx(x, **kwargs)

    def egrad2rgrad(self, x: torch.Tensor, u: torch.Tensor, **kwargs) -> torch.Tensor:
        return self.base.egrad2rgrad(x, u, **kwargs)

    def transp(
        self, x: torch.Tensor, y: torch.Tensor, v: torch.Tensor, **kwargs
    ) -> torch.Tensor:
        return self.base.transp(x, y, v, **kwargs)

    def random(self, *size, dtype=None, device=None, **kwargs) -> torch.Tensor:
        return self.base.random(*size, dtype=dtype, device=device, **kwargs)