import numpy as np
from skimage.exposure import match_histograms


def match_hist_all(all_ds_avg_ch):
    ref = all_ds_avg_ch[0][0]
    all_ds_avg_ch_matched = []
    for ds_avg_ch in all_ds_avg_ch:
        ds_avg_ch_matched = []
        for i in range(len(ds_avg_ch)):
            ds_avg_ch_matched.append(match_histograms(ds_avg_ch[i], ref))
        all_ds_avg_ch_matched.append(ds_avg_ch_matched)
        
    return all_ds_avg_ch_matched

def make_rgb_img(img1, img2):
    img_rgb = np.zeros((img1.shape[0], img1.shape[1], 3))
    # normalise images to 0-1
    img1_norm = (img1 - np.min(img1)) / (np.max(img1) - np.min(img1))
    img2_norm = (img2 - np.min(img2)) / (np.max(img2) - np.min(img2))
    img_rgb[:, :, 0] = img1_norm
    img_rgb[:, :, 1] = img2_norm
    
    return img_rgb

def saturate_perc(img_rgb, sat_perc=99):
    img_rgb = np.clip(img_rgb, 0, np.percentile(img_rgb, sat_perc)) 
    img_rgb = (img_rgb / np.max(img_rgb) * 255).astype(np.uint8)
    return img_rgb

def get_all_wind_mean_img(all_ds_mean_img, all_ds_centroids, all_pl_match_mat, nrn_id, plane_idx=0, win_size=64):

    all_wind_mean_img = []
    for i in range(len(all_ds_mean_img)):
        centroid_ids = all_pl_match_mat[plane_idx][nrn_id,:]
        cent = all_ds_centroids[i][plane_idx][centroid_ids[i]]
        mean_img = all_ds_mean_img[i][plane_idx]

        # pad mean image with mean and shift the centroid appropriately
        # mean_img = np.pad(mean_img, ((win_size,win_size),(win_size,win_size)))
        mean_img = np.pad(mean_img, ((win_size,win_size),(win_size,win_size)), mode='constant', constant_values=0)
        cent = cent + win_size
    
        wind_mean_img = mean_img[int(cent[0]-win_size/2):int(cent[0]+win_size/2), int(cent[1]-win_size/2):int(cent[1]+win_size/2)]


        all_wind_mean_img.append(wind_mean_img)

    return all_wind_mean_img