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

cuda_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>

__global__ void iou_tanh_kernel(const float* __restrict__ b1, 
                                const float* __restrict__ b2, 
                                float* __restrict__ out, 
                                int n) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        int offset = idx * 4;

        float b1_x1 = b1[offset + 0];
        float b1_y1 = b1[offset + 1];
        float b1_x2 = b1[offset + 2];
        float b1_y2 = b1[offset + 3];

        float b2_x1 = b2[offset + 0];
        float b2_y1 = b2[offset + 1];
        float b2_x2 = b2[offset + 2];
        float b2_y2 = b2[offset + 3];

        float inter_x1 = fmaxf(b1_x1, b2_x1);
        float inter_y1 = fmaxf(b1_y1, b2_y1);
        float inter_x2 = fminf(b1_x2, b2_x2);
        float inter_y2 = fminf(b1_y2, b2_y2);

        float inter_w = fmaxf(0.0f, inter_x2 - inter_x1);
        float inter_h = fmaxf(0.0f, inter_y2 - inter_y1);
        float inter_area = inter_w * inter_h;

        float area1 = fmaxf(0.0f, b1_x2 - b1_x1) * fmaxf(0.0f, b1_y2 - b1_y1);
        float area2 = fmaxf(0.0f, b2_x2 - b2_x1) * fmaxf(0.0f, b2_y2 - b2_y1);

        float union_area = area1 + area2 - inter_area;
        float iou = inter_area / (union_area + 1e-6f);

        out[idx] = tanhf(iou);
    }
}

torch::Tensor iou_tanh_cuda(torch::Tensor box1, torch::Tensor box2) {
    int n = box1.size(0);
    auto out = torch::empty({n}, box1.options());

    const int block_size = 256;
    int grid_size = (n + block_size - 1) / block_size;

    iou_tanh_kernel<<<grid_size, block_size>>>(
        box1.data_ptr<float>(),
        box2.data_ptr<float>(),
        out.data_ptr<float>(),
        n
    );

    return out;
}
"""

cpp_source = "torch::Tensor iou_tanh_cuda(torch::Tensor box1, torch::Tensor box2);"

module = load_inline(
    name="iou_tanh_ext",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["iou_tanh_cuda"],
    verbose=False,
    with_cuda=True
)


class ModelNew(nn.Module):
    def __init__(self):
        super(ModelNew, self).__init__()
        self.op = module

    def forward(self, box1, box2):
        res = self.op.iou_tanh_cuda(box1.contiguous(), box2.contiguous())
        return res.mean()