import mlx.core as mx


def nearest_interpolate(x, size=None, scale_factor=None):
    """
    Nearest neighbor interpolation that exactly matches PyTorch's behavior.
    """
    # Get input dimensions
    batch_size, channels, in_h, in_w = x.shape

    # Calculate output dimensions
    if size is not None:
        out_h, out_w = size
    elif scale_factor is not None:
        if isinstance(scale_factor, (int, float)):
            scale_h = scale_w = scale_factor
        else:
            scale_h, scale_w = scale_factor
        out_h, out_w = int(in_h * scale_h), int(in_w * scale_w)
    else:
        raise ValueError("Either size or scale_factor must be specified")

    # Create dimensions tensor
    dims = mx.array([batch_size, channels, in_h, in_w, out_h, out_w], dtype=mx.int32)

    # Reshape input tensor to 1D for kernel processing
    x_flat = x.reshape(-1)
    input_dtype = x.dtype
    if input_dtype != mx.float32:
        x_flat = x_flat.astype(mx.float32)

    # Metal kernel source that matches PyTorch's coordinate calculation
    source = """
        uint x_out = thread_position_in_grid.x;
        uint y_out = thread_position_in_grid.y;
        uint bc_idx = thread_position_in_grid.z;

        int batch_size = dims[0];
        int channels = dims[1];
        int in_h = dims[2];
        int in_w = dims[3];
        int out_h = dims[4];
        int out_w = dims[5];

        if (x_out >= (uint)out_w || y_out >= (uint)out_h || bc_idx >= (uint)(batch_size * channels))
            return;

        int c = bc_idx % channels;
        int b = bc_idx / channels;

        // PyTorch's coordinate calculation for nearest neighbor
        // This matches: torch.nn.functional.interpolate(..., mode='nearest')
        float scale_h = float(in_h) / float(out_h);
        float scale_w = float(in_w) / float(out_w);

        // PyTorch uses floor for nearest neighbor coordinate mapping
        int y_in = int(floor(float(y_out) * scale_h));
        int x_in = int(floor(float(x_out) * scale_w));

        // Clamp to bounds
        y_in = max(0, min(y_in, in_h - 1));
        x_in = max(0, min(x_in, in_w - 1));

        int input_offset = ((b * channels + c) * in_h + y_in) * in_w + x_in;
        int output_offset = ((b * channels + c) * out_h + y_out) * out_w + x_out;

        output[output_offset] = input[input_offset];
    """

    # Create and run kernel
    kernel = mx.fast.metal_kernel(
        name="nearest_interpolation",
        input_names=["input", "dims"],
        output_names=["output"],
        source=source,
    )

    threadgroup = get_optimal_threadgroup(out_w, out_h)
    outputs = kernel(
        inputs=[x_flat, dims],
        grid=(out_w, out_h, batch_size * channels),
        threadgroup=threadgroup,
        output_shapes=[(batch_size * channels * out_h * out_w,)],
        output_dtypes=[mx.float32],
    )

    result = outputs[0].reshape(batch_size, channels, out_h, out_w)
    if input_dtype != mx.float32:
        result = result.astype(input_dtype)

    return result


def bicubic_interpolate(
    x, size=None, scale_factor=None, align_corners=False, antialias=False
):
    """
    Bicubic interpolation using MLX's built-in interpolate function.

    Args:
        x: MLX tensor of shape [B, C, H, W]
        size: Tuple of (out_h, out_w) or None
        scale_factor: Float or tuple of (scale_h, scale_w) or None
        align_corners: Whether to align corners
        antialias: Whether to apply antialiasing

    Returns:
        Interpolated MLX tensor
    """
    # Get input dimensions
    batch_size, channels, in_h, in_w = x.shape

    # Calculate output dimensions
    if size is not None:
        out_h, out_w = size
        scale_h, scale_w = out_h / in_h, out_w / in_w
    elif scale_factor is not None:
        if isinstance(scale_factor, (int, float)):
            scale_h = scale_w = scale_factor
        else:
            scale_h, scale_w = scale_factor
        out_h, out_w = int(in_h * scale_h), int(in_w * scale_w)
    else:
        raise ValueError("Either size or scale_factor must be specified")

    # Calculate antialiasing parameters
    # PyTorch uses support = 2.0 for bicubic when antialiasing
    support = 2.0
    antialias_flag = 1.0 if (antialias and (scale_h < 1.0 or scale_w < 1.0)) else 0.0

    # When downsampling with antialias, PyTorch expands the filter support
    if antialias and scale_h < 1.0:
        filter_scale_h = 1.0 / scale_h
    else:
        filter_scale_h = 1.0

    if antialias and scale_w < 1.0:
        filter_scale_w = 1.0 / scale_w
    else:
        filter_scale_w = 1.0

    # Create parameters tensor
    params = mx.array(
        [
            scale_h,
            scale_w,
            1.0 if align_corners else 0.0,
            antialias_flag,
            filter_scale_h,
            filter_scale_w,
            support,
        ],
        dtype=mx.float32,
    )

    # Create dimensions tensor
    dims = mx.array([batch_size, channels, in_h, in_w, out_h, out_w], dtype=mx.int32)

    # Reshape input tensor to 1D for kernel processing
    x_flat = x.reshape(-1)

    # Convert to float32 for processing if needed
    input_dtype = x.dtype
    if input_dtype != mx.float32:
        x_flat = x_flat.astype(mx.float32)

    header = """
        // Bicubic kernel function
        float cubic_kernel(float x) {
            float absx = fabs(x);
            float absx2 = absx * absx;
            float absx3 = absx2 * absx;

            const float a = -0.5f;

            if (absx <= 1.0f) {
                return (a + 2.0f) * absx3 - (a + 3.0f) * absx2 + 1.0f;
            } else if (absx < 2.0f) {
                return a * absx3 - 5.0f * a * absx2 + 8.0f * a * absx - 4.0f * a;
            }
            return 0.0f;
        }

        // Antialiased bicubic kernel - scales the support region for downsampling
        float cubic_kernel_antialias(float x, float scale) {
            // When downsampling, we need to integrate over a wider region
            // This matches PyTorch's antialiasing behavior
            return cubic_kernel(x / scale);
        }
    """

    # Metal kernel source code with antialiasing support
    source = """
        // Get thread position
        uint x_out = thread_position_in_grid.x;
        uint y_out = thread_position_in_grid.y;
        uint bc_idx = thread_position_in_grid.z;

        // Extract dimensions
        int batch_size = dims[0];
        int channels = dims[1];
        int in_h = dims[2];
        int in_w = dims[3];
        int out_h = dims[4];
        int out_w = dims[5];

        // Extract parameters
        float scale_h = params[0];
        float scale_w = params[1];
        bool align_corners = params[2] > 0.5f;
        bool use_antialias = params[3] > 0.5f;
        float filter_scale_h = params[4];
        float filter_scale_w = params[5];
        float support = params[6];

        // Check bounds
        if (x_out >= (uint)out_w || y_out >= (uint)out_h || bc_idx >= (uint)(batch_size * channels))
            return;

        // Calculate batch and channel indices
        int c = bc_idx % channels;
        int b = bc_idx / channels;

        // Calculate input coordinates
        float x_in, y_in;

        if (align_corners && out_w > 1 && out_h > 1) {
            x_in = float(x_out) * (in_w - 1) / (out_w - 1);
            y_in = float(y_out) * (in_h - 1) / (out_h - 1);
        } else {
            // PyTorch's default coordinate mapping
            x_in = ((float(x_out) + 0.5f) / float(out_w)) * float(in_w) - 0.5f;
            y_in = ((float(y_out) + 0.5f) / float(out_h)) * float(in_h) - 0.5f;
        }

        // Calculate the support region based on antialiasing
        float support_h = use_antialias ? support * filter_scale_h : support;
        float support_w = use_antialias ? support * filter_scale_w : support;

        // Calculate the range of input pixels to sample
        int y_start = int(floor(y_in - support_h)) + 1;
        int y_end = int(floor(y_in + support_h)) + 1;
        int x_start = int(floor(x_in - support_w)) + 1;
        int x_end = int(floor(x_in + support_w)) + 1;

        // Clamp to valid range
        y_start = max(0, y_start);
        y_end = min(in_h, y_end);
        x_start = max(0, x_start);
        x_end = min(in_w, x_end);

        // Perform bicubic interpolation with antialiasing
        float result = 0.0f;
        float weight_sum = 0.0f;

        for (int y_pos = y_start; y_pos < y_end; y_pos++) {
            float dy = float(y_pos) - y_in;
            float wy = use_antialias ?
                cubic_kernel_antialias(dy, filter_scale_h) :
                cubic_kernel(dy);

            for (int x_pos = x_start; x_pos < x_end; x_pos++) {
                float dx = float(x_pos) - x_in;
                float wx = use_antialias ?
                    cubic_kernel_antialias(dx, filter_scale_w) :
                    cubic_kernel(dx);

                float weight = wy * wx;

                // Calculate input tensor offset
                int input_offset = ((b * channels + c) * in_h + y_pos) * in_w + x_pos;

                // Add weighted contribution
                result += input[input_offset] * weight;
                weight_sum += weight;
            }
        }

        // Normalize by weight sum
        if (weight_sum > 1e-8f) {
            result /= weight_sum;
        }

        // Calculate output tensor offset
        int output_offset = ((b * channels + c) * out_h + y_out) * out_w + x_out;

        // Assign the result to output
        output[output_offset] = result;
    """

    # Create the kernel
    kernel = mx.fast.metal_kernel(
        name="bicubic_interpolation_antialias",
        input_names=["input", "dims", "params"],
        output_names=["output"],
        source=source,
        header=header,
    )

    # Run the kernel
    threadgroup = get_optimal_threadgroup(out_w, out_h)
    outputs = kernel(
        inputs=[x_flat, dims, params],
        grid=(out_w, out_h, batch_size * channels),
        threadgroup=threadgroup,
        output_shapes=[(batch_size * channels * out_h * out_w,)],
        output_dtypes=[mx.float32],
    )

    # Reshape output back to 4D tensor and convert back to original dtype
    result = outputs[0].reshape(batch_size, channels, out_h, out_w)
    if input_dtype != mx.float32:
        result = result.astype(input_dtype)

    return result


def grid_sample(x, grid):
    """
    Grid sample using MLX's built-in interpolate function.
    Args:
        x: MLX tensor of shape [B, C, H, W]
        grid: MLX tensor of shape [B, gN, gM, 2]

    Returns:
        Interpolated MLX tensor
    """

    assert x.ndim == 4, "`x` must be 4D."
    assert grid.ndim == 4, "`grid` must be 4D."

    B, _, _, C = x.shape
    _, gN, gM, D = grid.shape
    out_shape = (B, gN, gM, C)

    assert D == 2, "Last dim of `grid` must be size 2."

    source = """
        uint elem = thread_position_in_grid.x;
        int H = x_shape[1];
        int W = x_shape[2];
        int C = x_shape[3];
        int gH = grid_shape[1];
        int gW = grid_shape[2];

        int w_stride = C;
        int h_stride = W * w_stride;
        int b_stride = H * h_stride;

        uint grid_idx = elem / C * 2;
        float ix = ((grid[grid_idx] + 1) * W - 1) / 2;
        float iy = ((grid[grid_idx + 1] + 1) * H - 1) / 2;

        int ix_nw = floor(ix);
        int iy_nw = floor(iy);

        int ix_ne = ix_nw + 1;
        int iy_ne = iy_nw;

        int ix_sw = ix_nw;
        int iy_sw = iy_nw + 1;

        int ix_se = ix_nw + 1;
        int iy_se = iy_nw + 1;

        T nw = (ix_se - ix)    * (iy_se - iy);
        T ne = (ix    - ix_sw) * (iy_sw - iy);
        T sw = (ix_ne - ix)    * (iy    - iy_ne);
        T se = (ix    - ix_nw) * (iy    - iy_nw);

        int batch_idx = elem / C / gH / gW * b_stride;
        int channel_idx = elem % C;
        int base_idx = batch_idx + channel_idx;

        T I_nw = x[base_idx + iy_nw * h_stride + ix_nw * w_stride];
        T I_ne = x[base_idx + iy_ne * h_stride + ix_ne * w_stride];
        T I_sw = x[base_idx + iy_sw * h_stride + ix_sw * w_stride];
        T I_se = x[base_idx + iy_se * h_stride + ix_se * w_stride];

        I_nw = iy_nw >= 0 && iy_nw <= H - 1 && ix_nw >= 0 && ix_nw <= W - 1 ? I_nw : 0;
        I_ne = iy_ne >= 0 && iy_ne <= H - 1 && ix_ne >= 0 && ix_ne <= W - 1 ? I_ne : 0;
        I_sw = iy_sw >= 0 && iy_sw <= H - 1 && ix_sw >= 0 && ix_sw <= W - 1 ? I_sw : 0;
        I_se = iy_se >= 0 && iy_se <= H - 1 && ix_se >= 0 && ix_se <= W - 1 ? I_se : 0;

        out[elem] = nw * I_nw + ne * I_ne + sw * I_sw + se * I_se;
    """

    kernel = mx.fast.metal_kernel(
        name="grid_sample",
        input_names=["x", "grid"],
        output_names=["out"],
        source=source,
    )

    outputs = kernel(
        inputs=[x, grid],
        template=[("T", x.dtype)],
        output_shapes=[out_shape],
        output_dtypes=[x.dtype],
        grid=(mx.prod(mx.array(out_shape)), 1, 1),
        threadgroup=(256, 1, 1),
    )
    return outputs[0]


def get_optimal_threadgroup(out_w, out_h):
    # Calculate optimal threadgroup dimensions based on output dimensions

    # Maximum threadgroup size for most Metal GPUs
    # This could be made more dynamic with Metal API queries if needed
    MAX_THREADS_PER_GROUP = 1024
    MAX_THREADS_PER_DIM = 1024

    # Start with a reasonable default size for 2D workloads
    default_threadgroup = (32, 32, 1)

    try:
        # Don't create threadgroups larger than the work dimensions
        max_width = min(MAX_THREADS_PER_DIM, out_w)
        max_height = min(MAX_THREADS_PER_DIM, out_h)

        # Find largest power of 2 that fits within our dimensions
        width = 2 ** (max_width.bit_length() - 1)
        if width > max_width:
            width = width // 2

        height = 2 ** (max_height.bit_length() - 1)
        if height > max_height:
            height = height // 2

        # Ensure we don't exceed maximum threads per threadgroup
        while width * height > MAX_THREADS_PER_GROUP:
            # Reduce the larger dimension first
            if width >= height:
                width = width // 2
            else:
                height = height // 2

        # Ensure minimum size for efficiency
        width = max(8, width)
        height = max(8, height)

        return (width, height, 1)

    except Exception:
        # Return safe defaults if calculation fails
        return default_threadgroup