#!/usr/bin/python
#Script to simulate Petri nets

import sys
import argparse
from xml.dom import minidom
#from XMLparse import parse_pnml_input_file
import copy
from copy import deepcopy
import random
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.backends.backend_pdf import PdfPages
from operator import add

def parse_pnml_input_file(pnml):
    transitions = []
    places = {}
    arc_weights_in = {}
    arc_weights_out = {}
    
    doc = minidom.parse(pnml)
    
    #Read in start places
    nodes = doc.getElementsByTagName('place')
    for node in nodes:
        if node.attributes.has_key('id'):
            plac = node.attributes['id'].value
            plac = str(plac)
            #values = 0
            try:
                values = node.getElementsByTagName('text')[0].firstChild.nodeValue
                values = float(values)
            except:
                values = 0
            places[plac] = values
    
    #Read in transitions
    nodes = doc.getElementsByTagName('transition')
    for node in nodes:
        if node.attributes.has_key('id'):
            transition = node.attributes['id'].value
            transition = str(transition)
            transitions.append(transition)
    
    #Read in arcs
    nodes = doc.getElementsByTagName('arc')
    for node in nodes:
        if node.attributes.has_key('source'):
            sources = node.attributes['source'].value
            sources = str(sources)
            try:
                values = node.getElementsByTagName('text')[0].firstChild.nodeValue
                values = float(values)
            except:
                values = 1
            if sources in places:
                if node.attributes.has_key('target'):
                    targets = node.attributes['target'].value
                    targets = str(targets)
                    if targets not in arc_weights_in:
                        arc_weights_in[targets] = {}
                        arc_weights_in[targets][sources] = values
                    else:
                        arc_weights_in[targets][sources] = values
            
            elif sources in transitions:
                if node.attributes.has_key('target'):
                    targets = node.attributes['target'].value
                    targets = str(targets)
                    if sources not in arc_weights_out:
                    	arc_weights_out[sources] = {}
                    	arc_weights_out[sources][targets] = values
                    else:
                    	arc_weights_out[sources][targets] = values
    
    return(transitions, places, arc_weights_in, arc_weights_out)

#SINGLE STEP
def single_step(transitions, places, arc_weights_in, arc_weights_out):
    places_temp1 = deepcopy(places)
    places_temp2 = deepcopy(places)
    transition_valid = "TRUE"
    transitions_rand = copy.copy(transitions)
    
    while len(transitions_rand) > 0:
        random.shuffle(transitions_rand)
        t = transitions_rand[0]
        for p in arc_weights_in[t]:
            if not places_temp2[p] - arc_weights_in[t][p] >= 0:
                transition_valid = "FALSE"
        if transition_valid == "TRUE":
            for p in arc_weights_in[t]:
                if places_temp2[p] - arc_weights_in[t][p] >= 1:
                    transition_valid = "TRUE"
                else:
                    transition_valid = "FALSE"
                places_temp2[p] = places_temp2[p] - arc_weights_in[t][p]
                places_temp1[p] = places_temp1[p] - arc_weights_in[t][p]
            for p in arc_weights_out[t]:
                places_temp1[p] = places_temp1[p] + arc_weights_out[t][p]
        elif transition_valid == "FALSE":
            transitions_rand.pop(0)
        transition_valid = "TRUE"
    return places_temp1

#SINGLE SIMULATION
def single_sim(num_steps, transitions, places, arc_weights_in, arc_weights_out, output):
    token_step_sim = deepcopy(places)
    for place in token_step_sim:
        token_step_sim[place] = []
    for place in places:
        token_step_sim[place].append(float(places[place]))
    i = 0
    while i < num_steps:
        places = single_step(transitions, places, arc_weights_in, arc_weights_out)
        i += 1
        for place in places:
            token_step_sim[place].append(float(places[place]))
    for place in output:
        output[place].append(token_step_sim[place])
    token_step_sim = {}
    return output

#MULTIPLE SIMULATIONS
def multi_sim(num_sim, num_steps, transitions, places, arc_weights_in, arc_weights_out, output):
    token_multi_sim = []
    i = 0
    while i < num_sim:
        output = single_sim(num_steps, transitions, places, arc_weights_in, arc_weights_out, output)
        i += 1
    return output


#MILTIPLE WEIGHTS
def multi_weight_sim(per, num_sim, num_steps, transitions, places, arc_weights_in, arc_weights_out):
    output = deepcopy(places)
    for key in output:
        output[key] = []
    
    CTNNB1_avg = []
    CTNNB1_std = []
    prot = ''
    
    if per == "WNT" or per == "GSK3":
        prot = per
        i = 0
        limit = 5
        while i <= limit:
            places[prot] = i
            output = multi_sim(num_sim, num_steps, transitions, places, arc_weights_in, arc_weights_out, output)
            CTNNB1 = []
            j = 0
            while j < len(output["CTNNB1"]):
                CTNNB1.append([sum(x) for x in zip(output["CTNNB1"][j],output["DCCTNNB1"][j],output["TCFLEFCTNNB1"][j])])
                j += 1
            CTNNB1_avg.append(np.mean(CTNNB1, axis=0))
            CTNNB1_std.append(0.5*(np.std(CTNNB1, axis=0)))
            for key in output:
                output[key] = []
            i += 1
    
    if per == "WNTfb":
        prot = "WNT"
        i = 0
        limit = 5
        k = 0
        while k <= 1:
            while i <= limit:
                places[prot] = i
                output = multi_sim(num_sim, num_steps, transitions, places, arc_weights_in, arc_weights_out, output)
                CTNNB1 = []
                j = 0
                while j < len(output["CTNNB1"]):
                    CTNNB1.append([sum(x) for x in zip(output["CTNNB1"][j],output["DCCTNNB1"][j],output["TCFLEFCTNNB1"][j])])
                    j += 1
                CTNNB1_avg.append(np.mean(CTNNB1, axis=0))
                CTNNB1_std.append(0.5*(np.std(CTNNB1, axis=0)))
                for key in output:
                    output[key] = []
                i += 1
            arc_weights_out["t11"]["AXIN"] = 0.15
            i = 0
            k += 1
    
    elif per == "APC":
        weigths = [0, 0.05, 0.1, 0.2]
        for i in weigths:
            arc_weights_out["t5"]["DC"] = i
            arc_weights_out["t5"]["AXIN"] = 1-i
            arc_weights_out["t5"]["APC"] = 1-i
            arc_weights_out["t5"]["GSK3"] = 1-i
            arc_weights_out["t5"]["CK1"] = 1-i
            output = multi_sim(num_sim, num_steps, transitions, places, arc_weights_in, arc_weights_out, output)
            CTNNB1 = []
            j = 0
            while j < len(output["CTNNB1"]):
                CTNNB1.append([sum(x) for x in zip(output["CTNNB1"][j],output["DCCTNNB1"][j],output["TCFLEFCTNNB1"][j])])
                j += 1
            CTNNB1_avg.append(np.mean(CTNNB1, axis=0))
            CTNNB1_std.append(0.5*(np.std(CTNNB1, axis=0)))
            for key in output:
                output[key] = []
    
    xdata = range(0, num_steps+1, 1)
    plt.axis([0, num_steps, 0, 100])
    
    if per == "WNT":
        plt.plot(xdata, CTNNB1_avg[0], color='black', linewidth=2, linestyle='-', label='total BCAT')
        plt.plot(xdata, CTNNB1_avg[3], color='black', linewidth=2, linestyle='-', label='total BCAT')
        plt.fill_between(xdata, np.array(CTNNB1_avg[3])-np.array(CTNNB1_std[3]), np.array(CTNNB1_avg[3])+np.array(CTNNB1_std[3]), alpha=0.5, edgecolor='lightgrey', facecolor='lightgrey')
        plt.plot(xdata, CTNNB1_avg[4], color='black', linewidth=2, linestyle='-', label='total BCAT')
        plt.fill_between(xdata, np.array(CTNNB1_avg[4])-np.array(CTNNB1_std[4]), np.array(CTNNB1_avg[4])+np.array(CTNNB1_std[4]), alpha=0.5, edgecolor='lightgrey', facecolor='lightgrey')
        plt.plot(xdata, CTNNB1_avg[5], color='black', linewidth=2, linestyle='-', label='total BCAT')
        plt.fill_between(xdata, np.array(CTNNB1_avg[5])-np.array(CTNNB1_std[5]), np.array(CTNNB1_avg[5])+np.array(CTNNB1_std[5]), alpha=0.5, edgecolor='lightgrey', facecolor='lightgrey')
        plt.savefig('Figure_WNT.ps', format='ps')
        plt.show()
        plt.cla()
    if per == "WNTfb":
        plt.fill_between(xdata, CTNNB1_avg[3], CTNNB1_avg[9], facecolor='lightgrey', alpha=0.3, edgecolor='none')
        plt.plot(xdata, CTNNB1_avg[3], color='black', linewidth=2, linestyle='-')
        plt.plot(xdata, CTNNB1_avg[9], color='black', linewidth=2, linestyle=':')
        plt.savefig('Figure_WNT5fb.ps', format='ps')
        plt.show()
        plt.cla()
        plt.axis([0, num_steps, 0, 100])
        plt.fill_between(xdata, CTNNB1_avg[4], CTNNB1_avg[10], facecolor='lightgrey', alpha=0.3, edgecolor='none')
        plt.plot(xdata, CTNNB1_avg[4], color='black', linewidth=2, linestyle='-')
        plt.plot(xdata, CTNNB1_avg[10], color='black', linewidth=2, linestyle=':')
        plt.savefig('Figure_WNT4fb.ps', format='ps')
        plt.show()
        plt.cla()
        plt.axis([0, num_steps, 0, 100])
        plt.fill_between(xdata, CTNNB1_avg[5], CTNNB1_avg[11], facecolor='lightgrey', alpha=0.3, edgecolor='none')
        plt.plot(xdata, CTNNB1_avg[5], color='black', linewidth=2, linestyle='-')
        plt.plot(xdata, CTNNB1_avg[11], color='black', linewidth=2, linestyle=':')
        plt.savefig('Figure_WNT3fb.ps', format='ps')
        plt.show()
        plt.cla()
    elif per == "GSK3":
        plt.plot(xdata, CTNNB1_avg[0], color='black', linewidth=2, linestyle='-', label='total BCAT')
        plt.plot(xdata, CTNNB1_avg[1], color='black', linewidth=2, linestyle='-', label='total BCAT')
        plt.fill_between(xdata, np.array(CTNNB1_avg[1])-np.array(CTNNB1_std[1]), np.array(CTNNB1_avg[1])+np.array(CTNNB1_std[1]), alpha=0.5, edgecolor='none', facecolor='lightgrey')
        plt.plot(xdata, CTNNB1_avg[2], color='black', linewidth=2, linestyle='-', label='total BCAT')
        plt.fill_between(xdata, np.array(CTNNB1_avg[2])-np.array(CTNNB1_std[2]), np.array(CTNNB1_avg[2])+np.array(CTNNB1_std[2]), alpha=0.5, edgecolor='none', facecolor='lightgrey')
        plt.plot(xdata, CTNNB1_avg[3], color='black', linewidth=2, linestyle='-', label='total BCAT')
        plt.savefig('Figure_GSK3.ps', format='ps')
        plt.show()
        plt.cla()
    elif per == "APC":
        plt.plot(xdata, CTNNB1_avg[0], color='black', linewidth=2, linestyle='-', label='total BCAT')
        plt.fill_between(xdata, np.array(CTNNB1_avg[0])-np.array(CTNNB1_std[0]), np.array(CTNNB1_avg[0])+np.array(CTNNB1_std[0]), alpha=0.5, edgecolor='none', facecolor='lightgrey')
        plt.plot(xdata, CTNNB1_avg[1], color='black', linewidth=2, linestyle='-', label='total BCAT')
        plt.fill_between(xdata, np.array(CTNNB1_avg[1])-np.array(CTNNB1_std[1]), np.array(CTNNB1_avg[1])+np.array(CTNNB1_std[1]), alpha=0.5, edgecolor='none', facecolor='lightgrey')
        plt.plot(xdata, CTNNB1_avg[2], color='black', linewidth=2, linestyle='-', label='total BCAT')
        plt.fill_between(xdata, np.array(CTNNB1_avg[2])-np.array(CTNNB1_std[2]), np.array(CTNNB1_avg[2])+np.array(CTNNB1_std[2]), alpha=0.5, edgecolor='none', facecolor='lightgrey')
        plt.plot(xdata, CTNNB1_avg[3], color='black', linewidth=2, linestyle='-', label='total BCAT')
        plt.fill_between(xdata, np.array(CTNNB1_avg[3])-np.array(CTNNB1_std[3]), np.array(CTNNB1_avg[3])+np.array(CTNNB1_std[3]), alpha=0.5, edgecolor='none', facecolor='lightgrey')
        plt.savefig('Figure_APC.ps', format='ps')
        plt.show()
        plt.cla()
    for key in output:
        output[key] = []


def main():
    num_sim = 100
    num_steps = 100
    #per = "GSK3"
    #per = "WNT"
    #per = "WNTfb"
    per = "APC"
    parser = argparse.ArgumentParser(description='This is a Petri net simulation tool')
    parser.add_argument('-i','--infile', help='pnml input file of the model', required=False)
    args = vars(parser.parse_args())
    
    if not args['infile']: #default model
        transitions = ["t1","t2","t3","t4","t5","t6","t7","t8","t9","t10","t11"]
        places = {"WNT": 0, "FZD": 5, "LRP": 5, "WNTFZDLRP": 0, "DVL": 5, "WNTFZDLRPDVL": 0, "RC": 0, "Controller": 1, "AXIN": 5, "APC": 5, "CK1": 5, "GSK3": 5, "DC": 0, "DCCTNNB1": 0, "CTNNB1": 0, "ctnnb1": 1, "TCFLEF": 1, "TCFLEFCTNNB1": 0, "axin2":1}
        arc_weights_in = {"t1": {"WNT": 1, "FZD": 1,  "LRP": 1},
        "t2": {"WNTFZDLRP": 1, "DVL": 1},
        "t3": {"WNTFZDLRPDVL": 1, "AXIN": 1},
        "t4": {"RC": 0.1, "Controller": 1},
        "t5": {"AXIN": 1, "APC": 1, "CK1": 1, "GSK3": 1},
        "t6": {"CTNNB1": 1, "DC": 1},
        "t7": {"DCCTNNB1": 1},
        "t8": {"DCCTNNB1": 1},
        "t9": {"ctnnb1": 1},
        "t10": {"CTNNB1": 3, "TCFLEF": 1},
        "t11": {"TCFLEFCTNNB1": 1, "axin2":1}}
        arc_weights_out = {"t1": {"WNTFZDLRP": 1},
        "t2": {"WNTFZDLRPDVL": 1},
        "t3": {"RC": 1},
        "t4": {"WNTFZDLRPDVL": 0.1, "AXIN": 0.1, "Controller": 1},
        "t5": {"DC": 1},
        "t6": {"DCCTNNB1": 1},
        "t7": {"DC": 1},
        "t8": {"AXIN": 1, "APC": 1, "CK1": 1, "GSK3": 1},
        "t9": {"ctnnb1": 1, "CTNNB1":1},
        "t10": {"CTNNB1": 2, "TCFLEFCTNNB1": 1},
        "t11": {"CTNNB1": 1, "TCFLEF": 1, "AXIN": 0, "axin2": 1}}
    else:
    	(transitions, places, arc_weights_in, arc_weights_out) = parse_pnml_input_file(args['infile']) #parse model from pnml file
    
    multi_weight_sim(per, num_sim, num_steps, transitions, places, arc_weights_in, arc_weights_out)

if __name__ == '__main__':
    sys.exit(main())