"""Q51M Mixture-of-Depths, exactly matched to Q-MoE's total budget."""
import math

from .configuration_q50m import Q50MConfig


class QMoDConfig(Q50MConfig):
    model_type = "q_mod"

    def __init__(self, mod_capacity=0.125, mod_layers=None,
                 mod_intermediate_size=None, router_aux_coef=0.01, **kwargs):
        super().__init__(**kwargs)
        self.mod_capacity = float(mod_capacity)
        self.mod_layers = (list(range(1, self.num_hidden_layers, 2))
                           if mod_layers is None else list(mod_layers))
        # Five extra SwiGLU channels in each of five routed blocks, plus five
        # scalar routers: 5 * (3 * 512 * 5 + 512) = 40,960 parameters.
        self.mod_intermediate_size = (self.intermediate_size + 5
                                      if mod_intermediate_size is None else mod_intermediate_size)
        self.router_aux_coef = float(router_aux_coef)
        if not 0 < self.mod_capacity <= 1:
            raise ValueError("mod_capacity must be in (0, 1]")
        if not math.isfinite(self.router_aux_coef) or self.router_aux_coef < 0:
            raise ValueError("router_aux_coef must be finite and nonnegative")
        if (len(set(self.mod_layers)) != len(self.mod_layers)
                or any(not isinstance(i, int) or not 0 <= i < self.num_hidden_layers
                       for i in self.mod_layers)):
            raise ValueError("mod_layers must contain unique valid layer indices")
        if not isinstance(self.mod_intermediate_size, int) or self.mod_intermediate_size < 1:
            raise ValueError("mod_intermediate_size must be a positive integer")