import numpy as np
import scipy.sparse


def average_2xFarEdge(grid_size, design_bounds):
    """
    Start with 2D versions of each grid and later just use repmat to
     make the S-matrix for the full 3D grid. The 2D version of Sz
     consists of a bunch of 0-rows above and below a block-diagonal
     matrix. The blocks are a bunch of 0-rows surrounding an identity
     matrix. The 0-rows create the y and x regions of the
     simulation grid which are unaffected by the design vector.
    For Sx and Sy, we need to average nearby locations. For Sx, this just
     means altering the individual blocks to add a right-shifted I. For
     Sy, this means shifting the block-diagonal matrix down.
    For both Sx and Sy, the edge values of the design vector are not
     averaged with anything, so the edge looks like
          [..., e(end-1) + e(end), e(end) + e(end)]
      If this causes weird behavior, it's probably best to just add an
       extra row+column to the design vector, to make it look like
           [..., e(end-2) + e(end-1), e(end-1) + e(end)]
       where the new e(end) is only used in the expression for the edge
       cell.
    """
    design_area = np.squeeze(np.diff(design_bounds, axis=0)) + 1

    # Decide how much zero-padding is needed to create the full grid from the z
    # vector (i.e., how much smaller is the design area).
    numz_l = design_bounds[0, 0]
    numz_r = grid_size[0] - design_bounds[1, 0] - 1
    numz_t = design_bounds[0, 1]
    numz_b = grid_size[1] - design_bounds[1, 1] - 1

    zero_padding_l = scipy.sparse.csr_matrix((numz_l, design_area[0]))
    zero_padding_r = scipy.sparse.csr_matrix((numz_r, design_area[0]))
    zero_padding_t = scipy.sparse.csr_matrix((numz_t * grid_size[0],
                                              np.prod(design_area)))
    zero_padding_b = scipy.sparse.csr_matrix((numz_b * grid_size[0],
                                              np.prod(design_area)))

    # X grid averages pixels adjacent in x.
    I0 = scipy.sparse.identity(design_area[0], format='csr')

    I0_shift_x = scipy.sparse.hstack([I0[:, 1:], I0[:, :1]], format='csr')
    I0_shift_x[-1, 0] = 0
    I0_shift_x[-1, -1] = 1

    Sx_flat_block = scipy.sparse.vstack(
        (zero_padding_l, (I0 + I0_shift_x) / 2, zero_padding_r), format='csr')
    Sx_block_diag = scipy.sparse.block_diag([Sx_flat_block] * design_area[1])
    Sx_flat = scipy.sparse.vstack((zero_padding_t, Sx_block_diag,
                                   zero_padding_b))

    # Y grid averages pixels adjacent in y.
    Sy_flat_block = scipy.sparse.vstack(
        (zero_padding_l, I0, zero_padding_r), format='csr')
    Sy_flat_preshift = scipy.sparse.block_diag(
        [Sy_flat_block] * design_area[1], format='csr')

    Sy_flat_shift = scipy.sparse.hstack(
        (Sy_flat_preshift[:, design_area[0]:],
         Sy_flat_preshift[:, :design_area[0]]),
        format='csr')
    excess = Sy_flat_shift[:, 0:design_area[0]]
    # Undo circular part of shift.
    Sy_flat_shift[:, 0:design_area[0]] = 0
    # Double-up on stuff that doesn't have a second.
    Sy_flat_shift[:, -design_area[0]:] += excess

    Sy_flat = scipy.sparse.vstack(
        (zero_padding_t, (Sy_flat_preshift + Sy_flat_shift) / 2,
         zero_padding_b))

    # Z grid doesn't need to average anything (effects of design vector uniform
    # in z).
    Sz_flat = scipy.sparse.vstack((zero_padding_t, Sy_flat_preshift,
                                   zero_padding_b))

    # Repmat to make a full 3D grid.
    S = scipy.sparse.vstack([Sx_flat] * grid_size[2] +
                            [Sy_flat] * grid_size[2] + [Sz_flat] * grid_size[2])
    return S


def direct_lattice(grid_size,
                   design_bounds,
                   get_design_area: bool = False,
                   wrap_or_clip: list = [0, 0, 1]):
    """
    Creates a selection matrix that gives the user direct control of the Yee
    lattice for any planar / 2D design. This selection matrix is useful for
    properly rendering arbitrary polygons with anti-aliasing.

    This function assumes the design is in the x-y plane, with the
    z-direction normal to the plane. Now, consider an epsilon_z voxel in the
    Yee lattice, and look down on it in the z-direction. This voxel is broken
    into 4 quadrants by the overlapping epsilon_x and epsilon_y voxels. The
    selection matrix generated by direct_lattice assigns each of these 4
    quadrants to an independent design variable.

    'design_bounds' refers to the epsilon_z voxels in the x-y plane to be
    included in the selection matrix.
    """
    design_bounds[0][0] = np.max([design_bounds[0][0], 0])
    design_bounds[0][1] = np.max([design_bounds[0][1], 0])
    design_bounds[1][0] = np.min([design_bounds[1][0], grid_size[0] - 1])
    design_bounds[1][1] = np.min([design_bounds[1][1], grid_size[1] - 1])
    design_area = np.squeeze(np.diff(design_bounds, axis=0)) + 1

    x_ind = np.arange(0, 2 * design_area[0])
    y_ind = np.arange(0, 2 * design_area[1])

    x_ind, y_ind = np.meshgrid(x_ind, y_ind, indexing='ij')
    x_ind = x_ind.flatten(order='F')
    y_ind = y_ind.flatten(order='F')

    num_gridxy = np.prod(grid_size[0:2])
    num_design = np.prod(2 * design_area)

    mode = ['wrap', 'clip']
    # X-grid
    map_x = np.ravel_multi_index(
        [
            design_bounds[0][0] + (x_ind + 1) // 2 - 1,
            design_bounds[0][1] + y_ind // 2
        ],
        grid_size[0:2],
        order='F',
        mode=mode[wrap_or_clip[0]])
    Sx_flat = scipy.sparse.csr_matrix(
        (np.ones(len(map_x)), (map_x, np.array(np.arange(len(map_x))))),
        shape=(num_gridxy, num_design))

    # Y-grid
    map_y = np.ravel_multi_index(
        [
            design_bounds[0][0] + x_ind // 2,
            design_bounds[0][1] + (y_ind + 1) // 2 - 1
        ],
        grid_size[0:2],
        order='F',
        mode=mode[wrap_or_clip[1]])
    Sy_flat = scipy.sparse.csr_matrix(
        (np.ones(len(map_y)), (map_y, np.array(np.arange(len(map_y))))),
        shape=(num_gridxy, num_design))

    # Z-grid
    map_z = np.ravel_multi_index(
        [design_bounds[0][0] + x_ind // 2, design_bounds[0][1] + y_ind // 2],
        grid_size[0:2],
        order='F',
        mode=mode[wrap_or_clip[2]])
    Sz_flat = scipy.sparse.csr_matrix(
        (np.ones(len(map_z)), (map_z, np.array(np.arange(len(map_z))))),
        shape=(num_gridxy, num_design))

    # Normalize each row of matrix S so that the sum of each non-zero row is 1.
    def normalize_rows(S):
        row_sums = np.array(S.sum(axis=1)).squeeze()
        row_sums[row_sums != 0] = 1 / row_sums[row_sums != 0]
        num_rows, num_cols = S.shape
        return scipy.sparse.spdiags(row_sums, 0, num_rows, num_rows) @ S

    # Make sure each grid pixel has a weighting of 1.
    Sx_flat = normalize_rows(Sx_flat)
    Sy_flat = normalize_rows(Sy_flat)
    Sz_flat = normalize_rows(Sz_flat)

    # Extrude in z-direction to turn it into full 3D grid.
    Sx = scipy.sparse.vstack([Sx_flat] * grid_size[2])
    Sy = scipy.sparse.vstack([Sy_flat] * grid_size[2])
    Sz = scipy.sparse.vstack([Sz_flat] * grid_size[2])

    S = scipy.sparse.vstack((Sx, Sy, Sz))
    if get_design_area:
        return S, 2 * design_area
    else:
        return S
