# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: OpenMDW-1.1

from . import modeling_nemotron_h
from .configuration_nemotron_h_puzzle import NemotronHPuzzleConfig
from .modeling_nemotron_h import NemotronHForCausalLM, NemotronHBlock


class NemotronHPuzzleBlock(NemotronHBlock):
    def __init__(self, config: NemotronHPuzzleConfig, layer_idx: int):
        layer_config = config.get_nemotron_h_config_for_layer(layer_idx)
        super().__init__(config=layer_config, layer_idx=layer_idx)


class NemotronHPuzzleForCausalLM(NemotronHForCausalLM):
    """
    A child class of NemotronHForCausalLM to support heterogeneous layer configurations.

    This class uses monkey-patching to inject custom behavior into the parent class while maximizing
    code reuse and minimizing duplication. During `__init__`, it temporarily replaces the block layer
    class to use `NemotronHPuzzleBlock`, so that each layer is built from its own per-layer
    configuration (derived from the matching entry in `config.block_configs`).
    """

    config_class = NemotronHPuzzleConfig
    _no_split_modules = ["NemotronHPuzzleBlock"]

    def __init__(self, config):
        original_block_class = modeling_nemotron_h.NemotronHBlock

        try:
            modeling_nemotron_h.NemotronHBlock = NemotronHPuzzleBlock

            super().__init__(config)

        finally:
            modeling_nemotron_h.NemotronHBlock = original_block_class