import numpy as np 
from skimage.filters import threshold_otsu, threshold_minimum
from scipy.optimize import linear_sum_assignment

from track2p.match.utils import get_cost_mat, get_iou, init_all_pl_match_mat

# assigment of ROIs in each ref-reg pair

def get_all_ds_assign(track_ops, all_ds_all_roi_ref, all_ds_all_roi_reg):

    all_ds_assign = []
    all_ds_assign_thr = []
    all_ds_thr_met = []
    all_ds_thr = []

    for i in range(len(track_ops.all_ds_path)-1):
        print(f'Finding matches in ref-reg pair: {i+1}/{len(track_ops.all_ds_path)-1}')
        ds_assign = []
        ds_assign_thr = []
        ds_thr_met = []
        ds_thr = []
        for j in range(track_ops.nplanes):
            all_roi_ref = all_ds_all_roi_ref[i][j]
            all_roi_reg = all_ds_all_roi_reg[i][j]

            # 1) compute cost matrix (currently two methods available, see DefaultTrackOps)
            cost_mat, all_inds_ref_filt, all_inds_reg_filt = get_cost_mat(all_roi_ref, all_roi_reg, track_ops) 

            # 2) optimally assign pairs
            ref_ind_filt, reg_ind_filt = linear_sum_assignment(cost_mat)

            # 3) convert them to pre-filtered indices (these are the indices of the ROIs after iscell)
            ref_ind = all_inds_ref_filt[ref_ind_filt]
            reg_ind = all_inds_reg_filt[reg_ind_filt]

            # 4) for each matched pair (len(all_roi_ref)) compute thresholding metric (in this case IOU, the filtering will be done afterwards in the all-day assignment)
            thr_met = get_iou(all_roi_ref[:,:,ref_ind], all_roi_reg[:,:,reg_ind])
            thr_met_compute = thr_met[thr_met>0] if track_ops.thr_remove_zeros else thr_met # remove zeros for computing the threshold (otsu thresholding is squed 

            # 5) compute otsu threshold on thr_met
            if track_ops.thr_method == 'otsu':
                thr = threshold_otsu(thr_met_compute)
            elif track_ops.thr_method == 'min':
                thr = threshold_minimum(thr_met_compute)


            ds_assign.append([ref_ind, reg_ind])
            ds_assign_thr.append([ref_ind[thr_met>thr], reg_ind[thr_met>thr]])
            ds_thr_met.append(thr_met)
            ds_thr.append(thr)

        all_ds_assign.append(ds_assign)
        all_ds_assign_thr.append(ds_assign_thr)
        all_ds_thr_met.append(ds_thr_met)
        all_ds_thr.append(ds_thr)
        print(f'Done ref-reg pair: {i+1}/{len(track_ops.all_ds_path)-1}')

    return all_ds_assign, all_ds_assign_thr, all_ds_thr_met, all_ds_thr


# propagating matches across all days

def get_all_pl_match_mat(all_ds_all_roi_ref, all_ds_assign_thr, track_ops):

    all_pl_match_mat = init_all_pl_match_mat(all_ds_all_roi_ref, all_ds_assign_thr, track_ops)

    for i in range(track_ops.nplanes):
        pl_match_mat = all_pl_match_mat[i]
        # now for each row in the match matrix (each ROI in the ref recording) we need to find the match across all days, if there is none then we leave it as None

        for roi_idx in range(pl_match_mat.shape[0]): # roi_idx is the index on first session
            # if first column is none then we skip this row
            if pl_match_mat[roi_idx, 0] is None:
                continue
            # otherwise we find the match in the all_ds_assign_thr
            else:
                ref_roi_ds0 = pl_match_mat[roi_idx, 0]
                track_roi = np.array(ref_roi_ds0)
                for ds_ind in range(pl_match_mat.shape[1]-1):
                    matches = all_ds_assign_thr[ds_ind][i]
                    
                    ref_ind = matches[0]
                    reg_ind = matches[1]
                    
                    reg_ind_ind = np.where(ref_ind==track_roi.item())[0]

                    # if there is a match then we update the track_roi
                    if reg_ind_ind.size>0:
                        track_roi = reg_ind[reg_ind_ind]
                        pl_match_mat[roi_idx, ds_ind+1] = track_roi.item()
                        
                    # if there is no match then we stop tracking this ROI
                    else:
                        break
                    
        # compute how many ROIs are tracked across all days
        n_tracked = np.sum(np.all(pl_match_mat!=None, axis=1))
        print(f'Number of ROIs tracked in plane{i} across all days: {n_tracked}')
        track_ops.all_pl_match_mat = all_pl_match_mat
        track_ops.n_tracked = n_tracked
        
    return all_pl_match_mat