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_cuda(torch::Tensor anchor, torch::Tensor positive, torch::Tensor negative, double margin);

torch::Tensor triplet_forward(torch::Tensor anchor, torch::Tensor positive, torch::Tensor negative, double margin) {
    return triplet_forward_cuda(anchor, positive, negative, margin);
}
"""

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

#define THREADS 256

__global__ void triplet_kernel_vec(
    const float* __restrict__ anchor,
    const float* __restrict__ positive,
    const float* __restrict__ negative,
    double* __restrict__ sum_out,
    int64_t N,
    int64_t D,
    float margin)
{
    const int row = blockIdx.x;
    const int tid = threadIdx.x;
    if (row >= N) return;

    float acc_ap = 0.0f;
    float acc_an = 0.0f;
    const float eps = 1.0e-6f;
    const int64_t base = (int64_t)row * D;

    if ((D & 3) == 0) {
        const int64_t D4 = D >> 2;
        const float4* __restrict__ a4 = reinterpret_cast<const float4*>(anchor + base);
        const float4* __restrict__ p4 = reinterpret_cast<const float4*>(positive + base);
        const float4* __restrict__ n4 = reinterpret_cast<const float4*>(negative + base);

        for (int64_t j = tid; j < D4; j += THREADS) {
            float4 av = a4[j];
            float4 pv = p4[j];
            float4 nv = n4[j];

            float d0 = av.x - pv.x + eps;
            float e0 = av.x - nv.x + eps;
            acc_ap = fmaf(d0, d0, acc_ap);
            acc_an = fmaf(e0, e0, acc_an);

            float d1 = av.y - pv.y + eps;
            float e1 = av.y - nv.y + eps;
            acc_ap = fmaf(d1, d1, acc_ap);
            acc_an = fmaf(e1, e1, acc_an);

            float d2 = av.z - pv.z + eps;
            float e2 = av.z - nv.z + eps;
            acc_ap = fmaf(d2, d2, acc_ap);
            acc_an = fmaf(e2, e2, acc_an);

            float d3 = av.w - pv.w + eps;
            float e3 = av.w - nv.w + eps;
            acc_ap = fmaf(d3, d3, acc_ap);
            acc_an = fmaf(e3, e3, acc_an);
        }
    } else {
        for (int64_t j = tid; j < D; j += THREADS) {
            float av = anchor[base + j];
            float d = av - positive[base + j] + eps;
            float e = av - negative[base + j] + eps;
            acc_ap = fmaf(d, d, acc_ap);
            acc_an = fmaf(e, e, acc_an);
        }
    }

    __shared__ float sh_ap[THREADS];
    __shared__ float sh_an[THREADS];
    sh_ap[tid] = acc_ap;
    sh_an[tid] = acc_an;
    __syncthreads();

    for (int stride = THREADS / 2; stride >= 32; stride >>= 1) {
        if (tid < stride) {
            sh_ap[tid] += sh_ap[tid + stride];
            sh_an[tid] += sh_an[tid + stride];
        }
        __syncthreads();
    }

    if (tid < 32) {
        float v_ap = sh_ap[tid];
        float v_an = sh_an[tid];
        unsigned mask = 0xffffffffu;
        v_ap += __shfl_down_sync(mask, v_ap, 16);
        v_an += __shfl_down_sync(mask, v_an, 16);
        v_ap += __shfl_down_sync(mask, v_ap, 8);
        v_an += __shfl_down_sync(mask, v_an, 8);
        v_ap += __shfl_down_sync(mask, v_ap, 4);
        v_an += __shfl_down_sync(mask, v_an, 4);
        v_ap += __shfl_down_sync(mask, v_ap, 2);
        v_an += __shfl_down_sync(mask, v_an, 2);
        v_ap += __shfl_down_sync(mask, v_ap, 1);
        v_an += __shfl_down_sync(mask, v_an, 1);

        if (tid == 0) {
            float loss = sqrtf(v_ap) - sqrtf(v_an) + margin;
            if (loss > 0.0f) atomicAdd(sum_out, (double)loss);
        }
    }
}

__global__ void finalize_kernel(const double* __restrict__ sum_in, float* __restrict__ out, int64_t N) {
    out[0] = (float)(sum_in[0] / (double)N);
}

torch::Tensor triplet_forward_cuda(torch::Tensor anchor, torch::Tensor positive, torch::Tensor negative, double margin) {
    const int64_t N = anchor.size(0);
    const int64_t D = anchor.numel() / N;

    auto sum_opts = anchor.options().dtype(torch::kFloat64);
    torch::Tensor sum = torch::zeros({}, sum_opts);
    torch::Tensor out = torch::empty({}, anchor.options());

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    dim3 block(THREADS);
    dim3 grid((unsigned int)N);

    triplet_kernel_vec<<<grid, block, 0, stream>>>(
        anchor.data_ptr<float>(),
        positive.data_ptr<float>(),
        negative.data_ptr<float>(),
        sum.data_ptr<double>(),
        N, D, (float)margin);

    finalize_kernel<<<1, 1, 0, stream>>>(sum.data_ptr<double>(), out.data_ptr<float>(), N);
    return out;
}
"""

_triplet_ext = load_inline(
    name="triplet_margin_loss_inline_ext",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["triplet_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    with_cuda=True,
    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 _triplet_ext.triplet_forward(anchor, positive, negative, self.margin)