#!/usr/bin/env python
# coding: utf-8

# ### Define paths and animals for the analysis

# In[1]:


path = '/media/andrey/My Passport/GIN/Anesthesia_CA1/meta_data/meta_recordings_transition_state.xlsx'
path4results = '/media/andrey/My Passport/GIN/Anesthesia_CA1/validation/calcium_imaging_transition_state/' #To store transformation matrix
save_plots_path = '/media/andrey/My Passport/GIN/Anesthesia_CA1/validation/calcium_imaging_transition_state/'
log_file_path = save_plots_path + 'registration_logs.txt'

#animals_for_analysis = [8235,8237,8238] 
animals_for_analysis = ['F0','M0','M3']

# ### Align FOV's for all recordings

# In[2]:


repeat_calc = 1
silent_mode = False

#######################

import pandas as pd
import numpy as np
import os

import matplotlib.pyplot as plt

import sys
np.set_printoptions(threshold=sys.maxsize)

from pystackreg import StackReg

# Sobel filter (not used)
#from scipy import ndimage #https://docs.scipy.org/doc/scipy/reference/generated/scipy.ndimage.sobel.html

meta_data = pd.read_excel(path, engine='openpyxl')

#%% compute transformations matrices between recordings

recordings = meta_data['Number']

animals = animals_for_analysis

#print("Recordings: ", recordings)

#ALIGNMENT ALGORITHM

sr = StackReg(StackReg.AFFINE)


# ROI adjustment check for all recordings

from matplotlib import colors
Amap = colors.ListedColormap(['white','tab:blue', 'tab:red','purple'])

'''  
if not os.path.exists(save_plots_path + 'ROIAdjustmentCheck/'):
    os.makedirs(save_plots_path + 'ROIAdjustmentCheck/')
 
for animal in animals:
    if not os.path.exists(save_plots_path + 'ROIAdjustmentCheck/' + str(animal) + "/"):
        os.makedirs(save_plots_path + 'ROIAdjustmentCheck/' + str(animal) + "/")  
    
    tmats_loaded = np.load(path4results + 'StackReg/' + str(animal) + "_best_tmats" + '.npy')
    meta_animal = meta_data[meta_data['Mouse_all'] == animal]
    recordings = meta_animal['Number'] #,'blue', 'green','cyan'])
    print(recordings)
    images = np.zeros((512, 512, np.shape(meta_animal)[0]))

    for idx, recording in enumerate(recordings):
        stats = np.load(meta_data['Folder'][recording] +
                        str(meta_data['Subfolder'][recording]) +
                        str(int(meta_data['Recording idx'][recording])) +
                        '/suite2p/plane0/stat.npy',
                        allow_pickle=True)
        iscell = np.load(meta_data['Folder'][recording] +
                         str(meta_data['Subfolder'][recording]) +
                         str(int(meta_data['Recording idx'][recording])) +
                         '/suite2p/plane0/iscell.npy',
                         allow_pickle=True)
        stats = stats[iscell[:, 0].astype(bool)]

        for stats_neuron in stats:
                images[stats_neuron['ypix'], stats_neuron['xpix'], idx] += 1
        images[:,:,idx][images[:,:,idx]>1] = 1
        
    roi_corresp = np.zeros((512, 512))
 
    #for idx0 in range(np.shape(images)[2]):
        #for idx1 in range(idx0, np.shape(images)[2]):
    for idx0, recording0 in enumerate(recordings):
        for idx1, recording1 in enumerate(recordings):
           # roi_corresp = images[:,:,idx0] + 2*images[:,:,idx1]
            roi_corresp = images[:,:,idx0] + 2*sr.transform(images[:, :, idx1], tmats_loaded[idx0, idx1, :, :])
            file_title = str(idx0) + "_" + str(idx1) + "_" + meta_data['Condition'][recording0] + "_" + meta_data['Condition'][recording1]
            print(file_title)
            plt.imshow(roi_corresp[:,:],cmap=Amap) #'tab10'
            #plt.annotate('blue - reference / green - corrected / cyan - intersection', xy=(0, 0), xytext=(.8, 0), fontsize=12)
            plt.savefig(save_plots_path + "ROIAdjustmentCheck/" + str(animal) + '/' + file_title + "_ROI_correspondence.svg")
            #plt.show()
'''

if not os.path.exists(save_plots_path + 'ROIs/'):
    os.makedirs(save_plots_path + 'ROIs/')
   
for animal in animals:
    if not os.path.exists(save_plots_path + 'ROIs/' + str(animal) + "/"):
        os.makedirs(save_plots_path + 'ROIs/' + str(animal) + "/")  
    
    #tmats_loaded = np.load(path4results + 'StackReg/' + str(animal) + "_best_tmats" + '.npy')
    meta_animal = meta_data[meta_data['Mouse_all'] == animal]
    recordings = meta_animal['Number'] #,'blue', 'green','cyan'])
    print(recordings)
    images = np.zeros((512, 512, np.shape(meta_animal)[0]))

    for idx, recording in enumerate(recordings):
        stats = np.load(meta_data['Folder'][recording] +
                        str(meta_data['Subfolder'][recording]) +
                        str(int(meta_data['Recording idx'][recording])) +
                        '/suite2p/plane0/stat.npy',
                        allow_pickle=True)
        iscell = np.load(meta_data['Folder'][recording] +
                         str(meta_data['Subfolder'][recording]) +
                         str(int(meta_data['Recording idx'][recording])) +
                         '/suite2p/plane0/iscell.npy',
                         allow_pickle=True)
        #stats = stats[iscell[:, 0].astype(bool)]

        for stats_neuron in stats:
                images[stats_neuron['ypix'], stats_neuron['xpix'], idx] += 1
        images[:,:,idx][images[:,:,idx]>1] = 1
        
    roi_corresp = np.zeros((512, 512))

    for idx0, recording0 in enumerate(recordings):
    
            #roi_corresp = sr.transform(images[:, :, idx0], tmats_loaded[0, idx0, :, :])
            roi_corresp = images[:, :, idx0]
            #roi_corresp [roi_corresp > 0.5] = 1
            file_title = str(idx0) + "_" + meta_data['Condition'][recording0]
            print(file_title)
            plt.imshow(roi_corresp[:,:],cmap='binary') #'tab10'

            plt.savefig(save_plots_path + "ROIs/" + str(animal) + '/' + file_title + "_ROIs.svg")