import numpy as np
import utils

def layer_norm(x, gamma, beta, eps=1e-5):
    mu   = x.mean(-1, keepdims=True)
    var  = x.var( -1, keepdims=True)
    xn   = (x - mu) / np.sqrt(var + eps)
    return gamma * xn + beta

def multi_head_attention(x, mask, W_qkv, b_qkv, W_o, b_o, n_heads):
    B, T, D = x.shape
    d_k      = D // n_heads

    # print(f"multi_head_attention: x.shape={x.shape}, mask.shape={mask.shape if mask is not None else 'None'}, "
    #       f"W_qkv.shape={W_qkv.shape}, b_qkv.shape={b_qkv.shape}, "
    #       f"W_o.shape={W_o.shape}, b_o.shape={b_o.shape}, n_heads={n_heads}")
    
    # utils.save_cpp_array_2d('x', x[0], dtype='float')
    # utils.save_cpp_array_2d("w_qkv", W_qkv, dtype='float')
    # utils.save_cpp_array_1d("b_qkv", b_qkv, dtype='float')
    # utils.save_cpp_array_2d("w_o", W_o, dtype='float')
    # utils.save_cpp_array_1d("b_o", b_o, dtype='float')

    # 1) QKV projection
    # ONNX does:   x_T = x.transpose(1,0,2)  → MatMul(x_T, W_qkv) + b  
    # but since W_qkv is shape (D,3D), you can do it in batch form:
    qkv = x @ W_qkv + b_qkv        # → (B, T, 3D)
    # utils.save_cpp_array_2d('qkv', qkv[0], dtype='float')

    q, k, v = np.split(qkv, 3, axis=-1)

    # print(f"multi_head_attention: q.shape={q.shape}, k.shape={k.shape}, v.shape={v.shape}")
    # utils.save_cpp_array_2d('q', q[0], dtype='float')
    # utils.save_cpp_array_2d('k', k[0], dtype='float')
    # utils.save_cpp_array_2d('v', v[0], dtype='float')

    # 2) reshape for heads
    #    (B,T,D) → (B, H, T, d_k)
    q = q.reshape(B, T, n_heads, d_k).transpose(0,2,1,3)
    k = k.reshape(B, T, n_heads, d_k).transpose(0,2,1,3)
    v = v.reshape(B, T, n_heads, d_k).transpose(0,2,1,3)

    # print(f"multi_head_attention: q.shape={q.shape}, k.shape={k.shape}, v.shape={v.shape} after reshaping for heads")
    # utils.save_cpp_array_3d('qt', q[0], dtype='float')
    # utils.save_cpp_array_3d('kt', k[0], dtype='float')
    # utils.save_cpp_array_3d('vt', v[0], dtype='float')

    # 3) convert boolean mask → additive mask of -inf/0
    if mask is not None:
        # mask comes in as (B, T) float or bool
        if mask.ndim == 2:
            m = np.where(mask.astype(bool), -np.inf, 0.0).astype(np.float32)
            # broadcast to (B,1,1,T)
            mask = m[:,None,None,:]
        else:
            raise ValueError("Unexpected mask ndim")
        
    # print(f"multi_head_attention: mask.shape={mask.shape if mask is not None else 'None'} after conversion")
    # utils.save_cpp_array_1d('mask', mask[0][0][0] if mask is not None else np.array([]), dtype='float')


    # 4) scaled dot–product: ONNX splits the √d_k factor
    scale = d_k ** -0.25
    q = q * scale
    k = k * scale
    scores = q @ k.swapaxes(-2,-1)    # (B, H, T, T)
    if mask is not None:
        scores = scores + mask        # (B,H,T,T)
    # softmax over last axis
    e = np.exp(scores - scores.max(-1,keepdims=True))
    weights = e / e.sum(-1,keepdims=True)

    # print(f"multi_head_attention: scores.shape={scores.shape}, weights.shape={weights.shape} after softmax")
    # utils.save_cpp_array_3d('scores', scores[0], dtype='float')
    # utils.save_cpp_array_3d('weights', weights[0], dtype='float')

    # 5) attention output
    attn = weights @ v                # (B,H,T,d_k)
    # back to (B, T, D)
    concat = attn.transpose(0,2,1,3).reshape(B, T, D)

    # print(f"multi_head_attention: concat.shape={concat.shape} after attention output")
    # utils.save_cpp_array_2d('concat', concat[0], dtype='float')

    # 6) final projection via Gemm(transB=1)
    #    Y = concat @ W_o.T + b_o   
    result = concat @ W_o.T + b_o
    # print(f"multi_head_attention: result.shape={result.shape} after final projection")
    # utils.save_cpp_array_2d('result', result[0], dtype='float')
    return result

def position_wise_feed_forward(x, W1, b1, W2, b2):
    # ONNX uses ReLU
    h = x @ W1 + b1
    h = np.maximum(0, h)
    return h @ W2 + b2

def encoder_layer(x, mask, w):
    # --- 1) Self-Attention sublayer (post-norm input) ---
    # Note: x here is already the previous layer’s Norm2 output.
    mha_out  = multi_head_attention(
        x, mask,
        w['mha.W_qkv'], w['mha.b_qkv'],
        w['mha.W_o'],   w['mha.b_o'],
        w['n_heads']
    )
    y        = x + mha_out

    # print(f"encoder_layer: y.shape={y.shape} after self-attention")
    # utils.save_cpp_array_2d('y', y[0], dtype='float')
    # print(f"encoder_layer: norm1.gamma.shape={w['norm1.gamma'].shape}, norm1.beta.shape={w['norm1.beta'].shape}")
    # utils.save_cpp_array_1d('norm1_gamma', w['norm1.gamma'], dtype='float')
    # utils.save_cpp_array_1d('norm1_beta', w['norm1.beta'], dtype='float')

    y_norm   = layer_norm(y, w['norm1.gamma'], w['norm1.beta'])

    # print(f"encoder_layer: y_norm.shape={y_norm.shape} after layer norm")
    # utils.save_cpp_array_2d('y_norm', y_norm[0], dtype='float')

    # print(f"encoder_layer: ff.W1.shape={w['ff.W1'].shape}, ff.b1.shape={w['ff.b1'].shape}, ff.W2.shape={w['ff.W2'].shape}, ff.b2.shape={w['ff.b2'].shape}")
    # utils.save_cpp_array_2d('ff_W1', w['ff.W1'], dtype='float')
    # utils.save_cpp_array_1d('ff_b1', w['ff.b1'], dtype='float')
    # utils.save_cpp_array_2d('ff_W2', w['ff.W2'], dtype='float')
    # utils.save_cpp_array_1d('ff_b2', w['ff.b2'], dtype='float')
    # --- 2) Feed-Forward sublayer ---
    ff_out   = position_wise_feed_forward(
        y_norm,
        w['ff.W1'], w['ff.b1'],
        w['ff.W2'], w['ff.b2']
    )

    # utils.save_cpp_array_2d('ff_out', ff_out[0], dtype='float')
    # print(f"encoder_layer: ff_out.shape={ff_out.shape} after feed-forward")
    z        = y_norm + ff_out
    # print(f"encoder_layer: z.shape={z.shape} after adding feed-forward output")

    # print(f"encoder_layer: norm2.gamma.shape={w['norm2.gamma'].shape}, norm2.beta.shape={w['norm2.beta'].shape}")
    # utils.save_cpp_array_1d('norm2_gamma', w['norm2.gamma'], dtype='float')
    # utils.save_cpp_array_1d('norm2_beta', w['norm2.beta'], dtype='float')
    out      = layer_norm(z, w['norm2.gamma'], w['norm2.beta'])
    # print(f"encoder_layer: out.shape={out.shape} after final layer norm")
    # utils.save_cpp_array_2d('out', out[0], dtype='float')
    

    return out