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


def oscillation(fluxes):
    sum1 = np.empty((fluxes.shape[0], 3))
    sum2 = np.empty((fluxes.shape[0], 3))

    sum1[:,0] = fluxes[:,0] - 4*fluxes[:,1] + 3*fluxes[:,2]
    sum1[:,1] = -fluxes[:,1] + fluxes[:,3]
    sum1[:,2] = -3*fluxes[:,2] + 4*fluxes[:,3] - fluxes[:,4]

    for i in range(3):
        sum2[:,i] = fluxes[:,i] - 2*fluxes[:,i+1] + fluxes[:,i+2]

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

    return result


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

    raw_weights = np.empty((5,3))
    for i in range(5):
        for j in range(3):
            raw_weights[i,j] = linear[j]/(oscillation[i,j] + eps)**2

    weight_sum = raw_weights.sum(axis=1)
    weights = np.empty((5,3))
    for i in range(5):
        for j in range(3):
            weights[i,j] = raw_weights[i,j]/weight_sum[i]

    return weights


def flux_differences(f):
    w = np.array([-1, 3, -3, 1])
    res = np.empty((5, 3))
    for j in range(res.shape[1]):
        res[:,j] = f[:,j:j+4]@w
    return res


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(np.multiply(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 weno_flux(consistent, dissipation_pos, dissipation_neg):
    return consistent + dissipation_pos + dissipation_neg