import numpy as np
import matplotlib.pyplot as plt
from datasets.stas import get_cell_svd
from scipy.optimize import  minimize


def separate_cell_svds(sta, center=None, crop=True, crop_size=7):
    temporal_variances = np.var(sta, axis=0)
    if center is None:
        center = np.unravel_index(np.argmax(temporal_variances), sta.shape[1:])
    if crop:
        sta = sta[:, center[0] - crop_size:center[0] + crop_size, center[1] - crop_size:center[1] + crop_size]
    spatial_filter, temporal_filter = get_cell_svd(sta)
    temp_sign = get_filter_sign(temporal_filter)
    if temp_sign < 0:
        spatial_filter = spatial_filter * -1
        temporal_filter = temporal_filter * -1
    return sta, spatial_filter, temporal_filter, center


def get_filter_sign(f):
    max_value = np.max(f)
    min_value = np.min(f)
    if np.abs(min_value) > np.abs(max_value):
        return -1
    else:
        return 1


def crop_spat_filter(spat_filter, crop_size, center=None):
    if center is None:
        center = spat_filter.shape[0] // 2, spat_filter.shape[1] // 2
    spat_filter = spat_filter[center[0] - crop_size:center[0] + crop_size, center[1] - crop_size:center[1] + crop_size]
    return spat_filter


def dog_model(x, y, p):
    A1, x1, y1, sigma1x, sigma1y, rho1, A2, sigma2x, sigma2y, rho2 = p

    cov1 = np.array([[sigma1x ** 2, rho1 * sigma1x * sigma1y],
                     [rho1 * sigma1x * sigma1y, sigma1y ** 2]])

    cov2 = np.array([[sigma2x ** 2, rho2 * sigma2x * sigma2y],
                     [rho2 * sigma2x * sigma2y, sigma2y ** 2]])

    # Compute inverse and determinant of the covariance matrices
    cov1_inv = np.linalg.inv(cov1)
    cov2_inv = np.linalg.inv(cov2)

    # Compute the difference of Gaussians with correlations
    diff = np.array([x - x1, y - y1])

    gauss1 = A1 * np.exp(-0.5 * np.einsum('ij,ij->i', diff.T @ cov1_inv, diff.T))
    gauss2 = A2 * np.exp(-0.5 * np.einsum('ij,ij->i', diff.T @ cov2_inv, diff.T))

    gauss = gauss1.ravel() - gauss2.ravel()
    return gauss


def cost_function(p, x, y, data):
    model = dog_model(x, y, p)
    res = np.sum((model.ravel() - data.ravel()) ** 2)
    return res


def fit_dog(spat_kernel, wn=False):
    max_coord = np.unravel_index(np.argmax(np.abs(spat_kernel)), (spat_kernel.shape[0], spat_kernel.shape[1]))
    crop_size = spat_kernel.shape[0]
    x = np.arange(spat_kernel.shape[0])
    y = np.arange(spat_kernel.shape[1])
    x, y = np.meshgrid(x, y)

    try:
        sign = np.sign(spat_kernel[max_coord])
    except:
        sign = np.sign(spat_kernel[int(crop_size / 2), int(crop_size / 2)])

    bounds = [(0, float('inf')) if (sign > 0) else (-float('inf'), 0),  # Bound for A1
              (None, None),  # No bound for mu_x
              (None, None),  # No bound for mu_y
              (None, None),  # Bound for sigma1x, ensuring it's positive
              (None, None),  # Bound for sigma1y
              (0, None),  # Bound for rho1
              (0, 2) if (sign > 0) else (-2, 0),  # Bound for A2
              (None, None),  # Bound for sigma2x
              (None, None),
              (0, None)]  # Bound for rho2]


    initial_guess = [sign * 0.5, int(spat_kernel.shape[1] / 2), int(spat_kernel.shape[0] / 2), 1.2, 1.2, 0.00001,
                         sign * 0.3, 1.5, 1.5, 0.00001]

    popt = minimize(cost_function, initial_guess, method='L-BFGS-B', args=(x.ravel(), y.ravel(), spat_kernel.ravel()),
                    bounds=bounds, options={'eps': 10e-6, 'ftol': 1e-6})

    return popt


def get_fitted_dog(spat_kernel, params):
    height, width = spat_kernel.shape[0], spat_kernel.shape[1]
    x = np.arange(0, height)
    y = np.arange(0, width)
    x, y = np.meshgrid(x, y)

    # Generate the DoG function using popt parameters
    A1, x1, y1, sigma1x, sigma1y, rho1, A2, sigma2x, sigma2y, rho2 = params

    g1 = build_dog_kernel(A1, x, y, x1, y1, sigma1x, sigma1y, rho1)
    g2 = build_dog_kernel(A2, x, y, x1, y1, sigma2x, sigma2y, rho2)
    final = g1 - g2
    if not ((sigma1x * sigma1y < sigma2x * sigma2y) or (np.abs(A1) > np.abs(A2))):
        final = g2 - g1

    return final.reshape(height, width)


def plot_fitted_dog_1d(spat_kernel, final, ax):
    x1 = int(spat_kernel.shape[0] / 2)
    ax[0].plot(spat_kernel[:, x1] / np.max(np.abs(spat_kernel)), label='oritinal kernel', color='black')
    ax[0].plot(final[:, x1], label='fit dog', color='orange', linestyle='--')
    ax[0].set_title('1D gaussian')

    ax[1].plot(np.abs(np.fft.fftshift(np.fft.fft(spat_kernel[:, x1]))), label='original kernel', color='black')
    ax[1].plot(np.abs(np.fft.fftshift(np.fft.fft(final[:, x1]))), label='fit dog', color='orange', linestyle='--')
    ax[1].set_title('FFT')
    plt.legend()


def plot_fitted_dog(spat_kernel, final, ax):
    ax[0].imshow(spat_kernel / np.max(np.abs(spat_kernel)), cmap='gray', vmin=-1, vmax=1)
    ax[1].imshow(final, cmap='gray', vmin=-1, vmax=1)

    ax[0].grid(False)
    ax[1].grid(False)

    plt.show()


def build_dog_kernel(A, x, y, x1, y1, sigmax, sigmay, rho):
    cov = np.array([[sigmax ** 2, rho * sigmax * sigmay],
                    [rho * sigmax * sigmay, sigmay ** 2]])
    cov_inv = np.linalg.inv(cov)
    diff = np.array([x - x1, y - y1]).reshape((2, -1))
    exponent = -0.5 * np.einsum('im,ij,jm->m', diff, cov_inv, diff).reshape((len(x), len(y)))

    return A * np.exp(exponent)


def get_gauss_area(sigmax, sigmay, rho):
    cov_matrix = np.array([[sigmax ** 2, rho * sigmax * sigmay],
                           [rho * sigmax * sigmay, sigmay ** 2]])
    eigenvalues, _ = np.linalg.eig(cov_matrix)
    return 4 * np.pi * np.sqrt(eigenvalues[0]) * np.sqrt(eigenvalues[1]) * 30 * 30


def get_filter_size(spat_filter):
    popt = fit_dog(spat_filter)
    A1, mu_x, mu_y, sigma1x, sigma1y, rho1, A2, sigma2x, sigma2y, rho2 = popt.x
    size1 = get_gauss_area(sigma1x, sigma1y, rho1)
    size2 = get_gauss_area(sigma2x, sigma2y, rho2)

    if (size1 > size2) and (np.abs(A2) > np.abs(A1)):
        return size2, popt.x
    elif (size2 > size1) and (np.abs(A1) > np.abs(A2)):
        return size1, popt.x
    elif (np.abs(A2) > np.abs(A1)):
        return size2, popt.x
    else:
        return size1, popt.x


def get_filter_amplitude(spat_filter, params):
    fitted_filter = get_fitted_dog(spat_filter, params)
    fitted_filter = fitted_filter / np.max(np.abs(fitted_filter))
    amplitude = get_surround_amplitude(fitted_filter)
    return amplitude, fitted_filter


def get_surround_amplitude(spat_filter):
    sign = get_filter_sign(spat_filter)
    spat_filter = spat_filter / np.max(np.abs(spat_filter))
    if sign < 0:
        return np.max([0, np.max(spat_filter)])
    else:
        return np.abs(np.min(spat_filter))


def handle_filter(spat_filter, plot=False):
    size, params = get_filter_size(spat_filter)
    amplitude, fitted_filter = get_filter_amplitude(spat_filter, params)
    mu = params[1], params[2]
    if plot:
        fig, ax = plt.subplots(1, 2)
        ax[0].grid(False)
        ax[1].grid(False)
        plot_fitted_dog_1d(spat_filter, fitted_filter, ax)
        plt.show()
        fig, ax = plt.subplots(1, 2)
        ax[0].grid(False)
        ax[1].grid(False)
        plot_fitted_dog(spat_filter, fitted_filter, ax)
    return size, amplitude, mu


def center_filter(spat_filter, mu, plot=False):
    mu = int(mu[0]), int(mu[1])
    center = spat_filter.shape[0] // 2, spat_filter.shape[1] // 2
    row_shift = center[0] - mu[0]
    col_shift = center[1] - mu[1]

    shifted = np.roll(spat_filter, (col_shift, row_shift), axis=(0, 1))
    if plot:
        plt.imshow(spat_filter)
        plt.show()
        plt.imshow(shifted)
        plt.show()
    return shifted


def handle_temp_filter(temp_filter, plot=False, peak_height=4, normalize_temp=True):
    size = get_filter_width(temp_filter, peak_height=peak_height)
    peak, surr = measure_peaks(temp_filter, normalize=normalize_temp)
    return size, surr


def measure_peaks(temp_filter, normalize, plot=False):
    sign = get_filter_sign(temp_filter)
    if normalize:
        temp_filter = temp_filter / np.max(np.abs(temp_filter))
    if sign < 0:
        peak = np.min(temp_filter)
        surr = np.max(temp_filter)
        peak_index = np.argmin(temp_filter)
        surr_index = np.argmax(temp_filter)
    else:
        peak = np.max(temp_filter)
        surr = np.min(temp_filter)
        peak_index = np.argmax(temp_filter)
        surr_index = np.argmin(temp_filter)
    if plot:
        plt.plot(temp_filter)
        plt.scatter([peak_index, surr_index], [temp_filter[peak_index], temp_filter[surr_index]])
        plt.show()
    return np.abs(peak), np.abs(surr)


def get_filter_width(temp_filter, peak_height=4, plot=True):
    peak_index = np.argmax(np.abs(temp_filter))
    k_sign = np.sign(temp_filter[peak_index])
    if k_sign < 0:
        temp_filter *= -1
    mid_level = temp_filter[peak_index] / peak_height
    try:
        left_below = np.where(temp_filter[:peak_index] < mid_level)[0][-1]
        left_above = left_below + 1
        right_above = np.where(temp_filter[peak_index:] < mid_level)[0][0] + peak_index
        right_below = right_above - 1
    except Exception as e:
        print(e, 'Failed')
        return -1

    x = [left_below, left_above]
    y = [temp_filter[left_below], temp_filter[left_above]]
    l_edge = np.interp(mid_level, y, x)

    x = [right_above, right_below]
    y = [temp_filter[right_above], temp_filter[right_below]]
    r_edge = np.interp(mid_level, y, x)

    if plot:
        plt.plot(temp_filter)
        plt.scatter([left_below, right_below], [temp_filter[left_below], temp_filter[right_below]])
        plt.scatter([left_above, right_above], [temp_filter[left_above], temp_filter[right_above]])
        plt.scatter([l_edge, r_edge], [mid_level, mid_level])
        plt.show()
    distance = l_edge - r_edge
    return np.abs(distance)