import numpy as np

from pySDC.core.Hooks import hooks


class err_reduction_hook(hooks):

    def pre_iteration(self, step, level_number):
        """
        Routine called before iteration starts

        Args:
            step (pySDC.Step.step): the current step
            level_number (int): the current level number
        """
        super(err_reduction_hook, self).pre_iteration(step, level_number)

        L = step.levels[level_number]
        if step.status.iter == 2 and np.isclose(L.time + L.dt, 0.1):

            P = L.prob

            err = []
            for m in range(L.sweep.coll.num_nodes):
                uex = P.u_exact(L.time + L.dt * L.sweep.coll.nodes[m])
                err.append(abs(uex - L.u[m + 1]))
            err_full = max(err)
            self.add_to_stats(process=step.status.slot, time=L.time, level=L.level_index, iter=step.status.iter,
                              sweep=L.status.sweep, type='error_pre_iteration', value=err_full)
            # print(L.time, step.status.iter, err_full)

    def post_iteration(self, step, level_number):
        """
        Routine called after each iteration

        Args:
            step (pySDC.Step.step): the current step
            level_number (int): the current level number
        """
        super(err_reduction_hook, self).post_iteration(step, level_number)

        L = step.levels[level_number]

        if step.status.iter == 2 and np.isclose(L.time + L.dt, 0.1):

            P = L.prob

            err = []
            for m in range(L.sweep.coll.num_nodes):
                uex = P.u_exact(L.time + L.dt * L.sweep.coll.nodes[m])
                err.append(abs(uex - L.u[m + 1]))
            err_full = max(err)
            self.add_to_stats(process=step.status.slot, time=L.time, level=L.level_index, iter=step.status.iter,
                              sweep=L.status.sweep, type='error_post_iteration', value=err_full)
            # print(L.time, step.status.iter, err_full)
