import torch
import torch.nn as nn

class Model(nn.Module):
    """
    A model that computes Triplet Margin Loss for metric learning tasks.

    Parameters:
        margin (float): The margin between the positive and negative samples.
    """
    def __init__(self, margin=1.0):
        super(Model, self).__init__()
        self.loss_fn = torch.nn.TripletMarginLoss(margin=margin)

    def forward(self, anchor, positive, negative):
        return self.loss_fn(anchor, positive, negative)

batch_size = 32768
input_shape = (8192,)
dim = 1

def get_inputs():
    scale = torch.rand(())
    return [torch.rand(batch_size, *input_shape)*scale, torch.rand(batch_size, *input_shape), torch.rand(batch_size, *input_shape)]
    
def get_init_inputs():
    return [1.0]  # Default margin

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

_cpp_src = r'''
#include <torch/extension.h>
torch::Tensor triplet_forward(torch::Tensor anchor, torch::Tensor positive, torch::Tensor negative, double margin);
'''

_cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>

__inline__ __device__ float warp_reduce_sum(float v) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        v += __shfl_down_sync(0xffffffff, v, offset);
    }
    return v;
}

template<int BLOCK>
__global__ void triplet_loss_kernel(const float* __restrict__ a,
                                    const float* __restrict__ p,
                                    const float* __restrict__ n,
                                    float* __restrict__ out,
                                    int64_t B, int64_t D, float margin) {
    int row = blockIdx.x;
    int tid = threadIdx.x;
    const int64_t base = (int64_t)row * D;
    float sp = 0.0f;
    float sn = 0.0f;
    const float eps = 1.0e-6f;

    int64_t D4 = D >> 2;
    const float4* __restrict__ a4 = reinterpret_cast<const float4*>(a + base);
    const float4* __restrict__ p4 = reinterpret_cast<const float4*>(p + base);
    const float4* __restrict__ n4 = reinterpret_cast<const float4*>(n + base);

    for (int64_t j = tid; j < D4; j += BLOCK) {
        float4 av = a4[j];
        float4 pv = p4[j];
        float4 nv = n4[j];
        float d;
        d = av.x - pv.x + eps; sp = fmaf(d, d, sp);
        d = av.y - pv.y + eps; sp = fmaf(d, d, sp);
        d = av.z - pv.z + eps; sp = fmaf(d, d, sp);
        d = av.w - pv.w + eps; sp = fmaf(d, d, sp);

        d = av.x - nv.x + eps; sn = fmaf(d, d, sn);
        d = av.y - nv.y + eps; sn = fmaf(d, d, sn);
        d = av.z - nv.z + eps; sn = fmaf(d, d, sn);
        d = av.w - nv.w + eps; sn = fmaf(d, d, sn);
    }

    sp = warp_reduce_sum(sp);
    sn = warp_reduce_sum(sn);

    __shared__ float sh_p[32];
    __shared__ float sh_n[32];
    int lane = tid & 31;
    int warp = tid >> 5;
    if (lane == 0) {
        sh_p[warp] = sp;
        sh_n[warp] = sn;
    }
    __syncthreads();

    if (warp == 0) {
        sp = sh_p[lane];
        sn = sh_n[lane];
        sp = warp_reduce_sum(sp);
        sn = warp_reduce_sum(sn);
        if (lane == 0) {
            float loss = sqrtf(sp) - sqrtf(sn) + margin;
            if (loss > 0.0f) {
                atomicAdd(out, loss / (float)B);
            }
        }
    }
}

torch::Tensor triplet_forward(torch::Tensor anchor, torch::Tensor positive, torch::Tensor negative, double margin) {
    auto B = anchor.size(0);
    auto D = anchor.numel() / B;
    auto out = torch::empty({}, anchor.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    cudaMemsetAsync(out.data_ptr<float>(), 0, sizeof(float), stream);
    constexpr int BLOCK = 1024;
    triplet_loss_kernel<BLOCK><<<B, BLOCK, 0, stream>>>(
        anchor.data_ptr<float>(), positive.data_ptr<float>(), negative.data_ptr<float>(),
        out.data_ptr<float>(), B, D, (float)margin);
    return out;
}
'''

_ext = load_inline(
    name='triplet_margin_loss_ext_opt',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['triplet_forward'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self, margin=1.0):
        super(ModelNew, self).__init__()
        self.margin = float(margin)

    def forward(self, anchor, positive, negative):
        return _ext.triplet_forward(anchor, positive, negative, self.margin)