# importing packages
import math
import numpy as np
import pandas as pd
import scipy.stats as stats
import sklearn
import imblearn
import matplotlib.pyplot as plt
import seaborn as sns
plt.style.use('ggplot')

from imblearn.over_sampling import RandomOverSampler
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import RandomForestClassifier
from sklearn.svm import LinearSVC
import xgboost as xgb
from sklearn import metrics
from imblearn.over_sampling import SMOTE
from imblearn.over_sampling import ADASYN
import math
from sklearn.model_selection import GridSearchCV
from sklearn.ensemble import RandomForestClassifier
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.model_selection import KFold
from sklearn.model_selection import cross_val_score
from sklearn.utils import shuffle
from sklearn.preprocessing import MinMaxScaler
from sklearn.calibration import CalibratedClassifierCV
from collections import Counter

class modelpipeline:
    def __init__(self):
        pass

    def run_model(self, df, varlist, response, standardize, sampletype, modelname, text, n_fold):
        # Remove any features not wanted based on varlist input and re-order based on varlist
        df = df[varlist]
        # We have to remove response from varlist - varlist_noresponse - as it is used later to subset out features
        # Refer to the for loop for the cross validation where X_train and X_test is created at the end of loop
        varlist_noresponse = []
        for col in varlist:
            if col != response:
                varlist_noresponse.append(col)

        if isinstance(n_fold, int) and n_fold > 1:
            # Initialize dictionary to store results
            self.store = {"accuracy": [], "actual_accuracy": [], "sensitivity": [], "specificity": [],
                          "precision": [], "f1": [], "auc": [], "pr_auc": [], "final": {}}

            # Split dataframes into 2, one for positive response and one for negative
            df_zero = df[df[response] == 0]
            df_one = df[df[response] == 1]

            # Shuffle dataframe for response=0 and =1 so that train-test will not be biased in case rows that are similar are placed side by side
            # Later on, we will reset the index and select by the index number by sections
            df_zero = shuffle(df_zero, random_state=42)
            df_one = shuffle(df_one, random_state=42)
            df_zero = df_zero.reset_index(drop=True)
            df_one = df_one.reset_index(drop=True)

            # Get the average number of records required for negative response and positive response for test records
            # Train records will then have all the other records not in the test records
            # n_fold is the number of folds for cross validation
            start_index_one = 0
            end_index_one = math.floor(df_one.shape[0]/n_fold)
            start_index_zero = 0
            end_index_zero = math.floor(df_zero.shape[0]/n_fold)

            for i in range(1,n_fold+1):
                if i != n_fold:
                    print('Getting TEST DF for response 1 from index ' + str(start_index_one) + ' to ' + str(end_index_one))
                    df_one_test = df_one.iloc[start_index_one:end_index_one]
                    print('Getting TRAIN DF for response 1 from index 0 to ' + str(start_index_one) + ' and from index ' + str(end_index_one) + ' to ' + str(df_one.shape[0]))
                    df_one_train = pd.concat([df_one.iloc[0:start_index_one],df_one.iloc[end_index_one:]], axis=0)
                    start_index_one += math.floor(df_one.shape[0]/n_fold)
                    end_index_one += math.floor(df_one.shape[0]/n_fold)

                    print('Getting TEST DF for response 0 from index ' + str(start_index_zero) + ' to ' + str(end_index_zero))
                    df_zero_test = df_zero.iloc[start_index_zero:end_index_zero]
                    print('Getting TRAIN DF for response 0 from index 0 to ' + str(start_index_zero) + ' and from index ' + str(end_index_zero) + ' to ' + str(df_zero.shape[0]))
                    df_zero_train = pd.concat([df_zero.iloc[0:start_index_zero],df_zero.iloc[end_index_zero:]], axis=0)
                    start_index_zero += math.floor(df_zero.shape[0]/n_fold)
                    end_index_zero += math.floor(df_zero.shape[0]/n_fold)

                else:
                    # Last section of split needs to reach the end of dataset
                    print('Getting TEST DF for response 1 from index ' + str(start_index_one) + ' to ' + str(df_one.shape[0]))
                    df_one_test = df_one.iloc[start_index_one:df_one.shape[0]]
                    print('Getting TRAIN DF for response 1 from index 0 to ' + str(start_index_one))
                    df_one_train = df_one.iloc[0:start_index_one]

                    # Last section of split needs to reach the end of dataset
                    print('Getting TEST DF for response 0 from index ' + str(start_index_zero) + ' to ' + str(df_zero.shape[0]))
                    df_zero_test = df_zero.iloc[start_index_zero:df_zero.shape[0]]
                    print('Getting TRAIN DF for response 0 from index 0 to ' + str(start_index_zero))
                    df_zero_train = df_zero.iloc[0:start_index_zero]

                # Combine the subsetted sections for negatives and postives for both train and test before oversampling
                df_train = pd.concat([df_one_train, df_zero_train], axis=0)
                df_test = pd.concat([df_one_test, df_zero_test], axis=0)
                # varlist_noresponse has the feature list X without Y while response is the Y
                # print(varlist_noresponse)
                X_train = df_train[varlist_noresponse]
                # print('Check X train vars after combining pds')
                # print(X_train.columns.values)
                y_train = df_train[response]
                X_test = df_test[varlist_noresponse]
                y_test = df_test[response]

                if standardize == True:
                    scaling = MinMaxScaler(feature_range=(-1,1)).fit(X_train)
                    X_train = scaling.transform(X_train)
                    X_test = scaling.transform(X_test)
                    X_train = pd.DataFrame(X_train, columns=varlist_noresponse)
                    X_test = pd.DataFrame(X_test, columns=varlist_noresponse)

                if sampletype == 'smote':
                    X_train, X_test, y_train, y_test = sampling.smote_oversample(X_train, X_test, y_train, y_test, response)
                elif sampletype == 'adasyn':
                    X_train, X_test, y_train, y_test = sampling.adasyn_oversample(X_train, X_test, y_train, y_test, response)
                elif sampletype == 'naive':
                    X_train, X_test, y_train, y_test = sampling.naive_oversample(X_train, X_test, y_train, y_test, response)
                else:
                    # Convert all DF to numpy array for model building later
                    X_train = X_train.values
                    y_train = y_train.values
                    X_test = X_test.values
                    y_test = y_test.values

                # Build model in current fold/iteration and get accuracy, sensitivity, specificity, precision, f1, auc
                self.store = self.build_model(X_train, X_test, y_train, y_test, text, modelname, i, n_fold, self.store)

                # test model with all actual fraud results
                if standardize == True:
                    df_acc = pd.concat([pd.DataFrame(scaling.transform(df[varlist_noresponse]),columns=varlist_noresponse),df[response]],axis=1)
                    # print(df)
                    self.store['actual_accuracy'].append(evaluate.actual_acc(df_acc, self.store['model'], response))
                else:
                    self.store['actual_accuracy'].append(evaluate.actual_acc(df, self.store['model'], response))

            # Before results are returned, get average of all evaluation metrics and store in store['final'] section
            self.store['final']['accuracy'] = self.avg(self.store['accuracy'])
            self.store['final']['sensitivity'] = self.avg(self.store['sensitivity'])
            self.store['final']['specificity'] = self.avg(self.store['specificity'])
            self.store['final']['precision'] = self.avg(self.store['precision'])
            self.store['final']['f1'] = self.avg(self.store['f1'])
            self.store['final']['auc'] = self.avg(self.store['auc'])
            self.store['final']['pr_auc'] = self.avg(self.store['pr_auc'])
            self.store['final']['actual_accuracy'] = self.avg(self.store['actual_accuracy'])

            print('Final Results of ' + str(n_fold) + ' fold CV:')
            print(self.store['final'])
            return self.store

        else:
            print('n fold must be an integer greater than 1')
            return self.store

    def build_model(self, X_train, X_test, y_train, y_test, text, modelname, i, n_fold, store):
        if modelname == 'LogisticRegression':
            model = LogisticRegression(max_iter=300, C=0.8, solver='liblinear')
            model.fit(X_train,y_train)
        elif modelname == 'XGBoost':
            model = xgb.XGBClassifier(seed=42, nthread=1, max_depth=math.ceil(math.sqrt(X_train.shape[1])),
                                      n_estimators=100, random_state=42)
            model.fit(X_train, y_train, eval_set=[(X_train, y_train), (X_test, y_test)], verbose=5)
        elif modelname == 'XGBoostminus1':
            # XGBoost with one less depth
            model = xgb.XGBClassifier(seed=42, nthread=1, max_depth=math.ceil(math.sqrt(X_train.shape[1])-1),
                                      n_estimators=100, random_state=42)
            model.fit(X_train, y_train, eval_set=[(X_train, y_train), (X_test, y_test)], verbose=5)
        elif modelname == 'XGBoostplus1':
            # XGBoost with one more depth
            model = xgb.XGBClassifier(seed=42, nthread=1, max_depth=math.ceil(math.sqrt(X_train.shape[1]))+1,
                                      n_estimators=100, random_state=42)
            model.fit(X_train, y_train, eval_set=[(X_train, y_train), (X_test, y_test)], verbose=5)
        elif modelname == 'XGBoostplus3':
            # XGBoost with 3 more depth
            model = xgb.XGBClassifier(seed=42, nthread=1, max_depth=math.ceil(math.sqrt(X_train.shape[1]))+3,
                                      n_estimators=100, random_state=42)
            model.fit(X_train, y_train, eval_set=[(X_train, y_train), (X_test, y_test)], verbose=5)

        # Use Linear SVC instead of sklearn svm.SVC as the former as way faster processing speed
        # However, LinearSVC does not have .predict_proba function to get probability of response 1
        # Hence, we need to use CalibratedClassifier that provides .predict_proba functionality
        # On the bad side, it has it's own CV, so we put 10 fold CV to minimize the dataset loss due to train-test fold
        # Ideally, we should use the older pipeline code that does not do customized k-fold CV (refer to SCI16. Jupyter Notebook)
        elif modelname == 'SVM_Linear':
            model = LinearSVC(random_state=42) # default C=1 (regularization parameter)
            model = CalibratedClassifierCV(model, cv=10)
            model.fit(X_train,y_train)
        elif modelname == 'SVM_Linear2':
            model = LinearSVC(random_state=42, C=2)
            model = CalibratedClassifierCV(model, cv=10)
            model.fit(X_train,y_train)
        elif modelname == 'SVM_Linear0.5':
            model = LinearSVC(random_state=42, C=0.5)
            model = CalibratedClassifierCV(model, cv=10)
            model.fit(X_train,y_train)
        elif modelname == 'SVM_Linear0.3':
            model = LinearSVC(random_state=42, C=0.3)
            model = CalibratedClassifierCV(model, cv=10)
            model.fit(X_train,y_train)
        elif modelname == 'RandomForest':
            treedepth = math.ceil(math.sqrt(X_train.shape[1]))
            model = RandomForestClassifier(random_state=42, max_depth=treedepth, n_estimators=100)
            model.fit(X_train,y_train)
        elif modelname == 'RandomForestminus1':
            treedepth = math.ceil(math.sqrt(X_train.shape[1]))-1
            model = RandomForestClassifier(random_state=42, max_depth=treedepth, n_estimators=100)
            model.fit(X_train,y_train)
        elif modelname == 'RandomForestminus2':
            treedepth = math.ceil(math.sqrt(X_train.shape[1]))-2
            model = RandomForestClassifier(random_state=42, max_depth=treedepth, n_estimators=100)
            model.fit(X_train,y_train)
        elif modelname == 'RandomForestplus2':
            treedepth = math.ceil(math.sqrt(X_train.shape[1]))+2
            model = RandomForestClassifier(random_state=42, max_depth=treedepth, n_estimators=100)
            model.fit(X_train,y_train)
        elif modelname == 'RandomForestplus4':
            treedepth = math.ceil(math.sqrt(X_train.shape[1]))+4
            model = RandomForestClassifier(random_state=42, max_depth=treedepth, n_estimators=100)
            model.fit(X_train,y_train)
        else:
            # Parameters based on gridsearchcv of modelname = logistic regresion
            # Leave parameter blank for modelname to run this instance of logistic regression
            model = LogisticRegression(C=0.8, max_iter=300, solver='liblinear')
            model.fit(X_train,y_train)

        y_predict = model.predict(X_test)
        y_predictprob = model.predict_proba(X_test)[:, 1]
        store = evaluate.model_results(y_test, y_predict, y_predictprob, text, store, i, n_fold)

        # Store model for usage in measuring actual accuracy of fraud cases
        store['model'] = model
        print("Iteration " + str(i) + " out of " + str(n_fold) + " of CV for model fitting and obtaining results is complete!")
        print("\n")
        return store

    def avg(self, array):
        return sum(array) / len(array)

class sampling:
    def __init__(self):
        pass
    @staticmethod
    def naive_oversample(X_train, X_test, y_train, y_test, response):
        ros = RandomOverSampler(random_state=42)
        X_train, y_train = ros.fit_resample(X_train, y_train)
        # train test split keeps X_test and y_test as pd series, oversampler converts X_train, y_train to numpy
        # Convert all to numpy array for XGBoost to not have bugs
        X_test = X_test.values
        y_test = y_test.values
        print("Oversampling is complete!")
        return X_train, X_test, y_train, y_test

    @staticmethod
    def smote_oversample(X_train, X_test, y_train, y_test, response):
        X_train, y_train = SMOTE().fit_resample(X_train, y_train)
        # train test split keeps X_test and y_test as pd series, oversampler converts X_train, y_train to numpy
        # Convert all to numpy array for XGBoost to not have bugs
        X_test = X_test.values
        y_test = y_test.values
        print("Number of Xs and Ys for SMOTE:")
        print(sorted(Counter(y_train).items()))
        print("Oversampling is complete!")
        return X_train, X_test, y_train, y_test

    @staticmethod
    def adasyn_oversample(X_train, X_test, y_train, y_test, response):
        X_train, y_train = ADASYN().fit_resample(X_train, y_train)
        # train test split keeps X_test and y_test as pd series, oversampler converts X_train, y_train to numpy
        # Convert all to numpy array for XGBoost to not have bugs
        X_test = X_test.values
        y_test = y_test.values
        print("Number of Xs and Ys for ADASYN:")
        print(sorted(Counter(y_train).items()))
        print("Oversampling is complete!")
        return X_train, X_test, y_train, y_test

class evaluate:
    def __init__(self):
        pass

    @staticmethod
    def model_results(y_test, y_predict, y_predictprob, text, store, i, n_fold):
        cm = metrics.confusion_matrix(y_test, y_predict)
        print(cm)
        RFC_CM = pd.DataFrame(cm, ['Actual 0', 'Actual 1'], ['Predict 0', 'Predict 1'])
        sns.heatmap(RFC_CM, annot=True, annot_kws={"size": 16}, cmap='Greens', linewidths=1, fmt='g')# font size
        sns.set(font_scale=1.4)#for label size
        plt.title("Confusion Matrix for " + text)

        # fix for mpl bug that cuts off top/bottom of seaborn viz
        b, t = plt.ylim()
        b += 0.5
        t -= 0.5
        plt.ylim(b, t)
        plt.figure(1,figsize=(4,4))
        plt.show()

        accuracy = metrics.accuracy_score(y_test, y_predict)
        # print('Accuracy: ' + str(accuracy))
        sensitivity = cm[1][1] / (cm[1][1] + cm[1][0])
        recall = sensitivity
        # print('Sensitivity: ' + str(sensitivity))
        specificity = cm[0][0] / (cm[0][0] + cm[0][1])
        # print('Specificity: ' + str(specificity))
        precision = cm[1][1] / (cm[1][1] + cm[0][1])
        # print('Precision: ' + str(precision))
        f1 = 2 * (recall * precision)/(recall + precision)
        # print('f1 score: ' + str(f1))
        auc, pr_auc = evaluate.ROC(y_test, y_predictprob, text, i, n_fold)

        store['accuracy'].append(accuracy)
        store['sensitivity'].append(sensitivity)
        store['specificity'].append(specificity)
        store['precision'].append(precision)
        store['f1'].append(f1)
        store['auc'].append(auc)
        store['pr_auc'].append(pr_auc)

        return store

#     @staticmethod
#     def ROC(y_test, y_predictprob, text):
#         # IMPORTANT: first argument is true values, second argument is predicted probabilities
#         auc = metrics.roc_auc_score(y_test, y_predictprob)
#         # print("AUC value is: " + str(auc))
#         fpr, tpr, thresholds = metrics.roc_curve(y_test, y_predictprob)
#         # print("AUC value is also: " + str(metrics.auc(fpr, tpr)))
#         plt.plot(fpr, tpr)
#         plt.xlim([0.0, 1.0])
#         plt.ylim([0.0, 1.0])
#         plt.title('ROC curve for ' + text)
#         plt.xlabel('False Positive Rate (1 - Specificity)')
#         plt.ylabel('True Positive Rate (Sensitivity)')
#         plt.grid(True)
#         return auc

    @staticmethod
    def ROC(y_test, y_predictprob, text, i, n_fold):
        # IMPORTANT: first argument is true values, second argument is predicted probabilities
        auc = metrics.roc_auc_score(y_test, y_predictprob)
        # print("AUC value is: " + str(auc))
        print("AUC value is: " + str(auc))
        fpr, tpr, thresholds = metrics.roc_curve(y_test, y_predictprob)
        # print("AUC value is also: " + str(metrics.auc(fpr, tpr)))
        # Calculate precision and recall for each threshold
        precision, recall, _ = metrics.precision_recall_curve(y_test, y_predictprob)
        pr_auc = metrics.auc(recall, precision)
        # Only show ROC-AUC graph and PR-AUC graph on last iteration as they look very similar
        # The full results can be obtained in the results section
        if n_fold == i:
            fullgraph = plt.figure(1,figsize=(10,20))
            plt.style.use('ggplot')
            ROCAUC_plot = fullgraph.add_subplot(211)
            ROCAUC_plot.plot(fpr, tpr, color='blue')
            ROCAUC_plot.set_title('ROC curve for ' + text)
            ROCAUC_plot.set_xlabel('False Positive Rate (1 - Specificity)')
            ROCAUC_plot.set_ylabel('True Positive Rate (Sensitivity)')
            ROCAUC_plot.set_xlim([0.0, 1.0])
            ROCAUC_plot.set_ylim([0.0, 1.0])
            ROCAUC_plot.grid(True)
            PRAUC_plot = fullgraph.add_subplot(212)
            PRAUC_plot.plot(precision, recall, color='purple')
            PRAUC_plot.set_title('Precision-Recall curve for ' + text)
            PRAUC_plot.set_xlabel('Recall')
            PRAUC_plot.set_ylabel('Precision')
            PRAUC_plot.set_xlim([0.0, 1.0])
            PRAUC_plot.set_ylim([0.0, 1.0])
            PRAUC_plot.grid(True)
        return auc, pr_auc

    @staticmethod
    def actual_acc(df, model, response):
        allpositive = df[df[response] == 1]
        x_positive = allpositive.drop([response], axis=1)
        y_positive = allpositive[response]
        # Convert to numpy array due to XGBoost model.predict not working well for pandas
        x_positive = x_positive.values
        y_positive = y_positive.values
        y_pospredict = model.predict(x_positive)
        accuracy_positive = metrics.accuracy_score(y_positive, y_pospredict)
        # print("Accuracy with all fraud results is " + str(accuracy_positive * 100) + "%")
        return accuracy_positive


import numpy as np

import torch

from scipy.stats import wasserstein_distance

from generate.metrics import ca_metrics, amino_acid_metrics, secondary_sequence_metrics

class MMD:
    """
    Compute the Maximum Mean Discrepancy (MMD) between the predicted graph and a set of target ones.

    Graph statistics include amino acid sequence information and secondary structure sequence information.

    Source
    ------
    https://torchdrift.org/notebooks/note_on_mmd.html

    Attributes
    ----------
    pred_graph : (torch.Tensor, torch.Tensor)
        The predicted graph consisting of the predicted node features and the predicted distance matrix.
    target_graphs : list of (torch.Tensor, torch.Tensor)
        The target graphs, where each entry consists of the node features and the distance matrix.
    median_subset : int
        The size of the subset to use to compute `sigma`.
    """

    def __init__(self, pred_graph, target_graphs, median_subset=100):
        """
        Initialize the class.

        Parameters
        ----------
        pred_graph : (torch.Tensor, torch.Tensor)
            The predicted graph consisting of the predicted node features and the predicted distance matrix.
        target_graphs : list of (torch.Tensor, torch.Tensor)
            The target graphs, where each entry consists of the node features and the distance matrix.
        median_subset : int, optional
            The size of the subset to use to compute `sigma`. The default is 100.
        """
        self.pred_graph = pred_graph
        self.target_graphs = target_graphs
        self.median_subset = median_subset

    def _wasserstein_kernel(self, x, y, sigma):
        x_flat, y_flat = x.detach().view(-1), y.detach().view(-1)
        return torch.exp(wasserstein_distance(x_flat, y_flat)/2*(sigma**2))

    def _mmd(self, x, y):
        # Get number of samples
        n, m = x.shape[0], y.shape[0]
        # Compute sigma
        dists = torch.pdist(torch.cat([x.detach(), y.detach()], dim=0)[:,None])
        sigma = dists[:self.median_subset].median()/2
        # Compute the mmd
        xx, yy, xy = 0, 0, 0
        for i in range(n):
            for j in range(n):
                xx += self._wasserstein_kernel(x[i], x[j], sigma)
        for i in range(m):
            for j in range(m):
                yy += self._wasserstein_kernel(y[i], y[j], sigma)
        for i in range(n):
            for j in range(m):
                xy += self._wasserstein_kernel(x[i], y[j], sigma)
        return xx + yy - 2*xy

    def compare_graphs(self):
        """
        Compare the graphs.

        Returns
        -------
        scores : numpy.ndarray
            The comparison scores.
        """
        scores = []
        # Compute MMD score of the predicted graph
        x_aa = amino_acid_metrics(self.pred_graph[0])
        x_ss = secondary_sequence_metrics(self.pred_graph[0])
        x = torch.cat((x_aa,x_ss))
        for target_graph in self.target_graphs:
            y_aa = amino_acid_metrics(target_graph[0])
            y_ss = secondary_sequence_metrics(target_graph[0])
            y = torch.cat((y_aa,y_ss))
            # Compare
            scores.append(self._mmd(x, y))
        scores = np.array(scores)
        return scores

class QCP:
    """
    Compute the superimposition RMSD of the predicted atoms against validation ones.

    The idea of this control is that if the overall RMSD is low, then the generated protein is likely to be
    realistic.

    Attributes
    ----------
    pred_coords : numpy.ndarray
        The coordinates of the generated C-alpha atoms.
    target_proteins : list of Bio.PDB.Structure
        The list of target proteins.
    """

    def __init__(self, pred_coords, target_proteins):
        """
        Initialize the class.

        Parameters
        ----------
        pred_coords : torch.tensor
            The coordinates of the generated C-alpha atoms. The input should be a PyTorch tensor, which it then
            converted to a Numpy array.
        target_proteins : list of Bio.PDB.Structure
            The list of target proteins.
        """
        self.pred_coords = pred_coords
        self.target_proteins = target_proteins

    def _get_target_ca_atoms_coords(self):
        target_coords_list = []
        for protein in self.target_proteins:
            coords = []
            for model in protein:
                for chain in model:
                    for residue in chain:
                        if 'CA' in residue:
                            coords.append(residue['CA'].get_coord())
            target_coords_list.append(np.array(coords))
        return target_coords_list

    def _get_shifted_coords(self, coords_long, coords_short):
        shifted_coords, i = [], 0
        while i + len(coords_short) <= len(coords_long):
            shifted_coords.append(coords_long[i:i+len(coords_short)])
            i += 1
        return shifted_coords

    def superimpose(self):
        """
        Compute the superimposition.

        Returns
        -------
        numpy.ndarray
            The array of RMSD scores (each i-th entry corresponds to the comparison of the generated atoms with the
            i-th protein)
        """
        scores = []
        target_coords_list = self._get_target_ca_atoms_coords()
        for target_coords in target_coords_list:
            if len(target_coords) == len(self.pred_coords):
                scores.append(ca_metrics(target_coords, self.pred_coords))
            elif len(target_coords) > len(self.pred_coords):
                shifted_coords = self._get_shifted_coords(target_coords, self.pred_coords)
                rmsds = []
                for scs in shifted_coords:
                    rmsds.append(ca_metrics(scs, self.pred_coords))
                scores.append(np.min(np.array(rmsds)))
            else:
                shifted_coords = self._get_shifted_coords(self.pred_coords, target_coords)
                rmsds = []
                for scs in shifted_coords:
                    rmsds.append(ca_metrics(target_coords, scs))
                scores.append(np.min(np.array(rmsds)))
        scores = np.array(scores)
        return scores


import itertools
from functools import partial
import numpy as np
import matplotlib.pyplot as plt
import scipy.stats

# --- Preliminaries ---

# - Partitioning techniques -
class BestSolutionHolder:
    """
        Keeps in memory the best solution of a partitioning technique,
        is used in the function below.
    """
    def __init__(self, crit, n_pos, n_neg, frontier, orient,
                 left_n_right=False, n_pos_tot=None, n_neg_tot=None):
        self.crit = crit
        self.left_n_right = left_n_right
        if left_n_right:
            if n_pos_tot is None or n_neg_tot is None:
                raise ValueError("Please indicate the total number of positives or negatives.")
            self.n_pos_tot = n_pos_tot
            self.n_neg_tot = n_neg_tot

        self.best_val = None
        self.best_front = None
        self.best_orient = None

        self.calls = 0

        self.save_sol(n_pos, n_neg, frontier, orient)

    def save_sol(self, n_pos, n_neg, frontier, orient):
        """Saves a proposition as solution if it is better than before."""
        self.calls += 1

        val = self.crit(n_pos, n_neg)

        if self.best_val is None or val > self.best_val:
            self.best_val = val
            self.best_front = frontier
            self.best_orient = orient

        if self.left_n_right:
            val = self.crit(self.n_pos_tot - n_pos, self.n_neg_tot - n_neg)

            if self.best_val is None or val > self.best_val:
                self.best_val = val
                self.best_front = frontier
                self.best_orient = not orient

def bipart_partition(X, Y, criterion, epsilon=0.001):
    """
    Takes values in R associated to a class in {-1, +1}, a criterion that depends solely
    on the number of positive or negative values in a partition and finds the partition
    that maximizes the criterion.
    """
    # crit = lambda np, nn : np - nn
    n = X.shape[0]
    n_pos_tot = (Y == +1).sum()
    n_neg_tot = n - n_pos_tot
    XY = np.vstack([X, Y]).transpose()
    sorted_XY = XY[np.argsort(XY[:, 0])]
    n_pos, n_neg = 0, 0
    best_value = BestSolutionHolder(criterion, n_pos, n_neg, sorted_XY[0, 0]-epsilon, False,
                                    left_n_right=True, n_pos_tot=n_pos_tot, n_neg_tot=n_neg_tot)
    i = 0
    for xy, xyp in zip(sorted_XY, sorted_XY[1:]):
        x, y = xy
        xp, yp = xyp
        frontier = (x + xp)/2
        # assumed the density of X continuous hence xp > x
        n_pos += int(y == +1)
        n_neg += int(y == -1)
        best_value.save_sol(n_pos, n_neg, frontier, False)
        i += 1

    frontier = xp + epsilon
    n_pos += int(yp == +1)
    n_neg += int(yp == -1)
    best_value.save_sol(n_pos, n_neg, frontier, False)
    return best_value.best_front, best_value.best_val, best_value.best_orient

# - Plotting utils -

def plot_eta(a_collection, param_fun, n_points=100):
    """
    Plots a family of functions defined on [0,1] indexed by a
    for a set of possible values for a.
    """
    plt.figure(figsize=(4, 4))

    x_ax = np.linspace(0., 1., n_points)

    for a in a_collection:
        eta_fun = partial(param_fun, a=a)
        plt.plot(x_ax, list(map(eta_fun, x_ax)),
                 label="$a = {:0.2f}$".format(a))

    plt.title(r"The function $\eta_a$ for different a")
    plt.xlabel("x")
    plt.ylabel(r"$\eta_a$(x)")
    plt.grid()
    plt.legend()
    plt.show()

def plot_distribution(a_collection, gen_datafun):
    """
    Plots the distributions induced by posterior probability functions defined on [0,1]
    indexed by a for a set of possible values for a.
    """
    plt.figure(figsize=(8, 10))

    for i, a in enumerate(a_collection):
        plt.subplot(5, 2, i + 1)

        X, Y = gen_datafun(n=500, a=a)
        plt.grid()
        plt.hist(X[Y == -1], bins=30, color="red", normed=True, alpha=0.5)
        plt.hist(X[Y == 1], bins=30, color="green", normed=True, alpha=0.5)
        plt.title("Sample distribution, $a = {:2.2f}$".format(a))
    plt.tight_layout()
    plt.show()

def GetSlope(df_int, quant_val=0.9):
    """
    Gets the slope corresponding to the log of a quantile of
    the generalization error regressed by the logarithm of n, number of data points.
    Parameters:
    * df_int: Dataframe containing the columns "n" and "gen_error"
        with enough entries for each value of "n".
    * quant_val: quantile value that we choose.
    Returns:
    * constant, slope: values in R.
    """
    groupby_quant = df_int.groupby("n").quantile(quant_val)
    vals_med = groupby_quant["gen_error"].values
    ns = df_int["n"].unique()
    reg = scipy.stats.linregress(np.log(ns), np.log(vals_med))
    return np.exp(reg.intercept), reg.slope

def boxplot_slopes(df, quant=0.5, ylim=[10**(-4), 10**(0)]):
    """
    Does a boxplot for each a of the results, to show the different generalization speeds.
    Parameters:
    * df: Dataframe containing the columns "a", "n" and "gen_error"
        with enough entries for each values of "n","a".
    * quant_val: quantile value that we choose to regress the empirical generalization speed.
    """
    plt.figure(figsize=(20, 12))

    x_n_plots = len(df["a"].unique())//2 + 1
    for i, a in enumerate(df["a"].unique()):
        df_int = df[df["a"] == a]
        ns = df_int["n"].unique()
        bxplt_vals = [df_int[df_int["n"] == n_val]["gen_error"].values
                      for n_val in df_int["n"].unique()]

        correct_width = [5] + list(map(lambda x: x / 4, ns[1:]))
        plt.subplot(x_n_plots, 4, i + 1)
        plt.boxplot(bxplt_vals, positions=ns,
                    widths=correct_width)
        plt.xscale("log")
        plt.yscale("log")

        reg_med_const, reg_med_slope = GetSlope(df_int, quant_val=quant)

        plt.plot(ns, reg_med_const * np.power(ns, reg_med_slope), '-b')

        theo_slope = -1. / (2. - a)
        plt.plot(ns, 10**(theo_slope * np.log10(ns)), '-g')

        plt.title("Experiment $a = {:1.1f}$\nSlope med (blue): ${:0.2f}$\n".format(
            a, reg_med_slope) + "Theoretical bound slope (green) : ${:0.2f}$".format(
                theo_slope))
        plt.ylim(ylim)

        plt.xlabel("$n$")
        plt.ylabel("Regret")
        plt.grid()

    plt.tight_layout()

    plt.show()

def slopes_quants(df, quant_values=np.linspace(0.7, 0.9, 5)):
    """
    Compares the theoretical generalization slope and the empirical generalization slope.
    Parameters:
    * df: Dataframe containing the columns "a", "n" and "gen_error"
        with enough entries for each values of "n","a".
    * quant_values: quantile values that we choose to regress the empirical generalization speed.
    """
    plt.figure(figsize=(8, 15))
    a_collection = df["a"].unique()

    theo_slopes = [-1. / (2. - a) for a in a_collection]
    emp_vals = []

    for i, quant in enumerate(quant_values):
        reg_med_slopes = list()
        for a in a_collection:
            df_int = df[df["a"] == a]
            reg_med_slopes.append(GetSlope(df_int, quant_val=quant)[1])
        emp_vals.append(reg_med_slopes)
        plt.subplot(5, 2, i + 1)
        plt.plot(theo_slopes, reg_med_slopes, "bo")
        plt.plot([-1, -0.5], [-1, -0.5], color="red")
        plt.xlabel("Theoretical slopes")
        plt.ylabel("Experimental slopes\n(quantile at ${:2.2f}$)".format(quant))
        plt.grid()

    plt.tight_layout()

    plt.show()

def grid_count(grid, X):
    """Counts the number of elements on a grid."""
    del_elems = X < grid[0]
    remaining_X = X[~del_elems]
    res = list()
    for thre in grid[1:]:
        del_elems = remaining_X < thre
        remaining_X = X[~del_elems]
        res.append(del_elems.sum())
    return res

def plot_select_eta(a_collection, eta_fun_of_a, datagen_fun_of_a, n=1000,
                    n_points=100, n_bins=20):
    """
        Plots the empirical distribution of the data, as well as the theoretical
        distribution of it.
    """
    x_ax = np.linspace(0., 1., n_points)

    grid = np.linspace(0., 1., n_bins+1)
    sizebin = grid[1]-grid[0]
    plt.figure(figsize=(12, 6))
    for i, a in enumerate(a_collection):
        ax1 = plt.subplot(1, 2, i + 1)

        eta_fun = lambda x, a=a: 2*eta_fun_of_a(x, a=a)

        X, Y = datagen_fun_of_a(n=n, a=a)
        gridXpos = grid_count(grid, X[Y == +1])
        gridXneg = grid_count(grid, X[Y == -1])
        ax2 = ax1.twinx()
        ax2.bar(grid[0:n_bins], gridXpos, width=sizebin, align="edge",
                color="green", alpha=0.5, label="class 1")
        ax2.bar(grid[0:n_bins], gridXneg, width=sizebin, align="edge",
                bottom=gridXpos, color="red", alpha=0.5, label="class 2")
        ax1.plot(x_ax, list(map(eta_fun, x_ax)),
                 label=r"$\mu_1(x)$", color="black", linewidth=3)

        lines, labels = ax1.get_legend_handles_labels()
        lines2, labels2 = ax2.get_legend_handles_labels()
        ax2.legend(lines + lines2, labels + labels2, loc="lower right") # "upper left") #
        if i == 0:
            ax1.set_ylabel(r"$\mu_1(x)$")
        else:
            ax2.set_ylabel(r"$P_n(X \in $ bin $  \; | \; Y = i)$, $i \in \{1,2\}$",
                           labelpad=20)
        plt.title("$a = " + str(a)+"$")
        # if i == 1:
        ax1.set_xlabel("$x$")
    plt.tight_layout()

def plot_possible_alpha_values(min_max_alpha, a_collection, m, alpha, sup_inc_simrank=False):
    """
    Plot the possible alpha values.
    """
    plt.figure(figsize=(4, 3))
    zipped_int_alpha = list(map(lambda a: min_max_alpha(a, m), a_collection))
    inf_alphas = [v[0] for v in zipped_int_alpha]
    sup_alphas = [v[1] for v in zipped_int_alpha]
    plt.plot(a_collection, inf_alphas, label="min for $\\alpha$ ($C>0$)")
    plt.plot(a_collection, sup_alphas, label="max for $\\alpha$ ($C<1/2$)")
    plt.plot(a_collection, [alpha]*len(a_collection), label="chosen $\\alpha$ ")
    if sup_inc_simrank: # Only for the similarity ranking case.
        sup_inc = [0.5 - 0.5*np.power(np.abs(2*m-1), (1-a)/a) for a in a_collection]
        plt.plot(a_collection, sup_inc, label=r"$\eta$ increasing")
    plt.xlabel("$a$")
    plt.ylabel("$\\alpha$")
    plt.title("Limitations on the possible value for $\\alpha$")
    plt.legend()
    plt.grid()
    plt.show()

def plot_emp_mammen(gendata_fun_of_a, eta_fun_of_a, a_collection, n_obs=1000, simrank=False):
    """
    Tries and show visually what the Mammen-Tsybakov assumption means.
    """
    plt.figure(figsize=(3, 3))

    for a in a_collection:
        X, _ = gendata_fun_of_a(n=n_obs, a=a)

        if simrank:
            all_pairs = itertools.combinations(X, 2)
            all_vals = list(map(lambda x, a=a: abs(eta_fun_of_a(x[0], x[1], a=a)-0.5), all_pairs))
        else:
            all_vals = [np.abs(eta_fun_of_a(x, a=a) - 0.5) for x in X]

        plt.hist(all_vals, bins=50, alpha=0.5, cumulative=True,
                 normed=True, label="$a = {:0.2f}$".format(a))

    plt.title(r"The distribution of $|\eta-1/2|$ for different $a$")
    plt.xlabel(r"$|\eta-1/2|$")
    plt.ylabel(r"$P ( X \le t)$")
    plt.legend()
    plt.grid()
    plt.show()


<filename>phone_display_demo.py
"""
Copyright (C) 2011-2012 <NAME>

Permission is hereby granted, free of charge, to any person obtaining a copy of
this software and associated documentation files (the "Software"), to deal in
the Software without restriction, including without limitation the rights to
use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
of the Software, and to permit persons to whom the Software is furnished to do
so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.

Displays either day or nighttime traffic image processing in a mock-up UI
based on the HTC Desire smartphone.
"""
import numpy as np
import scipy
import scipy.ndimage as ndimg
from collections import deque
from copy import *
import PIL
import ImageOps
import pylab
import cv2
import os
import fnmatch
import sys
import pymorph
import night
import day
import argparse

class PhoneDemo(object):
    """ Object to run the phone UI demo. """

    TYPE_DAY = "DAY"
    TYPE_NIGHT = "NIGHT"
    HISTORY_FRAMES = 600

    class DayProcessor(object):
        """ Object used to process day sequences. """

        GMM_K = 3
        GMM_NUM_FRAMES = 25
        GMM_W_INIT = 0.1
        GMM_VAR_INIT = 20
        GMM_MAHA_THRESH = 3
        MASK_OVERLAY_ALPHA = 0.4

        def __init__(self, rgb):

            assert(rgb.dtype == 'uint8')

            self._gmm = day.GaussianMixtureModelUV(self.GMM_K, rgb.shape,
                                                   self.GMM_NUM_FRAMES,
                                                   self.GMM_W_INIT,
                                                   self.GMM_VAR_INIT,
                                                   self.GMM_MAHA_THRESH)
            self._ycbcr = np.zeros(rgb.shape, dtype='uint8')
            self._mask = np.zeros(rgb.shape[:2], dtype='uint8')
            self._red_mask = np.zeros(rgb.shape, dtype='uint8')
            self._rgb_red_masked = np.zeros(rgb.shape, dtype='uint8')
            self._process_count = 0

        def next(self, rgb):
            """ Process the next file and return the results. """

            # Do GMM steps.
            self._gmm.rgb2ycbcr(rgb, self._ycbcr)
            self._gmm.segment_cl(self._mask)
            self._gmm.update_cl(self._ycbcr)
            # Save total pixels in foreground.
            fg_pixel_count = np.sum(self._mask)
            # Pull alpha and render red overlay
            # (channels are reversed RGB = BGR).
            self._red_mask[:,:,2] = self._mask * 255
            self._rgb_red_masked[:,:] = \
                (self.MASK_OVERLAY_ALPHA * self._red_mask) + \
                ((1. - self.MASK_OVERLAY_ALPHA) * rgb)

            # Ignore the first GMM_NUM_FRAMES / 2 frames.
            self._process_count = self._process_count + 1
            if self._process_count > self.GMM_NUM_FRAMES / 2:
                return fg_pixel_count, self._rgb_red_masked
            else:
                return 0, self._rgb_red_masked

    class NightProcessor(object):
        """ Object used to process day sequences. """

        def __init__(self, rgb):

            pass

        def next(self, rgb):
            """ Process the next file and return the results. """

            def blackout_date_regions(image, blackout_rects):
                """ Black out specified regions. """

                for rect in blackout_rects:
                    image[rect[1]:rect[3], rect[0]:rect[2]] = 0

            # Do bright object detection.
            blackout_date_regions(rgb, night.BLACKOUT_RECTS)
            steps = night.bright_object_detection(rgb)
            # Return results (channels are reversed RGB = BGR).
            label_img = pymorph.overlay(steps['luminance'].astype('uint8'),
                                        blue=steps['detect_dilate'])
            return steps['bright_blob_count'], label_img

    def __init__(self):

      # Initialize plotting parameters.
      self._history_raw = deque()
      self._history_filtered = deque()
      self._max_sample = 0.001
      self._ui = PhoneDisplay()
      self._filter_exp = 0.1
      self._sample_exp_filter = 0.

    def run_sequence(self, type, filepath, seq_range=None, filter_exp=None):
        """ Run a TYPE_DAY or TYPE_NIGHT sequence. """

        QUIT_KEY_CODES = [ 27, 113, 81 ]
        PAUSE_KEY_CODES = [ 32, 112, 80 ]

        def pause():
            """ Poll input until the pause key is pressed. """

            while True:
                key = cv2.waitKey(100)
                if PAUSE_KEY_CODES.count(key) > 0:
                    break

        def bound_queue_push(val, q, maxlen=None):
            """ Push to bounded queue. """

            q.append(val)
            if maxlen is not None and len(q) > maxlen:
                q.popleft()

        assert(type == self.TYPE_DAY or type == self.TYPE_NIGHT)

        # TODO(reissb) -- The history frames and filtering need to become
        #   parameterized in some way. The history frames is fixed by the
        #   camera framerate. The filtering is fixed by the required
        #   detection sensitivity.
        if filter_exp is not None:
            self._filter_exp = filter_exp
        else:
            self._filter_exp = 0.1

        # Clear state.
        self._ui.clear()
        self._history_raw = deque()
        self._history_filtered = deque()
        self._max_sample = 0.001
        self._sample_exp_filter = 0.

        # Extract command-line parameters. This is the name of one file in the
        # series.
        path, filename = os.path.split(filepath)
        file_name, file_ext = os.path.splitext(os.path.basename(filename))
        series_name_end = file_name.rindex('_')
        series_name = file_name[:series_name_end]
        print "Processing image series {0} in path {1}.".format(series_name,
                                                                path)
        files_in_path = os.listdir(path)
        series_pattern = series_name + '_[0-9]*' + file_ext
        print "Processing files matching pattern {0}.".format(series_pattern)
        series_suffixes = [int(os.path.splitext(fn)[0].split('_')[-1]) \
                           for fn in files_in_path \
                           if fnmatch.fnmatch(fn, series_pattern)]
        series_suffixes.sort()
        num_files = len(series_suffixes)
        print "Found {0} files in image series {1}.".format(num_files,
                                                            series_name)
        # Check for limited range.
        if seq_range is not None:
            assert(seq_range[1] > seq_range[0] and seq_range[0] >= 0)
            print "Filtering series range [{},{}).".format(seq_range[0],
                                                           seq_range[1])
            series_suffixes = np.array(series_suffixes)
            f = (series_suffixes >= seq_range[0]) * \
                (series_suffixes < seq_range[1])
            series_suffixes = np.sort(series_suffixes * f)
            remove_count = len(series_suffixes) - np.sum(f)
            series_suffixes = np.delete(series_suffixes, range(remove_count))

        # Load first file and process.
        series_filename = series_name + '_' + str(series_suffixes[0]) + \
                          file_ext
        rgb = ndimg.imread(os.path.join(path, series_filename))
        # Initilaize the processor.
        type_processor = self.DayProcessor(rgb) if type is self.TYPE_DAY \
                         else self.NightProcessor(rgb)
        # Process the files.
        quit_flag = False
        process_count = 0
        history_n = int(self.HISTORY_FRAMES / \
                        (self._ui.history_frame_count - 1))
        for suffix in series_suffixes:
            # Process the next file.
            series_filename = series_name + '_' + str(suffix) + file_ext
            print "Processing file {0}.".format(series_filename)
            rgb = ndimg.imread(os.path.join(path, series_filename))
            sample_raw, display_img = type_processor.next(rgb)
            self._sample_exp_filter = \
                ((1. - self._filter_exp) * self._sample_exp_filter) + \
                (self._filter_exp * sample_raw)
            # Put new samples on queues.
            bound_queue_push(sample_raw,
                             self._history_raw, self.HISTORY_FRAMES)
            bound_queue_push(self._sample_exp_filter,
                             self._history_filtered, self.HISTORY_FRAMES)
            # Update UI.
            self._max_sample = max(self._max_sample,
                                   self._sample_exp_filter * 1.1)
            ybound = (0, self._max_sample)
            plot_img = self.plot_history(self._history_raw,
                                         self._history_filtered,
                                         ybound)
            self._ui.set_main_video_frame(display_img)
            self._ui.set_plot(plot_img)
            # Space history frames evenly over interval.
            if 0 == (process_count % history_n):
                self._ui.push_history_frame(display_img)
            process_count = process_count + 1
            # Show UI.
            cv2.imshow("Phone Display", self._ui.ui_image)
            key = cv2.waitKey(1)
            # Catch escape or 'q' or 'Q':
            if QUIT_KEY_CODES.count(key) > 0:
                quit_flag = True
                break
            # Catch spacebar, 'p' or 'P':
            if PAUSE_KEY_CODES.count(key) > 0:
                pause()

        # Cleanup GUI on complete.
        if not quit_flag:
            cv2.waitKey(-1)
        cv2.destroyAllWindows()

    @staticmethod
    def plot_history(raw, filtered, ybound):
        """ Make plot of raw and history and return as image. """

        p = pylab.subplot('111')
        p.clear()
        p.figure.set_size_inches(4, 3);
        p.plot(raw, '.r')
        p.plot(filtered, '-b')
        p.axes.set_ybound(ybound)
        p.figure.canvas.draw()
        buf = np.fromstring(p.figure.canvas.tostring_rgb(), dtype='uint8')
        h, w = p.figure.canvas.get_width_height()
        buf.shape = (w, h, 3)
        return buf

class PhoneDisplay(object):

    # Dictionary of UI resources.
    RESOURCES = {
        # Name of the phone UI image.
        "UI_BASE": "ui_base.jpg"
        }
    # Coordinates for the phone image display area.
    UI_LAYOUT = {
        "MARGIN": 15,
        "DISPLAY_RECT": { "UpperLeft": (255, 59), "LowerRight": (1106, 596) },
        "MAIN_VIDEO_SIZE": (360, 480),
        "PLOT_SIZE": (240, 320),
        "HISTORY_FRAME_COUNT": 5,
        "CLEAR_COLOR": np.array([60, 85, 45])
        }

    def __init__(self):
        """ Setup phone UI. """

        # Load UI base resource and set slice of display area.
        self._ui_base = ndimg.imread(self.RESOURCES["UI_BASE"])
        up_lt = self.UI_LAYOUT["DISPLAY_RECT"]["UpperLeft"]
        lw_rt = self.UI_LAYOUT["DISPLAY_RECT"]["LowerRight"]
        self._ui_display_area = self._ui_base[up_lt[1]:lw_rt[1] + 1,
                                              up_lt[0]:lw_rt[0] + 1]
        self._ui_display_area[:,:] = self.UI_LAYOUT["CLEAR_COLOR"]
        self._ui_display_area_size = self._ui_display_area.shape[:2]
        self._ui_main_video_size = self.UI_LAYOUT["MAIN_VIDEO_SIZE"]
        self._plot_size = self.UI_LAYOUT["PLOT_SIZE"]

        margin = self.UI_LAYOUT["MARGIN"]
        # Get main video frame area.
        vid_frm_x1 = self._ui_display_area_size[1] - margin
        vid_frm_x0 = vid_frm_x1 - self._ui_main_video_size[1]
        vid_frm_y0 = margin
        vid_frm_y1 = vid_frm_y0 + self._ui_main_video_size[0]
        self._ui_main_video_frame = self._ui_display_area[vid_frm_y0:
                                                          vid_frm_y1,
                                                          vid_frm_x0:
                                                          vid_frm_x1]
        # Get plot area.
        plt_frm_x0 = margin
        plt_frm_x1 = plt_frm_x0 + self._plot_size[1]
        plt_frm_y0 = margin
        plt_frm_y1 = plt_frm_y0 + self._plot_size[0]
        self._ui_plot_frame = self._ui_display_area[plt_frm_y0: plt_frm_y1,
                                                    plt_frm_x0: plt_frm_x1]
        # Compute history frame areas.
        his_frm_count = self.UI_LAYOUT["HISTORY_FRAME_COUNT"]
        his_frm_wid = int((self._ui_display_area_size[1] -
                          ((his_frm_count + 1)  * margin)) / his_frm_count)
        his_frm_ht = int((3./4.) * his_frm_wid)
        self._ui_history_frame_size = (his_frm_ht, his_frm_wid)
        his_frm_y0 = (2 * margin) + self._ui_main_video_size[0]
        his_frm_x0_fn = lambda n: margin + ((margin + his_frm_wid) * n)
        his_frm_x1_fn = lambda n: (margin + his_frm_wid) * (n + 1)
        self._ui_history_frames = map(
            lambda n: self._ui_display_area[
                his_frm_y0: his_frm_y0 + his_frm_ht,
                his_frm_x0_fn(n):his_frm_x1_fn(n)],
            range(self.UI_LAYOUT["HISTORY_FRAME_COUNT"]))

    def clear_display_area(self, color=np.array([0, 0, 0])):
        """ Clear UI base display area to given color. """

        self._ui_display_area[:,:] = color

    def set_main_video_frame(self, frame):
        """ Set the main video frame in the UI layout. """

        h, w = self._ui_main_video_size
        img = np.array(ImageOps.fit(PIL.Image.fromarray(frame), (w, h)))
        self._ui_main_video_frame[:,:] = img

    def push_history_frame(self, frame):
        """ Push a frame to the top of the history images. """

        # Shift back.
        for n in range(len(self._ui_history_frames) - 1):
            self._ui_history_frames[n][:,:] = self._ui_history_frames[n+1][:,:]
        # Update.
        h, w = self._ui_history_frame_size
        img = np.array(ImageOps.fit(PIL.Image.fromarray(frame), (w, h)))
        self._ui_history_frames[-1][:,:] = img

    def set_plot(self, plot):
        """ Set the plot image in the UI layout. """

        h, w = self._plot_size
        img = np.array(ImageOps.fit(PIL.Image.fromarray(plot), (w, h)))
        self._ui_plot_frame[:,:] = img

    def clear(self):
        """ Reset the UI. """

        ui_base = ndimg.imread(self.RESOURCES["UI_BASE"])
        self._ui_base[:,:,:] = ui_base
        self._ui_display_area[:,:] = self.UI_LAYOUT["CLEAR_COLOR"]

    def get_ui_image(self):
        return self._ui_base
    def get_history_frame_count(self):
        return len(self._ui_history_frames)

    ui_image = property(get_ui_image, doc="The main display image.")
    history_frame_count = property(get_history_frame_count,
                                   doc="Count of history frames.")

def main():
    #  Parse arguments for
    #    PhoneDemo.run_sequence(type, filepath, seq_range, filter_exp):
    parser = argparse.ArgumentParser(description='Run a UI demo of ' +
                                                 'image-based traffic ' +
                                                 'analysis algorithms.')
    parser.add_argument('SEQUENCE_TYPE', type=str, choices=('D', 'N'),
                        help='day or night image type')
    parser.add_argument('SEQUENCE_IMAGE_PATH', type=str,
                        help='path to an image within sequence')
    parser.add_argument('-r', '--range', default=None, nargs=2, type=int,
                        help='range of frames to process as in \'-r 0 100\'')
    parser.add_argument('-e', '--filter_exp', default=None, type=float,
                        help='exponential filter strength')
    args = parser.parse_args(sys.argv[1:])
    demo = PhoneDemo()
    demo.run_sequence(PhoneDemo.TYPE_DAY if 'D' == args.SEQUENCE_TYPE
                                         else PhoneDemo.TYPE_NIGHT,
                      args.SEQUENCE_IMAGE_PATH, args.range, args.filter_exp)

if __name__ == "__main__":
    main()


<reponame>loicmiller/policy-analysis<gh_stars>0
###############################################################################
# Imports

import sys
import argparse # Argument parser

from mgtoolkit.library import *
import networkx as nx

from sympy import to_dnf
from sympy.parsing.sympy_parser import parse_expr

import random # Random indices of list

import os

###############################################################################
# General utility

# Exit the program
def terminate_app(code):
    print("Exiting program...")
    sys.exit(code)

###############################################################################
# Argument parser

class Range(object):
    def __init__(self, start, end):
        self.start = start
        self.end = end
    def __eq__(self, other):
        return self.start <= other <= self.end

def get_parser():
    # Get parser for command line arguments
    parser = argparse.ArgumentParser(description="Workflow specification to Rego", formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("--version", action="version", version='%(prog)s 1.0')
    parser.add_argument("-v", "--verbose", action="count", default=0, help="increase output verbosity")
    parser.add_argument("workflow", type=str, metavar="FILE", help="workflow to generate policy from")
    parser.add_argument("-e", "--error-rate", type=float, choices=[Range(0.0, 1.0)], metavar="ERROR_RATE", default=0.0, help="rate of errors in the generated workflow")

    return parser

###############################################################################
# Functions

###############################################################################
# Main

def main(verbose, workflow):
    global glob_verbose
    glob_verbose = verbose

    print("\n\n###############################################################################")
    print("Loading workflow specification from file")
    print("###############################################################################")

    with open(workflow, 'r') as workflow_file:
        workflow_edges = workflow_file.readlines()
        workflow_edges = [(set(src.lstrip('{').rstrip('}').split(', ')), set(dst.lstrip('{').rstrip('}').split(', ')), attributes) for src, dst, attributes in (edge.rstrip().split(';') for edge in workflow_edges)]

    if glob_verbose >= 1:
        print("Edges")
        for edge in workflow_edges:
            print(edge)

    print("\n\n###############################################################################")
    print("Turning workflow graph into metagraph")
    print("###############################################################################")

    workflow_variables_set = set()
    workflow_propositions_set = set()
    workflow_edges_set = []

    # Simplify boolean expressions (Use simpy) https://stackoverflow.com/questions/52416781/how-to-simplify-these-boolean-statements
    for src, dst, attributes in workflow_edges:
        if glob_verbose >= 2:
            print("Edge: {} {} {}".format(src, dst, attributes))

        # Add src and dst to variable set if they are not present yet
        workflow_variables_set.update(src)
        workflow_variables_set.update(dst)

        # Parse policy into expression for simpy
        if attributes:
            edge_policy = parse_expr(attributes)
            if glob_verbose >= 2:
                print("Edge policy: {}".format(edge_policy))

            # Convert policy to Disjunctive Normal Form (DNF)
            # I think we don't want to simplify the expression for the comparison
            # since it is not simplified in the metagraph generated from the policy
            # https://en.wikipedia.org/wiki/Disjunctive_normal_form
            # https://docs.sympy.org/latest/modules/logic.html
            # https://docs.sympy.org/latest/modules/parsing.html
            edge_policy_dnf = to_dnf(edge_policy, simplify=False)
            if glob_verbose >= 2:
                print("DNF: {}".format(edge_policy_dnf))

            # Metagraph nodes
            # Each element in metagraph_nodes is the proposition part of a node in the metagraph
            metagraph_nodes = str(edge_policy_dnf).split("|")
            if glob_verbose >= 2:
                print("Metagraph nodes: {}".format(metagraph_nodes))

            # Policy elements in nodes
            # Each element is a part of the propositions_set
            for node_propositions in metagraph_nodes:
                policy_elements = node_propositions.split('&')
                policy_elements = [policy_element.strip().lstrip('(').rstrip(')') for policy_element in policy_elements] # Remove leading and trailing whitespaces, plus leading and trailing parentheses

                # Add policy_elements to propositions_set
                for index, policy_element in enumerate(policy_elements):
                    # Add ')' back for equalities
                    if 'Eq' in policy_element:
                        policy_element = policy_element + ')'
                        policy_elements[index] = policy_elements[index] + ')'
                    workflow_propositions_set.add(policy_element)
                workflow_edges_set.append(Edge(src, dst, attributes=policy_elements))

                if glob_verbose >= 2:
                    print("Policy elements: {}".format(policy_elements))

            if glob_verbose >= 2:
                print("\n")
        else:
            workflow_edges_set.append(Edge(src, dst, attributes=""))

    if glob_verbose >= 4:
        print("Variables set: {}".format(workflow_variables_set))
        print("Propositions set: {}\n".format(workflow_propositions_set))
        print("Metagraph edges: {}\n".format(workflow_edges_set))

    # Create workflow metagraph
    print("Creating workflow metagraph")
    workflow_metagraph = ConditionalMetagraph(workflow_variables_set, workflow_propositions_set)
    workflow_metagraph.add_edges_from(workflow_edges_set)

    if glob_verbose >= 4:
        print("Policy metagraph\n{}\n".format(repr(workflow_metagraph)))

    if glob_verbose >= 4:
        print("Workflow metagraph edges")
        print("{} {}".format("INVERTEX", "OUTVERTEX"))
        for edge in workflow_metagraph.edges:
            print("{} {}".format(list(edge.invertex), list(edge.outvertex)))

    return workflow_metagraph

if __name__ == '__main__':
    print("\n\n###############################################################################")
    print("Getting arguments")
    print("###############################################################################")

    parser = get_parser() # Create a parser
    args = parser.parse_args() # Parse arguments
    print(args)

    # Call main
    main(args.verbose, args.workflow)

    terminate_app(0)

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


<gh_stars>1-10
################################################################################
#
# Copyright (c) 2009 The MadGraph5_aMC@NLO Development team and Contributors
#
# This file is a part of the MadGraph5_aMC@NLO project, an application which
# automatically generates Feynman diagrams and matrix elements for arbitrary
# high-energy processes in the Standard Model and beyond.
#
# It is subject to the MadGraph5_aMC@NLO license which should accompany this
# distribution.
#
# For more information, visit madgraph.phys.ucl.ac.be and amcatnlo.web.cern.ch
#
################################################################################
from madgraph.core import base_objects
"""Methods and classes to import v4 format model files."""

import fractions
import logging
import os
import re

from madgraph import InvalidCmd, MG4DIR, ReadWrite

import madgraph.core.color_algebra as color
import madgraph.iolibs.files as files
import madgraph.iolibs.save_load_object as save_load_object

import madgraph.various.misc as misc

from madgraph.core.base_objects import Particle, ParticleList
from madgraph.core.base_objects import Interaction, InteractionList

logger = logging.getLogger('madgraph.import_v4')

#===============================================================================
# import_v4model
#===============================================================================
def import_model(model_path, mgme_dir = MG4DIR, absolute=True):
    """create a model from a MG4 model directory."""

    # Check for a valid directory
    model_path_old = model_path
    model_path = find_model_path(model_path, mgme_dir, absolute)

    files_list = [os.path.join(model_path, 'particles.dat'),\
                  os.path.join(model_path, 'interactions.dat')]

    for filepath in files_list:
        if not os.path.isfile(filepath):
            if not absolute:
                raise InvalidCmd,  "%s directory is not a valid v4 model" % \
                                                                    (model_path)
            else:
                return import_model(model_path_old, mgme_dir, False)

    # use pickle files if defined
    if files.is_uptodate(os.path.join(model_path, 'model.pkl'), files_list):
        model = save_load_object.load_from_file( \
                                          os.path.join(model_path, 'model.pkl'))
        if model.has_key('version_tag') and model.get('version_tag') == os.path.realpath(model_path) + str(misc.get_pkg_info()):
            return model, model_path

    model = base_objects.Model()
    model.set('particles',files.read_from_file( \
                                  os.path.join(model_path, 'particles.dat'),
                                  read_particles_v4))

    model.set('interactions',files.read_from_file( \
                                  os.path.join(model_path, 'interactions.dat'),
                                  read_interactions_v4,
                                  model['particles']))

    model.set('name', os.path.split(model_path)[-1])

    # save in a pickle files to fasten future usage
    if ReadWrite:
        try:
            save_load_object.save_to_file(os.path.join(model_path, 'model.pkl'), model)
        except Exception:
            logger.warning("fail to write %s. This is perfectly fine will just prevent speed boost in future load of this model" %\
                           os.path.join(model_path, 'model.pkl'))
    return model, model_path

def find_model_path(model_path, mgme_dir, absolute=True):
    """Find the path to the model, starting with path model_path."""

    # treat simple case (model_path is a valid path/ mgme_dir doesn't exist)
    if os.path.isdir(model_path) and absolute:
        return model_path
    elif mgme_dir and os.path.isdir(os.path.join(mgme_dir, 'models',
                                                 model_path + "_v4")):
        model_path = os.path.join(mgme_dir, 'models', model_path + "_v4")
    elif mgme_dir and os.path.isdir(os.path.join(mgme_dir, 'Models', model_path)):
        model_path = os.path.join(mgme_dir, 'Models', model_path)
    elif not mgme_dir:
        error_text = "Path %s is not a valid pathname\n" % model_path
        error_text += "and no MG_ME installation detected in order to search in Models"
        raise InvalidCmd(error_text)

    # Try to build the valid path
    path_possibilities = [os.path.join(mgme_dir, 'Models', model_path),
                         os.path.join(mgme_dir, 'models', model_path + "_v4"),
                         os.path.join(mgme_dir, 'models', model_path)
                         ]

    for path in path_possibilities:
        if os.path.exists(path) and \
                        not os.path.exists(os.path.join(path, 'particles.py')):
            return path

    # No valid path found
    raise InvalidCmd("Path %s is not a valid pathname" % model_path)

#===============================================================================
# read_particles_v4
#===============================================================================
def read_particles_v4(fsock):
    """Read a list of particle from stream fsock, using the old v4 format"""

    spin_equiv = {'s': 1,
                  'f': 2,
                  'v': 3,
                  't': 5}

    color_equiv = {'s': 1,
                   't': 3,
                   '6': 6,
                   'o': 8}

    line_equiv = {'d': 'dashed',
                  's': 'straight',
                  'w': 'wavy',
                  'c': 'curly'}

    logger.info('load particles')

    mypartlist = ParticleList()

    for line in fsock:
        mypart = Particle()

        if line.find("MULTIPARTICLES") != -1:
            break # stop scanning if old MULTIPARTICLES tag found

        line = line.split("#", 2)[0] # remove any comment
        line = line.strip() # makes the string clean

        if line != "":
            values = line.split()
            if len(values) != 9:
                # Not the right number tags on the line
                raise ValueError, \
                    "Unvalid initialization string:" + line
            else:
                try:
                    mypart.set('name', values[0].lower())
                    mypart.set('antiname', values[1].lower())

                    if mypart['name'] == mypart['antiname']:
                        mypart['self_antipart'] = True

                    if values[2].lower() in spin_equiv.keys():
                        mypart.set('spin',
                                   spin_equiv[values[2].lower()])
                    else:
                        raise ValueError, "Invalid spin %s" % \
                                values[2]

                    if values[3].lower() in line_equiv.keys():
                        mypart.set('line',
                                   line_equiv[values[3].lower()])
                    else:
                        raise ValueError, \
                                "Invalid line type %s" % values[3]

                    mypart.set("mass", values[4])
                    mypart.set("width", values[5])

                    if values[6].lower() in color_equiv.keys():
                        mypart.set('color',
                                   color_equiv[values[6].lower()])
                    else:
                        raise ValueError, \
                            "Invalid color rep %s" % values[6]

                    #mypart.set("texname", values[7])
                    mypart.set("pdg_code", int(values[8]))

                    mypart.set('charge', 0.)
                    #mypart.set('antitexname', mypart.get('texname'))

                except (Particle.PhysicsObjectError, ValueError), why:
                    logger.warning("Warning: %s, particle ignored" % why)
                else:
                    mypartlist.append(mypart)

    return mypartlist

#===============================================================================
# read_interactions_v4
#===============================================================================
def read_interactions_v4(fsock, ref_part_list):
    """Read a list of interactions from stream fsock, using the old v4 format.
    Requires a ParticleList object as an input to recognize particle names."""

    logger.info('load interactions')
    myinterlist = InteractionList()

    if not isinstance(ref_part_list, ParticleList):
        raise ValueError, \
            "Object %s is not a valid ParticleList" % repr(ref_part_list)

    for line in fsock:
        myinter = Interaction()

        line = line.split("#", 2)[0] # remove any comment
        line = line.strip() # makes the string clean

        if line != "": # skip blank
            values = line.split()
            part_list = ParticleList()

            try:
                for str_name in values:
                    curr_part = ref_part_list.get_copy(str_name.lower())
                    if isinstance(curr_part, Particle):
                        # Look at the total number of strings, stop if
                        # anyway not enough, required if a variable name
                        # corresponds to a particle! (eg G)
                        if len(values) >= 2 * len(part_list) + 1:
                            part_list.append(curr_part)
                        else: break
                    # also stops if string does not correspond to
                    # a particle name
                    else: break

                if len(part_list) < 3:
                    raise Interaction.PhysicsObjectError, \
                        "Vertex with less than 3 known particles found."

                # Flip part/antipart of first part for FFV, FFS, FFT vertices
                # according to v4 convention
                spin_array = [part['spin'] for part in part_list]
                if spin_array[:2] == [2, 2] and \
                   not part_list[0].get('self_antipart'):
                    part_list[0]['is_part'] = not part_list[0]['is_part']

                myinter.set('particles', part_list)

                # Give color structure
                # Order particles according to color
                # Don't consider singlets
                color_parts = sorted(enumerate(part_list), lambda p1, p2:\
                                     p1[1].get_color() - p2[1].get_color())
                color_ind = [(i, part.get_color()) for i, part in \
                             color_parts if part.get_color() !=1]
                colors = [c for i,c in color_ind]
                ind = [i for i,c in color_ind]

                # Set color empty by default
                myinter.set('color', [])
                if not colors:
                    # All color singlets - set empty
                    pass
                elif colors == [-3, 3]:
                    # triplet-triplet-singlet coupling
                    myinter.set('color', [color.ColorString(\
                        [color.T(ind[1], ind[0])])])
                elif colors == [8, 8]:
                    # octet-octet-singlet coupling
                    my_cs = color.ColorString(\
                        [color.Tr(ind[0], ind[1])])
                    my_cs.coeff = fractions.Fraction(2)
                    myinter.set('color', [my_cs])
                elif colors == [-3, 3, 8]:
                    # triplet-triplet-octet coupling
                    myinter.set('color', [color.ColorString(\
                        [color.T(ind[2], ind[1], ind[0])])])
                elif colors == [8, 8, 8]:
                    # Triple glue coupling
                    my_color_string = color.ColorString(\
                        [color.f(ind[0], ind[1], ind[2])])
                    my_color_string.is_imaginary = True
                    myinter.set('color', [my_color_string])
                elif colors == [-3, 3, 8, 8]:
                    my_cs1 = color.ColorString(\
                        [color.T(ind[2], ind[3], ind[1], ind[0])])
                    my_cs2 = color.ColorString(\
                        [color.T(ind[3], ind[2], ind[1], ind[0])])
                    myinter.set('color', [my_cs1, my_cs2])
                elif colors == [8, 8, 8, 8]:
                    # 4-glue coupling
                    cs1 = color.ColorString([color.f(0, 1, -1),
                                                   color.f(2, 3, -1)])
                    #cs1.coeff = fractions.Fraction(-1)
                    cs2 = color.ColorString([color.f(2, 0, -1),
                                                   color.f(1, 3, -1)])
                    #cs2.coeff = fractions.Fraction(-1)
                    cs3 = color.ColorString([color.f(1, 2, -1),
                                                   color.f(0, 3, -1)])
                    #cs3.coeff = fractions.Fraction(-1)
                    myinter.set('color', [cs1, cs2, cs3])
#   The following line are expected to be correct but not physical validations
#    have been performed. So we keep it commented for the moment.
#                elif colors == [3, 3, 3]:
#                    my_color_string = color.ColorString(\
#                        [color.Epsilon(ind[0], ind[1], ind[2])])
#                    myinter.set('color', [my_color_string])
#                elif colors == [-3, -3, -3]:
#                    my_color_string = color.ColorString(\
#                        [color.EpsilonBar(ind[0], ind[1], ind[2])])
#                    myinter.set('color', [my_color_string])
                else:
                    logger.warning(\
                        "Color combination %s not yet implemented." % \
                        repr(colors))

                # Set the Lorentz structure. Default for 3-particle
                # vertices is empty string, for 4-particle pair of
                # empty strings
                myinter.set('lorentz', [''])

                pdg_codes = sorted([part.get_pdg_code() for part in part_list])

                # WWWW and WWVV
                if pdg_codes == [-24, -24, 24, 24]:
                    myinter.set('lorentz', ['WWWW'])
                elif spin_array == [3, 3, 3, 3] and \
                             24 in pdg_codes and - 24 in pdg_codes:
                    myinter.set('lorentz', ['WWVV'])

                # gggg
                if pdg_codes == [21, 21, 21, 21]:
                    myinter.set('lorentz', ['gggg1', 'gggg2', 'gggg3'])

                # go-go-g
                # Using the special fvigox routine provides the minus
                # sign necessary for octet Majorana-vector interactions
                if spin_array == [2, 2, 3] and colors == [8, 8, 8] and \
                   part_list[0].get('self_antipart') and \
                   part_list[1].get('self_antipart'):
                    myinter.set('lorentz', ['go'])

                # If extra flag, add this to Lorentz
                if len(values) > 3 * len(part_list) - 4:
                    myinter.get('lorentz')[0] = \
                                      myinter.get('lorentz')[0]\
                                      + values[3 * len(part_list) - 4].upper()

                # Use the other strings to fill variable names and tags

                # Couplings: special treatment for 4-vertices, where MG4 used
                # two couplings, while MG5 only uses one (with the exception
                # of the 4g vertex, which needs special treatment)
                # DUM0 and DUM1 are used as placeholders by FR, corresponds to 1
                if len(part_list) == 3 or \
                   values[len(part_list) + 1] in ['DUM', 'DUM0', 'DUM1']:
                    # We can just use the first coupling, since the second
                    # is a dummy
                    myinter.set('couplings', {(0, 0):values[len(part_list)]})
                    if myinter.get('lorentz')[0] == 'WWWWN':
                        # Should only use one Helas amplitude for electroweak
                        # 4-vector vertices with FR. I choose W3W3NX.
                        myinter.set('lorentz', ['WWVVN'])
                elif values[len(part_list)] in ['DUM', 'DUM0', 'DUM1']:
                    # We can just use the second coupling, since the first
                    # is a dummy
                    myinter.set('couplings', {(0, 0):values[len(part_list)+1]})
                elif pdg_codes == [21, 21, 21, 21]:
                    # gggg
                    myinter.set('couplings', {(0, 0):values[len(part_list)],
                                              (1, 1):values[len(part_list)],
                                              (2, 2):values[len(part_list)]})
                elif myinter.get('lorentz')[0] == 'WWWW':
                    # Need special treatment of v4 SM WWWW couplings since
                    # MG5 can only have one coupling per Lorentz structure
                    myinter.set('couplings', {(0, 0):\
                                              'sqrt(' +
                                              values[len(part_list)] + \
                                             '**2+' + \
                                              values[len(part_list) + 1] + \
                                              '**2)'})
                else: #if myinter.get('lorentz')[0] == 'WWVV':
                    # Need special treatment of v4 SM WWVV couplings since
                    # MG5 can only have one coupling per Lorentz structure
                    myinter.set('couplings', {(0, 0):values[len(part_list)] + \
                                             '*' + \
                                              values[len(part_list) + 1]})
                    #raise Interaction.PhysicsObjectError, \
                    #    "Only FR-style 4-vertices implemented."

                # SPECIAL TREATMENT OF COLOR
                # g g sq sq (two different color structures, same Lorentz)
                if spin_array == [3, 3, 1, 1] and colors == [-3, 3, 8, 8]:
                    myinter.set('couplings', {(0, 0):values[len(part_list)],
                                              (1, 0):values[len(part_list)]})

                # Coupling orders - needs to be fixed
                order_list = values[2 * len(part_list) - 2: \
                                    3 * len(part_list) - 4]

                def count_duplicates_in_list(dupedlist):
                    """return a dictionary with key the element of dupeList and
                    with value the number of times that they are in this list"""
                    unique_set = set(item for item in dupedlist)
                    ret_dict = {}
                    for item in unique_set:
                        ret_dict[item] = dupedlist.count(item)
                    return ret_dict

                myinter.set('orders', count_duplicates_in_list(order_list))

                myinter.set('id', len(myinterlist) + 1)

                myinterlist.append(myinter)

            except Interaction.PhysicsObjectError, why:
                logger.error("Interaction ignored: %s" % why)

    return myinterlist

#===============================================================================
# read_proc_card.dat (mg4 format)
#===============================================================================
def read_proc_card_v4(fsock):
    """A simple function reading the files in fsock and returning a
    ProcCardv4Reader object. This function authorize to have the same syntax as
    for the other files treatment"""

    reader = ProcCardv4Reader(fsock)
    return reader

class ParticleError(InvalidCmd):
    """ A class to carch the error"""
    pass

class WrongFileFormat(InvalidCmd):
    """A specific class error for wrong V4 proc_card"""
    pass

class ProcCardv4Reader(object):
    """read a proc_card.dat in the mg4 format and creates the equivalent routine
    for mg5"""

    #tag in the proc_card.dat which split the proc_card content

    # line pattern (remove comment at the end of the line)
    pat_line = re.compile(r"""^\s*(?P<info>[^\#]*?)\s*(\#|$)""", re.DOTALL)

    def __init__(self, fsock):
        """init the variable"""

        self.process = [] # List of ProcessInfo
        self.model = ""   # name of the model
        self.multipart = [] # list of the mg4 definition of multiparticle
        self.particles_name = set() # set of authorize particle name
        self.couplings_name = set() # set of mandatory couplings
        self.process_path = os.path.realpath(os.path.join(
                                        os.path.dirname(fsock.name), os.pardir))

        # Reading the files and store the information in string format.
        self.analyze_v4_proc_card(fsock)

    def analyze_v4_proc_card(self, fsock):
        """read the file and fullfill the variable with mg4 line"""

        proc_card = fsock.read()

        # store process information
        process_open = False

        process_re = re.search(\
            r"^# Begin\s+PROCESS.*?^(?P<process>.*)^# End\s+PROCESS",
            proc_card, re.MULTILINE|re.DOTALL)

        if not process_re:
            raise WrongFileFormat('No valid Begin...End PROCESS tags')

        model_re = re.search(\
            r"^# Begin\s+MODEL.*?^(?P<model>.+?)(\s+|$)^# End\s+MODEL",
            proc_card, re.MULTILINE|re.DOTALL)

        if not model_re:
            raise WrongFileFormat('No valid Begin...End MODEL tags')

        multiparticles_re = re.search(\
            r"^# Begin\s+MULTIPARTICLES.*?^(?P<multiparticles>.*)^# End\s+MULTIPARTICLES",
            proc_card, re.MULTILINE|re.DOTALL)

        if not multiparticles_re:
            raise WrongFileFormat('No valid Begin...End MULTIPARTICLES tags')

        process_lines = process_re.group('process').split('\n')

        for line in process_lines:
            # an 'end_coup' stop the current process,
            #    'done' finish the list of process
            analyze_line = self.pat_line.search(line)
            if analyze_line:
                data = analyze_line.group('info') #skip the comment
                if not data:
                    continue
                if not process_open and 'done' not in data:
                    process_open = True
                    self.process.append(ProcessInfo(data))
                elif 'end_coup' in data:
                    process_open = False
                elif 'done' not in data:
                    self.process[-1].add_coupling(data)

        self.model = model_re.group('model')

        multiparticles_lines = multiparticles_re.group('multiparticles').split('\n')

        for line in multiparticles_lines:
            analyze_line = self.pat_line.search(line)
            if analyze_line:
                line = analyze_line.group('info') #skip the comment
                if not line:
                    continue
                data = line.split()
                self.particles_name.add(data[0].lower())
                self.multipart.append(line)

    def extract_command_lines(self, model):
        """Return the MG5 command line corresponding to this proc_card
        the MG5 command import model is skipped (since the model should be
        loaded -it is one of the argument-)"""

        # extract useful information of the model
        self.extract_info_from_model(model)

        # use the model information for the splitting in particles of the mg4
        #process line.
        for process in self.process:
            process.analyze_process(self.particles_name)

        #Now we are in position to write the lines call
        lines = []
        #first write the lines associate to the multiparticls definition
        if self.multipart:
            lines.append('# Define multiparticle labels')
        for multipart in self.multipart:
            data = self.separate_particle(multipart, self.particles_name)
            lines.append('define ' + ' '.join(data))

        # secondly define the lines associate with diagram
        if self.process:
            lines.append('# Specify process(es) to run')
        for i, process in enumerate(self.process):
            if i == 0:
                lines.append('generate %s' % \
                                  process.mg5_process_line(self.couplings_name))
            else:
                lines.append('add process %s' % \
                                  process.mg5_process_line(self.couplings_name))

        #finally export the madevent output
        lines.append('# Output processes to MadEvent directory')
        lines.append('output -f')

        return lines

    def extract_info_from_model(self, model):
        """ creates the self.particles_name (list of all valid name)
            and self.couplings_name (list of all couplings)"""

        # add in self.particles_name (it contains normally the mulpart name
        #already) all the valid name of particle of the model
        for particle in model['particles']:
            self.particles_name.add(particle['name'])
            self.particles_name.add(particle['antiname'])

        # add in self.couplings_name the couplings name of the model
        for interaction in model['interactions']:
            for coupling in interaction['orders'].keys():
                self.couplings_name.add(coupling)

    @staticmethod
    def separate_particle(line, possible_str):
        """ for a list of concatanate variable return a list of particle name"""

        line = line.lower() # Particle name are not case sensitive
        out = []            # list of the particles
        # The procedure to find particles is the following
        #  - check if the combination of 4 string form a valid particle name
        #    if it is, move of 4 characters and check for the next particles.
        #    if not try with 3, 2, 1
        #    if still not -> exit.

        pos = 0        # current starting position
        old_pos = -1   # check that we don't have infinite loop
        line += '     '  #add 4 blank for security
        while pos < len(line) - 4:
            #Check for infinite loop
            if pos == old_pos:
                logging.error('Invalid particle name: %s' % \
                              line[pos:pos + 4].rstrip())
                raise ParticleError('Invalid particle name %s' %
                                     line[pos:pos + 4].rstrip())
            old_pos = pos
            # check for pointless character
            if line[pos] in [' ', '\n', '\t']:
                pos += 1
                continue

            # try to find a match at 4(then 3/2/1) characters
            for i in range(4, 0, -1):
                if line[pos:pos + i] in possible_str:
                    out.append(line[pos:pos + i])
                    pos = pos + i
                    break

        return out

class ProcessInfo(object):
    """This is the basic object for storing process information"""

    def __init__(self, line):
        """Initialize information"""

        self.particles = [] # list tuple (level, particle)
        self.couplings = {} # coupling -> max_order
        self.decays = []    # ProcessInfo of the decays
        self.tag = ''       # tag of the process
        self.s_forbid = []  # list of particles forbids in s channel
        self.forbid = []     # list of particles forbids
        self.line = line    # initialization line

        self.is_mg5_valid = False
        #some shortcut
        self.separate_particle = ProcCardv4Reader.separate_particle

    def analyze_process(self, particles_name):
        """Add a line information
            two format are possible (decay chains or not)
            pp>h>WWj /a $u @3
            pp>(h>WW)j /a $u @3
        """

        line = self.line
        #extract the tag
        if '@' in line:
            split = line.split('@')
            line = split[0]
            self.tag = split[1]

        # check if we have a MG5 format
        if '/mg5/' in line:
            self.line = line.replace('/mg5/','')
            self.is_mg5_valid = True
            return
        if ',' in line or '=' in line:
            self.is_mg5_valid = True
            return

        # extract (S-)forbidden particle
        pos_forbid = line.find('/')
        pos_sforbid = line.find('$')

        # Select the restrictions (pos is -1 if not defined)
        #and remove the restrictions from the line
        if pos_forbid != -1 and pos_sforbid != -1:
            if  pos_forbid > pos_sforbid :
                self.forbid = self.separate_particle(line[pos_forbid + 1:], \
                                                                 particles_name)
                self.s_forbid = self.separate_particle(\
                               line[pos_sforbid + 1:pos_forbid], particles_name)
                line = line[:min(pos_forbid, pos_sforbid)]
            else:
                self.forbid = self.separate_particle(\
                               line[pos_forbid + 1:pos_sforbid], particles_name)
                self.s_forbid = self.separate_particle(line[pos_sforbid + 1:], \
                                                           particles_name)
                line = line[:min(pos_forbid, pos_sforbid)]
        # Same but if they are no S-forbidden particles
        elif pos_forbid != -1:
            self.forbid = self.separate_particle(line[pos_forbid + 1:], \
                                                                 particles_name)
            line = line[:pos_forbid]
        # Same but if they are no forbidden particles
        elif pos_sforbid != -1:
            self.s_forbid = self.separate_particle(line[pos_sforbid + 1:], \
                                                                 particles_name)
            line = line[:pos_sforbid]

        # Deal with decay chains, returns lines whitout the decay (and treat
        #the different decays.
        if '(' in line:
            line = self.treat_decay_chain(line, particles_name)

        #define the level of each particle
        level_content = line.split('>')
        for level, data in enumerate(level_content):
            particles = self.separate_particle(data, particles_name)
            if particles:
                [self.particles.append((level, name)) for name in particles]

    def treat_decay_chain(self, line, particles_name):
        """Split the information of the decays into a tree of ProcessInfo."""

        level = 0 #depth of the decay chain
        out_line = '' # core process
        for character in line:
            if character == '(':
                level += 1
                if level == 1:
                    decay_line = "" # initialize a new decay info
                else:
                    decay_line += '('
                continue
            elif character == ')':
                level -= 1
                if level == 0: #store the information
                    self.decays.append(ProcessInfo(decay_line))
                    self.decays[-1].add_restrictions(self.forbid, self.s_forbid,
                                                                 None)
                    self.decays[-1].analyze_process(particles_name)
                    out_line += decay_line[:decay_line.find('>')]
                else:
                    decay_line += ')'
                continue
            elif level:
                decay_line += character
            else:
                out_line += character
        return out_line

    def add_coupling(self, line):
        """Add the coupling information to the process"""
        data = line.split('=')
        self.couplings[data[0]] = int(data[1])

    def add_restrictions(self, forbid, s_forbid, couplings):
        """Associate some restriction to this diagram"""

        self.forbid = forbid
        self.s_forbid = s_forbid
        self.couplings = couplings

    def mg5_process_line(self, model_coupling):
        """Return a valid mg5 format for this process """

        if self.is_mg5_valid:
            return self.line

        text = ''
        # Write the process
        cur_level = 0
        for level, particle in self.particles:
            if level > cur_level:
                text += '> '
                cur_level += 1
            text += '%s ' % particle

        # Write the constraints
        if self.s_forbid:
            text += '$ ' + ' '.join(self.s_forbid) + ' '
        if self.forbid:
            text += '/ ' + ' '.join(self.forbid) + ' '

        #treat decay_chains
        for decay in self.decays:
            decay_text = decay.mg5_process_line(model_coupling)
            if ',' in decay_text:
                text = text.rstrip() + ', (%s) ' % decay_text.strip()
            else:
                text = text.rstrip() + ', %s ' % decay_text.strip()

        # write the tag
        if self.tag:
            text += '@%s ' % self.tag

        if self.couplings:
            if not self.tag:
                text += '@0 '
            #write the rules associate to the couplings
            text += self.mg5_couplings_line(model_coupling, len(self.particles))

        return text.rstrip()

    def mg5_couplings_line(self, model_coupling, nb_part):
        """Return the assignment of coupling for this process"""

        out = ''
        for coupling in model_coupling:
            if self.couplings.has_key(coupling):
                # Need coupling for all cases, since might be decay chain
                out += '%s=%s ' % (coupling, self.couplings[coupling])
            else:
                # if not define put to zero (mg4 default)
                out += '%s=0 ' % coupling

        return out