import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

N, D = 32, 128

assert D % 4 == 0, "Embedding dimension D must be a multiple of 4 for vectorization"


class ModelNew(nn.Module):

    def __init__(self, margin=1.0, swap=False):
        super().__init__()
        self.margin = float(margin)
        self.swap = swap
        self.block_size = 256
        self._compile_cuda_kernel()

    def _compile_cuda_kernel(self):
        cpp_source = """
        #include <torch/extension.h>

        torch::Tensor triplet_forward_cuda(
            torch::Tensor anchor,
            torch::Tensor positive,
            torch::Tensor negative,
            float margin,
            bool swap,
            int N,
            int D);
        """

        cuda_source = f"""
        #include <cuda_runtime.h>
        #include <cmath>

        #define BLOCK_SIZE {self.block_size}
        #define WARP_SIZE 32

        // Warp 归约工具
        __inline__ __device__ float warp_reduce_sum(float val) {{
            #pragma unroll
            for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {{
                val += __shfl_down_sync(0xffffffff, val, offset);
            }}
            return val;
        }}


        __global__ void triplet_l2_kernel(
            const float* __restrict__ anchor,
            const float* __restrict__ positive,
            const float* __restrict__ negative,
            float* __restrict__ output,
            float margin,
            bool swap,
            int D_vec  // D / 4
        ) {{
            const int n_idx = blockIdx.x; 
            const int tid = threadIdx.x;


            const int offset = n_idx * D_vec * 4;
            const float4* a_ptr = reinterpret_cast<const float4*>(anchor + offset);
            const float4* p_ptr = reinterpret_cast<const float4*>(positive + offset);
            const float4* n_ptr = reinterpret_cast<const float4*>(negative + offset);


            float sum_sq_ap = 0.0f;
            float sum_sq_an = 0.0f;
            float sum_sq_pn = 0.0f; 


            for (int i = tid; i < D_vec; i += BLOCK_SIZE) {{
                float4 a = __ldg(&a_ptr[i]);
                float4 p = __ldg(&p_ptr[i]);
                float4 n = __ldg(&n_ptr[i]);


                float4 diff_ap, diff_an, diff_pn;

                diff_ap.x = a.x - p.x; diff_ap.y = a.y - p.y; diff_ap.z = a.z - p.z; diff_ap.w = a.w - p.w;
                diff_an.x = a.x - n.x; diff_an.y = a.y - n.y; diff_an.z = a.z - n.z; diff_an.w = a.w - n.w;


                sum_sq_ap += diff_ap.x*diff_ap.x + diff_ap.y*diff_ap.y + diff_ap.z*diff_ap.z + diff_ap.w*diff_ap.w;
                sum_sq_an += diff_an.x*diff_an.x + diff_an.y*diff_an.y + diff_an.z*diff_an.z + diff_an.w*diff_an.w;

                if (swap) {{
                    diff_pn.x = p.x - n.x; diff_pn.y = p.y - n.y; diff_pn.z = p.z - n.z; diff_pn.w = p.w - n.w;
                    sum_sq_pn += diff_pn.x*diff_pn.x + diff_pn.y*diff_pn.y + diff_pn.z*diff_pn.z + diff_pn.w*diff_pn.w;
                }}
            }}


            __shared__ float shared_data[32][3];

            int lane = tid % WARP_SIZE;
            int wid = tid / WARP_SIZE;


            sum_sq_ap = warp_reduce_sum(sum_sq_ap);
            sum_sq_an = warp_reduce_sum(sum_sq_an);
            if (swap) sum_sq_pn = warp_reduce_sum(sum_sq_pn);


            if (lane == 0) {{
                shared_data[wid][0] = sum_sq_ap;
                shared_data[wid][1] = sum_sq_an;
                if (swap) shared_data[wid][2] = sum_sq_pn;
            }}
            __syncthreads();


            if (wid == 0) {{
                sum_sq_ap = (tid < blockDim.x / WARP_SIZE) ? shared_data[lane][0] : 0.0f;
                sum_sq_an = (tid < blockDim.x / WARP_SIZE) ? shared_data[lane][1] : 0.0f;
                sum_sq_pn = (tid < blockDim.x / WARP_SIZE && swap) ? shared_data[lane][2] : 0.0f;

                sum_sq_ap = warp_reduce_sum(sum_sq_ap);
                sum_sq_an = warp_reduce_sum(sum_sq_an);
                if (swap) sum_sq_pn = warp_reduce_sum(sum_sq_pn);


                if (tid == 0) {{
                    // 开根号得到 L2 距离 (加上 epsilon 防止梯度爆炸通常在backward处理，前向计算通常加个极小值)
                    float dist_ap = sqrtf(sum_sq_ap + 1e-8f);
                    float dist_an = sqrtf(sum_sq_an + 1e-8f);

                    if (swap) {{
                        float dist_pn = sqrtf(sum_sq_pn + 1e-8f);
                        if (dist_pn < dist_an) {{
                            dist_an = dist_pn;
                        }}
                    }}

                    // loss = max(d_ap - d_an + margin, 0)
                    float loss = fmaxf(dist_ap - dist_an + margin, 0.0f);
                    output[n_idx] = loss;
                }}
            }}
        }}

        torch::Tensor triplet_forward_cuda(
            torch::Tensor anchor,
            torch::Tensor positive,
            torch::Tensor negative,
            float margin,
            bool swap,
            int N,
            int D) 
        {{
            anchor = anchor.contiguous();
            positive = positive.contiguous();
            negative = negative.contiguous();

            auto output = torch::empty({{N}}, anchor.options());

            int D_vec = D / 4;

            dim3 blocks(N);
            dim3 threads(BLOCK_SIZE);

            triplet_l2_kernel<<<blocks, threads>>>(
                anchor.data_ptr<float>(),
                positive.data_ptr<float>(),
                negative.data_ptr<float>(),
                output.data_ptr<float>(),
                margin,
                swap,
                D_vec
            );

            return output;
        }}
        """

        self.op = load_inline(
            name='triplet_loss_cuda_v1',
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=['triplet_forward_cuda'],
            extra_cuda_cflags=['-O3', '--use_fast_math'],
            verbose=False
        )

    def forward(self, a: torch.Tensor, p: torch.Tensor, n: torch.Tensor) -> torch.Tensor:

        if not a.is_cuda: a = a.cuda()
        if not p.is_cuda: p = p.cuda()
        if not n.is_cuda: n = n.cuda()

        N, D = a.shape

        losses = self.op.triplet_forward_cuda(a, p, n, self.margin, self.swap, N, D)

        return losses.mean()