"""WavLM wrapper for frame-rate experiments.

This keeps the normal SpeechBrain/HuggingFace WavLM path, but exposes two
small knobs used by alignment papers: changing the final feature-extractor
stride and changing the expected input sampling rate for waveform upsampling.
"""

from __future__ import annotations

import logging
from functools import reduce
from operator import mul

from speechbrain.lobes.models.huggingface_transformers.wavlm import WavLM


logger = logging.getLogger(__name__)


def _is_empty(value) -> bool:
    if value is None:
        return True
    return str(value).strip().lower() in {"", "none", "null", "false"}


def _as_bool(value) -> bool:
    if isinstance(value, bool):
        return value
    if value is None:
        return False
    return str(value).strip().lower() in {"1", "true", "yes", "y", "on"}


class WavLMFrameRate(WavLM):
    """WavLM with optional feature-extractor frame-rate edits."""

    def __init__(
        self,
        source,
        save_path,
        output_norm=False,
        freeze=False,
        freeze_feature_extractor=False,
        apply_spec_augment=False,
        output_all_hiddens=False,
        last_conv_stride=None,
        feature_extractor_sampling_rate=None,
        random_init=False,
    ):
        super().__init__(
            source=source,
            save_path=save_path,
            output_norm=output_norm,
            freeze=freeze,
            freeze_feature_extractor=freeze_feature_extractor,
            apply_spec_augment=apply_spec_augment,
            output_all_hiddens=output_all_hiddens,
        )
        if _as_bool(random_init):
            self.random_initialize()
        if not _is_empty(feature_extractor_sampling_rate):
            self.set_feature_extractor_sampling_rate(int(feature_extractor_sampling_rate))
        if not _is_empty(last_conv_stride):
            self.set_last_conv_stride(int(last_conv_stride))

    def random_initialize(self) -> None:
        """Reset HuggingFace WavLM weights while keeping the loaded architecture."""
        hf_model = getattr(self, "model", None)
        init_fn = getattr(hf_model, "_init_weights", None)
        if hf_model is None or init_fn is None:
            raise RuntimeError("Could not find HuggingFace WavLM model initializer")
        hf_model.apply(init_fn)
        if hasattr(hf_model, "tie_weights"):
            hf_model.tie_weights()
        logger.info(
            "WavLM random_init=True: reset HuggingFace WavLM weights from architecture config"
        )

    def set_feature_extractor_sampling_rate(self, sampling_rate: int) -> None:
        """Allow the data pipeline to feed resampled waveform to WavLM."""
        if hasattr(self, "feature_extractor"):
            self.feature_extractor.sampling_rate = int(sampling_rate)
        logger.info("WavLM feature_extractor.sampling_rate=%s", sampling_rate)

    def get_last_conv_stride(self) -> tuple[int, ...]:
        feature_extractor = getattr(self.model, "feature_extractor", None)
        conv_layers = getattr(feature_extractor, "conv_layers", None)
        if not conv_layers:
            raise RuntimeError("Could not find WavLM feature_extractor.conv_layers")
        conv = getattr(conv_layers[-1], "conv", None)
        if conv is None:
            raise RuntimeError("Could not find the final WavLM convolution layer")
        return tuple(int(x) for x in conv.stride)

    def set_feature_extractor_frozen(self, frozen: bool) -> None:
        """Toggle trainability of the raw CNN feature extractor after init/load."""
        feature_extractor = getattr(self.model, "feature_extractor", None)
        if feature_extractor is None:
            raise RuntimeError("Could not find WavLM feature_extractor")
        for param in feature_extractor.parameters():
            param.requires_grad = not bool(frozen)
        self.freeze_feature_extractor = bool(frozen)
        logger.info("WavLM feature extractor frozen=%s", int(bool(frozen)))

    def estimated_frame_shift_ms(self, sampling_rate: int | None = None) -> float | None:
        """Estimate the frame shift in milliseconds from current conv strides."""
        config = getattr(self.model, "config", None)
        conv_stride = getattr(config, "conv_stride", None) if config is not None else None
        if conv_stride is None:
            return None
        if sampling_rate is None:
            sampling_rate = getattr(getattr(self, "feature_extractor", None), "sampling_rate", None)
        if not sampling_rate:
            return None
        ratio = int(reduce(mul, conv_stride, 1))
        return 1000.0 * float(ratio) / float(sampling_rate)

    def set_last_conv_stride(self, stride: int) -> None:
        """Change the final CNN stride, e.g. 2 -> 1 for roughly 10 ms frames."""
        if stride <= 0:
            raise ValueError(f"last_conv_stride must be positive, got {stride}")
        feature_extractor = getattr(self.model, "feature_extractor", None)
        conv_layers = getattr(feature_extractor, "conv_layers", None)
        if not conv_layers:
            raise RuntimeError("Could not find WavLM feature_extractor.conv_layers")
        last_layer = conv_layers[-1]
        conv = getattr(last_layer, "conv", None)
        if conv is None:
            raise RuntimeError("Could not find the final WavLM convolution layer")

        old_stride = tuple(conv.stride)
        conv.stride = (stride,)

        config = getattr(self.model, "config", None)
        if config is not None and hasattr(config, "conv_stride"):
            conv_stride = list(config.conv_stride)
            old_config_stride = list(conv_stride)
            conv_stride[-1] = int(stride)
            config.conv_stride = conv_stride
            try:
                config.inputs_to_logits_ratio = int(reduce(mul, conv_stride, 1))
            except Exception:
                logger.debug("Could not update inputs_to_logits_ratio", exc_info=True)
            logger.info(
                "WavLM final conv stride changed: layer %s -> %s, config %s -> %s",
                old_stride,
                tuple(conv.stride),
                old_config_stride,
                conv_stride,
            )
        else:
            logger.info("WavLM final conv stride changed: %s -> %s", old_stride, tuple(conv.stride))