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

import sys
import logging

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

import utilities as u


@pytest.mark.xfail
@pytest.mark.parametrize(("gen_fluxes_str,char_fluxes_pos_str,char_fluxes_neg_str,"
    "combined_metric,R_str,flux_expected_str"), [
        ("4 11.2 8 8 46.4,1 2.6 1 1 7.1",
            ("1.09071563 1.23015152 1.23015152 7.52305259 0.232956271,"
                "0.467376796 -0.6627417 -0.6627417 1.47953026 0.312527304"),
            ("-0.168354897 -0.0585786438 -0.0585786438 -0.727493464 -0.306026299,"
                "-0.0672231577 0.248528137 0.248528137 -0.107250611 -0.374562227"),
            1.0,
            ("1 0 0 0.45781246 0.45781246,"
                "1.58578644 0 0 1.43309957 0.0188860081,"
                "1.58578644 1.41421356 0 0.725992789 0.725992789,"
                "1.58578644 0 1.41421356 0.725992789 0.725992789,"
                "3.77207794 2.24264069 2.24264069 5.57860029 3.3359596"),
            "4.35371022 12.2479485 8.99522344 8.99522344 51.3903927"),
        ("-1 2.6 1 1 -7.1,-4 11.2 8 8 -46.4",
            ("0.0672231577 0.248528137 0.248528137 0.374562227 0.107250611,"
                "0.168354897 -0.0585786438 -0.0585786438 0.306026299 0.727493464"),
            ("-0.467376796 -0.6627417 -0.6627417 -0.312527304 -1.47953026,"
                "-1.09071563 1.23015152 1.23015152 -0.232956271 -7.52305259"),
            1.0,
            ("1 0 0 0.45781246 0.45781246,"
                "-1.58578644 0 0 -0.0188860081 -1.43309957,"
                "-1.58578644 1.41421356 0 -0.725992789 -0.725992789,"
                "-1.58578644 0 1.41421356 -0.725992789 -0.725992789,"
                "3.77207794 -2.24264069 -2.24264069 3.3359596 5.57860029"),
            "-4.35371022 12.2479485 8.99522344 8.99522344 -51.3903927")
    ])
def test_weno_flux_uniform_grid(
        ctx_factory, gen_fluxes_str, char_fluxes_pos_str, char_fluxes_neg_str,
        combined_metric, R_str, flux_expected_str):
    prg = u.get_weno_program_with_root_kernel("weno_flux")
    queue = u.get_queue(ctx_factory)

    nvars = 5

    gen_fluxes = u.expand_to_6(
            u.transposed_array_from_string(gen_fluxes_str))
    char_fluxes_pos = u.expand_to_6(
            u.transposed_array_from_string(char_fluxes_pos_str))
    char_fluxes_neg = u.expand_to_6(
            u.transposed_array_from_string(char_fluxes_neg_str))
    R = u.array_from_string(R_str)
    flux_dev = u.empty_array_on_device(queue, nvars)

    prg(queue, nvars=nvars,
            generalized_fluxes=gen_fluxes,
            characteristic_fluxes_pos=char_fluxes_pos,
            characteristic_fluxes_neg=char_fluxes_neg,
            combined_frozen_metrics=combined_metric,
            R=R,
            flux=flux_dev)

    flux_expected = u.array_from_string(flux_expected_str)
    u.compare_arrays(flux_dev.get(), flux_expected)


@pytest.mark.parametrize(("gen_fluxes_str,consistent_expected_str"), [
        ("4 11.2 8 8 46.4,1 2.6 1 1 7.1",
            "2.5 6.9 4.5 4.5 26.75"),
        ("-1 2.6 1 1 -7.1,-4 11.2 8 8 -46.4",
            "-2.5 6.9 4.5 4.5 -26.75")
    ])
def test_consistent_part_uniform_grid(
        ctx_factory, gen_fluxes_str, consistent_expected_str):
    prg = u.get_weno_program_with_root_kernel("consistent_part")
    queue = u.get_queue(ctx_factory)

    nvars = 5

    gen_fluxes = u.expand_to_6(
            u.transposed_array_from_string(gen_fluxes_str))
    consistent_dev = u.empty_array_on_device(queue, nvars)

    prg(queue, nvars=nvars,
            generalized_fluxes=gen_fluxes,
            consistent=consistent_dev)

    consistent_expected = u.array_from_string(consistent_expected_str)
    u.compare_arrays(consistent_dev.get(), consistent_expected)


@pytest.mark.slow
@pytest.mark.parametrize(("states_str,fluxes_str,R_inv_str,wavespeeds_str,"
        "fluxes_pos_expected_str,fluxes_neg_expected_str"), [
    ("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",
        ("0.367521364 0.265894836 0.265894836 0.265894836 -0.167673798,"
        "-1.12132034 0 0.707106781 0 0,"
        "-1.12132034 0 0 0.707106781 0,"
        "-0.430558632 0.416709665 -0.290397116 -0.290397116 0.183124984,"
        "1.81208206 -0.997503897 -0.290397116 -0.290397116 0.183124984"),
        "2.2 2.2 2.2 3.84632925 0.55367075",
        ("1.09071563 1.23015152 1.23015152 7.52305259 0.232956271,"
        "0.467376796 -0.6627417 -0.6627417 1.47953026 0.312527304"),
        ("-0.168354897 -0.0585786438 -0.0585786438 -0.727493464 -0.306026299,"
        "-0.0672231577 0.248528137 0.248528137 -0.107250611 -0.374562227")),
    ("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",
        ("0.367521364 -0.265894836 -0.265894836 -0.265894836 -0.167673798,"
        "1.12132034 0 0.707106781 0 0,"
        "1.12132034 0 0 0.707106781 0,"
        "1.81208206 0.997503897 0.290397116 0.290397116 0.183124984,"
        "-0.430558632 -0.416709665 0.290397116 0.290397116 0.183124984"),
        "2.2 2.2 2.2 0.55367075 3.84632925",
        ("0.0672231577 0.248528137 0.248528137 0.374562227 0.107250611,"
        "0.168354897 -0.0585786438 -0.0585786438 0.306026299 0.727493464"),
        ("-0.467376796 -0.6627417 -0.6627417 -0.312527304 -1.47953026,"
        "-1.09071563 1.23015152 1.23015152 -0.232956271 -7.52305259")),
    ("2 4 8 12 64,1 1 2 3 11", "4 11.2 16 24 134.4,1 2.6 2 3 12.6",
        ("-1.41187463 0.217276117 0.434552233 0.65182835 -0.13701474,"
        "-2.24264069 0 0.707106781 0 0,"
        "-3.36396103 0 0 0.707106781 0,"
        "1.79265641 0.444598298 -0.525016967 -0.78752545 0.165538358,"
        "4.03529709 -0.969615265 -0.525016967 -0.78752545 0.165538358"),
        "2.2 2.2 2.2 3.84632925 0.55367075",
        ("1.1162114 2.46030304 3.69045456 7.98692488 0.429010898,"
        "0.260735274 -1.3254834 -1.9882251 2.05242289 0.501788756"),
        ("-0.148282372 -0.117157288 -0.175735931 -0.889325254 -0.200040418,"
        "-0.00948821371 0.497056275 0.745584412 -0.430637881 -0.314318322")),
    ])
def test_flux_splitting_uniform_grid(
        ctx_factory, states_str, fluxes_str, R_inv_str, wavespeeds_str,
        fluxes_pos_expected_str, fluxes_neg_expected_str):
    prg = u.get_weno_program_with_root_kernel("split_characteristic_fluxes")
    queue = u.get_queue(ctx_factory)

    nvars = 5

    states = u.expand_to_6(u.transposed_array_from_string(states_str))
    fluxes = u.expand_to_6(u.transposed_array_from_string(fluxes_str))
    R_inv = u.array_from_string(R_inv_str)
    wavespeeds = u.array_from_string(wavespeeds_str)
    fluxes_pos_dev = u.empty_array_on_device(queue, nvars, 6)
    fluxes_neg_dev = u.empty_array_on_device(queue, nvars, 6)

    prg(queue, nvars=nvars,
            generalized_states_frozen=states,
            generalized_fluxes_frozen=fluxes,
            R_inv=R_inv,
            wavespeeds=wavespeeds,
            characteristic_fluxes_pos=fluxes_pos_dev,
            characteristic_fluxes_neg=fluxes_neg_dev)

    fluxes_pos_expected = u.expand_to_6(
        u.transposed_array_from_string(fluxes_pos_expected_str))
    u.compare_arrays(fluxes_pos_dev.get(), fluxes_pos_expected)

    fluxes_neg_expected = u.expand_to_6(
        u.transposed_array_from_string(fluxes_neg_expected_str))
    u.compare_arrays(fluxes_neg_dev.get(), fluxes_neg_expected)


@pytest.mark.slow
@pytest.mark.parametrize("lam_pointwise_str,lam_roe_str,lam_expected_str", [
    ("1 2 3 4 5,2 4 6 8 10", "1.5 3 4.5 6 7.5", "2.2 4.4 6.6 8.8 11"),
    ("1 2 3 4 5,-2 -4 -6 -8 -10", "1.5 3 4.5 6 7.5", "2.2 4.4 6.6 8.8 11"),
    ("1 2 3 4 5,-2 -4 -6 -8 -10", "3 6 9 12 15", "3.3 6.6 9.9 13.2 16.5"),
    ("1 2 3 4 5,2 4 6 8 10", "-3 -6 -9 -12 -15", "3.3 6.6 9.9 13.2 16.5"),
    ("3 2 9 4 5,2 6 6 12 10", "-1 -4 -3 -8 -15", "3.3 6.6 9.9 13.2 16.5")
    ])
def test_lax_wavespeeds(
        ctx_factory, lam_pointwise_str, lam_roe_str, lam_expected_str):
    prg = u.get_weno_program_with_root_kernel("lax_wavespeeds")
    queue = u.get_queue(ctx_factory)

    nvars = 5

    lam_pointwise = u.expand_to_6(u.transposed_array_from_string(lam_pointwise_str))
    lam_roe = u.array_from_string(lam_roe_str)
    lam_dev = u.empty_array_on_device(queue, nvars)

    prg(queue, nvars=nvars, lambda_pointwise=lam_pointwise,
            lambda_roe=lam_roe, wavespeeds=lam_dev)

    lam_expected = u.array_from_string(lam_expected_str)
    u.compare_arrays(lam_dev.get(), lam_expected)


@pytest.mark.slow
@pytest.mark.parametrize("states_str,direction,lam_expected_str", [
    ("2 4 4 4 20,1 1 1 1 5.5", "x",
        "2 2 2 3.49666295 0.503337045,1 1 1 2.49666295 -0.496662955"),
    ("2 4 4 4 20,1 1 1 1 5.5", "y",
        "2 2 2 3.49666295 0.503337045,1 1 1 2.49666295 -0.496662955"),
    ("2 4 4 4 20,1 1 1 1 5.5", "z",
        "2 2 2 3.49666295 0.503337045,1 1 1 2.49666295 -0.496662955"),
    ("1 -1 -1 -1 5.5,2 -4 -4 -4 20", "x",
        "-1 -1 -1 0.496662955 -2.49666295,-2 -2 -2 -0.503337045 -3.49666295"),
    ("1 -1 -1 -1 5.5,2 -4 -4 -4 20", "y",
        "-1 -1 -1 0.496662955 -2.49666295,-2 -2 -2 -0.503337045 -3.49666295"),
    ("1 -1 -1 -1 5.5,2 -4 -4 -4 20", "z",
        "-1 -1 -1 0.496662955 -2.49666295,-2 -2 -2 -0.503337045 -3.49666295"),
    ("2 4 8 12 64,1 1 2 3 11", "x",
        "2 2 2 3.49666295 0.503337045,1 1 1 2.49666295 -0.496662955"),
    ("2 4 8 12 64,1 1 2 3 11", "y",
        "4 4 4 5.49666295 2.50333705,2 2 2 3.49666295 0.503337045"),
    ("2 4 8 12 64,1 1 2 3 11", "z",
        "6 6 6 7.49666295 4.50333705,3 3 3 4.49666295 1.503337045"),
    ("1 -1 -2 -3 11,2 -4 -8 -12 64", "x",
        "-1 -1 -1 0.496662955 -2.49666295,-2 -2 -2 -0.503337045 -3.49666295"),
    ("1 -1 -2 -3 11,2 -4 -8 -12 64", "y",
        "-2 -2 -2 -0.503337045 -3.49666295,-4 -4 -4 -2.50333705 -5.49666295"),
    ("1 -1 -2 -3 11,2 -4 -8 -12 64", "z",
        "-3 -3 -3 -1.50333705 -4.49666295,-6 -6 -6 -4.50333705 -7.49666295")
    ])
def test_pointwise_eigenvalues_ideal_gas(
        ctx_factory, states_str, direction, lam_expected_str):
    prg = u.get_weno_program_with_root_kernel("pointwise_eigenvalues")
    queue = u.get_queue(ctx_factory)

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

    states = u.expand_to_6(u.transposed_array_from_string(states_str))
    lam_dev = u.empty_array_on_device(queue, nvars, 6)

    prg(queue, nvars=nvars, d=dirs[direction],
            states=states, lambda_pointwise=lam_dev)

    lam_expected = u.expand_to_6(u.transposed_array_from_string(lam_expected_str))
    u.compare_arrays(lam_dev.get(), lam_expected)


@pytest.mark.slow
@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"),
    ("1 -1 -2 -3 11,2 -4 -8 -12 64", "-1 2.6 2 3 -12.6,-4 11.2 16 24 -134.4", "x"),
    ("1 -1 -2 -3 11,2 -4 -8 -12 64", "-2 2 5.6 6 -25.2,-8 16 35.2 48 -268.8", "y"),
    ("1 -1 -2 -3 11,2 -4 -8 -12 64", "-3 3 6 10.6 -37.8,-12 24 48 75.2 -403.2", "z")
    ])
def test_roe_uniform_grid_ideal_gas(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, R_inv):
        d_state = states[:, 1] - states[:, 0]
        u.compare_arrays(R@(R_inv@d_state), d_state)

    def check_roe_property(states, fluxes, R, R_inv, lam):
        d_state = states[:, 1] - states[:, 0]
        d_flux = fluxes[:, 1] - fluxes[:, 0]

        temp = R_inv@d_state
        temp = np.multiply(lam, temp)
        u.compare_arrays(R@temp, d_flux)

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

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

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

    R_dev = u.empty_array_on_device(queue, nvars, nvars)
    R_inv_dev = u.empty_array_on_device(queue, nvars, nvars)
    lam_dev = u.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=R_inv_dev, lambda_roe=lam_dev)

    R = R_dev.get()
    R_inv = R_inv_dev.get()
    lam = lam_dev.get()

    print(lam)

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


@pytest.mark.slow
def test_matvec(ctx_factory):
    prg = u.get_weno_program_with_root_kernel("mult_mat_vec")
    queue = u.get_queue(ctx_factory)

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

    c = u.empty_array_on_device(queue, 10)

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

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


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

    queue = u.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 = u.get_weno_program()
    prg = u.transform_weno_for_gpu(prg)

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

    if write_code:
        u.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__])