#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# -----------------------------------------------------------------------------------------------------------
# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
# CANN Open Software License Agreement Version 2.0 (the "License").
# Please refer to the License for details. You may not use this file except in compliance with the License.
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
# See LICENSE in the root of the software repository for the full text of the License.
# -----------------------------------------------------------------------------------------------------------

"""NPU callback functions"""

import os
import weakref
from tensorflow.python import keras
from tensorflow.python.ops import state_ops
from tensorflow.python.ops import control_flow_ops
import tensorflow as tf
from npu_device.distribute import hccl

# mapping id(var) to weakref.ref(var), which id(var) return memory address of var
broadcast_registry = weakref.WeakValueDictionary()


def broadcast_helper(variables, root_rank=0):
    """Broadcast trainable variables, and register to avoid repetitive processing"""
    candicates = []
    for var in variables:
        if hasattr(var, '_cast_dtype') and getattr(var, "_cast_dtype") != var.dtype:
            continue
        if id(var) not in broadcast_registry:
            candicates.append(var)
    if candicates:
        hccl.broadcast(candicates)
        for value in candicates:
            broadcast_registry[id(value)] = value


def broadcast_keras_model(model, root_rank=0):
    """Broadcast trainable variables of keras Model"""
    if not isinstance(model, tf.keras.Model):
        return model
    if model.built:
        broadcast_helper(model.trainable_variables, root_rank)
    return model


def get_ranksize():
    if os.getenv("CM_WORKER_SIZE") is not None and os.getenv("RANK_SIZE") is not None:
        raise ValueError("RANK_SIZE and CM_WORK_SIZE cannot be configured at the same time")
    rank_size = os.getenv('RANK_SIZE') if os.getenv("RANK_SIZE") is not None else os.getenv('CM_WORKER_SIZE', '1')
    return rank_size


class NPUBroadcastGlobalVariablesCallback(keras.callbacks.Callback):
    """
    Keras Callback that will broadcast all global variables from root rank
    to all other processes during initialization.
    This is necessary to ensure consistent initialization of all workers when
    training is started with random weights or restored from a checkpoint.
    """

    def __init__(self, root_rank):
        """
        Construct a new BroadcastGlobalVariablesCallback that will broadcast all
        global variables from root rank to all other processes during initialization.
        Args:
            root_rank: Rank that will send data, other ranks will receive data.
        """
        super(NPUBroadcastGlobalVariablesCallback, self).__init__()
        self.root_rank = root_rank
        self.broadcast_done = False

        def on_batch_begin(self, batch, logs=None):
            """This function is called when every batch begins"""
            if self.broadcast_done:
                return

            rank_size = get_ranksize()
            if int(rank_size) > 1:
                broadcast_helper(self.model.trainable_variables, self.root_rank)

            self.broadcast_done = True


class NpuBroadcastScopeContext(object):
    """
    A context manager, used to record and broadcast trainable varibles in its
    scope automatically.
    """
    _is_in_scope = False
    _is_broadcast = False

    def __init__(self, org_scope_ctx, model_trainable_variables):
        self.org_scope_ctx = org_scope_ctx
        self.enter_scope_again_count = 0
        self.model_trainable_variables = model_trainable_variables

        def _variable_broadcast_creator(next_creator, **kwargs):
            var = next_creator(**kwargs)
            if var.trainable:
                self.model_trainable_variables.append(var)
            return var

        # all var created in scope ctx will auto-broadcast
        self.broadcast_scope = tf.variable_creator_scope(
            _variable_broadcast_creator)

    def __enter__(self):
        if not NpuBroadcastScopeContext._is_in_scope:
            NpuBroadcastScopeContext._is_in_scope = True
            if self.broadcast_scope:
                self.broadcast_scope.__enter__()
        else:
            self.enter_scope_again_count += 1 # handle with reentry
        if self.org_scope_ctx:
            self.org_scope_ctx.__enter__()

    def __exit__(self, exception_type, exception_value, traceback):
        if self.org_scope_ctx:
            self.org_scope_ctx.__exit__(exception_type, exception_value, traceback)
        if not NpuBroadcastScopeContext._is_broadcast:
            NpuBroadcastScopeContext._is_broadcast = True
            broadcast_helper(self.model_trainable_variables)
        if self.enter_scope_again_count > 0:
            self.enter_scope_again_count -= 1
            return
        if self.broadcast_scope:
            self.broadcast_scope.__exit__(exception_type, exception_value, traceback)
        NpuBroadcastScopeContext._is_in_scope = False


def npu_broadcast_scope_wrapper(strategy):
    """ wrap strategy.scope with NpuBroadcastScopeContext """
    if not isinstance(strategy, tf.distribute.Strategy):
        return strategy

    org_scope = strategy.scope

    model_trainable_variables = []
    def _npu_broadcast_scope():
        org_scope_ctx = org_scope()
        return NpuBroadcastScopeContext(org_scope_ctx, model_trainable_variables)

    if not hasattr(strategy, '_npu_scope_wrapped'):
        strategy.scope = _npu_broadcast_scope
        setattr(strategy, '_npu_scope_wrapped', True)
    return strategy