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


# FIXME: add new file ideal_gas.py that has flux, pressure, etc. (any
# computation that depends on EOS

def pressure(state):
    # FIXME: for ideal gas stuff, make nvars, ndim, gamma, etc module level
    # constants
    # FIXME: also split out velocity, rho, energy into separate routines
    gamma = 1.4
    ndim = state.size - 2

    rho = state[0]
    energy = state[ndim+1]

    ke = 0.0
    for i in range(ndim):
        ke += (state[i+1]/rho)**2
    ke = 0.5*rho*ke

    return (gamma - 1)*(energy - ke)


def sound_speed(state):
    gamma = 1.4
    p = pressure(state)
    rho = state[0]

    return np.sqrt(gamma*p/rho)


def velocity(state):
    rho = state[0]
    return state[1:4]/rho


def flux(state):
    nvars = state.size
    ndim = nvars - 2
    result = u.empty_array(nvars, ndim)

    rho = state[0]
    vel = state[1:ndim+1]/rho
    energy = state[ndim+1]
    p = pressure(state)

    for i in range(ndim):
        result[:,i] = vel[i]*state
        result[i+1,i] += p
        result[ndim+1,i] += vel[i]*p

    return result


def pointwise_fluxes(states):
    nvars, npoints = states.shape
    ndim = 3
    result = u.empty_array(nvars, ndim, npoints)

    for k in range(npoints):
        result[:,:,k] = flux(states[:,k])

    return result


def roe_eigensystem(state_pair, frozen_metrics, direction):
    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, direction):
        c = sound_speed(state)
        c_norm = c*metric_norm[direction]
        vel = 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 np.array([lam(s, m, direction-1)
        for s, m in zip(states.T, metric_norm)]).T.copy(order="F")


def split_char_fluxes(states, wavespeeds, frozen_metrics, frozen_jacobian, R_inv):
    nvars, npoints = states.shape
    ndim = frozen_metrics.size

    fluxes = pointwise_fluxes(states)

    char_fluxes_pos = u.empty_array(nvars, npoints)
    char_fluxes_neg = u.empty_array(nvars, npoints)
    for k in range(npoints):
        generalized_fluxes = u.zero_array(nvars)
        for i in range(ndim):
            generalized_fluxes += frozen_metrics[i]*fluxes[:,i,k]

        generalized_states = states[:,k]/frozen_jacobian

        combination_pos = u.empty_array(nvars, nvars)
        combination_neg = u.empty_array(nvars, nvars)
        for i in range(nvars):
            char_fluxes_pos[i,k] = 0.5*np.dot(R_inv[i,:],
                    (generalized_fluxes + wavespeeds[i]*generalized_states))
            char_fluxes_neg[i,k] = 0.5*np.dot(R_inv[i,:],
                    (generalized_fluxes - wavespeeds[i]*generalized_states))

    return char_fluxes_pos, char_fluxes_neg


def oscillation(fluxes):
    sum1 = u.empty_array(fluxes.shape[0], 3)
    sum2 = u.empty_array(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 = u.empty_array(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 = u.empty_array(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 = u.empty_array(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