import numpy as np
import matplotlib.pyplot as plt
from numpy.linalg import norm
from numpy import dot

def scatter(k0, g, n):
    # assumes k_incident is in the -z direction
    # k0 is the scalar incident wave vector (magnitude)
    # g is a reciprocal vector of the diffractive surface
    # n is the normal vector to that surface

    z_hat = np.array([0, 0, 1])

    k0_perp = -k0 * n * dot(z_hat, n)

    # k0_par = -k0 * z_hat + k0 * n * dot(z_hat, n)
    # k0_par = k0 * (n * dot(z_hat, n) - z_hat)
    # k1_par = k0_par + g

    k1_par = k0 * (n * dot(z_hat, n) - z_hat) + g

    k1_perp = (k0**2 - norm(k1_par)**2)**0.5 * n

    k1 = k1_par + k1_perp

    return k1


k0 = 10.
n = np.array([1, 0, 1]) / 2**0.5
g = 3. * np.array([0, 1, 0])

k1 = scatter(k0, g, n)

print('k0: ', k0)
print('n: ', n)
print('g: ', g)
print('')

print('check perpendicularity, g dot n: ', dot(g, n))
print('')

print('k1: ', k1)
print('')

print('check conservation of energy norm(k1) = k0: ', norm(k1), k0)