""" Fucntions to make the selection matrix """
from typing import Optional, Tuple

import copy
import numpy as np
import scipy.sparse

from spins import fdfd_tools


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_deprecated(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


def calculate_design_bounds_yee(eps_bg: fdfd_tools.VecField,
                                eps_fg: fdfd_tools.VecField):
    """Returns the design bounds of the design region

    Args:
        eps_bg: background epsilon
        eps_fg: foreground epsilon

    Returns:
        design_bounds: 2x2 array: [[ind0_min, ind1_min],[ind0_max, ind1_max]]
    """
    # TODO: cleanup, also this only works for ext_dir = z

    # Threshold epsilon difference.
    threshold = 1e-6
    diff_grid = [np.abs(fg - bg) for fg, bg in zip(eps_bg, eps_fg)]
    design_grid = (np.array(diff_grid).sum(axis=0) > threshold)

    # Flatten the design grid in the z direction.
    design_grid_flat = (design_grid.sum(axis=2) > 0)

    # Project on the x and y axis and find the indices of the True entries.
    design_grid_x = (design_grid_flat.sum(axis=1) > 0)
    ind_x = [i for i, x in enumerate(design_grid_x) if x]
    design_grid_y = (design_grid_flat.sum(axis=0) > 0)
    ind_y = [i for i, x in enumerate(design_grid_y) if x]

    # Return if there is a design grid
    if ind_x and ind_y:
        design_bounds = [[ind_x[0], ind_y[0]], [ind_x[-1], ind_y[-1]]]
    else:
        raise ValueError("zero size design grid")

    return design_bounds


def create_selection_matrix_deprecated(
        eps_bg: fdfd_tools.VecField,
        eps_fg: fdfd_tools.VecField,
        design_bounds: Optional[Tuple[int, int]] = None,
        reduced: bool = False,
        return_param_dims: bool = False) -> np.ndarray:

    if not design_bounds:
        design_bounds = calculate_design_bounds_yee(eps_bg, eps_fg)

    dims = eps_bg[0].shape
    selection_mat, param_dims = direct_lattice_deprecated(
        dims, design_bounds, get_design_area=True)

    indices = np.arange(0, 3 * np.prod(dims))
    weights = scipy.sparse.csr_matrix(
        (fdfd_tools.vec(eps_fg) - fdfd_tools.vec(eps_bg), (indices, indices)))

    # Collapse direct lattice selection matrix by factor of 4.
    if reduced:
        real_design_dims = np.squeeze(np.diff(design_bounds, axis=0) + 1)
        x_ind = np.repeat(np.arange(real_design_dims[0]), 2)
        y_ind = np.repeat(np.arange(real_design_dims[1]), 2)
        x_ind, y_ind = np.meshgrid(x_ind, y_ind, indexing="ij")

        map_x = np.ravel_multi_index((x_ind, y_ind),
                                     real_design_dims,
                                     order="F")
        map_x = map_x.flatten(order="F")
        M = scipy.sparse.csr_matrix(
            (np.ones(len(map_x)), (map_x, np.arange(len(map_x)))),
            shape=(np.prod(real_design_dims), 4 * np.prod(real_design_dims)))
        selection_mat = selection_mat @ M.T

        param_dims = real_design_dims

    if return_param_dims:
        return weights @ selection_mat, param_dims
    return weights @ selection_mat


def calculate_design_bounds(eps_bg: fdfd_tools.VecField,
                            eps_fg: fdfd_tools.VecField):
    """Returns the design bounds of the design region on the direct lattice.

    The design bounds are given for direct lattice, which is twice the size of the Yee
    lattice.

    Args:
       eps_bg: Background epsilon.
       eps_fg: Foreground epsilon.

    Returns:
       A 2x2 array `[[ind0_min, ind1_min],[ind0_max, ind1_max]]` where `ind0_min` and
       `ind0_max` are the minimum and maximum x-indices for which the difference `eps_fg - eps_bg`
       is non-zero. Note that this definition implies the design length in x-direction is
       `ind0_max - ind0_min + 1`. `ind1_min` and `ind1_max` are defined analogously.
    """

    # Make doubled 2d yee grid mapping the differences between eps_bg and eps_fg.
    diff_grid = [np.abs(fg - bg).sum(axis=2) for fg, bg in zip(eps_bg, eps_fg)]
    des_grid = 3 * [None]
    for i in range(3):
        des_grid[i] = np.array([[v
                                 for v in val
                                 for c in (0, 1)]
                                for val in diff_grid[i].tolist()
                                for _ in (0, 1)])
    des_grid[0] = np.roll(des_grid[0], axis=1, shift=-1)
    des_grid[1] = np.roll(des_grid[1], axis=0, shift=-1)
    des_grid[2] = np.roll(des_grid[2], axis=0, shift=-1)
    des_grid[2] = np.roll(des_grid[2], axis=1, shift=-1)
    threshold = 1 / 1000 * np.max(des_grid)
    design_grid_flat = (np.array(des_grid) > threshold).all(axis=0)

    # Project on the x and y axis and find the indices of the True entries.
    ind_x = np.nonzero(design_grid_flat.any(axis=1))[0].tolist()
    ind_y = np.nonzero(design_grid_flat.any(axis=0))[0].tolist()

    # Return if there is a design grid
    if ind_x and ind_y:
        design_bounds = [[ind_x[0], ind_y[0]], [ind_x[-1], ind_y[-1]]]
    else:
        raise ValueError("zero size design grid")

    return design_bounds


def direct_lattice_to_yee_grid(shape):
    """Returns a matrix which maps the direct lattice on the yee grid.

    Args:
        shape: Shape of the yee grid.

    Returns:
        Transform matrix from direct lattice to yee grid.
    """

    ind_yee = np.arange(np.prod(shape[0:2])).reshape(shape[0:2], order='F')
    ind_yee = np.array(
        [[v
          for v in val
          for _ in (0, 1)]
         for val in ind_yee.tolist()
         for _ in (0,
                   1)])  # This doubles every element in the x and y direction.
    ind_yee = np.stack(
        [ind_yee + iz * np.prod(shape[0:2]) for iz in range(shape[2])], axis=2)
    ind_yee_direct = np.arange(4 * np.prod(shape[0:2])).reshape(
        2 * np.array(shape[0:2]), order='F')
    ind_yee_direct = np.stack(shape[2] * [ind_yee_direct], axis=2)
    direct_size = 4 * np.prod(shape[0:2])

    # x-grid
    ind_direct_shift = np.roll(
        ind_yee_direct, axis=1, shift=1).flatten(order='F')
    epsx_trans = scipy.sparse.csr_matrix(
        (0.25 * np.ones(ind_yee.size),
         (ind_yee.flatten(order='F'), ind_direct_shift.flatten(order='F'))),
        shape=(np.prod(shape), direct_size))
    # y-grid
    ind_direct_shift = np.roll(
        ind_yee_direct, axis=0, shift=1).flatten(order='F')
    epsy_trans = scipy.sparse.csr_matrix(
        (0.25 * np.ones(ind_yee.size),
         (ind_yee.flatten(order='F'), ind_direct_shift.flatten(order='F'))),
        shape=(np.prod(shape), direct_size))
    # z-grid
    ind_direct_shift = np.roll(ind_yee_direct, axis=0, shift=1)
    ind_direct_shift = np.roll(
        ind_direct_shift, axis=1, shift=1).flatten(order='F')
    epsz_trans = scipy.sparse.csr_matrix(
        (0.25 * np.ones(ind_yee.size),
         (ind_yee.flatten(order='F'), ind_direct_shift.flatten(order='F'))),
        shape=(np.prod(shape), direct_size))

    return scipy.sparse.vstack([epsx_trans, epsy_trans, epsz_trans])


def direct_lattice(grid_shape, design_bounds, get_design_area: bool = False):
    """
    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 indices on the direct lattice in the x-y plane
    to be included in the selection matrix.

    Args:
        grid_shape: Shape of the grid.
        design_bounds: Bounds on the direct lattice.
        get_design_area: Boolean inicating if the design_area is added to the
            output.

    Returns:
        struct2yee, design_area.
    """

    design_area = [
        design_bounds[1][0] - design_bounds[0][0] + 1,
        design_bounds[1][1] - design_bounds[0][1] + 1
    ]

    # Create structure to direct lattice.
    ind0, ind1 = np.meshgrid(
        np.arange(design_bounds[0][0], design_bounds[1][0] + 1),
        np.arange(design_bounds[0][1], design_bounds[1][1] + 1),
        indexing='ij')
    mapping = np.ravel_multi_index(
        (ind0.flatten(order='F'), ind1.flatten(order='F')),
        dims=2 * np.array(grid_shape[0:2]),
        order='F')
    struct2direct = scipy.sparse.csr_matrix(
        (np.ones(len(mapping)), (mapping, np.arange(len(mapping)))),
        shape=(4 * np.prod(grid_shape[0:2]), np.prod(design_area)))

    # Calculate the transformation from structure to yee.
    direct2yee = direct_lattice_to_yee_grid(grid_shape)
    direct2yee_norm = np.asarray((direct2yee@struct2direct).sum(axis=1)).squeeze()
    direct2yee_norm[direct2yee_norm == 0] = 1
    direct2yee = scipy.sparse.diags(1 / direct2yee_norm, offsets=0) @ direct2yee

    struct2yee = direct2yee @ struct2direct

    if get_design_area:
        return struct2yee, design_area
    return struct2yee


def create_selection_matrix(
        eps_bg: fdfd_tools.VecField,
        eps_fg: fdfd_tools.VecField,
        reduced: bool = False,
        return_param_dims: bool = False
    ) -> Tuple[scipy.sparse.csr_matrix, Tuple[int, int]]:
    """Creates a selection matrix.

    The selection `S` is defined such that given a structure vector `struct`,
    we have that `eps = eps_bg + S * struct`. Note that if `struct` is the
    ones vector, then `eps = eps_fg`.

    The design region is defined over the smallest rectangle that contains
    any non-zero entries in `eps_fg - eps_bg`.

    Additional notes:
    *   The structure vector is defined over fourths of a Yee cell (i.e. over
        the direct lattice), so the size of the vector should be twice that of
        the number of Yee cells in the region.
    *   The design region is assumed to be a 2D region in the xy-plane with
        the structure being extruded in the third dimension (z).

    Args:
        eps_bg: Background epsilon.
        eps_fg: Foreground epsilon.
        reduced: If `True`, the structure vector length is equal to the number
            of Yee cells inthe design region.
        return_param_dims: If `True`, return the size of the structure vector.

    Returns:
        If `return_param_dims` is `True`, a tuple `(S, design_area)` where `S`
        is the selection matrix and `design_area` is a tuple `(dims_x, dims_y)`
        that describes the number of elements in the 2D design area.
        If `return_param_dims` is `False`, only the selection matrix `S` is
        returned.
    """

    #TODO(logan): figure out if reduced is still needed.
    if reduced:
        return create_selection_matrix_deprecated(
            eps_bg=eps_bg,
            eps_fg=eps_fg,
            reduced=reduced,
            return_param_dims=return_param_dims)

    bounds = calculate_design_bounds(eps_bg, eps_fg)

    direct2yee, design_area = direct_lattice(
        grid_shape=eps_bg[0].shape, design_bounds=bounds, get_design_area=True)

    diff_eps = scipy.sparse.diags(
        fdfd_tools.vec(eps_fg) - fdfd_tools.vec(eps_bg), offsets=0)

    selection_matrix = diff_eps @ direct2yee

    if return_param_dims:
        return selection_matrix, design_area
    return selection_matrix
