import math
from functools import partial

import torch
from benchmark.utils import (
    benchmark_all,
    benchmark_backward,
    benchmark_combined,
    benchmark_forward,
)
from einops import rearrange

from .cauchy import cauchy_mult, cauchy_mult_keops, cauchy_mult_torch


def generate_data(batch_size, N, L, symmetric=True, device="cuda"):
    if not symmetric:
        v = torch.randn(
            batch_size,
            N,
            dtype=torch.complex64,
            device=device,
            requires_grad=True,
        )
        w = torch.randn(
            batch_size,
            N,
            dtype=torch.complex64,
            device=device,
            requires_grad=True,
        )
        z = torch.randn(L, dtype=torch.complex64, device=device)
    else:
        assert N % 2 == 0
        v_half = torch.randn(
            batch_size, N // 2, dtype=torch.complex64, device=device
        )
        v = torch.cat([v_half, v_half.conj()], dim=-1).requires_grad_(True)
        w_half = torch.randn(
            batch_size, N // 2, dtype=torch.complex64, device=device
        )
        w = torch.cat([w_half, w_half.conj()], dim=-1).requires_grad_(True)
        z = torch.exp(1j * torch.randn(L, dtype=torch.float32, device=device))
    return v, z, w


if __name__ == "__main__":
    device = "cuda"
    bs = 1024
    N = 64
    L = 16384

    v, z, w = generate_data(bs, N, L, symmetric=True)
    v_half = v[:, : N // 2].clone().detach().requires_grad_(True)
    w_half = w[:, : N // 2].clone().detach().requires_grad_(True)

    repeat = 30
    benchmark_all(repeat, cauchy_mult_keops, v, z, w, desc="Cauchy mult keops")
    fn = partial(cauchy_mult, symmetric=False)
    benchmark_all(repeat, fn, v, z, w, desc="Cauchy mult")
    fn = partial(cauchy_mult, symmetric=True)
    benchmark_all(repeat, fn, v_half, z, w_half, desc="Cauchy mult symmetric")
