# Path: sockeye/train.py
# def fixed_param_names_from_strategy(config: model.ModelConfig,
#                                     params: Dict[str, torch.nn.parameter.Parameter],
#                                     strategy: str) -> List[str]:
#     """
#     Generate a fixed parameter list given a list of all parameter names and
#     a strategy.
#     """
#     # Number of encoder/decoder layers in model.
#     num_encoder_layers = config.config_encoder.num_layers
#     num_decoder_layers = config.config_decoder.num_layers
# 
#     def is_fixed(name: str) -> bool:
#         if strategy == C.FIXED_PARAM_STRATEGY_ALL_EXCEPT_DECODER:
#             # Any decoder layer.
#             return not name.startswith(C.DECODER_PREFIX)
#         if strategy == C.FIXED_PARAM_STRATEGY_ALL_EXCEPT_OUTER_LAYERS:
#             # First and last encoder and decoder layers.
#             first_encoder_prefix = f'{C.ENCODER_PREFIX}.layers.{0}'
#             last_encoder_prefix = f'{C.ENCODER_PREFIX}.layers.{num_encoder_layers - 1}'
#             first_decoder_prefix = f'{C.DECODER_PREFIX}.layers.{0}'
#             last_decoder_prefix = f'{C.DECODER_PREFIX}.layers.{num_decoder_layers - 1}'
#             return not (name.startswith(first_encoder_prefix) or
#                         name.startswith(last_encoder_prefix) or
#                         name.startswith(first_decoder_prefix) or
#                         name.startswith(last_decoder_prefix))
#         if strategy == C.FIXED_PARAM_STRATEGY_ALL_EXCEPT_EMBEDDINGS:
#             # Any type of learned embedding.
#             return not (name.startswith(C.SOURCE_EMBEDDING_PREFIX) or name.startswith(C.TARGET_EMBEDDING_PREFIX))
#         if strategy == C.FIXED_PARAM_STRATEGY_ALL_EXCEPT_OUTPUT_PROJ:
#             # Target output projection.
#             return not name.startswith(C.DEFAULT_OUTPUT_LAYER_PREFIX)
#         if strategy == C.FIXED_PARAM_STRATEGY_ALL_EXCEPT_FEED_FORWARD:
#             return not (name.endswith("ff.ff1.bias") or name.endswith("ff.ff1.weight") or
#                         name.endswith("ff.ff2.bias") or name.endswith("ff.ff2.weight"))
#         if strategy == C.FIXED_PARAM_STRATEGY_ENCODER_AND_SOURCE_EMBEDDINGS:
#             return name.startswith(C.ENCODER_PREFIX) or name.startswith(C.SOURCE_EMBEDDING_PREFIX)
#         if strategy == C.FIXED_PARAM_STRATEGY_ENCODER_HALF_AND_SOURCE_EMBEDDINGS:
#             if name.startswith(C.ENCODER_PREFIX):
#                 for i in range(num_encoder_layers // 2):
#                     if name.startswith(f"{C.ENCODER_PREFIX}.layers.{i}"):
#                         return True
#             return name.startswith(C.SOURCE_EMBEDDING_PREFIX)
#         raise ValueError("Unknown fixed parameter strategy: %s" % strategy)
# 
#     return [name for name in params if is_fixed(name)]


from unittest import mock
from sockeye.train import fixed_param_names_from_strategy
import pytest
import sockeye.constants as C
    'encoder.layers.1.pre_ff.layer_norm.beta',
    'encoder.layers.1.pre_ff.layer_norm.gamma',
    'encoder.layers.1.pre_self_attention.layer_norm.beta',
    'encoder.layers.1.pre_self_attention.layer_norm.gamma',
    'encoder.layers.1.self_attention.ff_in.weight',
    'encoder.layers.1.self_attention.ff_out.weight',
    'encoder.layers.2.ff.ff1.bias',
    'encoder.layers.2.ff.ff1.weight',
    'encoder.layers.2.ff.ff2.bias',
    'encoder.layers.2.ff.ff2.weight',
    'encoder.layers.2.pre_ff.layer_norm.beta',
    'encoder.layers.2.pre_ff.layer_norm.gamma',
    'encoder.layers.2.pre_self_attention.layer_norm.beta',
    'encoder.layers.2.pre_self_attention.layer_norm.gamma',
    'encoder.layers.2.self_attention.ff_in.weight',
    'encoder.layers.2.self_attention.ff_out.weight',
]


@pytest.mark.parametrize("param_names, strategy, expected_fixed_param_names", [
    (ALL_PARAMS, C.FIXED_PARAM_STRATEGY_ALL_EXCEPT_DECODER, ALL_EXCEPT_DECODER_PARAMS),
    (ALL_PARAMS, C.FIXED_PARAM_STRATEGY_ALL_EXCEPT_OUTER_LAYERS, ALL_EXCEPT_OUTER_LAYERS_PARAMS),
    (ALL_PARAMS, C.FIXED_PARAM_STRATEGY_ALL_EXCEPT_EMBEDDINGS, ALL_EXCEPT_EMBED_PARAMS),
    (ALL_PARAMS, C.FIXED_PARAM_STRATEGY_ALL_EXCEPT_OUTPUT_PROJ, ALL_EXCEPT_OUTPUT_PROJ_PARAMS),
])
def test_fixed_param_strategy(param_names, strategy, expected_fixed_param_names):
    config = mock.Mock()
    config.config_encoder.num_layers = NUM_LAYERS
    config.config_decoder.num_layers = NUM_LAYERS
    params = {name: None for name in ALL_PARAMS}

# Path: sockeye/train.py
# def fixed_param_names_from_strategy(config: model.ModelConfig,
#                                     params: Dict[str, torch.nn.parameter.Parameter],
#                                     strategy: str) -> List[str]:
#     """
#     Generate a fixed parameter list given a list of all parameter names and
#     a strategy.
#     """
#     # Number of encoder/decoder layers in model.
#     num_encoder_layers = config.config_encoder.num_layers
#     num_decoder_layers = config.config_decoder.num_layers
# 
#     def is_fixed(name: str) -> bool:
#         if strategy == C.FIXED_PARAM_STRATEGY_ALL_EXCEPT_DECODER:
#             # Any decoder layer.
#             return not name.startswith(C.DECODER_PREFIX)
#         if strategy == C.FIXED_PARAM_STRATEGY_ALL_EXCEPT_OUTER_LAYERS:
#             # First and last encoder and decoder layers.
#             first_encoder_prefix = f'{C.ENCODER_PREFIX}.layers.{0}'
#             last_encoder_prefix = f'{C.ENCODER_PREFIX}.layers.{num_encoder_layers - 1}'
#             first_decoder_prefix = f'{C.DECODER_PREFIX}.layers.{0}'
#             last_decoder_prefix = f'{C.DECODER_PREFIX}.layers.{num_decoder_layers - 1}'
#             return not (name.startswith(first_encoder_prefix) or
#                         name.startswith(last_encoder_prefix) or
#                         name.startswith(first_decoder_prefix) or
#                         name.startswith(last_decoder_prefix))
#         if strategy == C.FIXED_PARAM_STRATEGY_ALL_EXCEPT_EMBEDDINGS:
#             # Any type of learned embedding.
#             return not (name.startswith(C.SOURCE_EMBEDDING_PREFIX) or name.startswith(C.TARGET_EMBEDDING_PREFIX))
#         if strategy == C.FIXED_PARAM_STRATEGY_ALL_EXCEPT_OUTPUT_PROJ:
#             # Target output projection.
#             return not name.startswith(C.DEFAULT_OUTPUT_LAYER_PREFIX)
#         if strategy == C.FIXED_PARAM_STRATEGY_ALL_EXCEPT_FEED_FORWARD:
#             return not (name.endswith("ff.ff1.bias") or name.endswith("ff.ff1.weight") or
#                         name.endswith("ff.ff2.bias") or name.endswith("ff.ff2.weight"))
#         if strategy == C.FIXED_PARAM_STRATEGY_ENCODER_AND_SOURCE_EMBEDDINGS:
#             return name.startswith(C.ENCODER_PREFIX) or name.startswith(C.SOURCE_EMBEDDING_PREFIX)
#         if strategy == C.FIXED_PARAM_STRATEGY_ENCODER_HALF_AND_SOURCE_EMBEDDINGS:
#             if name.startswith(C.ENCODER_PREFIX):
#                 for i in range(num_encoder_layers // 2):
#                     if name.startswith(f"{C.ENCODER_PREFIX}.layers.{i}"):
#                         return True
#             return name.startswith(C.SOURCE_EMBEDDING_PREFIX)
#         raise ValueError("Unknown fixed parameter strategy: %s" % strategy)
# 
#     return [name for name in params if is_fixed(name)]


# Path: test/unit/test_fixed_param_strategy.py
from unittest import mock
from sockeye.train import fixed_param_names_from_strategy
import pytest
import sockeye.constants as C

    'encoder.layers.1.pre_ff.layer_norm.beta',
    'encoder.layers.1.pre_ff.layer_norm.gamma',
    'encoder.layers.1.pre_self_attention.layer_norm.beta',
    'encoder.layers.1.pre_self_attention.layer_norm.gamma',
    'encoder.layers.1.self_attention.ff_in.weight',
    'encoder.layers.1.self_attention.ff_out.weight',
    'encoder.layers.2.ff.ff1.bias',
    'encoder.layers.2.ff.ff1.weight',
    'encoder.layers.2.ff.ff2.bias',
    'encoder.layers.2.ff.ff2.weight',
    'encoder.layers.2.pre_ff.layer_norm.beta',
    'encoder.layers.2.pre_ff.layer_norm.gamma',
    'encoder.layers.2.pre_self_attention.layer_norm.beta',
    'encoder.layers.2.pre_self_attention.layer_norm.gamma',
    'encoder.layers.2.self_attention.ff_in.weight',
    'encoder.layers.2.self_attention.ff_out.weight',
]


@pytest.mark.parametrize("param_names, strategy, expected_fixed_param_names", [
    (ALL_PARAMS, C.FIXED_PARAM_STRATEGY_ALL_EXCEPT_DECODER, ALL_EXCEPT_DECODER_PARAMS),
    (ALL_PARAMS, C.FIXED_PARAM_STRATEGY_ALL_EXCEPT_OUTER_LAYERS, ALL_EXCEPT_OUTER_LAYERS_PARAMS),
    (ALL_PARAMS, C.FIXED_PARAM_STRATEGY_ALL_EXCEPT_EMBEDDINGS, ALL_EXCEPT_EMBED_PARAMS),
    (ALL_PARAMS, C.FIXED_PARAM_STRATEGY_ALL_EXCEPT_OUTPUT_PROJ, ALL_EXCEPT_OUTPUT_PROJ_PARAMS),
])
def test_fixed_param_strategy(param_names, strategy, expected_fixed_param_names):
    config = mock.Mock()
    config.config_encoder.num_layers = NUM_LAYERS
    config.config_decoder.num_layers = NUM_LAYERS
    params = {name: None for name in ALL_PARAMS}

    fixed_param_names = fixed_param_names_from_strategy(config, params, strategy)