import jax
import jax.numpy as jnp
from typing import Mapping, Protocol


def ravelize_function(f, pytree):
    """
    Takes a function that accepts a PyTree and a PyTree,
    and produces a function that accepts a flat array.
    """
    # note: ravel_pytree is only really safe when we
    # know all the dtypes are the same. See
    # https://jax.readthedocs.io/en/latest/_autosummary/jax.flatten_util.ravel_pytree.html
    # This is usually true in stats models
    _, unravel = jax.flatten_util.ravel_pytree(pytree)
    return lambda x: f(unravel(x))


class Shaped(Protocol):
    shape: tuple[int, ...]
    dtype: jnp.dtype


class ParameterConstraint:
    def __init__(self, shape=(), dtype=jnp.float32):
        # if (isinstance(shape, Iterable) and any(s < 0 for s in shape)) or (shape < 0):
            # raise ValueError("Shape dimensions must be non-negative")

        self.shape = shape
        self.dtype = dtype

    def __call__(self, x):
        return x

    def inverse(self, y):
        return y

    def jacobian(self, _):
        return 0.0


# simple alias: base class does transforms
real = ParameterConstraint


# basic example of a positive constraint
class positive(ParameterConstraint):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    def __call__(self, x):
        return jnp.exp(x)

    def inverse(self, y):
        return jnp.log(y)

    def jacobian(self, x):
        return jnp.sum(x)


# note: shape will need some care for something like a simplex,
# which should also include axis=... in its definition. The shape
# should be that of the unconstrained parameter.
class simplex(ParameterConstraint):
    def __init__(self, shape, axes=..., **kwargs):
        self.axes = axes
        shape_n = jnp.array(shape, dtype=int)
        shape_n = shape_n.at[axes,].set(shape_n[axes,] - 1)

        super().__init__(shape=tuple(shape_n), **kwargs)

    def __call__(self, x):
        raise NotImplementedError("Simplex constraint not yet implemented")

    def inverse(self, y):
        raise NotImplementedError("Simplex constraint not yet implemented")

    def jacobian(self, x):
        raise NotImplementedError("Simplex constraint not yet implemented")


def spec_to_pytree(
    parameter_spec: Mapping[str, Shaped],
) -> dict[str, jnp.ndarray]:
    """
    Turns a dictionary from parameter names to shapes/dtypes
    into a pytree of zeros with those shapes/dtypes.

    Note that the Shaped protocol is satisfied by ParameterConstraint
    objects, but also by something like a jax.numpy array.
    """
    return jax.tree.map(lambda x: jnp.zeros(x.shape, dtype=x.dtype), parameter_spec)

