# -*- coding: utf-8 -*-
"""
@author: Laura C. Kühle

d = detail coefficient (rename?)
other A (from M) = ? (Is it the same???)
A = basis_projection_left
M1 = wavelet_projection_left
phi = DG basis vector
psi = wavelet vector

TODO: Find better names for A, B, M1, and M2

"""
import numpy as np
import timeit


class UpdateScheme(object):
    def __init__(self, detector, limiter, init_cond, mesh, wave_speed, polynom_degree, num_grid_cells, final_time,
                 history_threshold, left_bound, right_bound):
        # Unpack positional arguments
        self._detector = detector
        self._limiter = limiter
        self._init_cond = init_cond
        self._mesh = mesh
        self._wave_speed = wave_speed
        self._polynom_degree = polynom_degree
        self._num_grid_cells = num_grid_cells
        self._final_time = final_time
        self._history_threshold = history_threshold
        self._left_bound = left_bound
        self._right_bound = right_bound

        self._reset()

    def get_name(self):
        return self.name

    def step(self, projection, cfl_number, current_time):
        current_projection, troubled_cells = self._apply_stability_method(projection, cfl_number)

        return current_projection, troubled_cells

    def _apply_stability_method(self, projection, cfl_number):
        return projection, []

    def _reset(self):
        # Set additional necessary fixed instance variables
        self.name = 'None'
        self._interval_len = self._right_bound-self._left_bound
        self._cell_len = self._interval_len / self._num_grid_cells

        # Set matrix A
        matrix = []
        for i in range(self._polynom_degree+1):
            new_row = []
            for j in range(self._polynom_degree+1):
                new_entry = -1.0
                if (j < i) & ((i+j) % 2 == 1):
                    new_entry = 1.0
                new_row.append(new_entry*np.sqrt((i+0.5) * (j+0.5)))
            matrix.append(new_row)
        self._A = np.array(matrix)  # former: inv_mass @ np.array(matrix)

        # Set matrix B
        matrix = []
        for i in range(self._polynom_degree+1):
            new_row = []
            for j in range(self._polynom_degree+1):
                new_entry = np.sqrt((i+0.5) * (j+0.5)) * (-1.0)**i
                new_row.append(new_entry)
            matrix.append(new_row)
        self._B = np.array(matrix)  # former: inv_mass @ np.array(matrix)

    def _apply_limiter(self, current_projection):
        troubled_cells = self._detector.get_cells(current_projection)

        new_projection = current_projection.copy()
        for cell in troubled_cells:
            new_projection[:,  cell] = self._limiter.apply(current_projection, cell)

        return new_projection, troubled_cells

    def _enforce_boundary_condition(self, current_projection):
        current_projection[:, 0] = current_projection[:, self._num_grid_cells]
        current_projection[:, self._num_grid_cells+1] = current_projection[:, 1]
        return current_projection


class SSPRK3(UpdateScheme):
    def __init__(self, detector, limiter, init_cond, mesh, wave_speed, polynom_degree, num_grid_cells, final_time,
                 history_threshold, left_bound, right_bound):
        super().__init__(detector, limiter, init_cond, mesh, wave_speed, polynom_degree, num_grid_cells, final_time,
                         history_threshold, left_bound, right_bound)

        # Set name of update scheme
        self.name = 'SSPRK3'

    # Override method of superclass
    def _apply_stability_method(self, projection, cfl_number):
        original_projection = projection

        current_projection = self._apply_first_step(original_projection, cfl_number)
        current_projection, __ = self._apply_limiter(current_projection)
        current_projection = self._enforce_boundary_condition(current_projection)

        current_projection = self._apply_second_step(original_projection, current_projection, cfl_number)
        current_projection, __ = self._apply_limiter(current_projection)
        current_projection = self._enforce_boundary_condition(current_projection)

        current_projection = self._apply_third_step(original_projection, current_projection, cfl_number)
        current_projection, troubled_cells = self._apply_limiter(current_projection)
        current_projection = self._enforce_boundary_condition(current_projection)

        return current_projection, troubled_cells

    def _update_right_hand_side(self, current_projection):
        # Initialize vector and set first entry to accommodate for ghost cell
        right_hand_side = [0]

        for j in range(self._num_grid_cells):
            right_hand_side.append(2*(self._A @ current_projection[:, j+1]
                                      + self._B @ current_projection[:, j]))

        # Set ghost cells to respective value
        right_hand_side[0] = right_hand_side[self._num_grid_cells]
        right_hand_side.append(right_hand_side[1])

        return np.transpose(right_hand_side)

    def _apply_first_step(self, original_projection, cfl_number):
        right_hand_side = self._update_right_hand_side(original_projection)
        return original_projection + (cfl_number*right_hand_side)

    def _apply_second_step(self, original_projection, current_projection, cfl_number):
        right_hand_side = self._update_right_hand_side(current_projection)
        return 1/4 * (3*original_projection + (current_projection + cfl_number*right_hand_side))

    def _apply_third_step(self, original_projection, current_projection, cfl_number):
        right_hand_side = self._update_right_hand_side(current_projection)
        return 1/3 * (original_projection + 2*(current_projection + cfl_number*right_hand_side))