#!/usr/bin/env python
# encoding: utf-8


"""
g_wave.py

Created by Chris Stevens 2022
Copyright (c) 2022 University of Canterbury. All rights reserved.
"""

import numpy as np
from functools import partial

from coffee.tslices import tslices
from coffee.system import System
from coffee.diffop.fd import ghost_point_processor

rt2 = np.sqrt(2.)
r2 = 1./np.sqrt(2.)

def bump(t):
	return np.sin(8.*t)**7.

class OneDAdvection(System):

	################################
	# Constructor
	################################

	def __init__(self, D, tau, CFL, char_speed):
		super(OneDAdvection, self).__init__()
		self.D          = D
		self.tau        = tau
		self.CFL        = CFL
		self.char_speed = char_speed
		self.name       = "1D Advection equation"
		
	def left(self, t):
		return bump(t)

	def evaluate(self, t, U, intStep = None):

		# Define useful variables
		Psi = U.data[0]

		x  = U.domain.axes[0]
		dx = U.domain.step_sizes[0]

		# Calculate spatial derivative
		DxPsi = self.D(Psi, dx)

		# Evaluate right hand side
		DtPsi = -self.char_speed*DxPsi

		# Impose boundary conditions
		new_derivatives, _ = U.communicate(
			partial(ghost_point_processor),
			data=np.array([
				DtPsi
			])
		)
		DtPsi = new_derivatives[0]

		pt_l       = self.D.penalty_boundary(dx, "left")
		pt_l_shape = pt_l.size

		C_left  = self.tau * pt_l * (-self.char_speed)

		l_Psi = self.left(t)

		b_data = U.external_slices()
		for dim, direction, d_slice in b_data:
			if direction == 1:
				pass
			else:
				DtPsi[:pt_l_shape]  -= C_left * (Psi[0] - l_Psi)

		print('t =', t)

		return tslices.TimeSlice([DtPsi], U.domain, time = t)

	def initial_data(self, t0, grid):

		x   = grid.meshes[0]
		Psi = np.zeros_like(x)

		tslice = tslices.TimeSlice([Psi], grid, time = t0)
		return tslice

	def timestep(self, U):

		return self.CFL * U.domain.step_sizes[0]