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 utilities as u
import ideal_gas as gas


def pointwise_fluxes(states):
    return np.array([gas.flux(s) for s in states.T])


def roe_eigensystem(state_pair, frozen_metrics, direction):
    # FIXME: startup for test suite is pretty slow due to this routine
    #   -- can we speed this up?

    nvars = state_pair.shape[0]
    ndim = frozen_metrics.shape[0]

    prg = u.get_weno_program_with_root_kernel("roe_eigensystem")
    queue = u.get_queue(cl._csc)

    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=direction,
            states=state_pair, metrics_frozen=frozen_metrics,
            R=R_dev, R_inv=R_inv_dev, lambda_roe=lam_dev)

    return R_dev.get(), R_inv_dev.get(), lam_dev.get()


def lambda_pointwise(states, metrics, direction):
    def lam(state, metric_norm):
        c = gas.sound_speed(state)
        c_norm = c*metric_norm[direction]
        vel = gas.velocity(state)[direction]

        result = np.repeat(vel, state.size)
        result[-2] += c_norm
        result[-1] -= c_norm

        return result

    metric_norm = np.sqrt((metrics**2).sum(axis=1))

    return u.transposed_array([lam(s, m) for s, m in zip(states.T, metric_norm)])


def wavespeeds(pointwise, roe):
    lam = np.c_[pointwise, roe]
    return 1.1*np.max(np.abs(lam), axis=1)


def split_char_fluxes(states, wavespeeds, frozen_metrics, frozen_jacobian, R_inv):
    def split(flux, state):
        generalized_fluxes = np.dot(flux, frozen_metrics)
        weighted_states = np.outer(wavespeeds, state/frozen_jacobian)

        return (0.5*np.sum(R_inv*(generalized_fluxes + weighted_states), axis=1),
                0.5*np.sum(R_inv*(generalized_fluxes - weighted_states), axis=1))

    fluxes = pointwise_fluxes(states)

    char_fluxes_pos, char_fluxes_neg = zip(
            *[split(f, s) for f, s in zip(fluxes, states.T)])

    return (u.transposed_array(char_fluxes_pos),
            u.transposed_array(char_fluxes_neg))


def oscillation(fluxes):
    coeffs1 = np.array([[1, -4, 3], [-1, 0, 1], [-3, 4, -1]])
    coeffs2 = np.array([1, -2, 1])

    indices = np.arange(3)[None,:] + np.arange(3)[:,None]

    sum1 = u.transposed_array(
            [np.dot(c, f) for c, f in zip(coeffs1, fluxes.T[indices])])
    sum2 = u.transposed_array([np.dot(coeffs2, f) for f in fluxes.T[indices]])

    return (1.0/4)*(sum1**2) + (13.0/12)*(sum2**2)


def weno_weights(oscillation, frozen_metric):
    linear = np.array([0.1, 0.6, 0.3])
    eps = 1e-6*frozen_metric

    raw_weights = linear[None,:]/(oscillation + eps)**2

    return raw_weights/raw_weights.sum(axis=1)[:,None]


def flux_differences(fluxes):
    w = np.array([-1, 3, -3, 1])
    indices = np.arange(3)[:,None] + np.arange(4)[None,:]

    return u.transposed_array([np.dot(w, f) for f in fluxes.T[indices]])


def combination_weighting(w):
    return np.array(
            [20*w[:,0] - 1, -10*(w[:,0] + w[:,1]) + 5, np.ones(w.shape[0])]
            ).T


def combine_fluxes(w, f):
    cw = combination_weighting(w)
    return np.sum(cw*f, axis=1)


def dissipation_part(R, char_fluxes, w, sign):
    flux_diff = flux_differences(char_fluxes)[:,::sign]
    flux_comb = combine_fluxes(w, flux_diff)

    return -sign*R@flux_comb/60


def consistent_part(fluxes):
    w = np.array([1.0, -8.0, 37.0, 37.0, -8.0, 1.0])/60.0
    return np.dot(fluxes, w)


def weno_flux(consistent, dissipation_pos, dissipation_neg):
    return consistent + dissipation_pos + dissipation_neg