import numpy as np

def calculate_boundaries(p, pold, dt, dx, c, b_width, method, order):
    # test for minimum boundary size
    if b_width<2:
        print('\n Error: Absorbing boundary must have at least size 2!')
    
    pbound = np.zeros(p.shape)
    non_corner = slice(b_width, -b_width)

    if method == 'clayton-engquist' and order == 1:
        # Left boundary
        pbound[non_corner, :b_width] = p[non_corner, :b_width] + dt / dx * c[non_corner, :b_width] * \
                                       (p[non_corner, 1:b_width+1] - p[non_corner, :b_width])
        # Right boundary
        pbound[non_corner, -b_width:] = p[non_corner, -b_width:] - dt / dx * c[non_corner, -b_width:] * \
                                       (p[non_corner, -b_width:] - p[non_corner, -b_width-1:-1])
        # Top boundary
        pbound[:b_width, non_corner] = p[:b_width, non_corner] + dt / dx * c[:b_width, non_corner] * \
                                        (p[1:b_width+1, non_corner] - p[:b_width, non_corner])
        # Bottom boundary
        pbound[-b_width:, non_corner] = p[-b_width:, non_corner] - dt / dx * c[-b_width:, non_corner] * \
                                        (p[-b_width:, non_corner] - p[-b_width-1:-1, non_corner])

    elif method == 'clayton-engquist' and order == 2:
        # Left boundary
        pzz = (p[b_width-1:-b_width-1, :b_width] - 2 * p[non_corner, :b_width] +
                     p[b_width+1:-b_width+1, :b_width]) / (dx ** 2)
        pxt = (p[non_corner, 1:b_width+1] - p[non_corner, :b_width] -
                     pold[non_corner, 1:b_width+1] + pold[non_corner, :b_width]) / (dx * dt)
        pbound[non_corner, :b_width] = 2*p[non_corner, :b_width] - pold[non_corner, :b_width]  \
                        + dt ** 2 * c[non_corner, :b_width] * pxt + dt**2 * (1/2 * c[non_corner, :b_width]**2) * pzz

        # Right boundary
        pzz = (p[b_width - 1:-b_width - 1, -b_width:] - 2 * p[non_corner, -b_width:] +
               p[b_width + 1:-b_width + 1, -b_width:]) / (dx ** 2)
        pxt = (p[non_corner, -b_width:] - p[non_corner, -b_width-1:-1] -
               pold[non_corner, -b_width:] + pold[non_corner, -b_width-1:-1]) / (dx * dt)
        pbound[non_corner, -b_width:] = 2 * p[non_corner, -b_width:] - pold[non_corner, -b_width:] \
                        - dt ** 2 * c[non_corner, -b_width:] * pxt + dt**2 * (1/2 * c[non_corner, -b_width:]**2) * pzz

        # Top boundary
        pxx = (p[:b_width, b_width - 1:-b_width - 1] - 2 * p[:b_width, non_corner] +
               p[:b_width, b_width + 1:-b_width + 1]) / (dx ** 2)
        pzt = (p[1:b_width + 1, non_corner] - p[:b_width, non_corner] -
               pold[1:b_width + 1, non_corner] + pold[:b_width, non_corner]) / (dx * dt)
        pbound[:b_width, non_corner] = 2 * p[:b_width, non_corner] - pold[:b_width, non_corner] \
                + dt ** 2 * c[:b_width, non_corner] * pzt + dt ** 2 * (1 / 2 * c[:b_width, non_corner] ** 2) * pxx

        # Bottom boundary
        pxx = (p[-b_width:, b_width - 1:-b_width - 1] - 2 * p[-b_width:, non_corner] +
               p[-b_width:, b_width + 1:-b_width + 1]) / (dx ** 2)
        pzt = (p[-b_width:, non_corner] - p[-b_width-1:-1, non_corner] -
               pold[-b_width:, non_corner] + pold[-b_width-1:-1, non_corner]) / (dx * dt)
        pbound[-b_width:, non_corner] = 2 * p[-b_width:, non_corner] - pold[-b_width:, non_corner] \
                - dt ** 2 * c[-b_width:, non_corner] * pzt + dt ** 2 * (1 / 2 * c[-b_width:, non_corner] ** 2) * pxx
    else:
        raise NotImplementedError

    # Top left corners
    pbound[:b_width, :b_width] = p[:b_width, :b_width] + dt / dx * c[:b_width, :b_width] / np.sqrt(2) * \
        ((p[:b_width, 1:b_width + 1] - p[:b_width, :b_width]) + (p[1:b_width+1, :b_width] - p[:b_width, :b_width]))
    # Top right corner
    pbound[:b_width, -b_width:] = p[:b_width, -b_width:] + dt / dx * c[:b_width, -b_width:] / np.sqrt(2) * \
        (- (p[:b_width, -b_width:] - p[:b_width, -b_width-1:-1]) + (p[1:b_width + 1, -b_width:] - p[:b_width, -b_width:]))
    # Bottom left corner
    pbound[-b_width:, :b_width] = p[-b_width:, :b_width] + dt / dx * c[-b_width:, :b_width] / np.sqrt(2) * \
            ((p[-b_width:, 1:b_width + 1] - p[-b_width:, :b_width]) -(p[-b_width:, :b_width] - p[-b_width-1:-1, :b_width]))
    # Bottom right corner
    pbound[-b_width:, -b_width:] = p[-b_width:, -b_width:] + dt / dx * c[-b_width:, -b_width:] / np.sqrt(2) * \
            (-(p[-b_width:, -b_width:] - p[-b_width:, -b_width-1:-1])-(p[-b_width:, -b_width:] - p[-b_width-1:-1, -b_width:]))

    return pbound

def set_rectengular_ring(matrix, d_border, value):
    # if d_border == 1: set ring at outer border
    matrix[(d_border-1, -d_border), slice(d_border-1, matrix.shape[1]-d_border+1)] = value
    matrix[slice(d_border-1, matrix.shape[0]-d_border+1), (d_border-1, -d_border)] = value

def get_sigma(matrix_shape, b_width, absorbing_function='linear'):
    sigma = np.ones(matrix_shape)
    for d_border, value in zip(range(1, b_width+2), np.linspace(0, 1, b_width+1)):
        set_rectengular_ring(sigma, d_border=d_border, value=value)

    if absorbing_function == 'linear':
        return sigma
    elif absorbing_function == 'cosine':
        # normalize from [0, 1] to [np.pi, 2*np.pi], then take cosine
        sigma = (sigma + 1)*np.pi
        sigma = (np.cos(sigma)+1)/2
    else:
        raise NotImplementedError

    return sigma
