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

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

"""
import numpy as np
import timeit


class UpdateScheme(object):
    def __init__(self, polynomial_degree, num_grid_cells, detector, limiter):
        # Unpack positional arguments
        self._polynomial_degree = polynomial_degree
        self._num_grid_cells = num_grid_cells
        self._detector = detector
        self._limiter = limiter

        self._reset()

    def _reset(self):
        # Set stiffness matrix
        matrix = []
        for i in range(self._polynomial_degree+1):
            new_row = []
            for j in range(self._polynomial_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._stiffness_matrix = np.array(matrix)  # former: inv_mass @ np.array(matrix)

        # Set boundary matrix
        matrix = []
        for i in range(self._polynomial_degree+1):
            new_row = []
            for j in range(self._polynomial_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._boundary_matrix = np.array(matrix)  # former: inv_mass @ np.array(matrix)

    def get_name(self):
        return self.__class__.__name__

    def step(self, projection, cfl_number):
        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 _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):
    # 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 _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))

    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._stiffness_matrix @ current_projection[:, j+1]
                                      + self._boundary_matrix @ 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)