import os
import glob
import scanpy as sc
import numpy as np
import pandas as pd
from scipy.stats import gaussian_kde
import seaborn as sns
import matplotlib.pyplot as plt
import time
import datetime
import pickle

from scipy.stats import zscore
from sklearn.linear_model import LogisticRegression
from sklearn import metrics
from scipy.stats import mannwhitneyu, tiecorrect, rankdata
from statsmodels.stats.multitest import multipletests

# settings
plt.rc('font', size = 9)
plt.rc('font', family='sans serif')
plt.rcParams['pdf.fonttype']=42
plt.rcParams['ps.fonttype']=42
plt.rcParams['legend.frameon']=False
plt.rcParams['axes.grid']=False
plt.rcParams['legend.markerscale']=0.5
plt.rcParams['savefig.dpi']=600
sns.set_style("ticks")

pfp = '/home/ngr4/project/collabs/grants/czi_rp_2103/results/'
with open('/home/ngr4/project/collabs/grants/czi_rp_2103/data/processed/rpczi.pkl', 'rb') as f:
    temp = pickle.load(f)
    f.close()
adata = temp['adata']

# standard recipe
# sc.pp.combat(adata)
sc.tl.pca(adata)
sc.pp.neighbors(adata, n_neighbors=30, n_pcs=50)
sc.tl.umap(adata)
sc.tl.leiden(adata)
sc.pl.umap(adata, color=['leiden', 'source'])

# define markers & cell types
## TODO (before final polishes): add canonical cell type markers for human lung from Table S1 https://www.biorxiv.org/content/10.1101/742320v2.full.pdf
## REF: https://www.ncbi.nlm.nih.gov/pmc/articles/PMC5135277/ (most lung tissue markers come from here)
cell_markers = {'Basal': ['KRT5', 'DAPL1', 'TP63'],
                'Basal (proliferating)': ['ADH7', 'MKI67', 'TOP2A', 'CDK1'],
                'Hillock': ['KRT4', 'KRT13'],
                'Club': [ 'KRT15', 'CD74','CXCL6'],
                'Ciliated': ['FOXJ1', 'CCDC153', 'CCDC113', 'MLF1', 'LZTFL1','TUBB1','TP73','CCDC78'],
                'Tuft': ['POU2F3', 'AVIL', 'MAFF','MIAT','NOS2'],
                'Ionocyte': ['FOXI1', 'CFTR',], # 'ASCL3' not found
                'Goblet': ['MUC5AC', 'MUC5B', 'SPDEF'],
                'Epithelial':['ABCA3','LPCAT1','NAPSA','SFTPB','SFTPC','SLC34A2'],
                'Neuroendocrine':['ACADSB','ADA','AFAP1','CPE'],
                'Dendritic':['ITGAX','CCR7','CD1A','CD207'], # 'LY75' not found
#                 'Macrophage':['CD68','CD14','CCL18','CD163'],
                'Endothelial':['CD34','PECAM1','VWF'],
                'Fibroblast':['THY1','CD36','PDGFRA','PTPN13'],
                'Tcell':['CD3E','CD3D','CD3G','CD8A','CD8B','CD4'],
                'Granulocyte':['CCR5','SMAD1','ITGAM'],
#                 'Alveolar':['SLC34A2','ABCA3','CD44'],
                'AT1':['SLC34A2','ABCA3','CD44','AGER','PDPN','CLIC5'],
                'AT2':['SLC34A2','ABCA3','CD44','SFTPB','SFTPC','SFTPD','MUC1'],
                'Myofibroblast':['ACTA2'],
                'Monocyte':['CD36','CD14','CD68'],
                'NK':['NCR1'],
                'Progenitor':['TM4SF1','CEACAM6'],
#                 'Neutrophil':['S100A9','S100A8','S100A12','VCAN','FCN1',
#                               'CSTA','TSPO','CD14','MNDA','CTSD','PLBD1'], # from Tianyang (Iwasaki lab) ORIGINAL
                # updated 051820
                'Eosinophil':['RNASE2','LGALS1','RETN','AC020656.1', # 'RNASE3' not found
                              'H1FX','SLC44A1','AL355922.1','RFLNB','SERPINB10'], # from Tianyang (Iwasaki lab) ORIGINAL
#                 'Macrophage':['S100A9','S100A8','FCGR3A','CD14','CD68','FCGR1A','MARCO','MSR1','MRC1','C1QB','C1QA','FABP4','APOC1','APOE','PPARG'],
#                 'Monocyte':['S100A9','S100A8','FCGR3A','CD14','CD68','FCGR1A','RNASE2','RNASE3','FCN1','TNFRSF1B','S100A12','VCAN','CCR2','SDS'],
#                 'Monocyte':['CCR2', 'FCN1', 'RNASE2', 'RNASE3', 'S100A12', 'SDS', 'TNFRSF1B', 'VCAN'], # no overlap btw Macrophage/Monocyte/Neutrophil
                'Monocyte':['CCR2', 'FCN1', 'RNASE2', 'S100A12', 'SDS', 'TNFRSF1B', 'VCAN'],
                'Macrophage':['APOC1', 'APOE', 'C1QA', 'C1QB', 'FABP4', 'MARCO', 'MRC1', 'MSR1', 'PPARG'], # no overlap btw Macrophage/Monocyte/Neutrophil
                'Neutrophil':['CEACAM1', 'CEACAM8', 'CSF3R', 'CXCR1', 'CXCR2', 'FCGR3B'], # no overlap btw Macrophage/Monocyte/Neutrophil
#                 'Neutrophil':['S100A9','S100A8','FCGR3A','CEACAM8','CXCR1','CXCR2','CEACAM1','FCGR3B','CSF3R'],
#                 'Eosinophil':['RNASE2','RNASE3','IL5RA','CCR3','EPX','PRG2','PRG3','PTGDR2','SIGLEC8','GATA2'], # don't use RNASE2/3 since they overlap
#                 'Eosinophil':['IL5RA','CCR3','PRG2','PTGDR2','SIGLEC8','GATA2'], # don't use RNASE2/3 since they overlap
#                 'Eosinophil':['IL5RA','CCR3','PRG2','PTGDR2','SIGLEC8','GATA2', 'EPO','CD9','RNASE3','RETN','H1FX','RFLNB'], # added EPO and CD9 <>
               }

# subset data to markers
genes = [g for k,v in cell_markers.items() for g in v]
x = pd.DataFrame(adata[:,genes].X, columns=genes)
x['cluster'] = adata.obs['leiden'].to_list()
add_pcs = True
if add_pcs:
    # add PCs?
    pcs = ['PC1','PC2']
    for i,pc in enumerate(pcs):
        x[pc] = adata.obsm['X_pca'][:,i]
    genes = genes + pcs

# standard scale
x.loc[:,genes] = zscore(x.loc[:,genes])

results = pd.DataFrame()
fname = 'covid3balfs'

verbose = True
tic = time.time()
counter = 0
ORthreshold = 0.9
total_iter = len(cell_markers.keys())*len(x['cluster'].unique())
new_markers = {}
print('Lasso logistic regression')
for i,ctype in enumerate(cell_markers.keys()):
    for j,cluster in enumerate(x['cluster'].unique()):
        if verbose:
            if counter % 50 == 0 and counter != 0:
                p_through = counter / total_iter
                toc = time.time() - tic
                print('  through {:.1f}-% in {:.2f}-s\t~{:.2f}-s remain'.format(100*p_through,toc,(toc/counter)*(total_iter-counter)))

        # binarize & subset
        y = (x['cluster']==cluster).astype(int)
        if add_pcs:
            X = x.loc[:,cell_markers[ctype]+pcs]
        else:
            X = x.loc[:,cell_markers[ctype]]

        # run default params (could add CV)
        ## results, solver='saga', time for ~25k cells: >>1min
        ## results, solver='lbfgs', time for ~25k cells: 14s
        ## results, solver='liblinear', time for ~25k cells: 25s
        model = LogisticRegression(max_iter=10000,
                                   penalty='l1',
                                   tol=1e-6,
                                   solver='liblinear') #n_jobs=-1 doesn't work for liblinear
        model.fit(X, y)

        status = 'OK'
        if any(np.exp(model.coef_)[0][:-len(pcs)] < ORthreshold):
            markers = [marker for i,marker in enumerate(cell_markers[ctype]) if i not in np.where(np.exp(model.coef_)[0][:-len(pcs)]<0.9)[0]]
            if len(markers) != 0:
                new_markers[ctype] = markers
                if add_pcs:
                    X = x.loc[:,markers+pcs]
                else:
                    X = x.loc[:,markers]
                model = LogisticRegression(max_iter=10000,
                                   penalty='l1',
                                   tol=1e-6,
                                   solver='liblinear') #n_jobs=-1 doesn't work for liblinear
                model.fit(X, y)
            else:
                status = 'No markers with ORs >= {}'.format(ORthreshold)
        else:
            markers = cell_markers[ctype]

        p1 = model.predict_proba(X)[:,1]
        fpr, tpr, thresholds = metrics.roc_curve(y, p1)
        optimal_idx = np.argmax(tpr-fpr)
        optimal_threshold = thresholds[optimal_idx]
        optimal_pred = (p1>optimal_threshold).astype(int)
        precision,recall,_ = metrics.precision_recall_curve(y, p1)
        auprc = metrics.auc(recall, precision)
        auroc = metrics.roc_auc_score(y,p1)
        ap = metrics.average_precision_score(y,p1)
        bs = metrics.brier_score_loss(y,p1)
        acc = metrics.accuracy_score(y,optimal_pred)

        # store results
        dt = pd.DataFrame({'ctype2pred':ctype,
                           'cluster':cluster,
                           'auroc':auroc,
                           'status':status,
                           'markers':[markers],
                           'ORs':np.exp(model.coef_).tolist(),
                           'ave_prec':ap,
                           'acc':acc,
                           'sensitivity':tpr[optimal_idx],
                           'specificity':1-fpr[optimal_idx]},
                          index=[0])
        results = results.append(dt, ignore_index=True)
        counter += 1
print('Classifiers done. Saving and plotting...')

top_per_ctype = pd.DataFrame()
top_n = 3
for ctype in results['ctype2pred'].unique():
    dt = results.loc[results['ctype2pred']==ctype,:]
    dt = dt.sort_values(by='auroc', ascending=False)
    top_per_ctype = top_per_ctype.append(dt.iloc[0:top_n,:], ignore_index=True)

top_per_cluster = pd.DataFrame()
top_n = 3
for cluster in results['cluster'].unique():
    dt = results.loc[results['cluster']==cluster,:]
    dt = dt.sort_values(by='auroc', ascending=False)
    top_per_cluster = top_per_cluster.append(dt.iloc[0:top_n,:], ignore_index=True)

if True:
    top_per_cluster.to_csv(os.path.join(pfp,'top_ctype_per_cluster_{}.csv'.format(fname)))

# plot init annotation
## taking top ctype per cluster
top1_per_cluster = pd.DataFrame()
for cluster in results['cluster'].unique():
    dt = results.loc[results['cluster']==cluster,:]
    dt = dt.sort_values(by='auroc', ascending=False)
    if True:
        # eliminate rows with poor status (no markers with OR>=threshold)
        dt = dt.loc[dt['status'] == 'OK',:]
    if dt.shape[0]==0:
        print('Cluster {} could not be annotated due to ORs of markers.'.format(cluster))
        continue
    top1_per_cluster = top1_per_cluster.append(dt.iloc[0:1,:], ignore_index=True)
ctype_annotation = {}
for cluster in top1_per_cluster['cluster']:
    ctype_annotation[cluster] = top1_per_cluster.loc[top1_per_cluster['cluster']==cluster,'ctype2pred'].values[0]
adata.obs['init_ctype'] = adata.obs['leiden'].astype(str)
adata.obs['init_ctype'] = adata.obs['init_ctype'].map(ctype_annotation)

## aesthetics
pal18=['#ee5264','#565656','#75a3b7','#ffe79e','#fac18a','#f1815f','#ac5861','#62354f','#2d284b','#f4b9b9','#c4bbaf',
               '#f9ebae','#aecef9','#aeb7f9','#f9aeae','#9c9583','#88bb92','#bde4a7','#d6e5e3']
cmap_ctype = {v:pal18[i] for i,v in enumerate(adata.obs['init_ctype'].unique())}

## plot
sc.pl.umap(adata, color=['leiden', 'init_ctype'])

if verbose:
    # print bad ones
    print('\nClusters hard to identify')
    print('-------------------------')

    underdetermined = top1_per_cluster.loc[top1_per_cluster['auroc'] <= 0.7, :]
    for i in range(underdetermined.shape[0]):
        print(underdetermined.iloc[i,:])

# save metadata
adata.obs.to_csv(os.path.join(pfp, 'metadata_{}'.format(datetime.datetime.now().strftime('%y%m%d'))))

import numpy,from scipy,from statsmodels


import numpy as np
import tensorflow as tf
from tensorflow.keras import layers as tfl

from common import utils
from layers import losses
from layers import mrcnn_layers as mrcnnl

class MaskRCNN(tf.keras.Model):

    def __init__(self, config: dict, *args, **kwargs):
        """
        MaskRCNN base class
        Args:
            config:    str, model config
            *args:
            **kwargs:
        """
        super().__init__(*args, **kwargs)
        self.config = config
        self.losses_dict = {}
        self.losses_tracker_dict = {}

    @property
    def metrics(self) -> list:
        # We list our `Metric` objects here so that `reset_states()` can be
        # called automatically at the start of each epoch or at the start of `evaluate()`.
        # If you don't implement this property, you have to call `reset_states()` yourself at the time of your choosing.
        return [x for x in self.losses_tracker_dict.values()]

    def compile(self, optimizer, losses_dict: dict = None, run_eagerly: bool = True, *args, **kwargs):
        super().compile(optimizer=optimizer, run_eagerly=run_eagerly, *args, **kwargs)
        # Update a dict of losses
        if losses_dict:
            self.losses_dict.update(losses_dict)
            # Add l2 regularization to losses
            self.losses_dict.update({'l2_regularizer': losses.L2RegLoss(model=self, config=self.config)})
            # Add losses tracker
            for loss_name in self.losses_dict.keys():
                self.losses_tracker_dict.update({loss_name: tf.keras.metrics.Mean(name=loss_name)})
            self.losses_tracker_dict.update({'loss_sum': tf.keras.metrics.Mean(name='loss_sum')})
        print(f'\nMaskRCNN Losses:\n' + ''.join([f'{name}: {loss}\n' for name, loss in self.losses_dict.items()]))

    def get_input(self, inputs: list) -> list:
        """
        Wrapper for input parsing
        Args:
            inputs:  batch_images,          (batch, w, h, 3)
                     batch_images_meta,     (batch, meta_shape)
                     batch_rpn_match,       (batch, anchors, 1)
                     batch_rpn_bbox,        (batch, rpn_train_anchors_per_image, 4)
                     batch_gt_class_ids,    (batch, max_gt_instances)
                     batch_gt_boxes,        (batch, max_gt_instances, 4)
                     batch_gt_masks,        (batch, w, h, max_gt_instances)
                     random_rois, optional  (batch, n_of_rois, 4), optional
        Returns: a list of prepared data
        """
        if self.config['use_rpn_rois']:
            return [*inputs, None]
        else:
            return inputs

    @staticmethod
    def check_loss_nan(name, value, assert_nans, dummy_value=0.0):
        """
        Check nan in loss
        Args:
            name:         loss name
            value:        loss value
            assert_nans:  raise assertion or not
            dummy_value:  value to be set instead of nan if `assert_nans` is False

        Returns: loss value

        """
        if np.any(np.isnan(value)):
            msg = f'Warning. Nan loss was found: {name}'
            if assert_nans:
                raise AssertionError(msg)
            else:
                print(f'{msg}. Changing it to {dummy_value}')
                value = tf.constant(dummy_value, dtype=tf.float32)
        return value

    def get_summary_loss(self, rpn_class_loss, rpn_bbox_loss, mrcnn_class_loss,
                         mrcnn_bbox_loss, mrcnn_mask_loss, l2_reg_loss, assert_nans=True):
        """
        Calculate summary weighted loss with weights stored in self.config['loss_weights']
        Args:
            rpn_class_loss:   How well the Region Proposal Network separates background with objects
            rpn_bbox_loss:    How well the RPN localize objects
            mrcnn_class_loss: How well the Mask RCNN localize objects
            mrcnn_bbox_loss:  How well the Mask RCNN recognize each class of object
            mrcnn_mask_loss:  How well the Mask RCNN segment objects
            l2_reg_loss:      L2 weights regularization
            assert_nans:      Raise an assertion if nan in loss happens

        Returns: summary loss
        """
        rpn_class_loss = self.check_loss_nan('rpn_class_loss', rpn_class_loss, assert_nans)
        rpn_bbox_loss = self.check_loss_nan('rpn_bbox_loss', rpn_bbox_loss, assert_nans)
        mrcnn_class_loss = self.check_loss_nan('mrcnn_class_loss', mrcnn_class_loss, assert_nans)
        mrcnn_bbox_loss = self.check_loss_nan('mrcnn_bbox_loss', mrcnn_bbox_loss, assert_nans)
        mrcnn_mask_loss = self.check_loss_nan('mrcnn_mask_loss', mrcnn_mask_loss, assert_nans)

        if not self.config['use_rpn_rois']:
            # Remove rpn losses if rpn_rois is not used
            loss_tuple = (mrcnn_class_loss, mrcnn_bbox_loss, mrcnn_mask_loss)
            loss_weights = np.array(self.config['loss_weights'])[[2, 3, 4]].tolist()
        elif self.config['tune_rpn_model_only'] and self.config['use_rpn_rois']:
            loss_tuple = (rpn_class_loss, rpn_bbox_loss)
            loss_weights = np.array(self.config['loss_weights'])[[0, 1]].tolist()
        else:
            loss_tuple = (rpn_class_loss, rpn_bbox_loss, mrcnn_class_loss, mrcnn_bbox_loss, mrcnn_mask_loss)
            loss_weights = self.config['loss_weights']

        loss = tf.tensordot(a=loss_tuple, b=tf.cast(loss_weights, 'float32'), axes=1)

        if l2_reg_loss:
            assert not np.any(np.isnan(l2_reg_loss))
            loss = tf.math.add(loss, l2_reg_loss)

        return loss

    def update_loss_tracker(self, rpn_class_loss_val, rpn_bbox_loss_val, mrcnn_class_loss_val,
                            mrcnn_bbox_loss_val, mrcnn_mask_loss_val, l2_reg_loss_val, loss) -> None:
        for loss_name, loss_value in zip(('rpn_class_loss', 'rpn_bbox_loss', 'mrcnn_class_loss', 'mrcnn_bbox_loss',
                                          'mrcnn_mask_loss', 'l2_regularizer', 'loss_sum'
                                          ),
                                         (rpn_class_loss_val, rpn_bbox_loss_val, mrcnn_class_loss_val,
                                          mrcnn_bbox_loss_val, mrcnn_mask_loss_val, l2_reg_loss_val, loss)
                                         ):
            if loss_value:
                self.losses_tracker_dict[loss_name].update_state(loss_value)

    def train_step(self, inputs: list) -> dict:

        batch_images, batch_images_meta, batch_rpn_match, batch_rpn_bbox, \
        batch_gt_class_ids, batch_gt_boxes, batch_gt_masks, random_rois = self.get_input(inputs)

        with tf.GradientTape() as tape:
            outputs = self(inputs)
            rpn_class_logits, rpn_class, rpn_bbox, mrcnn_class_logits, mrcnn_probs, mrcnn_bbox, mrcnn_mask, \
            rpn_rois, active_class_ids, output_rois, target_class_ids, target_mask, target_bbox = outputs

            rpn_class_loss_val = self.losses_dict['rpn_class_loss'].call(rpn_match=batch_rpn_match,
                                                                         rpn_class_logits=rpn_class_logits)
            rpn_bbox_loss_val = self.losses_dict['rpn_bbox_loss'].call(target_bbox=batch_rpn_bbox,
                                                                       rpn_match=batch_rpn_match,
                                                                       rpn_bbox=rpn_bbox)
            mrcnn_class_loss_val = self.losses_dict['mrcnn_class_loss'].call(target_class_ids=target_class_ids,
                                                                             pred_class_logits=mrcnn_class_logits,
                                                                             active_class_ids=active_class_ids)
            mrcnn_bbox_loss_val = self.losses_dict['mrcnn_bbox_loss'].call(target_bbox=target_bbox,
                                                                           target_class_ids=target_class_ids,
                                                                           pred_bbox=mrcnn_bbox)
            mrcnn_mask_loss_val = self.losses_dict['mrcnn_mask_loss'].call(target_masks=target_mask,
                                                                           target_class_ids=target_class_ids,
                                                                           pred_masks=mrcnn_mask)
            l2_reg_loss_val = self.losses_dict['l2_regularizer'].call()

            loss = self.get_summary_loss(rpn_class_loss_val, rpn_bbox_loss_val, mrcnn_class_loss_val,
                                         mrcnn_bbox_loss_val, mrcnn_mask_loss_val, l2_reg_loss_val)

        # Use the gradient tape to automatically retrieve the gradients of the trainable variables \
        # with respect to the loss.
        grads = tape.gradient(loss, self.trainable_weights)

        # Run one step of gradient descent by updating the value of the variables to minimize the loss.
        # If some grads do not exist< skip them
        # Check if is needed, None grads usually are for gamma, beta BatchNorm if l2 regularization skips them
        # optimizer.apply_gradients(zip(grads, model.trainable_weights))
        self.optimizer.apply_gradients(
            (grad, var) for (grad, var) in zip(grads, self.trainable_variables) if grad is not None)

        self.update_loss_tracker(rpn_class_loss_val, rpn_bbox_loss_val, mrcnn_class_loss_val,
                                 mrcnn_bbox_loss_val, mrcnn_mask_loss_val, l2_reg_loss_val, loss)

        return {m.name: m.result() for m in self.metrics}

    def test_step(self, inputs: list) -> dict:

        batch_images, batch_images_meta, batch_rpn_match, batch_rpn_bbox, \
        batch_gt_class_ids, batch_gt_boxes, batch_gt_masks, random_rois = self.get_input(inputs)

        outputs = self(inputs)
        rpn_class_logits, rpn_class, rpn_bbox, mrcnn_class_logits, mrcnn_probs, mrcnn_bbox, mrcnn_mask, \
        rpn_rois, active_class_ids, output_rois, target_class_ids, target_mask, target_bbox = outputs

        rpn_class_loss_val = self.losses_dict['rpn_class_loss'].call(rpn_match=batch_rpn_match,
                                                                     rpn_class_logits=rpn_class_logits)

        rpn_bbox_loss_val = self.losses_dict['rpn_bbox_loss'].call(target_bbox=batch_rpn_bbox,
                                                                   rpn_match=batch_rpn_match,
                                                                   rpn_bbox=rpn_bbox)

        mrcnn_class_loss_val = self.losses_dict['mrcnn_class_loss'].call(target_class_ids=target_class_ids,
                                                                         pred_class_logits=mrcnn_class_logits,
                                                                         active_class_ids=active_class_ids)

        mrcnn_bbox_loss_val = self.losses_dict['mrcnn_bbox_loss'].call(target_bbox=target_bbox,
                                                                       target_class_ids=target_class_ids,
                                                                       pred_bbox=mrcnn_bbox)

        mrcnn_mask_loss_val = self.losses_dict['mrcnn_mask_loss'].call(target_masks=target_mask,
                                                                       target_class_ids=target_class_ids,
                                                                       pred_masks=mrcnn_mask)

        # Calculate summary loss with weights stored in self.config['loss_weights']
        loss = self.get_summary_loss(rpn_class_loss_val, rpn_bbox_loss_val, mrcnn_class_loss_val,
                                     mrcnn_bbox_loss_val, mrcnn_mask_loss_val, None)

        self.update_loss_tracker(rpn_class_loss_val, rpn_bbox_loss_val, mrcnn_class_loss_val,
                                 mrcnn_bbox_loss_val, mrcnn_mask_loss_val, None, loss)

        return {m.name: m.result() for m in self.metrics}

class FMaskRCNN(MaskRCNN):
    def __init__(self, inputs: list, outputs: list, config: dict, name: str = 'mask_rcnn', *args, **kwargs):
        """
        A wrapper for functional tf.keras API
        Args:
            inputs:           Data inputs (x), list
            outputs           Data outputs (y), list
            config:           General MaskRCNN config, dict
            name:             Model name, string
            *args:
            **kwargs:
        """
        super(MaskRCNN, self).__init__(inputs=inputs, outputs=outputs, name=name, *args, **kwargs)
        self.config = config
        self.training = config['training']
        self.losses_dict = {}
        self.losses_tracker_dict = {}

class SMaskRCNN(MaskRCNN):

    def __init__(self, config: dict, name: str = 'mask_rcnn', *args, **kwargs):
        """
        A wrapper for subclassed tf.keras API
        Such complex subclassed model can not be saved in HDF5 for now, but is useful for research and debug.
        Args:
            config:   General MaskRCNN config, dict
            name:     Model name, string
            *args:
            **kwargs:
        """
        super(MaskRCNN, self).__init__(name=name, *args, **kwargs)
        self.config = config
        self.training = config['training']
        self.losses_dict = {}
        self.losses_tracker_dict = {}

        # Backbone model for features extraction
        self.input_image_wh = tf.constant(self.config['image_shape'][:2], dtype=tf.float32)
        self.backbone = mrcnnl.MaskRCNNBackbone(config=config).build()

        # Upsampling layer
        self.upsampling_layer = mrcnnl.UpSamplingLayer(config=self.config)

        # Anchors
        self.anchors_layer = mrcnnl.AnchorsLayer(config=config, training=self.training)
        self.anchors = self.anchors_layer(None)

        # RPN model
        self.rpn_layer = mrcnnl.RPNLayer(anchors_per_location=len(config['rpn_anchor_ratios']),
                                         anchor_stride=config['rpn_anchor_stride'])

        if self.training:
            # Generate proposals
            # Proposals are [batch, N, (y1, x1, y2, x2)] in normalized coordinates and zero padded.
            self.proposal_count = config['post_nms_rois_training']
            # Class ID mask to mark class IDs supported by the dataset the image came from.
            self.parse_img_meta_layer = mrcnnl.ImageMetaLayer()
            # Generate detection targets
            # Subsamples proposals and generates target outputs for training
            # Note that proposal class IDs, gt_boxes, and gt_masks are zero
            # padded. Equally, returned rois and targets are zero padded.
            self.detection_target_layer = mrcnnl.DetectionTargetLayer(self.config)
            # Coordinates normalization layer
            self.norm_boxes_layer = mrcnnl.NormBoxesLayer()
        else:
            self.proposal_count = config['post_nms_rois_inference']
            self.detections_layer = mrcnnl.DetectionLayer(proposals=self.proposal_count,
                                                          batch_size=config['batch_size'],
                                                          images_per_gpu=config['images_per_gpu'],
                                                          detection_min_confidence=config['detection_min_confidence'],
                                                          detection_max_instances=config['detection_max_instances'],
                                                          detection_nms_threshold=config['detection_nms_threshold'],
                                                          bbox_std_dev=config['bbox_std_dev'], )
            self.detected_boxes_extraction = mrcnnl.DetectedBoxesExtraction()

        # Proposals layer
        self.proposal_layer = mrcnnl.ProposalLayer(proposal_count=self.proposal_count, config=config)

        # FPN classifier layer
        self.fpn_classifier = \
            mrcnnl.FPNClassifier(pool_size=self.config['pool_size'],
                                 fc_layers_size=self.config['fpn_cls_fc_layers_size'],
                                 num_classes=self.config['num_classes'],
                                 train_bn=self.config['train_bn'])

        # FPN mask layer
        self.fpn_mask_layer = \
            mrcnnl.FPNMaskLayer(pool_size=self.config['mask_pool_size'],
                                num_classes=self.config['num_classes'],
                                train_bn=self.config['train_bn'])

        # Build subclassed model
        self.call([np.random.uniform(size=x) for x in [(self.config['batch_size'], *self.config['image_shape']),
                                                       (self.config['batch_size'], *self.config['meta_shape']),
                                                       (self.config['batch_size'], self.anchors.shape[0], 1),
                                                       (self.config['batch_size'],
                                                        self.config['rpn_train_anchors_per_image'], 4),
                                                       (self.config['batch_size'], self.config['max_gt_instances']),
                                                       (self.config['batch_size'], self.config['max_gt_instances'], 4),
                                                       (self.config['batch_size'],
                                                        *self.config['image_shape'][:2],
                                                        self.config['max_gt_instances']),
                                                       ]
                   ]
                  )

    def call(self, inputs: list, **kwargs) -> tuple:

        if self.training:
            # Training input
            batch_images, batch_images_meta, batch_rpn_match, batch_rpn_bbox, \
            batch_gt_class_ids, batch_gt_boxes, batch_gt_masks, random_rois = self.get_input(inputs)

        else:
            # Inference input
            # Input for inference requires preprocessing
            batch_images, batch_images_meta = inputs

        # Extract c2, c3, c4, c5 features
        _, c2, c3, c4, c5 = self.backbone(batch_images)

        # rpn_feature_maps = [p2, p3, p4, p5, p6]
        # mrcnn_feature_maps = [p2, p3, p4, p5]
        rpn_feature_maps, mrcnn_feature_maps = self.upsampling_layer([c2, c3, c4, c5])

        # Loop through pyramid layers
        layer_outputs = []  # A list of lists
        for p in rpn_feature_maps:
            layer_outputs.append(self.rpn_layer(p))

        # Concatenate layer outputs
        # Convert from list of lists of level outputs to list of lists
        # of outputs across levels.
        # e.g. [[a1, b1, c1], [a2, b2, c2]] => [[a1, a2], [b1, b2], [c1, c2]]
        output_names = ["rpn_class_logits", "rpn_class", "rpn_bbox"]
        outputs = list(zip(*layer_outputs))
        outputs = [tfl.Concatenate(axis=1, name=n)(list(o))
                   for o, n in zip(outputs, output_names)]
        rpn_class_logits, rpn_class, rpn_bbox = outputs

        rpn_rois = self.proposal_layer([rpn_class, rpn_bbox, self.anchors])

        if self.training:
            batch_gt_boxes = self.norm_boxes_layer((batch_gt_boxes, self.input_image_wh))

            active_class_ids = self.parse_img_meta_layer(batch_images_meta)['active_class_ids']

            if self.config['use_rpn_rois']:
                target_rois = rpn_rois
            else:
                target_rois = self.norm_boxes_layer((random_rois, self.input_image_wh))

            rois, target_class_ids, target_bbox, target_mask = \
                self.detection_target_layer([target_rois, batch_gt_class_ids, batch_gt_boxes, batch_gt_masks])

            # Network heads
            mrcnn_class_logits, mrcnn_class, mrcnn_bbox = \
                self.fpn_classifier([rois, batch_images_meta, mrcnn_feature_maps])
            mrcnn_mask = self.fpn_mask_layer([rois, batch_images_meta, mrcnn_feature_maps])
            output_rois = tf.identity(rois)

            return rpn_class_logits, rpn_class, rpn_bbox, mrcnn_class_logits, mrcnn_class, mrcnn_bbox, mrcnn_mask, \
                   rpn_rois, active_class_ids, output_rois, target_class_ids, target_mask, target_bbox

        else:
            # Network heads
            mrcnn_class_logits, mrcnn_class, mrcnn_bbox = \
                self.fpn_classifier([rpn_rois, batch_images_meta, mrcnn_feature_maps])
            # Detections output is [batch, num_detections, (y1, x1, y2, x2, class_id, score)] in normalized coordinates
            detections = self.detections_layer([rpn_rois, mrcnn_class, mrcnn_bbox, batch_images_meta])
            # Create masks for detections
            detection_boxes = self.detected_boxes_extraction(detections)
            mrcnn_mask = self.fpn_mask_layer([detection_boxes, batch_images_meta, mrcnn_feature_maps])

            return detections, mrcnn_class, mrcnn_bbox, mrcnn_mask, rpn_rois, rpn_class, rpn_bbox

def mask_rcnn_functional(config: dict) -> tf.keras.Model:
    """
    Construct a model in functional API
    Args:
        config: General MaskRCNN config, dict
    Returns: tf.keras.Model:

    """
    # Prevent creating keras names with index increment.
    # It is important for weights setting from training to inference graph
    tf.compat.v1.reset_default_graph()
    tf.keras.backend.clear_session()

    # Inputs
    input_image = tfl.Input(shape=config['image_shape'], name="input_image")
    input_image_meta = tfl.Input(shape=config['meta_shape'], name="input_image_meta")
    input_image_wh = tf.constant(config['image_shape'][:2], dtype=tf.float32)

    if config['training']:
        print('[MaskRCNN] Training mode')
        # Regions Proposal Network Ground truth inputs
        input_rpn_match = tfl.Input(shape=[None, 1], name="input_rpn_match", dtype=tf.int32)
        input_rpn_bbox = tfl.Input(shape=[None, 4], name="input_rpn_bbox", dtype=tf.float32)

        # Detection GT (class IDs, bounding boxes, and masks)
        # 1. GT Class IDs (zero padded)
        input_gt_class_ids = tfl.Input(shape=[None], name="input_gt_class_ids", dtype=tf.int32)
        # 2. GT Boxes in pixels (zero padded)
        # [batch, MAX_GT_INSTANCES, (y1, x1, y2, x2)] in image coordinates
        input_gt_boxes = tfl.Input(shape=[None, 4], name="input_gt_boxes", dtype=tf.float32)
        # Normalize coordinates
        gt_boxes = mrcnnl.NormBoxesLayer()((input_gt_boxes, input_image_wh))
        # 3. GT Masks (zero padded)
        # [batch, height, width, MAX_GT_INSTANCES]
        if config['use_mini_masks']:
            input_gt_masks = tfl.Input(shape=[*config['mini_mask_shape'], None],
                                       name="input_gt_masks", dtype=tf.bool)
        else:
            input_gt_masks = tfl.Input(shape=[config['image_shape'][0], config['image_shape'][1], None],
                                       name="input_gt_masks", dtype=tf.bool)

        anchors = mrcnnl.AnchorsLayer(config=config, training=config['training'])(input_image)
        proposal_count = config['post_nms_rois_training']

    else:
        print('\n[MaskRCNN] Inference mode')
        # Anchors in normalized coordinates
        # input_anchors = tfl.Input(shape=[None, 4], name="input_anchors")
        anchors = mrcnnl.AnchorsLayer(config=config, training=config['training'])(input_image)  # input_anchors
        proposal_count = config['post_nms_rois_inference']

    print(f"""\n[MaskRCNN] Backbone architecture: {config['backbone']}""")
    backbone = mrcnnl.MaskRCNNBackbone(config=config).build()
    _, c2, c3, c4, c5 = backbone(input_image)

    # rpn_feature_maps =   [p2, p3, p4, p5, p6]
    # mrcnn_feature_maps = [p2, p3, p4, p5]
    rpn_feature_maps, mrcnn_feature_maps = mrcnnl.upsampling_graph(inputs=[c2, c3, c4, c5], config=config)

    # RPN Model
    rpn_model = mrcnnl.build_rpn_model(anchor_stride=config['rpn_anchor_stride'],
                                       anchors_per_location=len(config['rpn_anchor_ratios']),
                                       depth=config['top_down_pyramid_size'],
                                       training=config['training'],
                                       frozen=config['frozen_rpn_model']
                                       )

    # Loop through pyramid layers
    layer_outputs = []
    for p in rpn_feature_maps:
        layer_outputs.append(rpn_model([p]))

    # Concatenate layer outputs
    # Convert from list of lists of level outputs to list of lists
    # of outputs across levels.
    # e.g. [[a1, b1, c1], [a2, b2, c2]] => [[a1, a2], [b1, b2], [c1, c2]]
    output_names = ["rpn_class_logits", "rpn_class", "rpn_bbox"]
    outputs = list(zip(*layer_outputs))
    outputs = [tfl.Concatenate(axis=1, name=f'concat_{n}')(list(o))
               for o, n in zip(outputs, output_names)]
    rpn_class_logits, rpn_class, rpn_bbox = outputs

    # Proposal layer
    rpn_rois = mrcnnl.ProposalLayer(proposal_count=proposal_count, config=config)([rpn_class, rpn_bbox, anchors])

    if config['training']:
        # Class ID mask to mark class IDs supported by the dataset the image came from.
        active_class_ids = tfl.Lambda(lambda x: utils.parse_image_meta_graph(x)["active_class_ids"],
                                      name='parse_image_meta_graph')(input_image_meta)

        if config['use_rpn_rois']:
            target_rois = rpn_rois
        else:
            if config['post_nms_rois_training'] != config['random_rois']:
                raise ValueError('Set post_nms_rois_training == random_rois' +
                                 f"""for use_rpn_rois={config['use_rpn_rois']} option""")
            # Ignore predicted ROIs and use ROIs provided as an input.
            input_rois = tfl.Input(shape=[config['post_nms_rois_training'], 4], name="input_roi", dtype=np.float32)
            # Normalize coordinates
            target_rois = mrcnnl.NormBoxesLayer(name='norm_boxes_layer')((input_rois, input_image_wh))
        # Generate detection targets
        # Subsamples proposals and generates target outputs for training
        # Note that proposal class IDs, gt_boxes, and gt_masks are zero
        # padded. Equally, returned rois and targets are zero padded.
        rois, target_class_ids, target_bbox, target_mask = \
            mrcnnl.DetectionTargetLayer(config, name="detection_targets_layer")(
                [target_rois, input_gt_class_ids, gt_boxes, input_gt_masks])

        # Network Heads
        mrcnn_class_logits, mrcnn_probs, mrcnn_bbox = \
            mrcnnl.fpn_classifier_graph(inputs=[rois, input_image_meta, mrcnn_feature_maps],
                                        pool_size=config['pool_size'],
                                        fc_layers_size=config['fpn_cls_fc_layers_size'],
                                        num_classes=config['num_classes'],
                                        train_bn=config['train_bn'],
                                        batch_size=config['batch_size'],
                                        post_nms_rois_inference=config['post_nms_rois_inference'],
                                        training=config['training'],
                                        frozen=config['frozen_cls_head'],
                                        leaky_relu=config['cls_head_leaky_relu']
                                        )
        mrcnn_mask = mrcnnl.fpn_mask_graph(inputs=[rois, input_image_meta, mrcnn_feature_maps],
                                           pool_size=config['mask_pool_size'],
                                           num_classes=config['num_classes'],
                                           train_bn=config['train_bn'],
                                           frozen=config['frozen_mask_head'],
                                           leaky_relu=config['mask_head_leaky_relu']
                                           )

        # Model training inputs
        inputs = [input_image, input_image_meta, input_rpn_match, input_rpn_bbox,
                  input_gt_class_ids, input_gt_boxes, input_gt_masks]
        if not config['use_rpn_rois']:
            inputs.append(input_rois)

        # Model training outputs
        outputs = [rpn_class_logits, rpn_class, rpn_bbox,
                   mrcnn_class_logits, mrcnn_probs, mrcnn_bbox, mrcnn_mask,
                   rpn_rois, active_class_ids, rois, target_class_ids, target_mask, target_bbox]
        name = 'mask_rcnn_training'

    else:
        # Network Heads
        # Proposal classifier and BBox regressor heads
        mrcnn_class_logits, mrcnn_probs, mrcnn_bbox = \
            mrcnnl.fpn_classifier_graph(inputs=[rpn_rois, input_image_meta, mrcnn_feature_maps],
                                        pool_size=config['pool_size'],
                                        fc_layers_size=config['fpn_cls_fc_layers_size'],
                                        num_classes=config['num_classes'],
                                        train_bn=config['train_bn'],
                                        batch_size=config['batch_size'],
                                        post_nms_rois_inference=config['post_nms_rois_inference'],
                                        training=config['training'],
                                        frozen=False,
                                        leaky_relu=config['cls_head_leaky_relu']
                                        )

        # Detections output is [batch, num_detections, (y1, x1, y2, x2, class_id, score)] in normalized coordinates
        detections = mrcnnl.DetectionLayer(proposals=proposal_count,
                                           batch_size=config['batch_size'],
                                           images_per_gpu=config['images_per_gpu'],
                                           detection_min_confidence=config['detection_min_confidence'],
                                           detection_max_instances=config['detection_max_instances'],
                                           detection_nms_threshold=config['detection_nms_threshold'],
                                           bbox_std_dev=config['bbox_std_dev'],
                                           )(
            [rpn_rois, mrcnn_probs, mrcnn_bbox, input_image_meta])
        # Create masks for detections
        detection_boxes = mrcnnl.DetectedBoxesExtraction()(detections)
        mrcnn_mask = mrcnnl.fpn_mask_graph(inputs=[detection_boxes, input_image_meta, mrcnn_feature_maps],
                                           pool_size=config['mask_pool_size'],
                                           num_classes=config['num_classes'],
                                           train_bn=config['train_bn'],
                                           frozen=False,
                                           leaky_relu=config['mask_head_leaky_relu']
                                           )

        # Model inference inputs and outputs
        inputs = [input_image, input_image_meta]
        outputs = [detections, mrcnn_probs, mrcnn_bbox, mrcnn_mask, rpn_rois, rpn_class, rpn_bbox]
        name = 'mask_rcnn_inference'

    # Make model graph in a customized tf.keras.Model with overridden train and test steps
    model = FMaskRCNN(inputs=inputs, outputs=outputs, config=config, name=name)
    trainable_p = np.sum([np.prod(v.get_shape()) for v in model.trainable_weights])
    non_trainable_p = np.sum([np.prod(v.get_shape()) for v in model.non_trainable_weights])
    print('\n[MaskRCNN] Total params: {:,}'.format(trainable_p + non_trainable_p))
    print('\n[MaskRCNN] Trainable params: {:,}'.format(trainable_p))
    return model


import torch

from run_on_video.data_utils import ClipFeatureExtractor
from run_on_video.model_utils import build_inference_model
from utils.tensor_utils import pad_sequences_1d
from moment_detr.span_utils import span_cxw_to_xx
from utils.basic_utils import l2_normalize_np_array
import torch.nn.functional as F
import numpy as np

class MomentDETRPredictor:
    def __init__(self, ckpt_path, clip_model_name_or_path="ViT-B/32", device="cuda"):
        self.clip_len = 2  # seconds
        self.device = device
        print("Loading feature extractors...")
        self.feature_extractor = ClipFeatureExtractor(
            framerate=1/self.clip_len, size=224, centercrop=True,
            model_name_or_path=clip_model_name_or_path, device=device
        )
        print("Loading trained Moment-DETR model...")
        self.model = build_inference_model(ckpt_path).to(self.device)

    @torch.no_grad()
    def localize_moment(self, video_path, query_list):
        """
        Args:
            video_path: str, path to the video file
            query_list: List[str], each str is a query for this video
        """
        # construct model inputs
        n_query = len(query_list)
        video_feats = self.feature_extractor.encode_video(video_path)
        video_feats = F.normalize(video_feats, dim=-1, eps=1e-5)
        n_frames = len(video_feats)
        # add tef
        tef_st = torch.arange(0, n_frames, 1.0) / n_frames
        tef_ed = tef_st + 1.0 / n_frames
        tef = torch.stack([tef_st, tef_ed], dim=1).to(self.device)  # (n_frames, 2)
        video_feats = torch.cat([video_feats, tef], dim=1)
        assert n_frames <= 75, "The positional embedding of this pretrained MomentDETR only support video up " \
                               "to 150 secs (i.e., 75 2-sec clips) in length"
        video_feats = video_feats.unsqueeze(0).repeat(n_query, 1, 1)  # (#text, T, d)
        video_mask = torch.ones(n_query, n_frames).to(self.device)
        query_feats = self.feature_extractor.encode_text(query_list)  # #text * (L, d)
        query_feats, query_mask = pad_sequences_1d(
            query_feats, dtype=torch.float32, device=self.device, fixed_length=None)
        query_feats = F.normalize(query_feats, dim=-1, eps=1e-5)
        model_inputs = dict(
            src_vid=video_feats,
            src_vid_mask=video_mask,
            src_txt=query_feats,
            src_txt_mask=query_mask
        )

        # decode outputs
        outputs = self.model(**model_inputs)
        # #moment_queries refers to the positional embeddings in MomentDETR's decoder, not the input text query
        prob = F.softmax(outputs["pred_logits"], -1)  # (batch_size, #moment_queries=10, #classes=2)
        scores = prob[..., 0]  # * (batch_size, #moment_queries)  foreground label is 0, we directly take it
        pred_spans = outputs["pred_spans"]  # (bsz, #moment_queries, 2)
        _saliency_scores = outputs["saliency_scores"].half()  # (bsz, L)
        saliency_scores = []
        valid_vid_lengths = model_inputs["src_vid_mask"].sum(1).cpu().tolist()
        for j in range(len(valid_vid_lengths)):
            _score = _saliency_scores[j, :int(valid_vid_lengths[j])].tolist()
            _score = [round(e, 4) for e in _score]
            saliency_scores.append(_score)

        # compose predictions
        predictions = []
        video_duration = n_frames * self.clip_len
        for idx, (spans, score) in enumerate(zip(pred_spans.cpu(), scores.cpu())):
            spans = span_cxw_to_xx(spans) * video_duration
            # # (#queries, 3), [st(float), ed(float), score(float)]
            cur_ranked_preds = torch.cat([spans, score[:, None]], dim=1).tolist()
            cur_ranked_preds = sorted(cur_ranked_preds, key=lambda x: x[2], reverse=True)
            cur_ranked_preds = [[float(f"{e:.4f}") for e in row] for row in cur_ranked_preds]
            cur_query_pred = dict(
                query=query_list[idx],  # str
                vid=video_path,
                pred_relevant_windows=cur_ranked_preds,  # List([st(float), ed(float), score(float)])
                pred_saliency_scores=saliency_scores[idx]  # List(float), len==n_frames, scores for each frame
            )
            predictions.append(cur_query_pred)

        return predictions

def run_example():
    # load example data
    from utils.basic_utils import load_jsonl
    video_path = "run_on_video/example/RoripwjYFp8_60.0_210.0.mp4"
    query_path = "run_on_video/example/queries.jsonl"
    queries = load_jsonl(query_path)
    query_text_list = [e["query"] for e in queries]
    ckpt_path = "run_on_video/moment_detr_ckpt/model_best.ckpt"

    # run predictions
    print("Build models...")
    clip_model_name_or_path = "ViT-B/32"
    # clip_model_name_or_path = "tmp/ViT-B-32.pt"
    moment_detr_predictor = MomentDETRPredictor(
        ckpt_path=ckpt_path,
        clip_model_name_or_path=clip_model_name_or_path,
        device="cuda"
    )
    print("Run prediction...")
    predictions = moment_detr_predictor.localize_moment(
        video_path=video_path, query_list=query_text_list)

    # print data
    for idx, query_data in enumerate(queries):
        print("-"*30 + f"idx{idx}")
        print(f">> query: {query_data['query']}")
        print(f">> video_path: {video_path}")
        print(f">> GT moments: {query_data['relevant_windows']}")
        print(f">> Predicted moments ([start_in_seconds, end_in_seconds, score]): "
              f"{predictions[idx]['pred_relevant_windows']}")
        print(f">> GT saliency scores (only localized 2-sec clips): {query_data['saliency_scores']}")
        print(f">> Predicted saliency scores (for all 2-sec clip): "
              f"{predictions[idx]['pred_saliency_scores']}")

if __name__ == "__main__":
    run_example()


import concurrent.futures as cf
import json
import multiprocessing as mp
import os
import re
from functools import reduce
import numpy as np

from wfdb import rdann

from metrics.adapted_fixed_window_metric import AdPositivePredictiveValue, AdSensitivity, AdSpecificity
from metrics.metric import MeanError, join
from util.util import powerset

def read_evaluated_algorithms():
    algorithm_metrics = []
    for alg_subdir in os.listdir("algorithms"):
        metric_results = read_single_algorithm_results(alg_subdir)
        if metric_results:
            algorithm_metrics.append(metric_results)
    return algorithm_metrics

def read_single_algorithm_results(alg_subdir):
    m_file_path = os.path.join("algorithms", alg_subdir, "evaluation", "metrics.json")
    if os.path.isdir(os.path.join("algorithms", alg_subdir)) and os.path.exists(m_file_path):
        with open(m_file_path) as m_file:
            return json.load(m_file)

def evaluate_algorithm(alg_store):
    print("Evaluating...")
    acc_metrics = [MeanError, AdPositivePredictiveValue, AdSensitivity, AdSpecificity]
    metric_values = {}
    noise_names = ["em", "ma", "bw"]
    ps = list(powerset(noise_names))
    # need to be adapted to the available values
    for noise_comb in ps:
        noises = "_".join(noise_comb)
        metric_values[noises] = {}
        reg = "[a-zA-Z][0-9-_]+{}[0-9-_]+[.]atr".format(noises)
        loaded_data = list(read_ann_files(alg_store, reg))
        print("Loaded Eval Data")
        with cf.ProcessPoolExecutor(max_workers=len(acc_metrics)) as pool:
            computed_metrics = pool.map(evaluate_algorithm_with_metric,
                                        acc_metrics,
                                        list([loaded_data] * len(acc_metrics)))
        for met_idx, met_per_beat in enumerate(computed_metrics):
            for k in met_per_beat:
                met_per_beat[k] = met_per_beat[k].compute()
            metric_values[noises][acc_metrics[met_idx].__abbrev__] = met_per_beat

    print("Saving Evaluation Results to JSON")
    convert_metrics_to_json(alg_store, metric_values)

def read_ann_files(alg_store, file_name_regex=".*"):
    p = re.compile(file_name_regex)
    file_names = os.listdir(alg_store.groundtruth_dir())
    file_names = [f for f in file_names if p.match(f)]
    print(len(file_names), "file names")
    with cf.ProcessPoolExecutor(max_workers=10) as pool:
        prepared_ann_tuples = list(pool.map(read_ann_file, [alg_store] * len(file_names), file_names))
    return prepared_ann_tuples

def read_ann_file(alg_store, ann_file):
    pred_ann_file = os.path.join(alg_store.prediction_dir(), ann_file.rsplit('.')[0])
    gt_ann_file = os.path.join(alg_store.groundtruth_dir(), ann_file.rsplit('.')[0])
    ann_file_type = ann_file.split("_")[0]
    if os.path.exists(pred_ann_file + "." + ann_file.rsplit('.')[1]):
        pred_ann_ref = rdann(pred_ann_file, 'atr')
        gt_ann_ref = rdann(gt_ann_file, 'atr')
        pred_beats = pred_ann_ref.sample
        anno = gt_ann_ref.sample
        sampling_frequency = gt_ann_ref.fs
        gt_beat_type = gt_ann_ref.symbol[1]

        return create_evaluation_format(ann_file_type, anno, gt_beat_type, pred_beats, sampling_frequency)
    else:
        # occurs if algorithm does not output any annotation
        # print("No Output")
        gt_ann_ref = rdann(gt_ann_file, 'atr')
        anno = gt_ann_ref.sample
        pred_ann_ref_sample = []  # [[max(0, anno[1] - gt_ann_ref.fs)]]
        return [anno[1]], [gt_ann_ref.symbol[1]], pred_ann_ref_sample, gt_ann_ref.fs, ann_file_type

def create_evaluation_format(ann_file_type, anno, gt_beat_type, pred_beats, sampling_frequency):
    adapted_start = (anno[0] + anno[1]) / 2 if anno[0] != anno[1] else 0
    adapted_end = (anno[1] + anno[2]) / 2 if anno[1] != anno[2] else pred_beats[-1]
    samples = list(filter(lambda s: adapted_start <= s <= adapted_end, pred_beats))
    return [anno[1]], [gt_beat_type], samples, sampling_frequency, ann_file_type

def evaluate_algorithm_with_metric(metric, loaded_data):
    workers = mp.cpu_count() // 2
    with cf.ProcessPoolExecutor(max_workers=workers) as pool:
        calc_metrics = list(pool.map(match_for_metric_on_data_part,
                                     loaded_data,
                                     list([metric]*len(loaded_data)),
                                     chunksize=(len(loaded_data) // workers) + 1))

        collected_metrics = {}
        for symbol, metric in calc_metrics:
            collected_metrics.setdefault(symbol, []).append(metric)

    typed_metrics = {}
    for symbol, metrics in collected_metrics.items():
        red_metric = reduce(join, metrics)
        typed_metrics[symbol] = red_metric
    return typed_metrics

def match_for_metric_on_data_part(ann_tuple, metric_class):
    typed_metric = metric_class()
    typed_metric.match_annotations(*ann_tuple[0:4])
    return ann_tuple[4], typed_metric

def convert_metrics_to_json(alg_store, acc_metrics):
    for key in acc_metrics:
        acc_metrics[key]["Name"] = alg_store.alg_name

    with open(os.path.join(alg_store.evaluation_dir(), "metrics.json"), "w") as metrics_file:
        metrics_file.write(json.dumps(acc_metrics))


from __future__ import print_function
import argparse
import sys
import os.path
import os
from os import listdir
from os.path import isfile, join
import datetime

#os.environ["CUDA_VISIBLE_DEVICES"] = ""

from keras.preprocessing.image import ImageDataGenerator
from keras.applications import mobilenet,mobilenet_v2,densenet,inception_resnet_v2,inception_v3
from keras.layers import Flatten, Dense, Dropout,GlobalAveragePooling2D, Reshape, Conv2D, Activation, MaxPooling2D, BatchNormalization, Lambda
from keras.models import Model, Sequential,model_from_json
from keras.optimizers import Adam, Nadam
from keras.callbacks import TensorBoard,ModelCheckpoint,EarlyStopping
from keras.applications.imagenet_utils import preprocess_input
from keras.regularizers import l2
from keras.losses import *
import tensorflow as tf
import keras.backend.tensorflow_backend as KTF
from keras import backend as K

import numpy as np
np.random.seed(123)

activ_func,loss_func,class_mode,metric,monitor='softmax',categorical_crossentropy,'categorical','accuracy','val_acc'

input_height=input_width=224 #192 #128 #192 #224 #200

def count_of_files(dir):
    #print 'dir=',dir
    dirs = next(os.walk(dir))[1]
    file_count=0
    for d in dirs:
        #file_count+=len(next(os.walk(os.path.join(dir,d)))[2])
        directory=os.path.join(dir,d)
        file_count+=sum(1 for entry in listdir(directory) if isfile(join(directory,entry)))
    return file_count,len(dirs)

#net_model=mobilenet_v2
net_model=mobilenet
def get_net():
    if False:
        return net_model.MobileNetV2(alpha=1.4,include_top=False, weights='imagenet',input_shape=(input_height, input_width, 3))
    else:
        return net_model.MobileNet(include_top=False, weights='imagenet',input_shape=(input_height, input_width, 3))

def get_predictive_model(base_model,classes_num):
    last_model_layer = base_model.output
    x = GlobalAveragePooling2D(name='reshape_1')(last_model_layer)
    preds = Dense(classes_num, activation=activ_func, kernel_regularizer=l2(4e-5), name='predictions')(x)
    f_model = Model(base_model.input, preds)
    return f_model

def model_architecture(classes_num):
    #net_model = VGG16(weights='imagenet', include_top=False,
    #                     input_shape=(input_height, input_width, 3))
    #net_model = ResNet50(weights='imagenet', include_top=False,
    #                      input_shape=(input_height, input_width, 3))
    # net_model = Xception(weights='imagenet', include_top=False, input_shape=(input_height, input_width, 3))
    # net_model = VGG19(weights='imagenet', include_top=False,
    #                      input_shape=(input_height, input_width, 3))
    base_model = get_net()
    return get_predictive_model(base_model,classes_num),base_model

def freeze_session(session, keep_var_names=None, output_names=None, clear_devices=True):
    from tensorflow.python.framework.graph_util import convert_variables_to_constants
    graph = session.graph
    with graph.as_default():
        freeze_var_names = list(set(v.op.name for v in tf.global_variables()).difference(keep_var_names or []))
        output_names = output_names or []
        output_names += [v.op.name for v in tf.global_variables()]
        input_graph_def = graph.as_graph_def()
        if clear_devices:
            for node in input_graph_def.node:
                node.device = ''
        frozen_graph = convert_variables_to_constants(session, input_graph_def,
                                                      output_names, freeze_var_names)
        return frozen_graph

def load_model(modelpath):
    base_model = get_net()
    last_model_layer = base_model.output

    out = GlobalAveragePooling2D(name='reshape_1')(last_model_layer)
    cnn_model = Model(base_model.input, out)
    cnn_model.load_weights(modelpath)
    return cnn_model,base_model

from tensorflow.python.framework.graph_util import convert_variables_to_constants
def convert_to_tf(modelpath):
    K.set_learning_phase(0)
    config = tf.ConfigProto(device_count={'CPU':1,'GPU':0})
    KTF.set_session(tf.Session(config=config))
    filename,ext=os.path.splitext(os.path.basename(modelpath))
    print (filename, ext)
    if ext=='.hdf5':
        model,_ = model_architecture(9131)
        model.summary()
        model.load_weights(modelpath)
        out = model.get_layer('reshape_1').output
        print('out=',out)

        cnn_model = Model(model.input, out)
        #cnn_model.save('models/'+filename+'.h5')

        cnn_model.save_weights('weights.h5')
        model_json = cnn_model.to_json()
        with open("model.json", "w") as json_file:
            json_file.write(model_json)

        json_file = open('model.json', 'r')
        loaded_model_json = json_file.read()
        json_file.close()
        loaded_model = model_from_json(loaded_model_json)
        loaded_model.load_weights('weights.h5')

        loaded_model.save('models/'+filename+'.h5')
    else:
        saved_model = load_model(modelpath)[0]
        output_names=[saved_model.output.op.name]
        print (output_names)

        session=K.get_session()
        graph = session.graph
        keep_var_names=None
        clear_devices=True
        with graph.as_default():
            freeze_var_names = list(set(v.op.name for v in tf.global_variables()).difference(keep_var_names or []))
            output_names = output_names or []
            output_names += [v.op.name for v in tf.global_variables()]
            input_graph_def = graph.as_graph_def()
            if clear_devices:
                for node in input_graph_def.node:
                    node.device = ''
            frozen_graph = convert_variables_to_constants(session, input_graph_def,
                                                          output_names, freeze_var_names)
        tf.train.write_graph(frozen_graph, 'models', filename+'.pb', as_text=False)

if __name__ == '__main__':
    #config = tf.ConfigProto()
    #config.gpu_options.allocator_type = 'BFC'
    #KTF.set_session(tf.Session(config=config))

    parser = argparse.ArgumentParser(description="Training script for face recognition.")

    parser.add_argument("--datapath", type=str, default='/home/datasets/images/vgg_face_dataset/vggface-2/tf1', help="Path where image data is stored.")
    parser.add_argument("--modelpath", type=str, default='', help="Path where trained model is stored.")

    train_batch_size=val_batch_size= 32 #76 # 88 #104 #128

    args = parser.parse_args()
    if args.modelpath!='':
        convert_to_tf(args.modelpath)
        sys.exit()

    train_path=args.datapath + '/train/'
    val_path=args.datapath + '/val/'

    train_datagen = ImageDataGenerator(shear_range=0.3, #0.2
                                       rotation_range=10,
                                       zoom_range=0.2, #0.1
                                       width_shift_range=0.1,height_shift_range=0.1,
                                       horizontal_flip=True,preprocessing_function=preprocess_input)
    test_datagen = ImageDataGenerator(preprocessing_function=preprocess_input)

    target_size=(input_height, input_width)

    test_generator = test_datagen.flow_from_directory(
        val_path,
        target_size=target_size,
        batch_size=val_batch_size, class_mode=class_mode)

    train_generator = train_datagen.flow_from_directory(
        train_path,
        target_size=target_size,
        batch_size=train_batch_size, class_mode=class_mode)

    nb_train_samples=train_generator.samples
    classes_num=train_generator.num_classes
    nb_validation_samples=test_generator.samples
    print('after read',nb_train_samples,nb_validation_samples,classes_num)

    model_file=''
    model,base_model = model_architecture(classes_num)

    opt=Adam(lr=1e-3,decay=1e-5)
    model.compile(loss=loss_func, optimizer=opt, metrics=[metric])
    model.summary()

    #filepath="models/casia_mobilenet128"
    filepath="models/vgg2_mobilenet2_224"
    #filepath="models/gender_mobilenet192_tf15"
    #filepath='models/age_mobilenet192_tf1_regr'
    #filepath='models/age_simple_tf2'
    #saved_weights_filepath=filepath+'-03-0.64.hdf5'
    #model.load_weights(saved_weights_filepath)

    last_path='-{epoch:02d}-{val_acc:.2f}.hdf5'
    checkpoint = ModelCheckpoint(filepath+last_path, monitor=monitor, verbose=1, save_best_only=True, mode='auto')
    es=EarlyStopping(monitor='val_acc',patience=2)
    callbacks=[checkpoint,es]
    history = model.fit_generator(
        train_generator,
        steps_per_epoch=nb_train_samples // train_batch_size,
        epochs=16,
        validation_data=test_generator,
        validation_steps=nb_validation_samples // val_batch_size,
        callbacks=callbacks)
    model.save(filepath+'.hdf5')


#!/usr/bin/env python3
# coding=UTF-8

'''
    Atom
    ====
        Class Atom define.
'''

# Import Python Lib
from numpy import array

# Import PDBTools
from .NotProtein import __NotProtein
from .Math import Dis

################################################################################
# Class Atom
################################################################################

class Atom(__NotProtein):

    __slots__ = ('name', 'num', 'coord', 'alt', 'occ', 'tempF', 'ele', 'chg', 'owner')

    def __init__(self, atomName = '', atomNum = 0, atomCoord = array((0., 0., 0.)),
        atomAltLoc = '', atomOccupancy = '', atomTempFactor = '', atomElement = '',
        atomCharge = '', owner = None):

        self.name  = atomName
        self.num   = atomNum
        self.coord = atomCoord
        self.alt   = atomAltLoc
        self.occ   = atomOccupancy
        self.tempF = atomTempFactor
        self.ele   = atomElement
        self.chg   = atomCharge
        self.owner = owner

        if owner != None:
            owner.sub.append(self)

    def __repr__(self):

        return '<Atom object: %d %s %s, at 0x%X>' % (
            self.num, self.name, str(self.coord), id(self))

    __str__ = __repr__

    def __sub__(self, subAtomObj):

        return Dis(self.coord, subAtomObj.coord)

    def Copy(self):

        return Atom(self.name, self.num, self.coord.copy(), self.alt, self.occ,
            self.tempF, self.ele, self.chg)

    def Dumps(self):

        ownerResObj = self.owner

        if ownerResObj != None:

            resName = ownerResObj.name
            resNum  = ownerResObj.num
            resIns  = ownerResObj.ins

            ownerChainObj = ownerResObj.owner
            chainName = ownerChainObj.name if ownerChainObj != None else ''

        else:
            resName, resNum, resIns, chainName = '', 0, '', ''

        if self.name[0].isdigit() or len(self.name) == 4:
            atomName = '%-4s' % self.name
        else:
            atomName = ' %-3s' % self.name

        return 'ATOM  %5d %s%1s%3s %1s%4d%1s   %8.3f%8.3f%8.3f%6s%6s          %2s%2s\n' % (
            self.num, atomName, self.alt, resName, chainName, resNum, resIns,
            self.coord[0], self.coord[1], self.coord[2], self.occ, self.tempF,
            self.ele, self.chg)


# -*- coding: UTF-8 -*-
from __future__ import print_function
import argparse
import os
import random
import numpy as np
import torch
import torch.nn as nn
import torch.backends.cudnn as cudnn
import torch.optim as optim
import torchvision.transforms as transforms
import torchvision.utils as vutils
from torch.autograd import Variable
from utils.voc_datasets import VOC2007_TEST
from models.DarkNet import DarkNet
import torch.nn.init as init
import torch.nn.functional as F
from models.test_util import write_results
from PIL import Image, ImageDraw

parser = argparse.ArgumentParser(description='train hed model')
parser.add_argument('--batchSize', type=int, default=1, help='with batchSize=1 equivalent to instance normalization.')
parser.add_argument('--niter', type=int, default=15, help='number of epochs to train for')
parser.add_argument('--lr', type=float, default=1e-4, help='learning rate, default=0.0002')
parser.add_argument('--lr_decay', type=float, help='learning rate decay', default=0.1)
parser.add_argument('--beta1', type=float, default=0.5, help='beta1 for adam. default=0.5')
parser.add_argument('--outf', default='checkpoints/', help='folder to output images and model checkpoints')
parser.add_argument('--imgdir', default='/media/data2/wb/VOCdevkit/VOC2007/JPEGImages/', help='folder to output images and model checkpoints')

opt = parser.parse_args()
print(opt)

device = torch.device('cuda: 0')

try:
    os.makedirs(opt.outf)
except OSError:
    pass

imgdir = opt.imgdir
lr_decay_epoch = {10, 12}
lr_decay = opt.lr_decay
mean=[0.485, 0.456, 0.406]
std=[0.229, 0.224, 0.225]

mask = [0, 1, 2]
anchors = [10,13,  16,30,  33,23,  30,61,  62,45,  59,119,  116,90,  156,198,  373,326]
mask1 = [6, 7, 8]
mask2 = [3, 4, 5]
mask3 = [0, 1, 2]
anchors = [(anchors[i], anchors[i+1]) for i in range(0, len(anchors),2)]
anchors1 = [anchors[i] for i in mask1] # 52x52
anchors2 = [anchors[i] for i in mask2] # 26x26
anchors3 = [anchors[i] for i in mask3] # 13x13
anchors = [anchors1, anchors2, anchors3]

print (anchors)

confidence = 0.6
num_classes = 20
nms_thesh = 0.4

##########   DATASET   ###########
transform = transforms.Compose([
                transforms.ToTensor(),
                transforms.Normalize(mean,std)
            ])

###########   DATASET   ###########
data = VOC2007_TEST(transform=transform, size=416)
train_loader = torch.utils.data.DataLoader(dataset=data,
                                           batch_size=1,
                                           shuffle=True,
                                           num_workers=2)

net = DarkNet(anchors, num_classes)

def weights_init(m):
    if isinstance(m, nn.Conv2d):
        init.xavier_uniform_(m.weight.data)
        init.constant_(m.bias.data, 0)
    elif isinstance(m, nn.BatchNorm2d):
        m.weight.data.normal_(1.0, 0.02)
        m.bias.data.fill_(0)

net.apply(weights_init)

net.load_state_dict(torch.load('checkpoints/darknet_200.pth'))
net.to(device)

lr = opt.lr
optimizer = torch.optim.Adam(net.parameters(),lr=lr, betas=(opt.beta1, 0.999))

########### Testing   ###########

net.eval()
topilimage = transforms.ToPILImage()
for i, image in enumerate(train_loader):
    img = image[0]
    gt_boxes = image[1]
    gt_classes = image[2]
    name = image[3][0] + '.jpg'

    img = img.to(device)
    gt_boxes = gt_boxes.to(device)
    gt_classes = gt_classes.to(device)

    detections = net(img, 0, 0)
    # print (detections[:,:,:5])

    prediction = write_results(detections, confidence, num_classes, nms = True, nms_conf = nms_thesh)
    print (prediction[:,1:5])

    # resize to the original img
    # ori_im = topilimage(img.data[0].cpu())
    ori_im = Image.open(imgdir + name).convert('RGB')
    W, H = ori_im.size

    scale_H = float(H) / img.size(2)
    scale_W = float(W) / img.size(3)
    drawObject = ImageDraw.Draw(ori_im)
    bboxes = prediction[:,1:5].clamp(min=0.0, max=415.0)
    print (bboxes)
    bboxes[:,0:4:2] = bboxes[:,0:4:2] * scale_W
    bboxes[:,1:4:2] = bboxes[:,1:4:2] * scale_H

    for l in range(len(bboxes)):

        bbox = list(bboxes[l])
        # print (bbox)

        drawObject.rectangle(bbox, outline='red')
    ori_im.show()

    exit()


#!/usr/bin/env python
from __future__ import print_function
import os
from os.path import splitext, join, isfile, isdir, basename
import argparse
import numpy as np
# from scipy import misc, ndimage
import tensorflow.keras.backend as K
from tensorflow.keras.models import model_from_json, load_model
import tensorflow as tf
import layers_builder as layers
from glob import glob
from utils import utils
from tensorflow.python.keras.utils.generic_utils import CustomObjectScope
import cv2
import math
from PIL import Image
# -- Fix for macos, uncomment it
# import matplotlib
# matplotlib.use('TkAgg')
# --
import matplotlib.pyplot as plt
from ade20k_labels import ade20k_label_dict as ade20k_labels

from imageio import imread
# These are the means for the ImageNet pretrained ResNet
DATA_MEAN = np.array([[[123.68, 116.779, 103.939]]])  # RGB order

class PSPNet(object):
    """Pyramid Scene Parsing Network by Hengshuang Zhao et al 2017"""

    def __init__(self, nb_classes, resnet_layers, input_shape, weights):
        self.input_shape = input_shape
        self.num_classes = nb_classes

        json_path = join("weights", "keras", weights + ".json")
        h5_path = join("weights", "keras", weights + ".h5")
        if 'pspnet' in weights:
            if os.path.isfile(json_path) and os.path.isfile(h5_path):
                print("Keras model & weights found, loading...")
                with CustomObjectScope({'Interp': layers.Interp}):
                    with open(json_path) as file_handle:
                        self.model = model_from_json(file_handle.read())
                self.model.load_weights(h5_path)
            else:
                print("No Keras model & weights found, import from npy weights.")
                self.model = layers.build_pspnet(nb_classes=nb_classes,
                                                 resnet_layers=resnet_layers,
                                                 input_shape=self.input_shape)
                self.set_npy_weights(weights)
        else:
            print('Load pre-trained weights')
            self.model = load_model(weights)

    def predict(self, img, flip_evaluation=False):
        """
        Predict segementation for an image.
        Arguments:
            img: must be rowsxcolsx3
        """

        if img.shape[0:2] != self.input_shape:
            print(
                "Input %s not fitting for network size %s, resizing. You may want to try sliding prediction for better results." % (
                img.shape[0:2], self.input_shape))
            img = np.array(Image.fromarray(img).resize(size=self.input_shape))
            # img = misc.imresize(img, self.input_shape)

        img = img - DATA_MEAN
        img = img[:, :, ::-1]  # RGB => BGR
        img = img.astype('float32')

        probs = self.feed_forward(img, flip_evaluation)

        return probs

    def predict_sliding(self, full_img, flip_evaluation):
        """
        Predict on tiles of exactly the network input shape.
        This way nothing gets squeezed.
        """
        tile_size = self.input_shape
        classes = self.num_classes
        overlap = 1 / 3

        stride = math.ceil(tile_size[0] * (1 - overlap))
        tile_rows = max(int(math.ceil((full_img.shape[0] - tile_size[0]) / stride) + 1), 1)  # strided convolution formula
        tile_cols = max(int(math.ceil((full_img.shape[1] - tile_size[1]) / stride) + 1), 1)
        print("Need %i x %i prediction tiles @ stride %i px" % (tile_cols, tile_rows, stride))
        full_probs = np.zeros((full_img.shape[0], full_img.shape[1], classes))
        count_predictions = np.zeros((full_img.shape[0], full_img.shape[1], classes))
        tile_counter = 0
        for row in range(tile_rows):
            for col in range(tile_cols):
                x1 = int(col * stride)
                y1 = int(row * stride)
                x2 = min(x1 + tile_size[1], full_img.shape[1])
                y2 = min(y1 + tile_size[0], full_img.shape[0])
                x1 = max(int(x2 - tile_size[1]), 0)  # for portrait images the x1 underflows sometimes
                y1 = max(int(y2 - tile_size[0]), 0)  # for very few rows y1 underflows

                img = full_img[y1:y2, x1:x2]
                padded_img = self.pad_image(img, tile_size)
                plt.imshow(padded_img)
                plt.show()
                tile_counter += 1
                print("Predicting tile %i" % tile_counter)
                padded_prediction = self.predict(padded_img, flip_evaluation)
                prediction = padded_prediction[0:img.shape[0], 0:img.shape[1], :]
                count_predictions[y1:y2, x1:x2] += 1
                full_probs[y1:y2, x1:x2] += prediction  # accumulate the predictions also in the overlapping regions

        # average the predictions in the overlapping regions
        full_probs /= count_predictions
        # visualize normalization Weights
        # plt.imshow(np.mean(count_predictions, axis=2))
        # plt.show()
        return full_probs

    @staticmethod
    def pad_image(img, target_size):
        """Pad an image up to the target size."""
        rows_missing = target_size[0] - img.shape[0]
        cols_missing = target_size[1] - img.shape[1]
        padded_img = np.pad(img, ((0, rows_missing), (0, cols_missing), (0, 0)), 'constant')
        return padded_img

    def predict_multi_scale(self, img, flip_evaluation, sliding_evaluation, scales):
        """Predict an image by looking at it with different scales."""

        full_probs = np.zeros((img.shape[0], img.shape[1], self.num_classes))
        h_ori, w_ori = img.shape[:2]

        print("Started prediction...")
        for scale in scales:
            print("Predicting image scaled by %f" % scale)
            scaled_img = np.array(Image.fromarray(img).resize(size=(int(scale * img.shape[0]), int(scale * img.shape[1])), resample=Image.BILINEAR))
            # scaled_img = misc.imresize(img, size=scale, interp="bilinear")

            if sliding_evaluation:
                scaled_probs = self.predict_sliding(scaled_img, flip_evaluation)
            else:
                scaled_probs = self.predict(scaled_img, flip_evaluation)

            # scale probs up to full size
            # visualize_prediction(probs)
            probs = cv2.resize(scaled_probs, (w_ori, h_ori))
            full_probs += probs
        full_probs /= len(scales)
        print("Finished prediction...")

        return full_probs

    def feed_forward(self, data, flip_evaluation=False):
        assert data.shape == (self.input_shape[0], self.input_shape[1], 3)

        if flip_evaluation:
            print("Predict flipped")
            input_with_flipped = np.array(
                [data, np.flip(data, axis=1)])
            prediction_with_flipped = self.model.predict(input_with_flipped)
            prediction = (prediction_with_flipped[
                          0] + np.fliplr(prediction_with_flipped[1])) / 2.0
        else:
            prediction = self.model.predict(np.expand_dims(data, 0))[0]
        return prediction

    def set_npy_weights(self, weights_path):
        npy_weights_path = join("weights", "npy", weights_path + ".npy")
        json_path = join("weights", "keras", weights_path + ".json")
        h5_path = join("weights", "keras", weights_path + ".h5")

        print("Importing weights from %s" % npy_weights_path)
        weights = np.load(npy_weights_path, encoding='bytes').item()
        for layer in self.model.layers:
            print(layer.name)
            if layer.name[:4] == 'conv' and layer.name[-2:] == 'bn':
                mean = weights[layer.name.encode()][
                    'mean'.encode()].reshape(-1)
                variance = weights[layer.name.encode()][
                    'variance'.encode()].reshape(-1)
                scale = weights[layer.name.encode()][
                    'scale'.encode()].reshape(-1)
                offset = weights[layer.name.encode()][
                    'offset'.encode()].reshape(-1)

                self.model.get_layer(layer.name).set_weights(
                    [scale, offset, mean, variance])

            elif layer.name[:4] == 'conv' and not layer.name[-4:] == 'relu':
                try:
                    weight = weights[layer.name.encode()]['weights'.encode()]
                    self.model.get_layer(layer.name).set_weights([weight])
                except Exception as err:
                    biases = weights[layer.name.encode()]['biases'.encode()]
                    self.model.get_layer(layer.name).set_weights([weight,
                                                                  biases])
        print('Finished importing weights.')

        print("Writing keras model & weights")
        json_string = self.model.to_json()
        with open(json_path, 'w') as file_handle:
            file_handle.write(json_string)
        self.model.save_weights(h5_path)
        print("Finished writing Keras model & weights")

class PSPNet50(PSPNet):
    """Build a PSPNet based on a 50-Layer ResNet."""

    def __init__(self, nb_classes, weights, input_shape):
        PSPNet.__init__(self, nb_classes=nb_classes, resnet_layers=50,
                        input_shape=input_shape, weights=weights)

class PSPNet101(PSPNet):
    """Build a PSPNet based on a 101-Layer ResNet."""

    def __init__(self, nb_classes, weights, input_shape):
        PSPNet.__init__(self, nb_classes=nb_classes, resnet_layers=101,
                        input_shape=input_shape, weights=weights)

def update_weights_json_to_python36p(dataset):
    if dataset == 'cityscapes':
        model = layers.build_pspnet(nb_classes=19, resnet_layers=101, input_shape=(713, 713), activation='softmax')
        with open('weights/keras/pspnet101_cityscapes.json', 'w') as json_file:
            json_file.write(model.to_json())
    elif dataset == 'ade20k':
        model = layers.build_pspnet(nb_classes=150, resnet_layers=50, input_shape=(473, 473), activation='softmax')
        with open('weights/keras/pspnet50_ade20k.json', 'w') as json_file:
            json_file.write(model.to_json())

def main(args):
    # Handle input and output args
    images = sorted(glob(args.glob_path))[::args.glob_interval] if args.glob_path else [args.input_path, ]
    if args.glob_path:
        fn, ext = splitext(args.output_path)
        if ext:
            parser.error("output_path should be a folder for multiple file input")
        if not isdir(args.output_path):
            os.mkdir(args.output_path)

    # Predict
    os.environ["CUDA_VISIBLE_DEVICES"] = args.id

    sess = tf.Session()
    K.set_session(sess)

    with sess.as_default():
        print(args)
        if not args.weights:
            if "pspnet50" in args.model:
                pspnet = PSPNet50(nb_classes=150, input_shape=(473, 473),
                                  weights=args.model)
            elif "pspnet101" in args.model:
                if "cityscapes" in args.model:
                    pspnet = PSPNet101(nb_classes=19, input_shape=(713, 713),
                                       weights=args.model)
                if "voc2012" in args.model:
                    pspnet = PSPNet101(nb_classes=21, input_shape=(473, 473),
                                       weights=args.model)

            else:
                print("Network architecture not implemented.")
        else:
            # pspnet = PSPNet50(nb_classes=2, input_shape=(
            #     768, 480), weights=args.weights)
            pspnet = PSPNet50(nb_classes=2, input_shape=(
                473, 473), weights=args.weights)

        EVALUATION_SCALES = [1.0]
        if args.multi_scale:
            EVALUATION_SCALES = [0.5, 0.75, 1.0, 1.25, 1.5, 1.75]  # must be all floats! Taken from original paper

        for i, img_path in enumerate(images):
            # if '287799' not in img_path: continue
            print("Processing image {} / {}".format(i + 1, len(images)))
            img = imread(img_path, pilmode='RGB')
            cv2.fillPoly(img, [np.array([[700, 0], [799, 100], [799, 0]])], 0)
            cv2.fillPoly(img, [np.array([[0, 100], [100, 0], [0, 0]])], 0)

            probs = pspnet.predict_multi_scale(img, args.flip, args.sliding, EVALUATION_SCALES)
            # cm = np.argmax(probs, axis=2)
            pm = np.sort(probs, axis=2)[..., ::-1]
            cm = np.argsort(probs, axis=2)[..., ::-1]

            # pm = probs[cm[..., predlayer_i]]
            colored_class_image = utils.color_class_image(cm[..., 0], args.model)
            alpha_blended = 0.5 * colored_class_image + 0.5 * img

            if args.glob_path:
                input_filename, ext = splitext(basename(img_path))
                filename = join(args.output_path, input_filename)
            else:
                filename, ext = splitext(args.output_path)

            # cv2.imwrite(filename + "_seg_read" + ext, cm)
            # cv2.imwrite(filename + "_seg" + ext, colored_class_image)
            # cv2.imwrite(filename + "_probs" + ext, pm)
            # cv2.imwrite(filename + "_seg_blended" + ext, alpha_blended)

            sky_orig = np.where((colored_class_image != ade20k_labels['sky']).any(axis=-1), 255, 0).astype(np.uint8)
            mask_building = np.where((colored_class_image == ade20k_labels['building']).any(axis=-1))
            mask_tree = np.where((colored_class_image == ade20k_labels['tree']).any(axis=-1))
            sky_orig[mask_building[0], mask_building[1]] = 205
            sky_orig[mask_tree[0], mask_tree[1]] = 155
            sky_orig[int(sky_orig.shape[0] * 2 / 3):, ...] = 255
            fn_split = filename.split('/')
            cv2.imwrite('/'.join(fn_split[:-1] + ['/superpixels2/'] + fn_split[-1:]) + ext,
                        sky_orig.astype(np.uint8))
            # msk_border_np = \
            #     cv2.dilate(sky_orig.astype('uint8'), cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (50, 50)),
            #                iterations=1) - \
            #     cv2.erode(sky_orig.astype('uint8'), cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (100, 100)),
            #               iterations=1)
            # sky = np.where(msk_border_np, 128, sky_orig)
            # sky = np.where((colored_class_image == [180, 120, 120]).all(axis=-1) * (sky == 128), 255, sky).astype(np.uint8)
            # sky[int(sky.shape[0] * 2 / 3):, ...] = 255
            # fn_split = filename.split('/')
            # cv2.imwrite('/'.join(fn_split[:-1] + ['/trimap/'] + fn_split[-1:]) + ext, sky)
            # cv2.imwrite('/'.join(fn_split[:-1] + ['/image/'] + fn_split[-1:]) + ext, img[..., ::-1])
            # cv2.imwrite('/'.join(fn_split[:-1] + ['/seg_orig/'] + fn_split[-1:]) + ext,
            #             np.where(sky_orig[..., None], img[..., ::-1], img[..., ::-1] // 2))

            # cm2 = np.where(pm[..., 1] > pm[..., 0] / 10, cm[..., 1], cm[..., 0])
            # # colored_class_image2 = utils.color_class_image(cm2, args.model)
            # colored_class_image2 = utils.color_class_image(cm[...,1], args.model)
            # sky_orig2 = np.where((colored_class_image2 != [6, 230, 230]).any(axis=-1), 255, 0).astype(np.uint8)
            # sky_multilayer = np.dstack([sky_orig, sky_orig2, sky_orig2])
            # cv2.imwrite('/'.join(fn_split[:-1] + ['/superpixels/'] + fn_split[-1:]) + ext, sky_multilayer.astype(np.uint8))

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('-m', '--model', type=str, default='pspnet50_ade20k',
                        help='Model/Weights to use',
                        choices=['pspnet50_ade20k',
                                 'pspnet101_cityscapes',
                                 'pspnet101_voc2012'])
    parser.add_argument('-w', '--weights', type=str, default=None)
    # parser.add_argument('-i', '--input_path', type=str, default='example_images/frame-000000115873.png',                        help='Path the input image')
    # parser.add_argument('-g', '--glob_path', type=str, default=None,                         help='Glob path for multiple images')
    # parser.add_argument('-g', '--glob_path', type=str, default='/home/noam/Desktop/gimp_monodepth/img/*.png', help='Glob path for multiple images')
    parser.add_argument('-g', '--glob_path', type=str, default='/media/noam/Storage/semanticseg/datasets/rgo_videos/ColorCamera/20201029_Ginegar2/Images_col_MLRI_800x800/*.png', help='Glob path for multiple images')
    parser.add_argument('-gi', '--glob_interval', type=int, default=30, help='skip input files with defined interval')
    # parser.add_argument('-o', '--output_path', type=str, default='example_results/frame-000000115873.png', help='Path to output')
    parser.add_argument('-o', '--output_path', type=str, default='/media/noam/Storage/semanticseg/findskies/code/DNN/sky_autolabel/pspnet_FBAmatting/20201029_Ginegar2', help='Path to output')
    parser.add_argument('--id', default="0")
    parser.add_argument('--input_size', type=int, default=500)
    parser.add_argument('-s', '--sliding', action='store_true',
                        help="Whether the network should be slided over the original image for prediction.")
    parser.add_argument('-f', '--flip', action='store_true', default=False,
                        help="Whether the network should predict on both image and flipped image.")
    parser.add_argument('-ms', '--multi_scale', action='store_true', default=False,
                        help="Whether the network should predict on multiple scales.")

    args = parser.parse_args()

    # update_weights_json_to_python36p()
    # args.weights = 'weights/keras/pspnet50_ade20k.h5'
    main(args)


from typing import List
from agent.agent import Agent
from  env.env_cmd import CmdEnv
import copy
import random
import numpy as np
import time
from .tools import distance_matrix

'''
    y     [-150000, +150000]
    red_x [-150000, -125000]
    blue_x [125000,  150000]
'''
class Special():
    NE = np.array([125000,   125000])
    N = np.array([0,   125000])
    E = np.array([125000,   0])
    NW = np.array([-125000,  125000])
    SE = np.array([125000,  -125000])
    SW = np.array([-125000, -125000])

    init_X = 125000
    init_Y = 0
    init_Z = 9500

    init_speed_vip = 400
    init_speed_drone = 300 * 0.6

    Dive_Horizental = 10000 # 躲导弹时，下潜点的水平距离

def nearby(dir, L):
    dir = dir / (np.linalg.norm(dir)+1e-3) * L

class Drone():
    MaxSpeed = 300
    MinSpeed = 100
    FollowSpeed = 150
    MaxHeight = 10000
    EscapeHeight = 5000
    MinHeight = 2000
    MaxAcc = 2
    MaxOverload = 12

class Vip():
    MaxSpeed = 400
    MinSpeed = 150
    FollowSpeed = 150
    MaxHeight = 15000
    EscapeHeight = 5000
    MinHeight = 2000
    MaxAcc = 1
    MaxOverload = 6

class plane(object):
    def __init__(self, data=None) -> None:
        super().__init__()
        if data is not None:
            self.update_info(data)

    def update_info(self, data):
        self.data = data
        for key in data:
            setattr(self, key, data[key])

        if not hasattr(self, 'is_drone'):
            self.is_drone = (self.Type ==2)
            self.is_vip = (self.Type ==1)
            Ability = Drone if self.is_drone else Vip
            for key in Ability.__dict__:
                if '__' in key: continue
                setattr(self, key, getattr(Ability,key))

        # self

    def incoming_msid(self):
        return 0

    def step_init(self):
        self.under_attack = []
        self.state = ''

class NewAgent(Agent):
    def __init__(self, name, config):
        super(NewAgent, self).__init__(name, config["side"])
        self.my_uvas_infos = []
        self.my_manned_info = []
        self.my_allplane_infos = []
        self.enemy_uvas_infos = []
        self.enemy_manned_info = []
        self.enemy_allplane_infos = []
        self.enemy_missile_infos = []
        self.my_missile_infos = []
        self.attack_handle_enemy = {}

        self.n_uav = 4
        self.cmd_list = []
        self.state_recall = {}
        self.STATE = 'Full_Assult'
        self.individual_state = ['Assult']*self.n_uav

        self.my_planes = None

    def get_uav(self):
        return { id_:self.my_planes[id_] for id_ in self.my_planes if self.my_planes[id_].is_drone}

    def get_vip(self):
        return { id_:self.my_planes[id_] for id_ in self.my_planes if self.my_planes[id_].is_vip}

    def init_planes(self):
        self.my_planes = {
             p['ID']:plane(p)   for p in self.my_allplane_infos
        }
        return

    def reset(self, **kwargs):
        self.attack_handle_enemy.clear()
        pass

    def step(self,sim_time, obs_side, **kwargs) -> List[dict]:
        self.cmd_list = []
        self.process_observation_and_show(obs_side)
        if self.my_planes is None:
            self.init_planes()

        self.process_decision(sim_time, obs_side)
        return copy.deepcopy(self.cmd_list)

    def process_decision(self, time, obs_side):
        # self.show()
        if time == 3: self.init_pos(self.cmd_list); self.print_()
        elif time > 5: self.make_decision()
        return

    def make_decision(self):
        if self.STATE == 'Full_Assult':
            self.Full_Assult()

    def Full_Assult(self):
        '''
            cmd_list.append(CmdEnv.make_attackparam(plane_info['ID'], close_target['enemy_data']['ID'], 1))
            cmd_list.append();
            cmd_list.append(CmdEnv.make_followparam(plane['ID'], enemy_leader['ID'], data['move_max_speed'], data['move_max_acc'], data['move_max_g']))
        '''

        # my_uav_id = [ p['ID']   for p in self.my_uvas_infos    ]
        # op_uav_id = [ p['ID']   for p in self.enemy_uvas_infos ]

        my_uav = self.get_uav()
        my_vip = self.get_vip()

        for plane_info in self.my_allplane_infos:
            plane_id = plane_info['ID']
            self.my_planes[plane_id].step_init()
            self.my_planes[plane_id].update_info(plane_info)

        everything = self.my_allplane_infos + self.my_missile_infos + self.enemy_allplane_infos + self.enemy_missile_infos
        everything_pos = np.array([ np.array([p['X'],p['Y'],p['Z']])  for p in everything ])
        self.id2index = { p['ID']:i for i,p in enumerate(everything)}
        self.everything_dis = distance_matrix(everything_pos)

        def id2index(arr):
            return [self.id2index[id_] for id_ in arr]

        def get_dis(id1,id2):
            index1 = self.id2index[id1]
            index2 = self.id2index[id2]
            return self.everything_dis[index1, index2]

        if 'ID_Bind' not in self.state_recall:
            self.state_recall['ID_Bind'] = { p['ID']:self.enemy_uvas_infos[i]['ID']   for i,p in enumerate(self.my_uvas_infos) }

        # deal ms
        for ms in self.enemy_missile_infos:
            my_plane_targeted = ms['EngageTargetID']
            self.my_planes[my_plane_targeted].under_attack.append(ms)

        # state change
        for pid in my_uav:
            p = my_uav[pid]
            if len(p.under_attack) >= 1:
                p.state = 'escape'

        def follow_cmd(uav):
            self.cmd_list.append(
                CmdEnv.make_followparam(
                    uav.ID,
                    self.state_recall['ID_Bind'][uav.ID],
                    uav.FollowSpeed,
                    uav.MaxAcc,
                    uav.MaxOverload))

        def escape_cmd(uav):
            # Dst = nearby(Special.NE, Special.Dive_Hori/zental)
            goto_location = [
                {
                    "X": Special.E[0],
                    "Y": Special.E[1],
                    "Z": uav.EscapeHeight
                }
            ]
            self.cmd_list.append(
                CmdEnv.make_linepatrolparam(
                    uav.ID,
                    goto_location,
                    uav.MaxSpeed,
                    uav.MaxAcc,
                    uav.MaxOverload
                )
            )

        for pid in my_uav:
            p = my_uav[pid]
            if p.state == '':
                follow_cmd(p)
            elif p.state == 'escape':
                escape_cmd(p)
        return

    def print_(self):
        return

    def process_observation_and_show(self, obs_side):
        my_entity_infos = obs_side['platforminfos'] # 拿到己方阵营有人机、无人机在内的所有飞机信息
        if len(my_entity_infos) < 1:
            return
        my_manned_info = [] # 用以保存当前己方有人机信息
        my_uvas_infos = []  # 用以保存当前己方无人机信息
        my_allplane_infos = []    # 用以保存当前己方所有飞机信息
        for uvas_info in my_entity_infos:
            if uvas_info['ID'] != 0 and uvas_info['Availability'] > 0.0001: # 判断飞机是否可用 飞机的ID即为飞机的唯一编号 飞机的Availability为飞机当前生命值
                uvas_info["Z"] = uvas_info["Alt"]
                if uvas_info['Type'] == 1:
                    my_manned_info.append(uvas_info)
                if uvas_info['Type'] == 2:
                    my_uvas_infos.append(uvas_info)
                my_allplane_infos.append(uvas_info)
        if len(my_manned_info) < 1:
            return
        enemy_entity_infos = obs_side['trackinfos']   # 拿到敌方阵营的飞机信息,包括敌方有人机、无人机在内的所有飞机信息
        enemy_manned_info = []  # 用以保存当前敌方有人机信息
        enemy_uvas_infos = []   # 用以保存当前敌方无人机信息
        enemy_allplane_infos = []     # 用以保存当前敌方所有飞机信息
        for uvas_info in enemy_entity_infos:
            if uvas_info['ID'] != 0 and uvas_info['Availability'] > 0.0001:  # 判断飞机是否可用 飞机的ID即为飞机的唯一编号 飞机的Availability为飞机当前生命值
                uvas_info['Z'] = uvas_info['Alt']         # 飞机的 Alt 即为飞机的当前高度
                if uvas_info['Type'] == 1:               # 所有类型为 1 的飞机是 有人机
                    enemy_manned_info.append(uvas_info)  # 将有人机保存下来 一般情况，每方有人机只有1架
                if uvas_info['Type'] == 2:               # 所有类型为 2 的飞机是 无人机
                    enemy_uvas_infos.append(uvas_info)   # 将无人机保存下来 一般情况，每方无人机只有4架
                enemy_allplane_infos.append(uvas_info)         # 将己方所有飞机信息保存下来 一般情况，每方飞机实体总共5架
        my_allplane_maps = {}
        for input_entity in my_allplane_infos:
            my_allplane_maps[int(input_entity['ID'])] = input_entity
        missile_infos = obs_side['missileinfos']  # 拿到空间中已发射且尚未爆炸的导弹信息
        enemy_missile_infos = [] #  用以保存敌方已发射且尚未爆炸的导弹信息
        my_missile_infos = [] #  用以保存敌方已发射且尚未爆炸的导弹信息
        for missile_info in missile_infos:
            if missile_info['LauncherID'] in my_allplane_maps:#  判断导弹是否为己方导弹 导弹的LauncherID即为导弹的发射者
                missile_info["Z"] = missile_info["Alt"]     # 导弹的 Alt 即为导弹的当前高度
                my_missile_infos.append(missile_info)    # 保存敌方已发射且尚未爆炸的导弹信息
            if (missile_info['ID'] != 0 and missile_info['Availability'] > 0.0001): # 判断导弹是否可用 导弹的ID即为导弹的唯一编号 导弹的Availability为导弹当前生命值
                missile_info["Z"] = missile_info["Alt"]     # 导弹的 Alt 即为导弹的当前高度
                enemy_missile_infos.append(missile_info)    # 保存敌方已发射且尚未爆炸的导弹信息
        self.my_uvas_infos = my_uvas_infos                  # 保存当前己方无人机信息
        self.my_manned_info = my_manned_info                # 保存当前己方有人机信息
        self.my_allplane_infos = my_allplane_infos          # 保存当前己方所有飞机信息
        self.enemy_uvas_infos = enemy_uvas_infos            # 保存当前敌方无人机信息
        self.enemy_manned_info = enemy_manned_info          # 保存当前敌方有人机信息
        self.enemy_allplane_infos = enemy_allplane_infos    # 保存当前敌方所有飞机信息
        self.enemy_missile_infos = enemy_missile_infos      # 保存敌方已发射且尚未爆炸的导弹信息
        self.my_missile_infos = my_missile_infos            # 保存敌方已发射且尚未爆炸的导弹信息
        # if not hasattr(self, 'mcv'):
        #     from .VISUALIZE.mcom import mcom
        #     self.mcv = mcom(ip='127.0.0.1',
        #                 port=12084,
        #                 path='./TMP/v2d_logger/',
        #                 digit=16, rapid_flush=True, draw_mode='pyqtgraph')
        #     self.mcv.v2d_init()
        # id_ = 0
        # for uav in my_uvas_infos:
        #     self.mcv.v2dx('rec|%d|r|2000'%(id_), uav['X'], uav['Y'], 0);id_ += 1
        # for uav in my_manned_info:
        #     self.mcv.v2dx('cir|%d|r|4000'%(id_), uav['X'], uav['Y'], 0);id_ += 1
        # for uav in enemy_uvas_infos:
        #     self.mcv.v2dx('rec|%d|b|2000'%(id_), uav['X'], uav['Y'], 0);id_ += 1
        # for uav in enemy_manned_info:
        #     self.mcv.v2dx('cir|%d|b|4000'%(id_), uav['X'], uav['Y'], 0);id_ += 1
        # self.mcv.v2d_show()
        # time.sleep(0.1)

    def init_pos(self, cmd_list):
        '''
            y     [-150000, +150000]
            red_x [-150000, -125000]
            blue_x [125000,  150000]
        '''
        leader_original_pos = {}    # 用以初始化当前方的位置
        if self.name == "red":
            leader_original_pos = {"X": -Special.init_X, "Y": -Special.init_Y, "Z": Special.init_Z}
        else :
            leader_original_pos = {"X":  Special.init_X, "Y":  Special.init_Y, "Z": Special.init_Z}

        interval_distance = 5000   # 间隔 5000米排列
        for leader in self.my_manned_info: # 为己方有人机设置初始位置
            # CmdEnv.make_entityinitinfo 指令，可以将 飞机实体置于指定位置点，参数依次为 实体ID，X坐标，Y坐标，Z坐标，初始速度，初始朝向
            cmd_list.append(
                CmdEnv.make_entityinitinfo(leader['ID'], leader_original_pos['X'], leader_original_pos['Y'],
                leader_original_pos['Z'], Special.init_speed_vip, 45))

        #己方无人机在有人机的y轴上分别以9500的间距进行部署
        sub_index = 0  # 编号 用以在有人机左右位置一次排序位置点
        for sub in self.my_uvas_infos: # 为己方每个无人机设置初始位置
            sub_pos = copy.deepcopy(leader_original_pos)  # 深拷贝有人机的位置点
            if sub_index & 1 == 0: # 将当前编号放在有人机的一侧
                # CmdEnv.make_entityinitinfo 指令，可以将 飞机实体置于指定位置点，参数依次为 实体ID，X坐标，Y坐标，Z坐标，初始速度，初始朝向
                cmd_list.append(
                    CmdEnv.make_entityinitinfo(sub['ID'],
                    sub_pos['X'], sub_pos['Y'] + interval_distance,
                    sub_pos['Z'], Special.init_speed_drone, 45))
            else:                   # 将当前编号放在有人机的另一侧
                # CmdEnv.make_entityinitinfo 指令，可以将 飞机实体置于指定位置点，参数依次为 实体ID，X坐标，Y坐标，Z坐标，初始速度，初始朝向
                cmd_list.append(
                    CmdEnv.make_entityinitinfo(sub['ID'],
                    sub_pos['X'], sub_pos['Y'] - interval_distance,
                    sub_pos['Z'], Special.init_speed_drone, 45))
                interval_distance *= 2 # 编号翻倍
            sub_index += 1 # 编号自增