from pathlib import Path

import torch
from einops import rearrange

from structured_kernels import cauchy_mult_sym_bwd, cauchy_mult_sym_fwd

# try:
#     from cauchy_mult import cauchy_mult_sym_fwd, cauchy_mult_sym_bwd
# except ImportError:
#     from torch.utils.cpp_extension import load
#     current_dir = Path(__file__).parent.absolute()
#     cauchy_mult_extension = load(
#         name='cauchy_mult',
#         sources=[str(current_dir / 'cauchy.cpp'), str(current_dir / 'cauchy_cuda.cu')],
#         extra_cflags=['-g', '-march=native', '-funroll-loops'],
#         extra_cuda_cflags=['-O3', '-lineinfo', '--use_fast_math'],
#         extra_include_paths=str(current_dir),
#         build_directory=str(current_dir),
#         verbose=True
#     )
#     cauchy_mult_sym_fwd = cauchy_mult_extension.cauchy_mult_sym_fwd
#     cauchy_mult_sym_bwd = cauchy_mult_extension.cauchy_mult_sym_bwd


def cauchy_mult_torch(
    v: torch.Tensor, z: torch.Tensor, w: torch.Tensor, symmetric=True
) -> torch.Tensor:
    """
    v: (B, N)
    z: (L)
    w: (B, N)
    symmetric: whether to assume that v and w contain complex conjugate pairs, of the form
    [v_half, v_half.conj()] and [w_half, w_half.conj()]
    """
    if not symmetric:
        return (
            rearrange(v, "b n -> b 1 n")
            / (rearrange(z, "l -> l 1") - rearrange(w, "b n -> b 1 n"))
        ).sum(dim=-1)
    else:
        N = v.shape[-1]
        assert N % 2 == 0
        vv = rearrange(v[:, : N // 2], "b n -> b 1 n")
        zz = rearrange(z, "l -> l 1")
        ww = rearrange(w[:, : N // 2], "b n -> b 1 n")
        # return 2 * ((zz * vv.real - vv.real * ww.real - vv.imag * ww.imag)
        #             / (zz * zz - 2 * zz * ww.real + ww.abs().square())).sum(dim=-1)
        return (vv / (zz - ww) + vv.conj() / (zz - ww.conj())).sum(dim=-1)


def cauchy_mult_keops(v, z, w):
    from pykeops.torch import LazyTensor

    v_l = LazyTensor(rearrange(v, "b N -> b 1 N 1"))
    z_l = LazyTensor(rearrange(z, "L -> 1 L 1 1"))
    w_l = LazyTensor(rearrange(w, "b N -> b 1 N 1"))
    sub = (
        z_l - w_l
    )  # (b N L 1), for some reason it doesn't display the last dimension
    div = v_l / sub
    s = div.sum(dim=2, backend="GPU")
    return s.squeeze(-1)


def _cauchy_mult(v, z, w):
    return CauchyMultiplySymmetric.apply(v, z, w)


def cauchy_mult(v, z, w):
    """Wrap the cuda method to deal with shapes"""
    v, w = torch.broadcast_tensors(v, w)
    shape = v.shape
    # z_shape = z.shape
    # z = z.squeeze()
    assert len(z.shape) == 1

    v = v.contiguous()
    w = w.contiguous()
    z = z.contiguous()

    N = v.size(-1)
    assert w.size(-1) == N
    y = _cauchy_mult(v.view(-1, N), z, w.view(-1, N))
    y = y.view(*shape[:-1], z.size(-1))
    return y


class CauchyMultiplySymmetric(torch.autograd.Function):

    @staticmethod
    def forward(ctx, v, z, w):
        batch, N = v.shape
        supported_N_values = [
            1 << log_n for log_n in [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
        ]
        L = z.shape[-1]
        if not N in supported_N_values:
            raise NotImplementedError(
                f"Only support N values in {supported_N_values}"
            )
        max_L_value = 32 * 1024 * 64 * 1024
        if L > max_L_value:
            raise NotImplementedError(
                f"Only support L values <= {max_L_value}"
            )
        if not (v.is_cuda and z.is_cuda and w.is_cuda):
            raise NotImplementedError(f"Only support CUDA tensors")
        ctx.save_for_backward(v, z, w)
        return cauchy_mult_sym_fwd(v, z, w)

    @staticmethod
    def backward(ctx, dout):
        v, z, w = ctx.saved_tensors
        dv, dw = cauchy_mult_sym_bwd(v, z, w, dout)
        return dv, None, dw
