import numpy as np
import numpy.linalg as la
import pyopencl as cl
import pyopencl.array  # noqa
import pyopencl.tools  # noqa
import pyopencl.clrandom  # noqa
import loopy as lp  # noqa

import sys
import logging

import pytest
from pytest import approx
from pyopencl.tools import (  # noqa
        pytest_generate_tests_for_pyopencl
        as pytest_generate_tests)

from utilities import *


@pytest.mark.parametrize("states_str,fluxes_str,direction", [
    ("2 4 4 4 20,1 1 1 1 5.5", "4 11.2 8 8 46.4,1 2.6 1 1 7.1", "x"),
    ("2 4 4 4 20,1 1 1 1 5.5", "4 8 11.2 8 46.4,1 1 2.6 1 7.1", "y"),
    ("2 4 4 4 20,1 1 1 1 5.5", "4 8 8 11.2 46.4,1 1 1 2.6 7.1", "z"),
    ("1 -1 -1 -1 5.5,2 -4 -4 -4 20", "-1 2.6 1 1 -7.1,-4 11.2 8 8 -46.4", "x"),
    ("1 -1 -1 -1 5.5,2 -4 -4 -4 20", "-1 1 2.6 1 -7.1,-4 8 11.2 8 -46.4", "y"),
    ("1 -1 -1 -1 5.5,2 -4 -4 -4 20", "-1 1 1 2.6 -7.1,-4 8 8 11.2 -46.4", "z"),
    ("2 4 8 12 64,1 1 2 3 11", "4 11.2 16 24 134.4,1 2.6 2 3 12.6", "x"),
    ("2 4 8 12 64,1 1 2 3 11", "8 16 35.2 48 268.8,2 2 5.6 6 25.2", "y"),
    ("2 4 8 12 64,1 1 2 3 11", "12 24 48 75.2 403.2,3 3 6 10.6 37.8", "z")
    ])
def test_roe_uniform_grid(ctx_factory, states_str, fluxes_str, direction):
    def identity_matrix(n):
        return np.identity(n).astype(np.float32).copy(order="F")

    def check_roe_identity(states, R, Rinv):
        dState = states[:,1] - states[:,0]
        compare_arrays(R@(Rinv@dState), dState)

    def check_roe_property(states, fluxes, R, Rinv, lam):
        dState = states[:,1] - states[:,0]
        dFlux = fluxes[:,1] - fluxes[:,0]

        temp = Rinv@dState
        temp = np.multiply(lam, temp)
        compare_arrays(R@temp, dFlux)

    prg = get_weno_program_with_root_kernel("roe_eigensystem")
    queue = get_queue(ctx_factory)

    nvars = 5
    ndim = 3
    dirs = {"x" : 1, "y" : 2, "z" : 3}

    states = transposed_array_from_string(states_str)
    fluxes = transposed_array_from_string(fluxes_str)
    metrics_frozen = identity_matrix(ndim)

    R_dev = empty_array_on_device(queue, nvars, nvars)
    Rinv_dev = empty_array_on_device(queue, nvars, nvars)
    lam_dev = empty_array_on_device(queue, nvars)

    prg(queue, nvars=nvars, ndim=ndim, d=dirs[direction],
            states=states, metrics_frozen=metrics_frozen,
            R=R_dev, R_inv=Rinv_dev, lambda_roe=lam_dev)

    R = R_dev.get()
    Rinv = Rinv_dev.get()
    lam = lam_dev.get()

    print(lam)

    check_roe_identity(states, R, Rinv)
    check_roe_property(states, fluxes, R, Rinv, lam)


def test_matvec(ctx_factory):
    prg = get_weno_program_with_root_kernel("mult_mat_vec")
    queue = get_queue(ctx_factory)

    a = random_array_on_device(queue, 10, 10)
    b = random_array_on_device(queue, 10)

    c = empty_array_on_device(queue, 10)

    prg(queue, alpha=1.0, a=a, b=b, c=c)

    compare_arrays(a.get()@b.get(), c.get())


@pytest.mark.slow
def test_compute_flux_derivatives(ctx_factory):
    prg = get_weno_program()

    queue = get_queue(ctx_factory)
    prg = prg.copy(target=lp.PyOpenCLTarget(queue.device))

    lp.auto_test_vs_ref(prg, ctx_factory(), warmup_rounds=1,
            parameters=dict(ndim=3, nvars=5, nx=16, ny=16, nz=16))


@pytest.mark.slow
def test_compute_flux_derivatives_gpu(ctx_factory, write_code=False):
    prg = get_weno_program()
    prg = transform_weno_for_gpu(prg)

    queue = get_queue(ctx_factory)
    prg = prg.copy(target=lp.PyOpenCLTarget(queue.device))
    prg = lp.set_options(prg, no_numpy=True)

    if write_code:
        write_target_device_code(prg)

    lp.auto_test_vs_ref(prg, ctx_factory(), warmup_rounds=1,
            parameters=dict(ndim=3, nvars=5, nx=16, ny=16, nz=16))


# This lets you run 'python test.py test_case(cl._csc)' without pytest.
if __name__ == "__main__":
    if len(sys.argv) > 1:
        logging.basicConfig(level="INFO")
        exec(sys.argv[1])
    else:
        pytest.main([__file__])