from itertools import product

import cupy as cp
import numpy as np
from cupy import testing
from cupyx.profiler import benchmark

from cupyx.scipy.spatial.delaunay_2d._kernels import DELAUNAY_MODULE
from cupyx.scipy.interpolate import CloughTocher2DInterpolator as GPUCt
from scipy.interpolate import CloughTocher2DInterpolator as CPUCt

from tqdm import tqdm

DELAUNAY_MODULE.compile()

# input_sizes = [5, 10, 100, 1000, 5000]
input_sizes = [100, 1000, 5000, 10000, 50000, 100000, 500000]
# input_sizes = [50000]
# input_sizes = [100, 1000, 5000, 10000, 50000]
modules = [(cp, GPUCt, 'CuPy'), (np, CPUCt, 'SciPy')]
kwargs = []
# kwargs = [('order', [1, 2, 4, 6, 10])]
# kwargs = [('extrapolate', [True])]
kwargs = list(product(*[list(product([x[0]], x[1])) for x in kwargs]))
measurement = {}


def compute_circumcenter(tri_points, xp):
    D = (tri_points[:, 0, 0] * (tri_points[:, 1, 1] - tri_points[:, 2, 1]) +
         tri_points[:, 1, 0] * (tri_points[:, 2, 1] - tri_points[:, 0, 1]) +
         tri_points[:, 2, 0] * (tri_points[:, 0, 1] - tri_points[:, 1, 1]))
    D = 2 * D

    out = xp.empty((tri_points.shape[0], 2), dtype=xp.float64)
    out[:, 0] = ((tri_points[:, 0] ** 2).sum(-1) * (
                    tri_points[:, 1, 1] - tri_points[:, 2, 1]) +
                 (tri_points[:, 1] ** 2).sum(-1) * (
                     tri_points[:, 2, 1] - tri_points[:, 0, 1]) +
                 (tri_points[:, 2] ** 2).sum(-1) * (
                     tri_points[:, 0, 1] - tri_points[:, 1, 1])) / D

    out[:, 1] = ((tri_points[:, 0] ** 2).sum(-1) * (
                    tri_points[:, 2, 0] - tri_points[:, 1, 0]) +
                 (tri_points[:, 1] ** 2).sum(-1) * (
                     tri_points[:, 0, 0] - tri_points[:, 2, 0]) +
                 (tri_points[:, 2] ** 2).sum(-1) * (
                     tri_points[:, 1, 0] - tri_points[:, 0, 0])) / D
    return out


def compute_random_points(tri_points, xp):
    bary = xp.empty((tri_points.shape[0], 3), xp.float64)
    s = testing.shaped_random((tri_points.shape[0],), np, np.float64,
                              scale=1.0)
    t = testing.shaped_random((tri_points.shape[0],), np, np.float64,
                              scale=1-s)
    bary[:, 0] = xp.asarray(s)
    bary[:, 1] = xp.asarray(t)
    bary[:, 2] = xp.asarray(1 - s - t)

    # breakpoint()

    return (xp.expand_dims(bary, -1) * tri_points).sum(1)


def gather_time(prof):
    cpu_time = prof.cpu_times.mean() * 1000
    gpu_time = prof.gpu_times.mean() * 1000
    return max(cpu_time, gpu_time)


def input_creation(xp, fn, size, **kwargs):
    # np.linspace()
    # # breakpoint()
    # x = xp.linspace(0, 10, int(np.sqrt(size[0])))
    # y = xp.linspace(0, 10, int(np.sqrt(size[0])))

    # noise_x = testing.shaped_random((int(np.sqrt(size[0]))**2,), xp, xp.float64)
    # noise_y = testing.shaped_random((int(np.sqrt(size[0]))**2,), xp, xp.float64, seed=1234)
    # x += noise_x
    # y += noise_y
    # X, Y = xp.meshgrid(x, y)
    # points = xp.c_[X.reshape(-1) + noise_x, Y.reshape(-1) + noise_y]
    points = testing.shaped_random(size + (2,), xp, xp.float64)
    values = testing.shaped_random(points.shape, xp, xp.float64)
    instance = fn(points, values)
    # breakpoint()
    # search = testing.shaped_random(size + (2,), xp, xp.float64, 5, seed=1111)
    search = compute_random_points(points[instance.tri.simplices], xp)
    # search = points[instance.tri.simplices].mean(1)
    return instance, search


def closure(fn, sig):
    return fn(sig)


for kwarg_comb in kwargs:
    kwarg_measurement = measurement.get(kwarg_comb, {})
    measurement[kwarg_comb] = kwarg_measurement
    kwarg_comb = dict(kwarg_comb)
    for size in tqdm(input_sizes):
        size_measurement = kwarg_measurement.get(size, {})
        kwarg_measurement[size] = size_measurement
        for xp, cls, _id in modules:
            b, x = input_creation(xp, cls, (size,), **kwarg_comb)
            prof = benchmark(closure, (b, x), n_repeat=100)
            size_measurement[_id] = gather_time(prof)


lines = []
for kwarg_comb in kwargs:
    kwarg_line = ', '.join([f'{x}={y}' for x, y in kwarg_comb])
    lines.append(f'### `{kwarg_line}`\n')
    lines.append('| Size | CuPy (ms) | SciPy (ms) | Speedup |')
    lines.append('|:----:|:---------:|:----------:|:-------:|')
    kwarg_measurement = measurement[kwarg_comb]
    for size in input_sizes:
        comp = f'{size}'
        size_measurement = kwarg_measurement[size]
        times = []
        for _, _, _id in modules:
            time = size_measurement[_id]
            times.append(time)
            time_info = f'{time:3f}'
            comp = f'{comp} | {time_info}'
        speedup = times[1] / times[0]
        comp = f'{comp} | {speedup:3f}'
        lines.append(f'| {comp} |')
    lines.append('\n')


print('\n'.join(lines))