# AUTOGENERATED! DO NOT EDIT! File to edit: ../source_nbs/lib_nbs/utils_challenge.ipynb.

# %% auto 0
__all__ = ['majority_filter', 'label_filter', 'label_continuous_to_list', 'label_list_to_continuous', 'array_to_df',
           'df_to_array', 'get_VIP', 'changepoint_assignment', 'changepoint_alpha_beta', 'jaccard_index',
           'single_changepoint_error', 'ensemble_changepoint_error', 'create_binary_segment',
           'jaccard_between_segments', 'segment_assignment', 'metric_anomalous_exponent',
           'metric_diffusion_coefficient', 'metric_diffusive_state', 'check_no_changepoints', 'segment_property_errors',
           'extract_ensemble', 'multimode_dist', 'distribution_distance', 'error_Ensemble_dataset',
           'check_prediction_length', 'separate_prediction_values', 'load_file_to_df', 'error_SingleTraj_dataset',
           'listdir_nohidden', 'codalab_scoring']

# %% ../source_nbs/lib_nbs/utils_challenge.ipynb 2
import numpy as np
from scipy.optimize import linear_sum_assignment
import pandas
from tqdm.auto import tqdm
import warnings

# %% ../source_nbs/lib_nbs/utils_challenge.ipynb 6
def majority_filter(seq, width):
    '''
    Given a vector, applies a majority filter of given width.
    
    Parameters
    ----------
    seq : list
        Vector to filter.
    width : int
        Size of the window in which the filter is applied.
    
    Returns
    -------
    list
        Filtered vector
    '''
    offset = width // 2
    seq = [0] * offset + seq
    return [max(set(a), key=a.count) 
        for a in (seq[i:i+width] for i in range(len(seq) - offset))]

def label_filter(label, 
                 window_size = 5, 
                 min_seg = 3):
    '''
    Given a vector of changing labels, applies a majority filter such that the minimum segment of a particular label is
    bigger than the minimum set segment.
    
    Parameters
    ----------
    label : list
        label vector to filter.
    window_size : int
        Size of the window in which the majority filter is applied.
    min_seg : int
        Minimum segment size after filtering.
    
    Returns
    -------
    np.array
        Filtered label vector
    '''
    
    if np.min(label) < 0:
        raise ValueError('This function only works with positive labels')
        
    # if there are no changes:
    if np.sum(label[1:] != label[:-1]) == 0:
        return label
    
    # define dummy vector with all zeros and ones
    values, dummy = np.unique(label, return_inverse = True)
    
    # check if there are segment smaller than minimal segment (min_seg)
    cp = np.argwhere(dummy[1:] != dummy[:-1])
    cp = np.append(0, cp)
    current_min = (cp[1:]-cp[:-1]).flatten().min()

    while (current_min < min_seg):

        filt = majority_filter(dummy.tolist(), width = window_size)
        filt = np.array(filt)
        
        # check if there are segment smaller than minimal segment (min_seg)
        cp = np.argwhere(filt[1:] != filt[:-1])
        
        # If all changepoints were eliminated
        if cp.size == 0:
            break
            
        cp = np.append(0, cp)
        current_min = (cp[1:]-cp[:-1]).flatten().min()
        
        if (dummy == filt).all():  
            # If all failed and still have segments smaller than min_seg
            seg_lengths = (cp[1:]-cp[:-1]).flatten().astype(int)
            seg_smaller = np.argwhere(seg_lengths < min_seg).flatten()
            
            # We go over each segment and we asign the values 'by hand'
            for idxsegs in seg_smaller:     
                if seg_lengths[idxsegs] == 1:
                    filt[(cp[idxsegs]+1)] = filt[cp[idxsegs]]     
                elif seg_lengths[idxsegs] == 2:      
                    filt[(cp[idxsegs]+1)] = filt[cp[idxsegs]]     
                    filt[(cp[idxsegs]+2)] = filt[cp[idxsegs]+3]                     
            
            dummy = filt
            break         
        dummy = filt
        
        
    # Check boundaries
    if dummy[0] != dummy[1] or dummy[1] != dummy[2]:
        dummy[:2] = dummy[2]
    if dummy[-2] != dummy[-3] or dummy[-1] != dummy[-2]:
        dummy[-3:] = dummy[-3]
        
    # reset to label values
    dummy_ret = np.zeros_like(dummy).astype(float)
    
    for idx, v in enumerate(values):
        dummy_ret[dummy == idx] = v

    return dummy_ret

# %% ../source_nbs/lib_nbs/utils_challenge.ipynb 16
def label_continuous_to_list(labs):
    ''' 
    Given an array of T x 2 labels containing the anomalous exponent and diffusion 
    coefficient at each timestep, returns 3 arrays, each containing the changepoints, 
    exponents and coefficient, respectively. 
    If labs is size T x 3, then we consider that diffusive states are given and also
    return those.
    
    Parameters
    ----------
    labs : array
        T x 2  or T x 3 labels containing the anomalous exponent, diffusion 
        and diffusive state.
        
    Returns
    -------
    tuple
        - First element is the list of change points
        - The rest are corresponding segment properties (order: alpha, Ds and states)        
        '''
    
    # Check if states were given
    are_states = False
    if labs.shape[1] == 3:
        are_states = True
        
    # Check in which variable there is changes
    CP = np.argwhere((labs[:-1, :] != labs[1:, :]).sum(1) != 0).flatten()+1 
    T = labs.shape[0] 

    alphas = np.zeros(len(CP)+1)
    Ds = np.zeros(len(CP)+1)
    if are_states: states = np.zeros(len(CP)+1)
        
    for idx, cp in enumerate(np.append(CP, T)):
        alphas[idx] = labs[cp-1, 0]
        Ds[idx] = labs[cp-1, 1]
        if are_states: states[idx] = labs[cp-1, 2]
    
    CP = np.append(CP, T)
    
    if are_states:
        return CP, alphas, Ds, states
    else:
        return CP, alphas, Ds    

# %% ../source_nbs/lib_nbs/utils_challenge.ipynb 20
def label_list_to_continuous(CP, label):
    '''
    Given a list of change points and the labels of the diffusion properties of the
    resulting segments, generates and array of continuous labels. The last change point
    indicates the array length.
    
    Parameters
    ----------
    CP : array, list
        list of change points. Last change point indicates label length.
    label : array, list
        list of segment properties
        
    Returns
    -------
    array
        Continuous label created from the given change points and segment properties
    '''    
    
    if isinstance(label, list):
        label = np.array(label)
    segs = create_binary_segment(CP[:-1], CP[-1])
    return (segs.transpose()*label).sum(1)

# %% ../source_nbs/lib_nbs/utils_challenge.ipynb 24
from .utils_trajectories import segs_inside_fov


def array_to_df(trajs, 
               labels,
               min_length = 10,
               fov_origin = [0,0], fov_length= 100.0, cutoff_length = 10):
    '''
    Given arrays for the position and labels of trajectories, creates a dataframe with that
    data. The function also applies the demanded FOV. If you don't want a field of view, chose a 
    FOV length bigger (smaller) that your maximum (minimum) trajectory position.
   
    Parameters
    ----------
    trajs : array 
        Trajectories to store in the df (dimension: T x N x 3)
    labels : array
        Labels to store in the df (dimension: T x N x 3)        
    fov_origin : tuple
        Bottom left point of the square defining the FOV.
    fov_length : float
        Size of the box defining the FOV.
    cutoff_length : int
        Minimum length of a trajectory inside the FOV to be considered in the output dataset.
    
    
    Returns
    -------
    tuple
        - df_in (dataframe): dataframe with trajectories
        - df_out (datafram): dataframe with labels 
    '''
    
    xs = []
    ys = []
    idxs = []   
    
    df_out = pandas.DataFrame(columns = ['traj_idx', 'Ds', 'alphas', 'states', 'changepoints']) 
    
    idx_t = 0
    for traj, l_alpha, l_D, l_s in zip(tqdm(trajs), labels[:, :, 0], labels[:, :, 1], labels[:, :, 2]):

        # Check FOV and 
        idx_inside_segments = segs_inside_fov(traj, fov_origin, fov_length, cutoff_length)

        if idx_inside_segments is not None:

            for idx_in in idx_inside_segments:            
                seg_x = traj[idx_in[0]:idx_in[1], 0]
                seg_y = traj[idx_in[0]:idx_in[1], 1]
                seg_alpha = l_alpha[idx_in[0]:idx_in[1]]
                seg_D = l_D[idx_in[0]:idx_in[1]]
                seg_state = l_s[idx_in[0]:idx_in[1]]

                # Filtering
                seg_alpha = label_filter(seg_alpha)
                seg_D = label_filter(seg_D)
                seg_state = label_filter(seg_state)
                
                
                # Stacking data of input dataframe
                xs += seg_x.tolist()
                ys += seg_y.tolist()
                idxs += (np.ones(len(seg_x))*idx_t).tolist()
                
                # Transforming to list of changepoints and physical properties
                merge = np.hstack((seg_alpha.reshape(seg_alpha.shape[0], 1),
                                   seg_D.reshape(seg_D.shape[0], 1),
                                   seg_state.reshape(seg_state.shape[0], 1)))
                
                CP, alphas, Ds, states = label_continuous_to_list(merge)
                
                # Saving each segment info in output dataframe
                df_out.loc[df_out.shape[0]] = [idx_t, Ds, alphas, states, CP]
                
                # Updating segment index
                idx_t += 1

    
    # Saving trajectories in Dataframe
    tr_to_df = np.vstack((idxs,
                          xs,
                          ys)).transpose()
    df_in = pandas.DataFrame(tr_to_df, columns = ['traj_idx', 'x', 'y'])  
    
    return df_in, df_out

# %% ../source_nbs/lib_nbs/utils_challenge.ipynb 29
def df_to_array(df, pad = -1):
    '''
    Transform a dataframe as the ones given in the ANDI 2 challenge (i.e. 4 columns:
    traj_idx, frame, x, y) into a numpy array. To deal with irregular temporal supports,
    we pad the array whenever the trajectory is not present.
    The output array has the typical shape of ANDI datasets: TxNx2
    
    Parameters
    ----------
    df : dataframe
        Dataframe with four columns 'traj_idx': the trajectory index, 'frame' the time frame and 
        'x' and 'y' the positions of the particle.
    pad : int
        Number to use as padding.
    
    Returns
    -------
    array
        Array containing the trajectories from the dataframe, with usual ANDI shape (TxNx2).
    
    
    '''

    max_T = int(df.frame.max()+1)
    num_part = int(df.iloc[-1].traj_idx)
    array_trajs = np.ones((max_T, num_part+1, 2))*pad

    for idx in np.unique(df.traj_idx).astype(int):

        df_part = df.loc[df.traj_idx == idx]

        array_trajs[df_part.frame.values.astype(int), idx, 0] = df_part.x.values
        array_trajs[df_part.frame.values.astype(int), idx, 1] = df_part.y.values
        
    return array_trajs

# %% ../source_nbs/lib_nbs/utils_challenge.ipynb 31
from scipy.spatial import distance


def get_VIP(array_trajs, num_vip = 5, min_distance = 2, pad = -1):
    '''
    Given an array of trajectories, finds the particles VIP particles that participants will
    need to characterize in the video trakcl.
    
    The function first finds the particles that exist at frame 0 (i.e. that their first value 
    is different from pad). Then, iterates over this particles to find num_vip that are at 
    distance > than min_distance in the first frame.
    
    Parameters
    ----------
    array_trajs : array
        Position of the trajectories that will be considered for the VIP search.
    num_vip : int
        Number of VIP particles to flag.
    min_distance : float
        Minimum distance between two VIP particles.
    pad : int
        Number used to indicate in the temporal support that the particle is outside of the FOV.
        
    Returns
    -------
    list
        List of indices of the chosen VIP particles
    
    '''
    
    candidates_vip = np.argwhere(array_trajs[0,:,0] != pad).flatten()
    if len(candidates_vip) < num_vip:
        raise ValueError('Number of VIP demanded is bigger than available particles.')

    elected = []
    count_while = 0
    while len(elected) < num_vip:

        elected = [np.random.choice(candidates_vip)]

        for c_idx in candidates_vip:
            if c_idx == elected[0]:
                continue
            if len(array_trajs[0, elected,:].shape) < 2:
                all_rest = np.expand_dims(array_trajs[0, elected,:], 0)
            else:
                all_rest = array_trajs[0, elected,:]

            dist = distance.cdist(np.expand_dims(array_trajs[0,c_idx,:], 0), all_rest, metric='euclidean').transpose()

            if dist.min() > 2:
                elected.append(c_idx)

            if len(elected) == num_vip:
                break


        count_while += 1
        if count_while > 100: 
            raise ValueError('Could not find suitable VIP particles. This is due to either having to few particles or them being too close')
            
    return elected


# %% ../source_nbs/lib_nbs/utils_challenge.ipynb 34
def changepoint_assignment(GT, preds):
    ''' 
    Given a list of groundtruth and predicted changepoints, solves the assignment problem via
    the Munkres algorithm (aka Hungarian algorithm) and returns two arrays containing the index of the
    paired groundtruth and predicted changepoints, respectively.
    
    The distance between change point is the Euclidean distance.
    
    Parameters
    ----------
    GT : list
        List of groundtruth change points.
    preds : list
        List of predicted change points.
    
    Returns
    -------
    tuple
        - tuple of two arrays, each corresponding to the assigned GT and pred changepoints
        - Cost matrix
    
    '''
    
    cost_matrix = np.zeros((len(GT), len(preds)))

    for idxg, gt in enumerate(GT):
        for idxp, pred in enumerate(preds):
            cost_matrix[idxg, idxp] = np.abs(gt-pred)
            
    return linear_sum_assignment(cost_matrix), cost_matrix

# %% ../source_nbs/lib_nbs/utils_challenge.ipynb 36
def changepoint_alpha_beta(GT, preds, threshold = 10):
    '''
    Calculate the alpha and beta measure of paired changepoints.
    Inspired from Supplemantary Note 3 in https://www.nature.com/articles/nmeth.2808
    
    Parameters
    ----------
    GT : list
        List of groundtruth change points.
    preds : list
        List of predicted change points.
    threshold : float
        Distance from which predictions are considered to have failed. They are then assigned this number.
    
    Returns
    -------
    tuple
        alpha, beta
        
    '''

    assignment, _ = changepoint_assignment(GT, preds)
    assignment = np.array(assignment)

    threshold = 10
    distance = np.abs(GT[assignment[0]] - preds[assignment[1]])
    distance[distance > threshold] = threshold
    distance = np.sum(distance)

    d_x_phi = threshold*len(GT)
    d_ybar_phi = max([0, (len(preds)-len(GT))*threshold])

    alpha = 1-distance/d_x_phi
    beta = (d_x_phi-distance)/(d_x_phi+d_ybar_phi)

    return alpha, beta

# %% ../source_nbs/lib_nbs/utils_challenge.ipynb 38
def jaccard_index(TP: int, # true positive
                  FP: int, # false positive
                  FN: int # false negative
                 )-> float: # Jaccard Index
    '''
    Given the true positive, false positive and false negative rates, calculates the Jaccard Index
    '''
    return TP/(TP+FP+FN)

# %% ../source_nbs/lib_nbs/utils_challenge.ipynb 39
def single_changepoint_error(GT, preds, threshold = 5):
    '''
    Given the groundtruth and predicted changepoints for a single trajectory, first solves the assignment problem between changepoints,
    then calculates the RMSE of the true positive pairs and the Jaccard index.
    
    Parameters
    ----------
    GT : list
        List of groundtruth change points.
    preds : list
        List of predicted change points.
    threshold : float
        Distance from which predictions are considered to have failed. They are then assigned this number.
    
    Returns
    -------
    tuple
        - TP_rmse: root mean square error of the true positive change points.
        - Jaccard Index of the ensemble predictions        
        
    '''
    
    assignment, _ = changepoint_assignment(GT, preds)
    assignment = np.array(assignment)
    
    TP, FP, FN = 0, 0, 0
    TP_rmse = []
    for p in assignment.transpose():
        
        if np.abs(GT[p[0]] - preds[p[1]]) < threshold:
            TP += 1
            TP_rmse.append((GT[p[0]] - preds[p[1]])**2)
        else:
            FP += 1
            FN += 1
    # Calculating RMSE
    TP_rmse = np.sqrt(np.mean(TP_rmse))

    # Checking false positive and missed events
    if len(preds) > len(GT):
        FP += len(preds) - len(GT)
    elif len(preds) < len(GT):
        FN += len(GT) - len(preds)
    
    return TP_rmse, jaccard_index(TP, FP, FN)

# %% ../source_nbs/lib_nbs/utils_challenge.ipynb 40
def ensemble_changepoint_error(GT_ensemble, pred_ensemble, threshold = 5):    
    ''' 
    Given an ensemble of groundtruth and predicted change points, iterates
    over each trajectory's changepoints. For each, it solves the assignment problem 
    between changepoints. Then, calculates the RMSE of the true positive pairs and
    the Jaccard index over the ensemble of changepoints (i.e. not the mean of them 
    w.r.t. to the trajectories)
    
    Parameters
    ----------
    GT_ensemble : list, array
        Ensemble of groutruth change points.
    pred_ensemble : list
        Ensemble of predicted change points.
    threshold : float
        Distance from which predictions are considered to have failed. They are then assigned this number.
    
    Returns
    -------
    tuple
        - TP_rmse: root mean square error of the true positive change points.
        - Jaccard Index of the ensemble predictions
    
    '''
    
    TP, FP, FN = 0, 0, 0
    TP_rmse = []
    
    for gt_traj, pred_traj in zip(GT_ensemble, pred_ensemble):
        
        assignment, _ = changepoint_assignment(gt_traj, pred_traj)
        assignment = np.array(assignment)
        
        for p in assignment.transpose():
            
            if np.abs(gt_traj[p[0]] - pred_traj[p[1]]) < threshold:
                TP += 1
                TP_rmse.append((gt_traj[p[0]] - pred_traj[p[1]])**2)
            else:
                FP += 1
                FN += 1    
                
        # Checking false positive and missed events
        if len(pred_traj) > len(gt_traj):
            FP += len(pred_traj) - len(gt_traj)
        elif len(pred_traj) < len(gt_traj):
            FN += len(gt_traj) - len(pred_traj)
                
    if TP+FP+FN == 0:
        wrn_str = f'No segments found in this dataset.'
        warnings.warn(wrn_str)
        return threshold, 0
        
    # Calculating RMSE
    TP_rmse = np.sqrt(np.mean(TP_rmse))

    
        
    return TP_rmse, jaccard_index(TP, FP, FN)

# %% ../source_nbs/lib_nbs/utils_challenge.ipynb 43
def create_binary_segment(CP: list, # list of changepoints
                          T: int # length of the trajectory
                         )-> list: # list of arrays with value 1 in the temporal support of the current segment.
    '''
    Given a set of changepoints and the lenght of the trajectory, create segments which are equal to one
    if the segment takes place at that position and zero otherwise.
    '''
    segments = np.zeros((len(CP)+1, T))
    CP = np.append(0, CP)
    for idx, (cp1, cp2) in enumerate(zip(CP[:-1], CP[1:])):
        segments[idx, cp1+1:cp2+1] = 1
    segments[-1, CP[-1]+1:] = 1
    segments[0, 0] = 1
    return segments

# %% ../source_nbs/lib_nbs/utils_challenge.ipynb 45
def jaccard_between_segments(gt, pred):
    '''
    Given two segments, calculates the Jaccard index between them by considering TP as correct labeling,
    FN as missed events and FP leftover predictions.
    
    Parameters
    ----------
    gt : array
        groundtruth segment, equal to one in the temporal support of the given segment, zero otherwise.
    pred : array
        predicted segment, equal to one in the temporal support of the given segment, zero otherwise.
    
    Returns
    -------
    float
        Jaccard index between the given segments.
    '''
    
    if len(gt) > len(pred):
        pred = np.append(pred, np.zeros(len(gt) - len(pred)))
    elif len(pred) > len(gt):                        
        gt = np.append(gt, np.zeros(len(pred) - len(gt)))
    
    
    tp = np.sum(np.logical_and(pred == 1, gt == 1))
    fp = np.sum(np.logical_and(pred == 1, gt == 0))
    fn = np.sum(np.logical_and(pred == 0, gt == 1))
    
    # special case for absence of changepoint
    if tp+fp+fn == 0: return 0    
    else: return jaccard_index(tp, fp, fn)

# %% ../source_nbs/lib_nbs/utils_challenge.ipynb 46
def segment_assignment(GT, preds, T:int = None):
    ''' 
    Given a list of groundtruth and predicted changepoints, generates a set of segments. Then constructs 
    a cost matrix by calculting the Jaccard Index between segments. From this cost matrix, we solve the 
    assignment  problem via the Munkres algorithm (aka Hungarian algorithm) and returns two arrays 
    containing the index of the groundtruth and predicted segments, respectively.
    
    If T = None, then we consider that GT and preds may have different lenghts. In that case, the end
    of the segments is the the last CP of each set of CPs.
    
    Parameters
    ----------
    GT : list
        List of groundtruth change points.
    preds : list
        List of predicted change points.
    T : int, None
        Length of the trajectory. If None, considers different GT and preds length.
    
    Returns
    -------
    tuple
        - tuple of two arrays, each corresponding to the assigned GT and pred changepoints
        - Cost matrix calculated via JI of segments   
    
    '''
   
    if T is not None:
        T_gt = T_pred = T
        # Check if the GT or predictions are a single integer or an empty array
        if isinstance(GT, int): GT = [GT]
        elif len(GT) == 0: GT = [T-1]

        if isinstance(preds, int): preds = [preds]
        elif len(preds) == 0: preds = [T-1]
    else:
        T_gt = GT[-1]
        if len(GT) > 1:
            GT = GT[:-1]            
            
        T_pred = preds[-1]
        if len(preds) > 1:
            preds = preds[:-1]
        
    
    
    seg_GT = create_binary_segment(GT, T_gt)
    seg_preds = create_binary_segment(preds, T_pred)
    
    cost_matrix = np.zeros((seg_GT.shape[0], seg_preds.shape[0]))

    for idxg, gt in enumerate(seg_GT):
        for idxp, pred in enumerate(seg_preds):
            cost_matrix[idxg, idxp] = 1-jaccard_between_segments(gt, pred)

    return linear_sum_assignment(cost_matrix), cost_matrix

# %% ../source_nbs/lib_nbs/utils_challenge.ipynb 56
from sklearn.metrics import mean_squared_log_error as msle, f1_score
from .models_phenom import models_phenom

def metric_anomalous_exponent(gt = None,
                              pred = None,
                              max_error = np.abs(models_phenom().bound_alpha[0]-models_phenom().bound_alpha[1])):
    ''' 
    Compute the mean absolute error (mae) between anomalous exponents.
    Checks the current bounds of anomalous exponents from models_phenom to calculate the maximum error.
    ''' 
    error = np.mean(np.abs(gt-pred))
    if error > max_error:
        return max_error
    else: 
        return error

def metric_diffusion_coefficient(gt = None, pred = None, 
                                 threshold_min = models_phenom().bound_D[0],                               
                                 max_error = msle([models_phenom().bound_D[0]],
                                                  [models_phenom().bound_D[1]])):
    ''' 
    Compute the mean squared log error (msle) between diffusion coefficients.
    Checks the current bounds of diffusion from models_phenom to calculate the maximum error. 
    ''' 
    
     # considering the presence of zeros and negatives
    pred = np.array(pred).copy(); gt = np.array(gt).copy()
    pred[pred <= threshold_min] = threshold_min
    gt[gt <= threshold_min] = threshold_min    
    # mean squared log error
    error = msle(gt, pred)
    
    if error > max_error:
        return max_error
    else: 
        return error

def metric_diffusive_state(gt = None, pred = None, max_error = False):
    ''' 
    Compute the F1 score between diffusive states. 
    ''' 
    return f1_score(gt.astype(int), pred.astype(int), average = 'micro')

# %% ../source_nbs/lib_nbs/utils_challenge.ipynb 60
def check_no_changepoints(GT_cp, GT_alpha, GT_D, GT_s,
                          preds_cp, preds_alpha, preds_D, preds_s,
                          T:bool|int = None):
    '''
    Given predicionts over changepoints and variables, checks if in both GT and preds there is an 
    absence of change point. If so, takes that into account to pair variables.
    
    Parameters
    ----------
    GT_cp : list, int, float
        Groundtruth change points
    GT_alpha : list, float
        Groundtruth anomalous exponent
    GT_D : list, float
        Groundtruth diffusion coefficient
    GT_s : list, float
        Groundtruth diffusive state
    preds_cp : list, int, float
        Predicted change points
    preds_alpha : list, float
        Predicted anomalous exponent
    preds_D : list, float
        Predicted diffusion coefficient
    preds_s : list, float
        Predicted diffusive state
    T : bool,int
        (optional) Length of the trajectories. If none, last change point is length.
    
    Returns
    -------
    tuple
        - False if there are change points. True if there were missing change points.
        - Next three are either all Nones if change points were detected, or paired exponents, 
        coefficient and states if some change points were missing.
    
    '''


    if isinstance(GT_cp, int) or isinstance(GT_cp, float):
        GT_cp = [GT_cp]
    if isinstance(preds_cp, int) or isinstance(preds_cp, float):
        preds_cp = [preds_cp]
        
    no_GT_cp = False; no_preds_cp = False
    # CP always contain the final point of the trajectory, hence minimal length is one
    if len(GT_cp) == 1: no_GT_cp = True
    if len(preds_cp) == 1: no_preds_cp = True       
        

    if no_GT_cp + no_preds_cp == 0:
        return False, None, None, None
    
    else:

        [row_ind, col_ind], _ = segment_assignment(GT_cp, preds_cp, T)   

        if no_GT_cp and not no_preds_cp:
            paired_alpha = np.array([[GT_alpha[0], preds_alpha[col_ind[0]]]])
            paired_D = np.array([[GT_D[0], preds_D[col_ind[0]]]])
            paired_s = np.array([[GT_s[0], preds_s[col_ind[0]]]])

        if no_preds_cp and not no_GT_cp:
            row_position = np.argwhere(col_ind == 0).flatten()[0]            
            paired_alpha = np.array([[GT_alpha[row_position], preds_alpha[col_ind[row_position]]]])
            paired_D = np.array([[GT_D[row_position], preds_D[col_ind[row_position]]]])
            paired_s = np.array([[GT_s[row_position], preds_s[col_ind[row_position]]]])
            
        if no_preds_cp and no_GT_cp: 
            paired_alpha = np.array([[GT_alpha[0], preds_alpha[0]]])
            paired_D = np.array([[GT_D[0], preds_D[0]]])
            paired_s = np.array([[GT_s[0], preds_s[0]]])
            

        return True, paired_alpha, paired_D, paired_s

# %% ../source_nbs/lib_nbs/utils_challenge.ipynb 61
def segment_property_errors(GT_cp, GT_alpha, GT_D, GT_s,
                            preds_cp, preds_alpha, preds_D, preds_s,
                            return_pairs = False,
                            T = None):
    '''
    Given predicionts over change points and the value of diffusion parameters in the generated
    segments, computes the defined metrics.
    
    Parameters
    ----------
    GT_cp : list, int, float
        Groundtruth change points
    GT_alpha : list, float
        Groundtruth anomalous exponent
    GT_D : list, float
        Groundtruth diffusion coefficient
    GT_s : list, float
        Groundtruth diffusive state
    preds_cp : list, int, float
        Predicted change points
    preds_alpha : list, float
        Predicted anomalous exponent
    preds_D : list, float
        Predicted diffusion coefficient
    preds_s : list, float
        Predicted diffusive state
    return_pairs : bool
        If True, returns the assigment pairs for each diffusive property.
    T : bool,int
        (optional) Length of the trajectories. If none, last change point is length.
    
    Returns
    -------
    tuple
        - if return_pairs = True, returns the assigned pairs of diffusive properties
        - if return_pairs = False, returns the errors for each diffusive property
    '''
    
    # Check cases in which changepoint where not detected or there were none in groundtruth
    no_change_point_case, paired_alpha, paired_D, paired_s = check_no_changepoints(GT_cp, GT_alpha, GT_D, GT_s,
                                                                                   preds_cp, preds_alpha, preds_D, preds_s, T)
   
    if not no_change_point_case:
        # Solve the assignment problem
        [row_ind, col_ind], _ = segment_assignment(GT_cp, preds_cp, T)
   
        # iterate over the groundtruth segments
        paired_alpha, paired_D, paired_s = [], [], []
        for idx_seg, (gt_alpha, gt_D) in enumerate(zip(GT_alpha, GT_D)):

            row_position = np.argwhere(row_ind == idx_seg).flatten()

            # if the GT segment was associated to a prediction
            if len(row_position) > 0:
                row_position = int(row_position)
                # alpha                
                gt_a_seg = GT_alpha[idx_seg]                
                pred_a_seg = preds_alpha[col_ind[row_position]]
                # d
                gt_d_seg = GT_D[idx_seg]
                pred_d_seg = preds_D[col_ind[row_position]]
                # state
                gt_s_seg = GT_s[idx_seg]
                pred_s_seg = preds_s[col_ind[row_position]]

                paired_alpha.append([gt_a_seg, pred_a_seg])
                paired_D.append([gt_d_seg, pred_d_seg])
                paired_s.append([gt_s_seg, pred_s_seg])

        paired_alpha, paired_D, paired_s = np.array(paired_alpha), np.array(paired_D), np.array(paired_s) 
    
    if return_pairs:
        return paired_alpha, paired_D, paired_s 
    else:
        error_alpha = metric_anomalous_exponent(paired_alpha[:,0], paired_alpha[:,1])
        error_D = metric_diffusion_coefficient(paired_D[:,0], paired_D[:,1])
        error_s = metric_diffusive_state(paired_s[:,0], paired_s[:,1])
        return error_alpha, error_D, error_s

# %% ../source_nbs/lib_nbs/utils_challenge.ipynb 69
from .models_phenom import models_phenom
def extract_ensemble(state_label, dic):
        ''' 
        Given an array of the diffusive state and a dictionary with the diffusion information,
        returns a summary of the ensemble properties for the current dataset.

        Parameters
        ----------
        state_label : array
            Array containing the diffusive state of the particles in the dataset.
            For multi-state and dimerization, this must be the number associated to the
            state (for dimerization, 0 is free, 1 is dimerized). For the rest, we follow
            the numeration of models_phenom().lab_state.
        dic : dict 
            Dictionary containing the information of the input dataset.
       
       Returns
       -------
       array
           Matrix containing the ensemble information of the input dataset. It has the following shape:
            |mu_alpha1      mu_alpha2     ... |
            |sigma_alpha1   sigma_alpha2  ... |
            |mu_D1          mu_D1         ... | 
            |sigma_D1       sigma_D2      ... |
            |counts_state1  counts_state2 ... |
        '''

        # Single state
        if dic['model'] == 'single_state': 
            ensemble = np.vstack((dic['alphas'][0],
                                   dic['alphas'][1],
                                   dic['Ds'][0],
                                   dic['Ds'][1],
                                   len(state_label)
                                   ))
        # Multi-state
        if dic['model'] == 'multi_state':
            states, counts = np.unique(state_label, return_counts=True)    
            # If the number of visited stated is not equal to the expected number of states
            if len(states) != dic['alphas'].shape[0]:
                states_corrected = np.ones(dic['alphas'].shape[0])
                counts_corrected = np.ones(dic['alphas'].shape[0])
                for s, c in zip(states, counts):
                    counts_corrected[int(s)] = c
            else: 
                counts_corrected = counts

            ensemble = np.vstack((dic['alphas'][:, 0],
                                   dic['alphas'][:, 1],
                                   dic['Ds'][:, 0],
                                   dic['Ds'][:, 1],
                                   counts_corrected
                                   ))

        # Immobile
        if dic['model'] == 'immobile_traps':
            counts = [len(state_label[state_label == models_phenom().lab_state.index('i')]),
                      len(state_label[state_label == models_phenom().lab_state.index('f')])]  
            ensemble = np.vstack(([0, dic['alphas'][0]],
                                   [0, dic['alphas'][1]],
                                   [0, dic['Ds'][0]],
                                   [0, dic['Ds'][1]],
                                   counts
                                   ))
        # dimerization    
        if dic['model'] == 'dimerization':
            counts = [len(state_label[state_label == 0]),
                      len(state_label[state_label == 1])]           
            ensemble = np.vstack((dic['alphas'][:, 0],
                                   dic['alphas'][:, 1],
                                   dic['Ds'][:, 0],
                                   dic['Ds'][:, 1],
                                   counts
                                   ))

        if dic['model'] == 'confinement':
            counts = [len(state_label[state_label == models_phenom().lab_state.index('f')]),
                      len(state_label[state_label == models_phenom().lab_state.index('c')])]   
            ensemble = np.vstack((dic['alphas'][:, 0],
                                   dic['alphas'][:, 1],
                                   dic['Ds'][:, 0],
                                   dic['Ds'][:, 1],
                                   counts
                                   ))
        return ensemble

# %% ../source_nbs/lib_nbs/utils_challenge.ipynb 70
import scipy.stats
def multimode_dist(params, weights, bound, x, normalized = False):
    '''
    Generates a multimodal distribution with given parameters.
    Also accounts for single mode if weight is float or int.
    
    Parameters
    ----------
    params : list
        Mean and variances of every mode.
    weights : list, float
        Weight of every mode. If float, we consider a single mode.
    bound : tuple
        Bounds (min, max) of the functions support.
    x : array
        Support upon which the distribution is created.
    normalize : bool
        If True, returns the normalized distribution.
    
    Returns
    -------
    array
        Value of the distribution in each point of the given support
        
    '''
    func = scipy.stats.truncnorm
    dist = np.zeros_like(x)
    lower, upper = bound 
   
    # If we have single state, change values to list to still
    # have a loop:
    if isinstance(weights, float) or isinstance(weights, int):
        params = [params]
        weights = [weights]
        
    for param, w in zip(params, weights):
        mean, var  = param  
        # introduce a cutoff to avoid nan when var = 0
        if var == 0: var = 1e-9
        unimodal = func.pdf(x,
                            (lower-mean)/np.sqrt(var),
                            (upper-mean)/np.sqrt(var),
                            loc = mean,
                            scale = np.sqrt(var))
        dist += w*unimodal
    if normalized:
        dist /= np.sum(dist)
    return dist

# %% ../source_nbs/lib_nbs/utils_challenge.ipynb 71
def distribution_distance(p:np.array, # distribution 1
                          q:np.array # distribution 2
                         )-> float:  # distance between distributions
    ''' Calculates mean absolute error between two distributions. '''
#     return np.sum(np.where(p != 0, p * np.log(p / q), 0))
    return np.abs(p-q).mean()

# %% ../source_nbs/lib_nbs/utils_challenge.ipynb 74
from .models_phenom import models_phenom

def error_Ensemble_dataset(true_data, pred_data, return_distributions = False):
    ''' 
    Calculates the ensemble metrics for the ANDI 2 challenge. The input are matrices of shape:
    
    | col1 (state 1) | col2 (state 2) | col3 (state 3) | ... |
    |:--------------:|:--------------:|:--------------:|:---:|
    | $\mu_a^1$      | $\mu_a^2$      | $\mu_a^3$      | ... |
    | $\sigma_a^1$   | $\sigma_a^2$   | $\sigma_a^3$   | ... |
    | $\mu_D^1$      | $\mu_D^2$      | $\mu_D^3$      | ... |        
    | $\sigma_D^1$   | $\sigma_D^2$   | $\sigma_D^3$   | ... |
    | $N_1$          | $N_2$          | $N_3$          | ... |
    
    Parameters
    ----------
    true_data : array
        Matrix containing the groundtruth data.
    pred_data : array
        Matrix containing the predicted data.
    return_distributions : bool
        If True, the function also outputs the generated distributions.
    
    Returns
    -------
    tuple
        - distance_alpha: distance between anomalous exponents
        - distance_D: distance between diffusion coefficients
        - dists (if asked): distributions of both groundtruth and predicted data.        
    
    '''
    
    dists = []
    for data in [true_data, pred_data]:
        
        if len(data.shape) > 1: # If we have more than one state
            alpha_info = np.delete(data, [2,3, -1], 0)
            d_info = data[2:-1,:]
            weights = data[-1,:]
            if weights.sum() > 1: weights /= weights.sum()
        else: # If single state
            alpha_info = data[:2]
            d_info = data[2:-1]
            weights = 1
            
        for idx, (var, bound) in enumerate(zip([alpha_info, d_info], 
                                               [models_phenom().bound_alpha, models_phenom().bound_D])):
            if idx == 0: x = np.linspace(bound[0], bound[1], 1000)
            else: x = np.logspace(np.log10(bound[0]), np.log10(bound[1]), 1000)
            dists.append(multimode_dist(var.T, weights, bound, x))
            
    # Distance between alpha dists
    distance_alpha = distribution_distance(dists[0], dists[2])
    distance_D = distribution_distance(dists[1], dists[3])
    
    if return_distributions:
        return distance_alpha, distance_D, dists
    else:
        return distance_alpha, distance_D

# %% ../source_nbs/lib_nbs/utils_challenge.ipynb 76
def check_prediction_length(pred):
    '''
    Given a trajectory segments prediction, checks whether it has C changepoints and C+1 segments properties values.
    As it must also contain the index of the trajectory, this is summarized by being multiple of 4. 
    In some cases, the user needs to also predict the final point of the trajectory. In this case, 
    we will have a residu of 1.
    '''
    if len(pred) % 4 == 0 or len(pred) % 4 == 1 :
        return True
    else: 
        return False

# %% ../source_nbs/lib_nbs/utils_challenge.ipynb 77
def separate_prediction_values(pred):
    '''
    Given a prediction over trjaectory segments, extracts the predictions for each segment property
    as well as the changepoint values.
    '''        
    Ds = pred[1::4]
    alphas = pred[2::4]
    states = pred[3::4]
    cp = pred[4::4]    
    return Ds, alphas, states, cp

# %% ../source_nbs/lib_nbs/utils_challenge.ipynb 78
def load_file_to_df(path_file, 
                    columns = ['traj_idx', 'Ds', 'alphas', 'states', 'changepoints']):
    '''
    Given the path of a .txt file, extract the segmentation predictions based on 
    the rules of the ANDI 2 challenge022
    '''

    with open(path_file) as f:
        lines_pred = f.read().splitlines()

    df = pandas.DataFrame(columns = columns)

    for line in lines_pred:
        # Extract values with comma separator and transform to float
        pred_traj = line.split(',')
        pred = [float(i) for i in pred_traj]
        
        # Check that prediction has the correct shape
        pred_correct = check_prediction_length(pred)
        
        # If correct size, then extract parameters and add it to dataframe
        if pred_correct:
            preds_D, preds_a, preds_s, preds_cp = separate_prediction_values(pred)

            current_row = df.shape[0]
            for param, pred_param in zip(columns, [pred[0], preds_D, preds_a, preds_s, preds_cp]):
                df.loc[current_row, param] = pred_param
                
    return df

# %% ../source_nbs/lib_nbs/utils_challenge.ipynb 83
def _get_error_bounds():
    '''
    Sets the current maximum errors we can do in the different diffusive properties.
    '''
    
    threshold_error_alpha = 2
    threshold_error_D = 1e5
    threshold_error_s = -1
    threshold_cp = 10
    return threshold_error_alpha, threshold_error_D, threshold_error_s, threshold_cp

# %% ../source_nbs/lib_nbs/utils_challenge.ipynb 84
def error_SingleTraj_dataset(df_pred, df_true, 
                              threshold_error_alpha = 2, max_val_alpha = 2, min_val_alpha = 0, 
                              threshold_error_D = 1e5, max_val_D = 1e6, min_val_D = 1e-6, 
                              threshold_error_s = -1,
                              threshold_cp = 10,
                              prints = True, disable_tqdm = False
                             ):
    '''
    Given two dataframes, corresponding to the predictions and true labels of a set
    of trajectories from the ANDI 2 challenge022, calculates the corresponding metrics
    Columns must be for both (no order needed):
    traj_idx | alphas | Ds | changepoints | states
    df_true must also contain a column 'T'.
    
    Parameters
    ----------
    df_pred : dataframe
        Predictions
    df_true : dataframe
        Groundtruth
    threshold_error_alpha : float
        (same for D, s, cp) Maximum possible error allowed. If bigger, it is substituted by this error.
    max_val_alpha : float
        (same for D, s, cp) Maximum value of the parameter.
    min_val_alpha : float
        (same for D, s, cp) Minimum value of the parameter.
    print : bool
        If True, prints the results.
    disable_tqdm : bool
        If True, disables the progress bar.
    
    Returns
    -------
    tuple
        - rmse_CP: root mean squared error change points
        - JI: Jaccard index change points
        - error_alpha: mean absolute error anomalous exponents
        - error_D: mean square log error diffusion coefficients
        - error_s: Jaccar index diffusive states
    
    '''
    # Initiate counting missing trajectories
    missing_traj = False
    
    # Deleter saving variables, just in case...
    try: del paired_alpha, paired_D, paired_s
    except: pass

    # for every trajectory, we stack paired segment properties. We also store changepoints info
    ensemble_pred_cp, ensemble_true_cp = [], []
    for t_idx in tqdm(df_true['traj_idx'].values, disable = disable_tqdm):
        
        traj_trues = df_true.loc[df_true.traj_idx == t_idx]

        traj_preds = df_pred.loc[df_pred.traj_idx == t_idx]    
        if traj_preds.shape[0] == 0:
            # If there is no trajectory, we give maximum error. To do so, we redefine predictions
            # and trues so that they give maximum error
            missing_traj += 1                       
            
            preds_cp, preds_alpha, preds_D, preds_s = [[10],
                                                       [0],
                                                       [1],
                                                       [0]]

            trues_cp, trues_alpha, trues_D, trues_s = [[10+threshold_cp],
                                                       [threshold_error_alpha],
                                                       [1+threshold_error_D],
                                                       [10]]
        
        else:      

            preds_cp, preds_alpha, preds_D, preds_s = [np.array(traj_preds.changepoints.values[0]).astype(int),
                                                       traj_preds.alphas.values[0],
                                                       traj_preds.Ds.values[0],
                                                       traj_preds.states.values[0]]

            trues_cp, trues_alpha, trues_D, trues_s = [np.array(traj_trues.changepoints.values[0]).astype(int),
                                                       traj_trues.alphas.values[0],
                                                       traj_trues.Ds.values[0],
                                                       traj_trues.states.values[0]]


        # Collecting changepoints for metric
        # In this metric, we don't want to enter the final point of the trajectory
        ensemble_pred_cp.append(preds_cp[:-1])
        ensemble_true_cp.append(trues_cp[:-1])        
        
        # collecting segment properties error after segment assignment
        pair_a, pair_d, pair_s = segment_property_errors(trues_cp, trues_alpha, trues_D, trues_s, 
                                                         preds_cp, preds_alpha, preds_D, preds_s,
                                                         return_pairs = True)
        

        
        try:
            paired_alpha = np.vstack((paired_alpha, pair_a))
            paired_D = np.vstack((paired_D, pair_d))
            paired_s = np.vstack((paired_s, pair_s))        
        except:
            paired_alpha = pair_a
            paired_D = pair_d
            paired_s = pair_s
               
    #### Calculate metrics from assembled properties   

    # checking for nans and problems in predictions
    wrong_alphas = np.argwhere(np.isnan(paired_alpha[:, 1]) | (paired_alpha[:, 1] > 2) | (paired_alpha[:, 1] < 0)).flatten()
    paired_alpha[wrong_alphas, 1] = paired_alpha[wrong_alphas, 0] + threshold_error_alpha

    wrong_ds = np.argwhere(np.isnan(paired_D[:, 1])).flatten()
    paired_D = np.abs(paired_D)
    paired_D[wrong_ds, 1] = paired_D[wrong_ds, 0] + threshold_error_D
    
    wrong_s = np.argwhere((paired_s[:, 1] > 4) | (paired_s[:, 1]<0))
    paired_s[wrong_s, 1] = threshold_error_s    
    
    # Changepoints
    rmse_CP, JI = ensemble_changepoint_error(ensemble_true_cp, ensemble_pred_cp, threshold = threshold_cp)

    # Segment properties
    error_alpha = metric_anomalous_exponent(paired_alpha[:,0], paired_alpha[:,1])
    error_D = metric_diffusion_coefficient(paired_D[:,0], paired_D[:,1])
    error_s = metric_diffusive_state(paired_s[:,0], paired_s[:,1])
    
    if prints:        
        print(f'Summary of metrics assesments:')
        if missing_traj is not False:            
            print(f'\n{missing_traj} missing trajectory/ies. ')           
        if rmse_CP == threshold_cp:
            print(f'No change points found. RMSE set to max ({threshold_cp})')
        print(f'\nChangepoint Metrics \nRMSE: {round(rmse_CP, 3)} \nJaccard Index: {round(JI, 3)}',
              f'\n\nDiffusion property metrics \nMetric anomalous exponent: {error_alpha} \nMetric diffusion coefficient: {error_D} \nMetric diffusive state: {error_s}')
              
              

    return rmse_CP, JI, error_alpha, error_D, error_s

# %% ../source_nbs/lib_nbs/utils_challenge.ipynb 96
import re
import sys
import os

# %% ../source_nbs/lib_nbs/utils_challenge.ipynb 97
def listdir_nohidden(path):
    for f in os.listdir(path):
        if not f.startswith(('.','_')):
            yield f

# %% ../source_nbs/lib_nbs/utils_challenge.ipynb 98
def codalab_scoring(input_dir , output_dir):
    '''
    Given an input directoy where predictions and groundtruths for the ANDI 2 challenge can be found,
    calculates metrics and outputs the results in a file in the given output directory.
    This code is prepared to be run in Codalab.
    '''
    
    # Error bounds
    threshold_error_alpha, threshold_error_D, threshold_error_s, threshold_cp = _get_error_bounds()
    
    ### Saving variables
    # Track 1 - Videos
    t1_ens = {'alpha': [],
              'D': []}

    t1_st = {'RMSE': [],
             'JI': [],
             'alpha': [],
             'D': [],
             'state': [],
             'num_traj': [],
             'num_traj_CP': []} # this last one takes into account no changepoint from single state

    # Track 2 - Trajectories
    t2_ens = {'alpha': [],
              'D': []}

    t2_st = {'RMSE': [],
             'JI': [],
             'alpha': [],
             'D': [],
             'state': [],
             'num_traj': [],
             'num_traj_CP': []} # this last one takes into account no changepoint from single state
    
    # Handling paths of input files
    submit_dir = os.path.join(input_dir, 'pred')
    truth_dir = os.path.join(input_dir, 'true')
    if not os.path.isdir(submit_dir):
        print( "%s doesn't exist", truth_dir)
        
    # Calculate metrics if directories exist
    if os.path.isdir(submit_dir) and os.path.isdir(truth_dir):
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

        # Extracts all files in reference directory
        true_files_list = sorted(list(listdir_nohidden(truth_dir)))#os.listdir(truth_dir)

        # Run over all files        
        missing_tracks = []
        for filename in tqdm(true_files_list):
            task = re.search('_(.+?)_labs', filename).group(1)
            exp = re.search('exp_(.+?)_', filename).group(1)
            fov = re.search('fov_(.+?).', filename).group(1)
            # check track and save found tracks
            track = int(filename[1]) 
            

            true_file = os.path.join(truth_dir, filename)
            corresponding_submission_file = os.path.join(submit_dir, filename)
            
            if not os.path.isfile(corresponding_submission_file):
                if track not in missing_tracks:
                    missing_tracks.append(track)
                if len(missing_tracks) == 2:
                    raise FileNotFoundError(f'Failed to find prediction files.')
                else:
                    continue
            
            # if not os.path.isfile(corresponding_submission_file) and missing_tracks == 1:
            #     raise FileNotFoundError(f'Failed to find prediction files.')
                # raise FileNotFoundError(f'Prediction file for: track {track}, task {task}, experiment {exp} and FOV {fov} not found.')
            
            # extract model
            if task == 'ens':
                model = np.genfromtxt(true_file, dtype='str', skip_footer=5)[1][:-1]
            else:
                file_ens = os.path.join(truth_dir, f't{track}_ens_labs_exp_{exp}_fov_{fov}.txt')
                model = np.genfromtxt(file_ens, dtype='str', skip_footer=5)[1][:-1]
            
            # Ensemble
            if task == 'ens':

                true = np.loadtxt(true_file, skiprows=1, delimiter = ';')
                pred = np.loadtxt(corresponding_submission_file, skiprows=1, delimiter = ';')

                mae_alpha, mae_D, = error_Ensemble_dataset(true_data = true,
                                                           pred_data = pred)

                if track == 1:
                    t1_ens['alpha'].append(mae_alpha)
                    t1_ens['D'].append(mae_D)
                if track == 2:                
                    t2_ens['alpha'].append(mae_alpha)
                    t2_ens['D'].append(mae_D)

            # Single trajectory
            if task == 'traj':    
                df_true = load_file_to_df(true_file)
                df_pred = load_file_to_df(corresponding_submission_file)

                rmse_CP, JI, error_alpha, error_D, error_s = error_SingleTraj_dataset(df_true = df_true, df_pred = df_pred, 
                                                                                      threshold_error_alpha = threshold_error_alpha,
                                                                                      threshold_error_D = threshold_error_D, 
                                                                                      threshold_error_s = threshold_error_s,
                                                                                      threshold_cp = threshold_cp,
                                                                                      prints = False, disable_tqdm = True)

                if track == 1:
                    # to avoid single state entering in CP metrics
                    if model != 'single_state': 
                        t1_st['RMSE'].append(rmse_CP)
                        t1_st['JI'].append(JI)
                        t1_st['num_traj_CP'].append(df_true.shape[0])
                        
                    t1_st['alpha'].append(error_alpha)
                    t1_st['D'].append(error_D)
                    t1_st['state'].append(error_s)                    
                    t1_st['num_traj'].append(df_true.shape[0])
                if track == 2:
                    # to avoid single state entering in CP metrics
                    if model != 'single_state': 
                        t2_st['RMSE'].append(rmse_CP)
                        t2_st['JI'].append(JI)
                        t2_st['num_traj_CP'].append(df_true.shape[0])
                    
                    t2_st['alpha'].append(error_alpha)
                    t2_st['D'].append(error_D)
                    t2_st['state'].append(error_s)        
                    t2_st['num_traj'].append(df_true.shape[0])

            # print(f'Track {track}, Task {task}, Exp {exp}, FOV {fov}: OK!')
       
    ### Saving data
    '''CHECK HOW TO DO THE MEAN!'''    
    # Define output file
    output_filename = os.path.join(output_dir, 'scores.txt')
    output_file = open(output_filename, 'w')

    # Single trajectory data
    # We define a variable that gives the worst values for each metric. This is applied
    # separetedly for every FOV
    worst_value_st = {'RMSE': threshold_cp,
                      'JI': 0,
                      'alpha': threshold_error_alpha,
                      'D': threshold_error_D,
                      'state': 0}
    # Run over all keys
    for key in t1_st: 
        
        # Compare results with 
        if key in ['RMSE', 'alpha', 'D']:            
            if key == 'RMSE': avg_against = 'num_traj_CP'
            else: avg_against = 'num_traj'
            
            if 1 not in missing_tracks:
                
                save_t1 = np.nanmin(np.vstack([t1_st[key],
                                               np.ones_like(t1_st[key])*worst_value_st[key]]),
                                    axis = 0)
                save_t1 = np.average(save_t1, axis = 0, weights = t1_st[avg_against])
                
            if 2 not in missing_tracks:                     
                save_t2 = np.nanmin(np.vstack([t2_st[key],
                                               np.ones_like(t2_st[key])*worst_value_st[key]]),
                                    axis = 0)            
                save_t2 = np.average(save_t2, axis = 0, weights = t2_st[avg_against])
            
        elif key in ['JI', 'state']:                        
            if key == 'JI': avg_against = 'num_traj_CP'
            else: avg_against = 'num_traj'
            
            if 1 not in missing_tracks:
                save_t1 = np.nanmax(np.vstack([t1_st[key],
                                               np.ones_like(t1_st[key])*worst_value_st[key]]),
                                    axis = 0)
                save_t1 = np.average(save_t1, axis = 0, weights = t1_st[avg_against])
                    
            if 2 not in missing_tracks: 
                save_t2 = np.nanmax(np.vstack([t2_st[key],
                                               np.ones_like(t2_st[key])*worst_value_st[key]]),
                                    axis = 0)           
                save_t2 = np.average(save_t2, axis = 0, weights = t2_st[avg_against])

        if 1 not in missing_tracks: 
            output_file.write('T1_st_'+ key +f' : {save_t1}\n')    
        if 2 not in missing_tracks: 
            output_file.write('T2_st_'+ key +f' : {save_t2}\n')

    ### Saving ensemble data
    '''WHAT ARE THE THRESHOLDS FOR THIS?'''
    worst_value_ens = {'alpha': 100,
                       'D': 100}

    for key in t1_ens: 
        if key == 'num_traj': continue
        
        if 1 not in missing_tracks: 
            save_t1 = np.nanmin(np.vstack([t1_ens[key],
                                               np.ones_like(t1_ens[key])*worst_value_ens[key]]),
                                    axis = 0).mean()
        if 2 not in missing_tracks: 
            save_t2 = np.nanmin(np.vstack([t2_ens[key],
                                               np.ones_like(t2_ens[key])*worst_value_ens[key]]),
                                    axis = 0).mean()

        if 1 not in missing_tracks: 
            output_file.write('T1_ens_'+ key +f' : {save_t1}\n')    
        if 2 not in missing_tracks: 
            output_file.write('T2_ens_'+ key +f' : {save_t2}\n')

    output_file.close()
