__copyright__ = "Copyright (C) 2019 Timothy A. Smith"

__license__ = """
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
"""

import numpy as np
import pytest
import utilities as u
import weno_reference_implementation as ref


# {{{ FluxDataSingle

class FluxDataSingle:
    # FIXME: can we set some of these constants from ref.gas?
    #   -- if all nvars references come from there, it's relatively easy to
    #   introduce a new gas with more (e.g. scalar) variables
    nvars = 5
    ndim = 3
    dirs = {"x": 1, "y": 2, "z": 3}

    def __init__(self, queue, states_str, direction):

        self.direction = self.dirs[direction]
        self.dir_internal = self.direction-1

        self.metrics = np.array([np.identity(self.ndim) for i in range(6)],
                dtype=np.float64, order="F")
        self.jacobians = np.repeat(1.0, 6)

        # FIXME: should be computed directly from the metrics and jacobians
        self.frozen_metrics = np.mean(self.metrics[2:4], axis=0)
        self.frozen_jacobian = np.mean(self.jacobians[2:4], axis=0)
        self.combined_frozen_metrics = 1.0

        # FIXME: Move array_from_string stuff outside FluxDataSingle
        #   -- just pass an array & have external utilities that generate
        #   Riemann, sine wave, etc. initial conditions
        # FIXME: Consider handling row swapping outside as well?
        # FIXME: Do we even need to swap rows?
        self.state_pair = self.swap_array_rows(
                u.transposed_array_from_string(states_str), self.dir_internal)
        self.states = u.expand_to_n(self.state_pair, 6)

        # FIXME: these should be generalized fluxes
        # FIXME: make a clear distinction between fluxes in physical and
        # generalized coordinates
        self.flux_pair = ref.pointwise_fluxes(
                self.state_pair)[:,:,self.dir_internal].T.copy(order="F")
        self.fluxes = ref.pointwise_fluxes(
                self.states)[:,:,self.dir_internal].T.copy(order="F")

        self.lam_pointwise = ref.lambda_pointwise(
                self.states, self.metrics, self.dir_internal)
        self.R, self.R_inv, self.lam_roe = ref.roe_eigensystem(
                queue, self.state_pair, self.frozen_metrics, self.direction)
        self.wavespeeds = ref.wavespeeds(self.lam_pointwise, self.lam_roe)

        self.char_fluxes_pos, self.char_fluxes_neg = ref.split_char_fluxes(
                self.states, self.wavespeeds,
                self.frozen_metrics[self.dir_internal], self.frozen_jacobian,
                self.R_inv)

        self.oscillation_pos = ref.oscillation(self.char_fluxes_pos)
        self.oscillation_neg = ref.oscillation(self.char_fluxes_neg[:,::-1])

        self.weno_weights_pos = ref.weno_weights(
                self.oscillation_pos, self.combined_frozen_metrics)
        self.weno_weights_neg = ref.weno_weights(
                self.oscillation_neg, self.combined_frozen_metrics)

        self.consistent = ref.consistent_part(self.fluxes)
        self.dissipation_pos = ref.dissipation_part(
                self.R, self.char_fluxes_pos, self.weno_weights_pos, 1)
        self.dissipation_neg = ref.dissipation_part(
                self.R, self.char_fluxes_neg, self.weno_weights_neg, -1)

        self.weno_flux = ref.weno_flux(
                self.consistent, self.dissipation_pos, self.dissipation_neg)

    def swap_array_rows(self, arr, d):
        p = self.permutation(d)
        arr[p, :] = arr[[1, 2, 3], :]
        return arr

    def permutation(self, d):
        return [(d+i) % 3 + 1 for i in range(3)]

# }}}

# {{{ FluxDataVector

# FIXME: is there a better way to divide responsibilities with these fixture classes?
class FluxDataVector:
    nvars = 5
    ndim = 3
    dirs = {"x": 1, "y": 2, "z": 3}
    halo = 3

    def __init__(self, nx, ny, nz, states_str, direction):
        self.direction = self.dirs[direction]
        self.dir_internal = self.direction-1

        self.nx = nx
        self.ny = ny
        self.nz = nz

        self.nxhalo = self.nx + 2*self.halo
        self.nyhalo = self.ny + 2*self.halo
        self.nzhalo = self.nz + 2*self.halo

        self.flux_dims = (self.nvars, self.nx, self.ny, self.nz)

        self.metrics = np.stack(
                [np.stack(
                    [np.stack(
                        [np.identity(self.ndim) for i in range(self.nxhalo)],
                        axis=-1) for j in range(self.nyhalo)],
                    axis=-1) for k in range(self.nzhalo)],
                axis=-1).copy(order="F")
        self.jacobians = np.ones((self.nxhalo, self.nyhalo, self.nzhalo),
                order="F")

        state_pair = self.swap_array_rows(
                u.transposed_array_from_string(states_str), self.direction)
        # FIXME: Move array_from_string stuff outside FluxDataSingle
        #   -- just pass an array & have external utilities that generate
        #   Riemann, sine wave, etc. initial conditions
        # FIXME: Consider handling row swapping outside as well?
        # FIXME: Do we even need to swap rows?
        self.state_pair = self.swap_array_rows(
                u.transposed_array_from_string(states_str), self.dir_internal)
        # NOTE: dimensions are nvars x nxhalo x nyhalo x nzhalo
        self.states = self.fill_from_pair()

        # NOTE: dimensions are nvars x nxhalo x nyhalo x nzhalo
        # FIXME: these should be generalized fluxes
        # FIXME: make a clear distinction between fluxes in physical and
        # generalized coordinates
        npoints = self.nxhalo*self.nyhalo*self.nzhalo
        flat_states = self.states.reshape((self.nvars, npoints))
        self.fluxes = ref.pointwise_fluxes(
                flat_states)[:,:,self.dir_internal].T.reshape(
                        (self.nvars, self.nxhalo, self.nyhalo, self.nzhalo)
                        ).copy(order="F")

        # FIXME: use reference implementation
        # NOTE: dimensions are nvars x nx x ny x nz
        self.flux_derivatives = np.zeros((self.nvars, self.nx, self.ny,
            self.nz), order="F")

    def swap_array_rows(self, arr, d):
        p = self.permutation(d)
        arr[p, :] = arr[[1, 2, 3], :]
        return arr

    def permutation(self, d):
        return [(d-1+i) % 3 + 1 for i in range(3)]

    def fill_from_pair(self):
        d = self.dir_internal
        nx_arr = np.array([self.nxhalo, self.nyhalo, self.nzhalo])
        result = u.expand_to_n(self.state_pair, nx_arr[d])

        for i in range(d):
            result = self.add_dimension(result, nx_arr[i])
            result = np.swapaxes(result, -2, -1)
        for i in range(d+1,self.ndim):
            result = self.add_dimension(result, nx_arr[i])

        return result.copy(order="F")

    def add_dimension(self, arr, n):
        return np.stack([arr for i in range(n)], axis=-1)

# }}}

@pytest.fixture(scope="session", params=[
    ("1 1 1 1 5.5,1 1 1 1 5.5", "x"),
    ("1 1 1 1 5.5,1 1 1 1 5.5", "y"),
    ("1 1 1 1 5.5,1 1 1 1 5.5", "z"),
    ("2 4 4 4 20,1 1 1 1 5.5", "x"),
    ("2 4 4 4 20,1 1 1 1 5.5", "y"),
    ("2 4 4 4 20,1 1 1 1 5.5", "z"),
    ("1 -1 -1 -1 5.5,2 -4 -4 -4 20", "x"),
    ("1 -1 -1 -1 5.5,2 -4 -4 -4 20", "y"),
    ("1 -1 -1 -1 5.5,2 -4 -4 -4 20", "z"),
    ("2 4 8 12 64,1 1 2 3 11", "x"),
    ("2 8 12 4 64,1 2 3 1 11", "y"),
    ("2 12 4 8 64,1 3 1 2 11", "z"),
    ("1 -1 -2 -3 11,2 -4 -8 -12 64", "x"),
    ("1 -2 -3 -1 11,2 -8 -12 -4 64", "y"),
    ("1 -3 -1 -2 11,2 -12 -4 -8 64", "z")
    ])
def flux_test_data_fixture(request, queue):
    return FluxDataSingle(queue, *request.param)


vector_data = {}

vector_data["Case flat:x"] = FluxDataVector(
        nx=6, ny=2, nz=2,
        states_str="1 1 1 1 5.5,1 1 1 1 5.5",
        direction="x")
vector_data["Case flat:y"] = FluxDataVector(
        nx=2, ny=6, nz=2,
        states_str="1 1 1 1 5.5,1 1 1 1 5.5",
        direction="y")
vector_data["Case flat:z"] = FluxDataVector(
        nx=2, ny=2, nz=6,
        states_str="1 1 1 1 5.5,1 1 1 1 5.5",
        direction="z")

vector_data["Case a:x"] = FluxDataVector(
        nx=6, ny=2, nz=2,
        states_str="2 4 4 4 20,1 1 1 1 5.5",
        direction="x")
vector_data["Case a:y"] = FluxDataVector(
        nx=2, ny=6, nz=2,
        states_str="2 4 4 4 20,1 1 1 1 5.5",
        direction="y")
vector_data["Case a:z"] = FluxDataVector(
        nx=2, ny=2, nz=6,
        states_str="2 4 4 4 20,1 1 1 1 5.5",
        direction="z")


@pytest.fixture(scope="session", params=[
    "Case flat:x", "Case flat:y", "Case flat:z"])
def cfd_test_data_fixture(request):
    return vector_data[request.param]


# vim: foldmethod=marker