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

# In[1]:


import cPickle as pickle
import networkx as nx


# In[2]:


train_cont_sc_1_3_graphs = []
train_cont_sc_1_3_labels = []
# labels:  linear = 0, lineage = 1


# In[3]:


from os import listdir
from os.path import isfile, join


# In[4]:


linear_continuity_files = [f for f in listdir("train-linear/continuity") if isfile(join("train-linear/continuity", f))]
lineage_continuity_files = [f for f in listdir("train-lineage/continuity") if isfile(join("train-lineage/continuity", f))]


# In[5]:


for f in linear_continuity_files:
    path = "train-linear/continuity/" + f
    try:
        g = nx.read_gml(path)
    except:
        print "bad file format: %s" % f
    train_cont_sc_1_3_graphs.append(g)
    train_cont_sc_1_3_labels.append(0)
    
for f in lineage_continuity_files:
    path = "train-lineage/continuity/" + f
    try:
        g = nx.read_gml(path)
    except:
        print "bad file format: %s" % f
    train_cont_sc_1_3_graphs.append(g)
    train_cont_sc_1_3_labels.append(1)


# In[6]:


pickle.dump(train_cont_sc_1_3_graphs,open("train-cont-sc-1-3-graphs.pkl",'wb'))
pickle.dump(train_cont_sc_1_3_labels,open("train-cont-sc-1-3-labels.pkl",'wb'))


# In[ ]: