from math import pi

import keras
from keras import ops

from bayesflow.types import Tensor
from bayesflow.utils import logging, jvp, find_network, expand_right_as, expand_right_to, layer_kwargs
from bayesflow.utils.serialization import deserialize, serializable, serialize

from ..inference_network import InferenceNetwork


@serializable("bayesflow.networks")
class StableConsistencyModel(InferenceNetwork):
    """(IN) Implements an sCM (simple, stable, and scalable Consistency Model) with continuous-time Consistency Training
    (CT) as described in [1]. The sampling procedure is taken from [2].

    [1] Lu, C., & Song, Y. (2024).
    Simplifying, Stabilizing and Scaling Continuous-Time Consistency Models
    arXiv preprint arXiv:2410.11081

    [2] Song, Y., Dhariwal, P., Chen, M. & Sutskever, I. (2023).
    Consistency Models. arXiv preprint arXiv:2303.01469
    """

    TIME_MLP_DEFAULT_CONFIG = {
        "widths": (256, 256, 256, 256, 256),
        "activation": "mish",
        "kernel_initializer": "he_normal",
        "residual": True,
        "dropout": 0.05,
        "spectral_normalization": False,
        "time_embedding_dim": 32,
        "merge": "concat",
        "norm": "layer",
    }

    WEIGHT_MLP_DEFAULT_CONFIG = {
        "widths": (256,),
        "activation": "mish",
        "kernel_initializer": "he_normal",
        "residual": False,
        "dropout": 0.05,
        "spectral_normalization": False,
    }

    EPS_WARN = 0.1

    def __init__(
        self,
        subnet: str | type | keras.Layer = "time_mlp",
        sigma: float = 1.0,
        subnet_kwargs: dict[str, any] = None,
        weight_mlp_kwargs: dict[str, any] = None,
        **kwargs,
    ):
        """Creates an instance of an sCM to be used for consistency training (CT).

        Parameters
        ----------
        subnet : str, type, or keras.Layer, optional, default="time_mlp"
            The neural network architecture used for the consistency model.
            If a string is provided, it should be a registered name (e.g., "time_mlp").
            If a type or keras.Layer is provided, it will be directly instantiated
            with the given ``subnet_kwargs``. Any subnet must accept a tuple of tensors (target, time, conditions).
        sigma : float, optional, default=1.0
            Standard deviation of the target distribution for the consistency loss.
            Controls the scale of the noise injected during training.
        subnet_kwargs : dict[str, any], optional, default=None
            Keyword arguments passed to the constructor of the chosen ``subnet``. For example, number of hidden units,
            activation functions, or dropout settings.
        weight_mlp_kwargs : dict[str, any], optional, default=None
            Keyword arguments for an auxiliary MLP used to generate weights within the consistency model. Typically,
            includes depth, hidden sizes, and non-linearity choices.
        **kwargs
            Additional keyword arguments passed to the parent ``InferenceNetwork`` initializer
            (e.g., ``name``, ``dtype``, or ``trainable``).
        """
        super().__init__(base_distribution="normal", **kwargs)

        subnet_kwargs = subnet_kwargs or {}
        if subnet == "time_mlp":
            subnet_kwargs = StableConsistencyModel.TIME_MLP_DEFAULT_CONFIG | subnet_kwargs
        self.subnet = find_network(subnet, **subnet_kwargs)

        self.subnet_projector = None

        weight_mlp_kwargs = weight_mlp_kwargs or {}
        weight_mlp_kwargs = StableConsistencyModel.WEIGHT_MLP_DEFAULT_CONFIG | weight_mlp_kwargs
        self.weight_fn = find_network("mlp", **weight_mlp_kwargs)

        self.weight_fn_projector = keras.layers.Dense(
            units=1, bias_initializer="zeros", kernel_initializer="zeros", name="weight_fn_projector"
        )

        self.sigma = sigma
        self.seed_generator = keras.random.SeedGenerator()

    @classmethod
    def from_config(cls, config, custom_objects=None):
        return cls(**deserialize(config, custom_objects=custom_objects))

    def get_config(self):
        base_config = super().get_config()
        base_config = layer_kwargs(base_config)

        config = {
            "subnet": self.subnet,
            "sigma": self.sigma,
        }

        return base_config | serialize(config)

    @staticmethod
    def _discretize_time(num_steps: int, rho: float = 3.5):
        t = keras.ops.linspace(0.0, pi / 2, num_steps)
        times = keras.ops.exp((t - pi / 2) * rho) * pi / 2
        times = keras.ops.concatenate([keras.ops.zeros((1,)), times[1:]], axis=0)

        # if rho is set too low, bad schedules can occur
        if times[1] > StableConsistencyModel.EPS_WARN:
            logging.warning("Warning: The last time step is large.")
            logging.warning(f"Increasing rho (was {rho}) or n_steps (was {num_steps}) might improve results.")
        return times

    def build(self, xz_shape, conditions_shape=None):
        if self.built:
            # building when the network is already built can cause issues with serialization
            # see https://github.com/keras-team/keras/issues/21147
            return

        self.base_distribution.build(xz_shape)

        self.subnet_projector = keras.layers.Dense(
            units=xz_shape[-1],
            bias_initializer="zeros",
            name="output_projector",
        )

        # construct input shape for subnet and subnet projector
        time_shape = tuple(xz_shape[:-1]) + (1,)  # same batch/sequence dims, 1 feature
        self.subnet.build((xz_shape, time_shape, conditions_shape))
        input_shape = self.subnet.compute_output_shape((xz_shape, time_shape, conditions_shape))
        self.subnet_projector.build(input_shape)

        # input shape for weight function and projector
        input_shape = (xz_shape[0], 1)
        self.weight_fn.build(input_shape)
        input_shape = self.weight_fn.compute_output_shape(input_shape)
        self.weight_fn_projector.build(input_shape)

    def _forward(self, x: Tensor, conditions: Tensor = None, **kwargs) -> Tensor:
        # Consistency Models only learn the direction from noise distribution
        # to target distribution, so we cannot implement this function.
        raise NotImplementedError("Consistency Models are not invertible")

    def _inverse(self, z: Tensor, conditions: Tensor = None, **kwargs) -> Tensor:
        """Generate random draws from the approximate target distribution
        using the multistep sampling algorithm from [2], Algorithm 1.

        Parameters
        ----------
        z           : Tensor
            Samples from a standard normal distribution
        conditions  : Tensor, optional, default: None
            Conditions for an approximate conditional distribution
        **kwargs    : dict, optional, default: {}
            Additional keyword arguments. Include `steps` (default: 15) to
            adjust the number of sampling steps.

        Returns
        -------
        x            : Tensor
            The approximate samples
        """
        steps = kwargs.get("steps", 15)
        rho = kwargs.get("rho", 3.5)

        # noise distribution has variance sigma
        x = keras.ops.copy(z) * self.sigma
        discretized_time = keras.ops.flip(self._discretize_time(steps, rho=rho), axis=-1)
        t = keras.ops.full((*keras.ops.shape(x)[:-1], 1), discretized_time[0], dtype=x.dtype)
        x = self.consistency_function(x, t, conditions=conditions)
        for n in range(1, steps):
            noise = keras.random.normal(keras.ops.shape(x), dtype=keras.ops.dtype(x), seed=self.seed_generator)
            x_n = ops.cos(t) * x + ops.sin(t) * noise
            t = keras.ops.full_like(t, discretized_time[n])
            x = self.consistency_function(x_n, t, conditions=conditions)
        return x

    def consistency_function(
        self,
        x: Tensor,
        t: Tensor,
        conditions: Tensor = None,
        training: bool = False,
    ) -> Tensor:
        """Compute consistency function at time t.

        Parameters
        ----------
        x           : Tensor
            Input vector
        t           : Tensor
            Vector of time samples in [0, pi/2]
        conditions  : Tensor
            The conditioning vector
        training    : bool
            Flag to control whether the inner network operates in training or test mode
        """
        subnet_out = self.subnet((x / self.sigma, t, conditions), training=training)
        f = self.subnet_projector(subnet_out)
        out = ops.cos(t) * x - ops.sin(t) * self.sigma * f
        return out

    def compute_metrics(
        self, x: Tensor, conditions: Tensor = None, stage: str = "training", **kwargs
    ) -> dict[str, Tensor]:
        base_metrics = super().compute_metrics(x, conditions=conditions, stage=stage)

        # $# Implements Algorithm 1 from [1]

        # training parameters
        p_mean = -1.0
        p_std = 1.6

        c = 0.1

        # generate noise vector
        z = keras.random.normal(keras.ops.shape(x), dtype=keras.ops.dtype(x), seed=self.seed_generator) * self.sigma

        # sample time
        tau = (
            keras.random.normal(keras.ops.shape(x)[:1], dtype=keras.ops.dtype(x), seed=self.seed_generator) * p_std
            + p_mean
        )
        t_ = ops.arctan(ops.exp(tau) / self.sigma)
        t = expand_right_as(t_, x)

        # generate noisy sample
        xt = ops.cos(t) * x + ops.sin(t) * z

        # calculate estimator for dx_t/dt
        dxtdt = ops.cos(t) * z - ops.sin(t) * x

        r = 1.0  # TODO: if consistency distillation training (not supported yet) is unstable, add schedule here

        def f_teacher(x, t):
            o = self.subnet((x, t, conditions), training=stage == "training")
            return self.subnet_projector(o)

        primals = (xt / self.sigma, t)
        tangents = (
            ops.cos(t) * ops.sin(t) * dxtdt,
            ops.cos(t) * ops.sin(t) * self.sigma,
        )

        teacher_output, cos_sin_dFdt = jvp(f_teacher, primals, tangents, return_output=True)
        teacher_output = ops.stop_gradient(teacher_output)
        cos_sin_dFdt = ops.stop_gradient(cos_sin_dFdt)

        # calculate output of the network
        subnet_out = self.subnet((xt / self.sigma, t, conditions), training=stage == "training")
        student_out = self.subnet_projector(subnet_out)

        # calculate the tangent
        g = -(ops.cos(t) ** 2) * (self.sigma * teacher_output - dxtdt) - r * ops.cos(t) * ops.sin(t) * (
            xt + self.sigma * cos_sin_dFdt
        )

        # apply normalization to stabilize training
        g = g / (ops.norm(g, axis=-1, keepdims=True) + c)

        # compute adaptive weights and calculate loss
        w = self.weight_fn_projector(self.weight_fn(expand_right_to(t_, 2)))

        D = ops.shape(x)[-1]

        loss = ops.mean(
            (ops.exp(w) / D)
            * ops.mean(
                ops.reshape(((student_out - teacher_output - g) ** 2), (ops.shape(teacher_output)[0], -1)), axis=-1
            )
            - w
        )

        return base_metrics | {"loss": loss}
