from typing import Optional, Callable, Tuple

import numpy as np
from .components import BlochMZI
from scipy.stats import multivariate_normal

from .config import NP_FLOAT


def clements_decomposition(unitary: np.ndarray, hadamard: bool=False,
                           pbar_handle: Callable = None) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """

    Args:
        unitary: unitary matrix :math:`U` to be decomposed into pairwise operators.
        hadamard: Whether to use Hadamard convention
        pbar_handle: Useful for larger matrices

    Returns:
        Clements decomposition of unitary matrix :math:`U` in terms of
        :math:`\\boldsymbol{\\theta}`, :math:`\\boldsymbol{\\phi}`, :math:`\\boldsymbol{\\gamma}`.

    """
    hadamard = hadamard
    unitary_hat = np.copy(unitary)
    units = unitary.shape[0]
    # odd and even layer dimensions
    theta_checkerboard = np.zeros_like(unitary, dtype=NP_FLOAT)
    phi_checkerboard = np.zeros_like(unitary, dtype=NP_FLOAT)
    iterator = pbar_handle(range(units - 1)) if pbar_handle else range(units - 1)
    for i in iterator:
        if i % 2 == 0:
            for j in range(i + 1):
                pairwise_index = i - j
                target_row, target_col = units - j - 1, i - j
                m, n = units - 1 - target_row, units - 1 - target_col
                theta = np.arctan(np.abs(
                    unitary_hat[target_row, target_col] / unitary_hat[target_row, target_col + 1])) * 2 + np.pi
                phi = np.angle(
                    unitary_hat[target_row, target_col] / unitary_hat[target_row, target_col + 1])
                mzi = BlochMZI(theta, phi, hadamard=hadamard, dtype=np.complex128)
                right_multiplier = mzi.givens_rotation(units=units, m=pairwise_index, i_factor=True)
                unitary_hat = unitary_hat @ right_multiplier.conj().T
                theta_checkerboard[m, n - 1] = theta
                phi_checkerboard[m, n - 1] = phi
        else:
            for j in range(i + 1):
                pairwise_index = units + j - i - 2
                target_row, target_col = units + j - i - 1, j
                m, n = target_row, target_col
                theta = np.arctan(np.abs(unitary_hat[target_row, target_col] / unitary_hat[target_row - 1, target_col])) * 2 + np.pi
                phi = np.angle(-unitary_hat[target_row, target_col] / unitary_hat[target_row - 1, target_col])
                mzi = BlochMZI(theta, phi, hadamard=hadamard, dtype=np.complex128)
                left_multiplier = mzi.givens_rotation(units=units, m=pairwise_index, i_factor=True)
                unitary_hat = left_multiplier @ unitary_hat
                theta_checkerboard[m, n] = theta
                phi_checkerboard[m, n] = phi
    theta_checkerboard = to_absolute_theta(theta_checkerboard).T
    phi_checkerboard = np.mod(phi_checkerboard, 2 * np.pi)
    return theta_checkerboard, phi_checkerboard, np.diag(unitary_hat)  # not quite the same as the model in MeshPhases


def to_stripe_array(nparray: np.ndarray, units: int):
    """
    Convert a numpy array of phase shifts of size (`num_layers`, `units`)
    or (`batch_size`, `num_layers`, `units`) into striped array for use in
    general feedforward mesh architectures.

    Args:
        nparray: phase shift values for all columns
        units: dimension the stripe array acts on (depends on parity)

    Returns:
        A general mesh stripe array arrangement that is of size (`units`, `num_layers`)
        or (`batch_size`, `units`, `num_layers`)
    """
    if len(nparray.shape) == 2:
        num_layers, _ = nparray.shape
        stripe_array = np.zeros((units - 1, num_layers), dtype=nparray.dtype)
        stripe_array[::2] = nparray.T
        stripe_array = np.vstack([stripe_array, np.zeros(shape=(1, num_layers))])
    else:
        num_layers, _, batch_size = nparray.shape
        stripe_array = np.zeros((units - 1, num_layers, batch_size), dtype=nparray.dtype)
        stripe_array[::2] = nparray.transpose((1, 0, 2))
        stripe_array = np.vstack([stripe_array, np.zeros(shape=(1, num_layers, batch_size))])
    return stripe_array


def to_absolute_theta(theta: np.ndarray) -> np.ndarray:
    theta = np.mod(theta, 2 * np.pi)
    theta[theta > np.pi] = 2 * np.pi - theta[theta > np.pi]
    return theta


def get_haar_diagonal_sequence(diagonal_length, parity_odd: bool=False):
    odd_nums = list(diagonal_length + 1 - np.flip(np.arange(1, diagonal_length + 1, 2), axis=0))
    even_nums = list(diagonal_length + 1 - np.arange(2, 2 * (diagonal_length - len(odd_nums)) + 1, 2))
    nums = np.asarray(odd_nums + even_nums)
    if parity_odd:
        nums = nums[::-1]
    return nums


def get_alpha_checkerboard(units: int, num_layers: int, include_off_mesh: bool=False, flipud=False):
    if units < num_layers:
        raise ValueError("Require units >= num_layers!")
    alpha_checkerboard = np.zeros((units - 1, num_layers))
    diagonal_length_to_sequence = [get_haar_diagonal_sequence(i, bool(num_layers % 2)) for i in range(1, num_layers + 1)]
    for i in range(units - 1):
        for j in range(num_layers):
            if (i + j) % 2 == 0:
                if i < num_layers and j > i:
                    diagonal_length = num_layers - np.abs(i - j)
                elif i > units - num_layers and j < i - units + num_layers:
                    diagonal_length = num_layers - np.abs(i - j - units + num_layers) - 1 * (units == num_layers)
                else:
                    diagonal_length = num_layers - 1 * (units == num_layers)
                alpha_checkerboard[i, j] = 1 if diagonal_length == 1 else \
                    diagonal_length_to_sequence[int(diagonal_length) - 1][min(i, j)]
            else:
                if include_off_mesh:
                    alpha_checkerboard[i, j] = 1
    # symmetrize the checkerboard
    if units != num_layers:
        alpha_checkerboard = (alpha_checkerboard + np.flipud(alpha_checkerboard)) / 2
    return alpha_checkerboard if not flipud else np.flipud(alpha_checkerboard)


def get_alpha_checkerboard_general(units: int, num_layers: int):
    alpha_checkerboards = [get_alpha_checkerboard(units, units, flipud=bool(n % 2 and units % 2))
                           for n in range(num_layers // units)]
    extra_layers = num_layers - num_layers // units * units
    if extra_layers < units:
        # partial checkerboard
        alpha_checkerboards.append(
            get_alpha_checkerboard(units, extra_layers, flipud=not num_layers // units % 2 and units % 2))
    return np.hstack(alpha_checkerboards)


def get_efficient_coarse_grain_block_sizes(units: int, tunable_layers_per_block: int=2, use_cg_sequence: bool=True):
    num_blocks = int(np.rint(np.log2(units)))
    sampling_frequencies = [2 ** (block_num + 1) for block_num in range(num_blocks - 1)]
    if use_cg_sequence:
        sampling_frequencies = 2 ** get_haar_diagonal_sequence(num_blocks - 1)
    tunable_block_sizes = [tunable_layers_per_block for _ in range(num_blocks - 1)]
    return np.asarray(tunable_block_sizes, dtype=np.int32), np.asarray(sampling_frequencies, dtype=np.int32)


def get_default_coarse_grain_block_sizes(units: int, use_cg_sequence: bool=True):
    num_blocks = int(np.rint(np.log2(units)))
    sampling_frequencies = [2 ** (block_num + 1) for block_num in range(num_blocks - 1)]
    if use_cg_sequence:
        sampling_frequencies = 2 ** get_haar_diagonal_sequence(num_blocks - 1)
    tunable_layer_rank = int(np.floor(units / num_blocks))
    tunable_layer_rank = tunable_layer_rank + 1 if tunable_layer_rank % 2 else tunable_layer_rank
    tunable_block_sizes = [tunable_layer_rank for _ in range(num_blocks - 1)]
    tunable_block_sizes.append(units - tunable_layer_rank * (num_blocks - 1))
    return np.asarray(tunable_block_sizes, dtype=np.int32), np.asarray(sampling_frequencies, dtype=np.int32)


def prm_permutation(units: int, tunable_block_sizes: np.ndarray,
                    sampling_frequencies: np.ndarray, butterfly: bool=False):
    grid_perms = [grid_permutation(units, tunable_block_size) for tunable_block_size in tunable_block_sizes]
    perms_to_concatenate = [grid_perms[0][0]]
    for idx, frequency in enumerate(sampling_frequencies):
        perm_prev = grid_perms[idx][-1]
        perm_next = grid_perms[idx + 1][0]
        perm = butterfly_layer_permutation(units, frequency) if butterfly else rectangular_permutation(units, frequency)
        glued_perm = glue_permutations(perm_prev, perm)
        glued_perm = glue_permutations(glued_perm, perm_next)
        perms_to_concatenate += [grid_perms[idx][1:-1], glued_perm]
    perms_to_concatenate.append(grid_perms[-1][1:])
    return np.vstack(perms_to_concatenate)


def butterfly_layer_permutation(units: int, frequency: int):
    if units % 2:
        raise NotImplementedError('Odd input dimension case not yet implemented.')
    frequency = frequency
    unpermuted_indices = np.arange(units)
    num_splits = units // frequency
    total_num_indices = num_splits * frequency
    unpermuted_indices_remainder = unpermuted_indices[total_num_indices:]
    permuted_indices = np.hstack(
        [np.hstack([i, i + frequency] for i in range(frequency)) + 2 * frequency * split_num
         for split_num in range(num_splits // 2)] + [unpermuted_indices_remainder]
    )
    return permuted_indices.astype(np.int32)


def rectangular_permutation(units: int, frequency: int):
    unpermuted_indices = np.arange(units)
    frequency_offset = np.empty((units,))
    frequency_offset[::2] = -frequency
    frequency_offset[1::2] = frequency
    permuted_indices = unpermuted_indices + frequency_offset
    for idx in range(units):
        if permuted_indices[idx] < 0:
            permuted_indices[idx] = -1 - permuted_indices[idx]
        if permuted_indices[idx] > units - 1:
            permuted_indices[idx] = 2 * units - 1 - permuted_indices[idx]
    return permuted_indices.astype(np.int32)


def grid_permutation(units: int, num_layers: int):
    ordered_idx = np.arange(units)
    split_num_layers = (num_layers - num_layers // 2, num_layers // 2)
    left_shift = np.roll(ordered_idx, -1, axis=0)
    right_shift = np.roll(ordered_idx, 1, axis=0)
    permuted_indices = np.zeros((num_layers, units))
    permuted_indices[::2] = np.ones((split_num_layers[0], 1)) @ left_shift[np.newaxis, :]
    permuted_indices[1::2] = np.ones((split_num_layers[1], 1)) @ right_shift[np.newaxis, :]
    if num_layers % 2:
        return np.vstack((ordered_idx.astype(np.int32),
                          permuted_indices[:-1].astype(np.int32),
                          ordered_idx.astype(np.int32)))
    return np.vstack((ordered_idx.astype(np.int32),
                      permuted_indices.astype(np.int32)))


def grid_viz_permutation(units: int, num_layers: int, flip: bool=False):
    ordered_idx = np.arange(units)
    split_num_layers = (num_layers - num_layers // 2, num_layers // 2)
    right_shift = np.roll(ordered_idx, 1, axis=0)
    permuted_indices = np.zeros((num_layers, units))
    if flip:
        permuted_indices[::2] = np.ones((split_num_layers[0], 1)) @ ordered_idx[np.newaxis, :]
        permuted_indices[1::2] = np.ones((split_num_layers[1], 1)) @ right_shift[np.newaxis, :]
    else:
        permuted_indices[::2] = np.ones((split_num_layers[0], 1)) @ right_shift[np.newaxis, :]
        permuted_indices[1::2] = np.ones((split_num_layers[1], 1)) @ ordered_idx[np.newaxis, :]
    return np.vstack((ordered_idx.astype(np.int32),
                      permuted_indices[:-1].astype(np.int32),
                      ordered_idx.astype(np.int32)))


def ordered_viz_permutation(units: int, num_layers: int):
    ordered_idx = np.arange(units)
    permuted_indices = np.ones((num_layers + 1, 1)) @ ordered_idx[np.newaxis, :]
    return permuted_indices.astype(np.int32)


def plot_complex_matrix(plt, matrix: np.ndarray):
    plt.figure(figsize=(15, 5), dpi=200)
    plt.subplot(131)
    plt.title('Absolute')
    plt.imshow(np.abs(matrix), cmap='hot')
    plt.colorbar(shrink=0.7)
    plt.subplot(132)
    plt.title('Real')
    plt.imshow(np.real(matrix), cmap='hot')
    plt.colorbar(shrink=0.7)
    plt.subplot(133)
    plt.title('Imag')
    plt.imshow(np.imag(matrix), cmap='hot')
    plt.colorbar(shrink=0.7)


def random_gaussian_batch(batch_size: int, units: int, covariance_matrix: Optional[np.ndarray]=None,
                          seed: Optional[int]=None) -> np.ndarray:
    if seed is not None:
        np.random.seed(seed)
    input_matrix = multivariate_normal.rvs(
        mean=np.zeros(units),
        cov=1 if not covariance_matrix else covariance_matrix,
        size=batch_size
    )
    random_phase = np.random.rand(batch_size, units).astype(dtype=NP_FLOAT) * 2 * np.pi
    return input_matrix * np.exp(1j * random_phase)


def glue_permutations(perm_idx_1: np.ndarray, perm_idx_2: np.ndarray):
    perm_idx = np.zeros_like(perm_idx_1)
    perm_idx[perm_idx_2] = perm_idx_1
    return perm_idx.astype(np.int32)


def inverse_permutation(permuted_indices: np.ndarray):
    inv_permuted_indices = np.zeros_like(permuted_indices)
    for idx, perm_idx in enumerate(permuted_indices):
        inv_permuted_indices[perm_idx] = idx
    return inv_permuted_indices


def pairwise_off_diag_permutation(units: int):
    ordered_idx = np.arange(units)
    perm_idx = np.zeros_like(ordered_idx)
    if units % 2:
        perm_idx[:-1][::2] = ordered_idx[1::2]
        perm_idx[1::2] = ordered_idx[:-1][::2]
        perm_idx[-1] = ordered_idx[-1]
    else:
        perm_idx[::2] = ordered_idx[1::2]
        perm_idx[1::2] = ordered_idx[::2]
    return perm_idx.astype(np.int32)


def butterfly_permutation(num_layers: int):
    ordered_idx = np.arange(2 ** num_layers)
    permuted_idx = np.vstack(
        [butterfly_layer_permutation(2 ** num_layers, 2 ** layer) for layer in range(num_layers)]
    ).astype(np.int32)
    return np.vstack((ordered_idx.astype(np.int32),
                      permuted_idx[1:].astype(np.int32),
                      ordered_idx.astype(np.int32)))


def neurophox_matplotlib_setup(plt):
    plt.rc('text', usetex=True)
    plt.rc('font', family='serif')
    # plt.rc('text', usetex=True)
    # plt.rc('font', **{'family': 'serif', 'serif': ['Charter']})
    # plt.rcParams['mathtext.fontset'] = 'dejavuserif'
    plt.rcParams.update({'text.latex.preamble': [r'\usepackage{siunitx}', r'\usepackage{amsmath}']})
