import tensorflow as tf
import math

# ===== 補助関数 =====
def extract_pairwise_multi_position_embedding(position_mat, emb_dim):
    """
    【スタブ】ペアごとの位置埋め込みを抽出する関数。
    position_mat: [num_fg_classes, num_rois, num_rois, 4]
    emb_dim: 出力チャネル数（ここでは fc_dim[0] と同じ）
    ※実際の実装では論文の記述に従った埋め込み処理を行う
    """
    # 例として、各位置情報 (4次元) に対して Dense 層で emb_dim 次元に変換する
    dense_layer = tf.keras.layers.Dense(units=emb_dim, name='pairwise_pos_embedding')
    # 入力をフラット化して処理後、元の形状に戻す
    shape = tf.shape(position_mat)
    num_fg = shape[0]
    num_rois = shape[1]
    pos_reshaped = tf.reshape(position_mat, [-1, 4])
    embedded = dense_layer(pos_reshaped)
    embedded = tf.reshape(embedded, [num_fg, num_rois, num_rois, emb_dim])
    return embedded

# ===== attention_module_multi_head =====
def attention_module_multi_head(roi_feat, position_embedding, non_gt_index,
                                fc_dim, feat_dim, 
                                dim=(1024, 1024, 1024), group=16, index=1):
    """
    Multi-head attention モジュール（通常版）
    
    Args:
      roi_feat: Tensor of shape [num_rois, feat_dim]
      position_embedding: Tensor of shape [1, emb_dim, num_rois, nongt_dim] （NCHW形式）
      non_gt_index: インデックス（None の場合は全 RoI を使用）
      fc_dim: 整数（通常 group と同じ値）
      feat_dim: roi_feat の次元（dim[2] と同じである必要がある）
      dim: (query_dim, key_dim, output_dim) のタプル
      group: head の数
      index: 識別用の番号（ネットワーク内での命名用）
      
    Returns:
      output: Tensor of shape [num_rois, dim[2]]
    """
    # 各 head ごとの次元
    dim_group = (dim[0] // group, dim[1] // group, dim[2] // group)
    
    # non_gt_index が指定されていれば対象 RoI を抽出
    if non_gt_index is None:
        nongt_roi_feat = roi_feat
    else:
        nongt_roi_feat = tf.gather(roi_feat, non_gt_index, axis=0)
    
    # ----- 位置特徴の変換 -----
    # position_embedding: [1, emb_dim, num_rois, nongt_dim]（NCHW）
    conv_layer = tf.keras.layers.Conv2D(filters=fc_dim, kernel_size=(1,1), strides=(1,1),
                                        padding='valid', data_format="channels_first",
                                        name='pair_pos_fc1_' + str(index))
    position_feat_1 = conv_layer(position_embedding)  # → [1, fc_dim, num_rois, nongt_dim]
    position_feat_1_relu = tf.nn.relu(position_feat_1)
    
    # 転置して [num_rois, fc_dim, nongt_dim, 1] にし、最後の次元を squeeze
    aff_weight = tf.transpose(position_feat_1_relu, perm=[2, 1, 3, 0])
    aff_weight = tf.squeeze(aff_weight, axis=-1)  # → [num_rois, fc_dim, nongt_dim]
    
    # ----- Multi-head のための Query, Key の生成 -----
    # Query: roi_feat から Dense 層で変換 → [num_rois, dim[0]]
    query_dense = tf.keras.layers.Dense(units=dim[0], name='query_' + str(index))
    q_data = query_dense(roi_feat)
    # reshape: [num_rois, group, dim_group[0]] → 転置して [group, num_rois, dim_group[0]]
    q_data_batch = tf.reshape(q_data, [-1, group, dim_group[0]])
    q_data_batch = tf.transpose(q_data_batch, perm=[1, 0, 2])
    
    # Key: nongt_roi_feat から Dense 層で変換 → [num_nongt, dim[1]]
    key_dense = tf.keras.layers.Dense(units=dim[1], name='key_' + str(index))
    k_data = key_dense(nongt_roi_feat)
    k_data_batch = tf.reshape(k_data, [-1, group, dim_group[1]])
    k_data_batch = tf.transpose(k_data_batch, perm=[1, 0, 2])
    
    # Value: nongt_roi_feat（変換なし）
    v_data = nongt_roi_feat  # shape: [num_nongt, feat_dim]
    
    # ----- Attention 計算 -----
    # 各 head ごとに Q と K の内積（batch_dot）を計算
    aff = tf.matmul(q_data_batch, k_data_batch, transpose_b=True)  # → [group, num_rois, nongt_dim]
    aff_scale = aff / math.sqrt(dim_group[1])
    # 転置して [num_rois, group, nongt_dim]
    aff_scale = tf.transpose(aff_scale, perm=[1, 0, 2])
    
    # fc_dim と group は同じ値である前提
    weighted_aff = tf.math.log(tf.maximum(aff_weight, 1e-6)) + aff_scale  # [num_rois, fc_dim, nongt_dim]
    aff_softmax = tf.nn.softmax(weighted_aff, axis=2)
    
    # reshape: [num_rois * fc_dim, nongt_dim]
    num_rois = tf.shape(roi_feat)[0]
    aff_softmax_reshape = tf.reshape(aff_softmax, [num_rois * fc_dim, -1])
    
    # Value と掛け合わせる: → [num_rois * fc_dim, feat_dim]
    output_t = tf.matmul(aff_softmax_reshape, v_data)
    # reshape → [num_rois, fc_dim * feat_dim, 1, 1]
    output_t = tf.reshape(output_t, [num_rois, fc_dim * feat_dim, 1, 1])
    
    # ----- 出力層：グループ付き 1x1 畳み込み -----
    conv_linear = tf.keras.layers.Conv2D(filters=dim[2], kernel_size=(1,1),
                                         groups=fc_dim, data_format="channels_first",
                                         name='linear_out_' + str(index))
    linear_out = conv_linear(output_t)
    output = tf.reshape(linear_out, [num_rois, dim[2]])
    return output

# ===== attention_module_nms_multi_head =====
def attention_module_nms_multi_head(roi_feat, position_mat, num_rois,
                                    dim=(1024, 1024, 1024), fc_dim=(64, 16),
                                    feat_dim=1024, group=16, index=1):
    """
    NMS 用 Multi-head attention モジュール
    
    Args:
      roi_feat: Tensor of shape [num_rois, num_fg_classes, feat_dim]
      position_mat: Tensor of shape [num_fg_classes, num_rois, num_rois, 4]
      num_rois: RoI 数
      dim: (query_dim, key_dim, output_dim)
      fc_dim: タプル（例: (64, 16)）; fc_dim[1] は group と一致すべき
      feat_dim: roi_feat の次元
      group: head の数
      index: ネットワーク内での識別番号
      
    Returns:
      output: Tensor of shape [num_rois, num_fg_classes, dim[2]]
      aff_softmax: 計算された attention 重み（後で解析等に利用可能）
    """
    # 各 head ごとの次元
    dim_group = (dim[0] // group, dim[1] // group, dim[2] // group)
    
    # roi_feat: [num_rois, num_fg_classes, feat_dim] → 転置して [num_fg_classes, num_rois, feat_dim]
    roi_feat = tf.transpose(roi_feat, perm=[1, 0, 2])
    # 平坦化: [num_fg_classes * num_rois, feat_dim]
    roi_feat_reshape = tf.reshape(roi_feat, [-1, feat_dim])
    
    # ----- 位置埋め込み -----
    # 補助関数を用いて，position_mat から [num_fg_classes, num_rois, num_rois, fc_dim[0]] の埋め込みを得る
    position_embedding = extract_pairwise_multi_position_embedding(position_mat, fc_dim[0])
    # 平坦化: [num_fg_classes * num_rois * num_rois, fc_dim[0]]
    position_embedding_reshape = tf.reshape(position_embedding, [-1, fc_dim[0]])
    
    # FullyConnected 層で変換 → [num_fg_classes * num_rois * num_rois, fc_dim[1]]
    fc_layer = tf.keras.layers.Dense(units=fc_dim[1], name='nms_pair_pos_fc1_' + str(index))
    position_feat_1 = fc_layer(position_embedding_reshape)
    # reshape → [num_fg_classes, num_rois, num_rois, fc_dim[1]]
    position_feat_1 = tf.reshape(position_feat_1, [-1, num_rois, num_rois, fc_dim[1]])
    aff_weight = tf.nn.relu(position_feat_1)
    # 転置して → [num_fg_classes, fc_dim[1], num_rois, num_rois]
    aff_weight = tf.transpose(aff_weight, perm=[0, 3, 1, 2])
    
    # ----- Multi-head attention の計算（クエリ・キー） -----
    # Query: roi_feat_reshape から Dense 層で → [num_fg_classes*num_rois, dim[0]]
    q_dense = tf.keras.layers.Dense(units=dim[0], name='nms_query_' + str(index))
    q_data = q_dense(roi_feat_reshape)
    # reshape → [num_fg_classes, num_rois, group, dim_group[0]] → 転置して [num_fg_classes, group, num_rois, dim_group[0]]
    q_data_batch = tf.reshape(q_data, [-1, num_rois, group, dim_group[0]])
    q_data_batch = tf.transpose(q_data_batch, perm=[0, 2, 1, 3])
    # 平坦化: [num_fg_classes * group, num_rois, dim_group[0]]
    q_data_batch = tf.reshape(q_data_batch, [-1, num_rois, dim_group[0]])
    
    # Key: 同様に
    k_dense = tf.keras.layers.Dense(units=dim[1], name='nms_key_' + str(index))
    k_data = k_dense(roi_feat_reshape)
    k_data_batch = tf.reshape(k_data, [-1, num_rois, group, dim_group[1]])
    k_data_batch = tf.transpose(k_data_batch, perm=[0, 2, 1, 3])
    k_data_batch = tf.reshape(k_data_batch, [-1, num_rois, dim_group[1]])
    
    # Value: roi_feat（[num_fg_classes, num_rois, feat_dim]）
    v_data = roi_feat
    
    # ----- Attention スコア計算 -----
    aff = tf.matmul(q_data_batch, k_data_batch, transpose_b=True)  # → [num_fg_classes * group, num_rois, num_rois]
    aff_scale = aff / math.sqrt(dim_group[1])
    
    # aff_weight の形状は [num_fg_classes, fc_dim[1], num_rois, num_rois] で，
    # 平坦化して [num_fg_classes * fc_dim[1], num_rois, num_rois]
    aff_weight_reshape = tf.reshape(aff_weight, [-1, num_rois, num_rois])
    weighted_aff = tf.math.log(tf.maximum(aff_weight_reshape, 1e-6)) + aff_scale
    aff_softmax = tf.nn.softmax(weighted_aff, axis=2)
    
    # reshape: [num_fg_classes, fc_dim[1] * num_rois, num_rois]
    # ※ fc_dim[1] == group と仮定
    num_fg = tf.shape(aff_weight)[0]
    aff_softmax_reshape = tf.reshape(aff_softmax, [num_fg, fc_dim[1] * num_rois, num_rois])
    
    # ----- 出力生成 -----
    # batch_dot: [num_fg_classes, fc_dim[1]*num_rois, feat_dim]
    output_t = tf.matmul(aff_softmax_reshape, v_data)
    # reshape → [num_fg_classes, fc_dim[1], num_rois, feat_dim]
    output_t_reshape = tf.reshape(output_t, [-1, fc_dim[1], num_rois, feat_dim])
    # 転置して → [fc_dim[1], feat_dim, num_rois, num_fg_classes]
    output_t_reshape = tf.transpose(output_t_reshape, perm=[1, 3, 2, 0])
    # reshape → [1, fc_dim[1]*feat_dim, num_rois, num_fg_classes]
    output_t_reshape = tf.reshape(output_t_reshape, [1, fc_dim[1] * feat_dim, num_rois, -1])
    
    # グループ付き 1×1 畳み込み（groups=fc_dim[1]）
    conv_linear = tf.keras.layers.Conv2D(filters=dim[2], kernel_size=(1,1),
                                         groups=fc_dim[1], data_format="channels_first",
                                         name='nms_linear_out_' + str(index))
    linear_out = conv_linear(output_t_reshape)
    # reshape: [dim[2], num_rois, num_fg_classes] → 転置して [num_rois, num_fg_classes, dim[2]]
    linear_out_reshape = tf.reshape(linear_out, [dim[2], num_rois, -1])
    output = tf.transpose(linear_out_reshape, perm=[1, 2, 0])
    
    return output, aff_softmax