import pickle
import tempfile

from collections import defaultdict, Counter
import networkx as nx
import scipy.stats

from networkx.drawing.nx_agraph import graphviz_layout, pygraphviz_layout
import os, sys

from synonymes.GeneOntology import GeneOntology
from utils.tmutils import normalize_gene_names

sys.path.insert(0, str(os.path.dirname("/mnt/d/dev/git/poreSTAT/")))

from porestat.utils.DataFrame import DataFrame, DataRow, ExportTYPE

from synonymes.mirnaID import miRNA, miRNAPART, miRNACOMPARISONLEVEL
from textdb.makeNetworkView import DataBasePlotter
from utils.cytoscape_grapher import CytoscapeGrapher

import numpy as np
import matplotlib.pyplot as plt

if __name__ == '__main__':

    cellObo = GeneOntology("/mnt/d/owncloud/data/miRExplore/obodir/meta_cells.obo")

    cellTypeName2Terms = {
        "EC": ["META:52"],
        "MC": ["META:148", "META:99"],
        "FC": ["CL:0000891"],
        "SMC": ["META:83"],
    }

    cellType2AccTerms = {}
    for cellT in cellTypeName2Terms:

        cellType2AccTerms[cellT] = set()

        for et in cellTypeName2Terms[cellT]:

            oboT = cellObo.getID(et)

            if oboT != None:
                cellType2AccTerms[cellT].add(et)
                for x in oboT.getAllChildren():
                    cellType2AccTerms[cellT].add(x.termid)

            else:
                print("No such obo term:", et)

    missRates = sorted([2/19, 5/17, 1/4, 5/12, 0/6, 2/8, 3/12, 1/23])

    print("missrate", sum(missRates)/len(missRates))

    print(scipy.stats.describe(missRates))
    print(np.median(missRates))
    #exit(0)

    largeInteractions = {
        'CCL9': ['miR-30d-3p', 'miR-3473c', 'let-7g-5p'],
        'CXCL5': ['miR-204-5p', 'let-7g-5p', 'miR-362-3p', 'miR-155-5p'],
        'CXCL1': ['miR-194-2-3p', 'miR-128-3p', 'miR-194-5p', 'miR-199b-5p', 'miR-467g', 'miR-122-5p'],
        'CXCL13': ['miR-122-5p'],
        'CXCL14': ['miR-301b-3p'],
        'CXCR2': ['let-7g-5p', 'let-7b-5p', 'let-7f-5p', 'let-7c-5p', 'let-7a-5p', 'let-7i-5p', 'miR-98-5p'],
        'CXCL7': ['let-7g-5p'],
        'CCL2': ['let-7a-5p', 'let-7b-5p', 'let-7f-5p', 'let-7c-5p', 'let-7g-5p', 'let-7i-5p', 'miR-181a-5p'],
        'CXCL9': ['miR-1935'],
        'CCL3': ['miR-30a-5p','miR-30b-5p','miR-30c-5p','miR-30d-5p','miR-30e-5p'],
        'CCL7': ['miR-181a-5p', 'miR-322-5p', 'miR-29a-5p', 'miR-29b-1-5p'],
        'CCL22': [  'miR-34a-5p'],
        'CXCL10': ['miR-503-3p', 'miR-186-5p'],
        'CCR5': ['miR-186-5p', 'miR-669j', 'miR-21-5p', 'miR-146a-5p', 'miR-150-5p', 'miR-146b-5p', 'miR-669k-3p', 'miR-142-3p', 'miR-34a-5p'],
        'CCL4': ['miR-27b-3p', 'miR-27a-3p', 'miR-21-3p', 'miR-467f'],
        'CX3CL1': ['miR-15a-5p', 'miR-322-5p', 'miR-706', 'miR-762', 'miR-665-3p', 'miR-758-3p', 'miR-381-3p'],
        'CXCR4': ['miR-381-3p', 'miR-21-3p', 'miR-467a-5p', 'miR-467h', 'miR-218-5p', 'miR-1a-3p', 'miR-181d-5p',
                  'miR-206-3p', 'miR-181b-5p', 'miR-9-5p', 'miR-132-3p', 'miR-25-3p', 'miR-467d-5p', 'miR-669k-3p',
                  'miR-146b-5p', 'miR-467b-5p', 'miR-467e-5p', 'miR-467f', 'miR-146a-5p'],
        'CCR7': ['let-7g-5p', 'miR-23b-3p', 'miR-669p-5p', 'miR-23a-5p', 'let-7e-5p', 'miR-669l-5p', 'miR-15a-5p', 'miR-467e-5p', 'miR-21-5p', 'miR-16-5p', 'let-7d-5p', 'miR-669n', 'miR-98-5p', 'let-7b-5p', 'let-7a-5p', 'let-7i-5p', 'let-7c-5p', 'miR-15b-5p', 'miR-467h'],
        'CXCL12': [
            'miR-532-5p', 'miR-130b-3p', 'miR-222-3p', 'miR-144-3p', 'miR-542-3p', 'miR-149-5p', 'miR-330-3p', 'miR-532-3p', 'miR-3470b', 'miR-125b-5p', 'miR-221-3p', 'miR-19b-3p', 'miR-301b-3p',
            'miR-34b-5p', 'miR-125a-3p', 'miR-126-3p', 'miR-16-1-3p', 'miR-882', 'miR-497-5p', 'miR-26a-5p', 'miR-124-3p', 'miR-26b-5p', 'miR-5620-3p', 'mIR-19a-3p', 'miR-130a-3p', 'miR-690',
            'miR-185-5p', 'miR-31-5p', 'miR-340-5p', 'miR-1843-5p', 'miR-466f-3p', 'miR-301a-3p', 'miR-101a-3p', 'miR-210-3p', 'miR-107-3p', 'miR-706', 'miR-23b-3p', 'miR-146a-5p', 'miR-467f',
            'miR-322-5p', 'miR-15a-5p', 'miR-29b-1-5p', 'let-7e-5p', 'miR-23a-3p', 'miR-338-3p', 'miR-103-3p', 'miR-362-3p', 'let-7g-5p', 'miR-155-5p', 'miR-140-5p', 'miR-122-5p', 'miR-22-3p', 'miR-3470a', 'let-7d-5p'
        ]

    }

    inflammatory_ec = {
        'CCL2': ['miR-495'],
        'PPARA': ['miR-21'],
        'THBS1': ['let-7g'],
        'TGFBR1': ['let-7g'],
        'SMAD2': ['let-7g'],
        'SIRT1': ['let-7g', 'miR-34a'],
        'SOCS5': ['miR-92a'],
        'KLF2': ['miR-92a'],
        'KLF4': ['miR-92a', 'miR-103'],
        'TRC': ['miR-10a'],
        'TAK1': ['miR-10a'],
        'KPNA3': ['miR-181b'],
        'TRAF6': ['miR-146'],
        'IRAK1': ['miR-146'],
        'HUR': ['miR-146'],
        'CXCL12': ['miR-126-3p'],
        'RGS16': ['miR-126-3p'],
        'ETS1': ['miR-155', 'miR-221', 'miR-222'],
        'TNPO2': ['miR-181b']
    }

    ccl2_macrophage = {
        'CCL2': ['miR-124a', 'miR-150'],
        'CHI3L1': ['miR-24'],
        'TLR4': ['miR-146a'],
        'TRAF6': ['miR-146'],
        'IRAK1': ['miR-146'],
        'AKT1': ['miR-342-5p'],
        'BCL6': ['miR-155'],
        'LPL': ['miR-467b']
    }

    normGeneSymbols = normalize_gene_names(path="/mnt/d/owncloud/data/miRExplore/obodir/" + "/hgnc_no_withdrawn.syn")

    andreouInts1= {'ICAM1': ['miR-223', 'miR-17-3p'], 'SELE': ['miR-31'], 'VCAM1': ['miR-126', 'miR-126-5p'], 'DLK1': ['miR-126', 'miR-126-5p'], 'JAMA': ['miR-143', 'miR-145'], 'SIRT1': ['miR-217'], 'NOS3': ['miR-155'], 'ETS1': ['miR-155', 'miR-126', 'miR-126-5p'], 'PPARA': ['miR-21'], 'TIMP3': ['miR-712'], 'MAP3K7': ['miR-10'], 'KPNA4': ['miR-181b'], 'TNPO2': ['miR-181b'], 'TRAF6': ['miR-146a', 'miR-146b'], 'IRAK1': ['miR-146a', 'miR-146b'], 'IRAK2': ['miR-146a', 'miR-146b'], 'KLF2': ['miR-143', 'miR-145', 'miR-126', 'miR-126-5p', 'miR-92a'], 'CXCL12': ['miR-126', 'miR-126-5p'], 'KLF4': ['miR-92a', 'miR-663'], 'SOCS5': ['miR-92a']}
    andreouInts2= {'IFNG': ['miR-29', 'miR-155'], 'CD40': ['miR-146a', 'miR-181a'], 'CD40L': ['miR-146a', 'miR-181a'], 'MMP2': ['miR-21', 'miR-21', 'miR-24', 'miR-29'], 'MMP3': ['miR-21', 'miR-24', 'miR-29'], 'MMP8': ['miR-21', 'miR-24', 'miR-29'], 'MMP9': ['miR-21', 'miR-24', 'miR-29'], 'MMP12': ['miR-21', 'miR-24', 'miR-29'], 'MMP13': ['miR-21', 'miR-24', 'miR-29'], 'MMP14': ['miR-21', 'miR-24', 'miR-29']}
    andreouInts3= {'MTTP': ['miR-30c'], 'ABCA1': ['miR-33', 'miR-302a'], 'ABCG1': ['miR-33'], 'CPT1A': ['miR-33'], 'KLF2': ['miR-92a'], 'KLF4': ['miR-92a', 'miR-145'], 'SOCS5': ['miR-92a'], 'DLK1': ['miR-126-5p'], 'RGS16': ['miR-126-3p'], 'BCL6': ['miR-155'], 'SOCS1': ['miR-155'], 'MAP3K10': ['miR-155'], 'KPNA4': ['miR-181b'], 'TNPO2': ['miR-181b'],'AKT1': ['miR-342-5p'], 'TIMP3': ['miR-712']}

    andreouInteractions1 = {}
    andreouInteractions2 = {}
    andreouInteractions3 = {}

    for gene in andreouInts1:
        andreouInteractions1[normGeneSymbols.get(gene, gene)] = andreouInts1[gene]

    for gene in andreouInts2:
        andreouInteractions2[normGeneSymbols.get(gene, gene)] = andreouInts2[gene]

    for gene in andreouInts3:
        andreouInteractions3[normGeneSymbols.get(gene, gene)] = andreouInts3[gene]

    networks = {}
    networks['large_chemokines'] = largeInteractions
    networks['large_chemokines_cv'] = largeInteractions
    networks['large_chemokines_athero'] = largeInteractions

    networks['andreou_fig1'] = andreouInteractions1
    networks['andreou_fig1_cv'] = andreouInteractions1
    networks['andreou_fig1_athero'] = andreouInteractions1

    networks['andreou_fig2'] = andreouInteractions2
    networks['andreou_fig2_cv'] = andreouInteractions2
    networks['andreou_fig2_athero'] = andreouInteractions2

    networks['andreou_table1'] = andreouInteractions3
    networks['andreou_table1_cv'] = andreouInteractions3
    networks['andreou_table1_athero'] = andreouInteractions3

    networks['inflammatory_ec'] = inflammatory_ec
    networks['inflammatory_ec_cv'] = inflammatory_ec
    networks['inflammatory_ec_athero'] = inflammatory_ec

    networks['macrophages'] = ccl2_macrophage
    networks['macrophages_cv'] = ccl2_macrophage
    networks['macrophages_athero'] = ccl2_macrophage
    networks['macrophages_all'] = ccl2_macrophage

    networkGraphs = {}
    makeStory = [
        ('macrophages', 'macrophages_cv', 'macrophages_athero', 'macrophages_all'),
        ('inflammatory_ec','inflammatory_ec_cv', 'inflammatory_ec_athero'),
        ('large_chemokines', 'large_chemokines_cv', 'large_chemokines_athero'),
        ('andreou_fig1', 'andreou_fig1_cv', 'andreou_fig1_athero'),
        ('andreou_fig2', 'andreou_fig2_cv', 'andreou_fig1_athero'),
        ('andreou_table1', 'andreou_table1_cv', 'andreou_table1_athero'),
    ]

    ignoreNetworks = []#['large_chemokines']

    networkRestrictions = {
        'inflammatory_ec':
            {
                'sentences': "false",
                "cells": [
                    {"group": "cells", "name": "endothelial cell", "termid": "META:52"}
                ]
            },
        'inflammatory_ec_cv':
            {
                'sentences': "false",
                "cells": [
                    {"group": "cells", "name": "endothelial cell", "termid": "META:52"}
                ],
                "disease": [
                    {'group': 'disease', 'termid': 'DOID:1287', 'name': 'cardiovascular system disease'},
                    {'group': 'disease', 'termid': 'DOID:2349', 'name': 'arteriosclerosis'}
                ]
            },

        'inflammatory_ec_athero':
            {
                'sentences': "false",
                "cells": [
                    {"group": "cells", "name": "endothelial cell", "termid": "META:52"}
                ],
                "disease": [
                    {'group': 'disease', 'termid': 'DOID:1936', 'name': 'atherosclerosis'}
                ]
            },
        'macrophages':
            {
                'sentences': "false",
                "cells": [
                    {"group": "cells", "name": "macrophage", "termid": "META:99"}
                ]
            },
        'macrophages_all':
            {
                'sentences': "false",
            },
        'macrophages_cv':
            {
                'sentences': "false",
                "cells": [
                    {"group": "cells", "name": "macrophage", "termid": "META:99"}
                ],
                "disease": [
                    {'group': 'disease', 'termid': 'DOID:1287', 'name': 'cardiovascular system disease'},
                    {'group': 'disease', 'termid': 'DOID:2349', 'name': 'arteriosclerosis'}
                ]
            },
        'macrophages_athero':
            {
                'sentences': "false",
                "cells": [
                    {"group": "cells", "name": "macrophage", "termid": "META:99"}
                ],
                "disease": [
                    {'group': 'disease', 'termid': 'DOID:1936', 'name': 'atherosclerosis'}
                ]
            },
        'large_chemokines': {
            'sentences': "false",
        },
        'large_chemokines_cv': {
            'sentences': "false",
            "disease": [
                {'group': 'disease', 'termid': 'DOID:1287', 'name': 'cardiovascular system disease'},
                {'group': 'disease', 'termid': 'DOID:2349', 'name': 'arteriosclerosis'}
            ]        },
        'large_chemokines_athero': {
            'sentences': "false",
            "disease": [
                {'group': 'disease', 'termid': 'DOID:1936', 'name': 'atherosclerosis'}
            ]},
        'andreou_fig1': {
            'sentences': "false",
            "cells": [
                {"group": "cells", "name": "endothelial cell", "termid": "META:52"}
            ],
        },
        'andreou_fig2': {
            'sentences': "false",
        },
        'andreou_table1': {
            'sentences': "false",
        },

        'andreou_fig1_cv':
            {
                'sentences': "false",
                "disease": [
                    {'group': 'disease', 'termid': 'DOID:1287', 'name': 'cardiovascular system disease'},
                    {'group': 'disease', 'termid': 'DOID:2349', 'name': 'arteriosclerosis'}
                ],
                "cells": [
                    {"group": "cells", "name": "endothelial cell", "termid": "META:52"}
                ],
            },
        'andreou_fig2_cv':
            {
                'sentences': "false",
                "disease": [
                    {'group': 'disease', 'termid': 'DOID:1287', 'name': 'cardiovascular system disease'},
                    {'group': 'disease', 'termid': 'DOID:2349', 'name': 'arteriosclerosis'}
                ]
            },
        'andreou_table1_cv':
            {
                'sentences': "false",
                "disease": [
                    {'group': 'disease', 'termid': 'DOID:1287', 'name': 'cardiovascular system disease'},
                    {'group': 'disease', 'termid': 'DOID:2349', 'name': 'arteriosclerosis'}
                ]
            },

        'andreou_fig1_athero':
            {
                'sentences': "false",
                "disease": [
                    {'group': 'disease', 'termid': 'DOID:1936', 'name': 'atherosclerosis'} #{'group': 'disease', 'termid': 'DOID:2349', 'name': 'arteriosclerosis'}
                ],
                "cells": [
                    {"group": "cells", "name": "endothelial cell", "termid": "META:52"}
                ],
            },
        'andreou_fig2_athero':
            {
                'sentences': "false",
                "disease": [
                    {'group': 'disease', 'termid': 'DOID:1936', 'name': 'atherosclerosis'}
                ]
            },
        'andreou_table1_athero':
            {
                'sentences': "false",
                "disease": [
                    {'group': 'disease', 'termid': 'DOID:1936', 'name': 'atherosclerosis'}
                ]
            },
    }

    restrictDF = DataFrame()
    restrictDF.addColumns(["Network", "Cells", "Disease", "Other"], "")

    for x in networkRestrictions:

        restricts = networkRestrictions[x]

        networkDRdict = defaultdict(str)
        networkDRdict["Network"] =  x.replace("_", " ")

        diseaseElems = []
        cellElems = []
        otherElems = []

        for restrictType in restricts:

            if restrictType == "sentences":
                continue

            if restrictType in ["disease"]:
                for elem in restricts[restrictType]:
                    diseaseElems.append( elem['name'] + " ("+elem['termid']+")")

            elif restrictType in ["cells"]:
                for elem in restricts[restrictType]:
                    cellElems.append( elem['name'] + " ("+elem['termid']+")")

            else:
                for elem in restricts[restrictType]:
                    otherElems.append( elem['name'] + " ("+elem['termid']+")")

        networkDRdict['Cells'] =  "\makecell[l]{" +"\\\\".join(sorted(cellElems)) + "}"
        networkDRdict['Disease'] = "\makecell[l]{" + "\\\\".join(sorted(diseaseElems)) + "}"
        networkDRdict['Other'] = "\makecell[l]{" + "\\\\".join(sorted(otherElems)) + "}"

        dr = DataRow.fromDict(networkDRdict)
        restrictDF.addRow(dr)

    print(restrictDF._makeLatex())

    mirna2cellOut = open("/mnt/d/yanc_network/important_networks.txt", 'w')

    def acceptEvidence(ev):

        return True

    figidx = 0

    for network in networks:
        figidx+= 1

        networkGraph = nx.Graph()

        if network in ignoreNetworks:
            continue

        #if not network in ['large_chemokines_athero']: #'large_chemokines', 'large_chemokines_cv',
        #    continue

        interactions = networks[network]
        acceptedInteractions = defaultdict(set)
        typeByGene = defaultdict(lambda: Counter())
        elemsByGene = defaultdict(lambda: defaultdict(set))

        allMirna = set()
        miStr2mirna = {}

        normalizedInteractions = defaultdict(set)

        for gene in interactions:

            for mirna in interactions[gene]:

                try:
                    oMirna = miRNA(mirna)
                    mirnaN = oMirna.getStringFromParts([miRNAPART.MATURE, miRNAPART.ID, miRNAPART.PRECURSOR])
                    mirna = mirnaN

                except:
                    pass

                allMirna.add(mirna)
                normalizedInteractions[gene].add(mirna)

        allTargetMirna = []

        for x in allMirna:
            try:

                oMirna = miRNA(x)

                allTargetMirna.append( oMirna )
                miStr2mirna[x] = oMirna

            except:
                pass

        requestData = None

        if network in networkRestrictions:
            requestData = networkRestrictions[network]
        else:
            requestData = {
                            'sentences': "false",
                            }

        requestData["gene"]= [x for x in interactions]
        #requestData['disease'] = [{'group': 'disease', 'termid': 'DOID:1287', 'name': 'cardiovascular system disease'}]

        graph, nodeCounter, edge2datasourceCount, jsonRes = DataBasePlotter.fetchGenes(requestData, minPMIDEvCount=0, minTgtCount=0, acceptEv=acceptEvidence)

        print(len(jsonRes['rels']))

        htmlDF = DataFrame()
        htmlDF.addColumns(['gene rel', 'gene', 'miRNA Group', 'miRNA', 'Original Network', 'PubMed', 'MIRECORD', 'MIRTARBASE', 'DIANA', 'Disease', 'Cells', 'GO'])

        foundGenes = set()

        interactionCountExpDB = 0
        interactionCountPubmed = 0
        interactionCountPubmedExpDB = 0

        interactionsWithPubmed = set()
        interactionsWithExpDB = set()

        interactionCountAccepted = 0
        interactionCountAdditional = 0
        interactionCountMissing = 0

        expDBEdge2Source = defaultdict(set)
        mirnaGeneSourceCounter = Counter()
        mirnaGeneIntSourceCounter = Counter()

        intToSources = defaultdict(set)

        simpleMirna2Genes = defaultdict(set)
        gene2simpleMirna = defaultdict(set)

        distinctMirna2Genes = defaultdict(set)
        gene2distinctMirna = defaultdict(set)

        simpleInteractions = set()
        distinctInteractions = set()

        pmidEvidences = set()

        for rel in jsonRes['rels']:

            orderedEdge = [None, None]

            if rel['ltype'] == "gene":
                orderedEdge[0] = rel['lid']
            elif rel['ltype'] == "mirna":
                orderedEdge[1] = rel['lid']

            if rel['rtype'] == "gene":
                orderedEdge[0] = rel['rid']
            elif rel['rtype'] == "mirna":
                orderedEdge[1] = rel['rid']

            if orderedEdge[0] == "TNPO2":
                orderedEdge[0] = "KPNA4"

            orderedEdges = set()

            if orderedEdge[1].startswith("microRNAS"):
                continue

            wasAccepted = False
            """
            for tMirna in allTargetMirna:

                if tMirna.accept(orderedEdge[1]):

                    wasAccepted = True

                    orderedEdges.add(
                        (orderedEdge[0], str(tMirna))
                    )
            """

            if not wasAccepted:
                orderedEdges.add(tuple(orderedEdge))

            for oEdge in orderedEdges:

                foundGenes.add(oEdge[0])
                origEdge = tuple(oEdge)

                try:
                    miObj = miRNA(oEdge[1])
                    miStr = miObj.getStringFromParts([miRNAPART.MATURE, miRNAPART.ID, miRNAPART.PRECURSOR])

                    selfObj = miRNA(miStr)

                    oEdge = list(oEdge)
                    oEdge[1] = miStr
                    oEdge = tuple(oEdge)
                except:
                    print("Ignoring edge for mirna", origEdge, oEdge)
                    continue
                    #pass

                edgeStatus = None

                allGeneMirna = normalizedInteractions[oEdge[0]]
                miAccepted = False

                allAcceptedStr = set()

                if network == "macrophages_athero":
                    print(oEdge)

                for strMirna in allGeneMirna:
                    miObj = miStr2mirna.get(strMirna, None)

                    if miObj == None:
                        continue

                    miObjAccepts = miObj.accept(oEdge[1], compLevel=miRNACOMPARISONLEVEL.PRECURSOR)
                    miAccepted = miAccepted or miObjAccepts

                    if "fig1_athero" in network:
                        if oEdge[0] == "IRAK1":
                            print(oEdge)
                            miObjAccepts = miObj.accept(oEdge[1], compLevel=miRNACOMPARISONLEVEL.PRECURSOR)

                    if miObjAccepts:
                        acceptedInteractions[oEdge[0]].add(strMirna)
                        edgeStatus = "accepted"

                        allAcceptedStr.add(strMirna)

                        #print(oEdge[0], oEdge[1], strMirna)

                if network == "macrophages_athero":
                    print(allAcceptedStr)

                if not miAccepted:
                    edgeStatus = "additional"

                if edgeStatus == 'accepted':
                    interactionCountAccepted += 1
                else:
                    interactionCountAdditional += 1

                # determine kindness of interaction
                edgeResources = edge2datasourceCount.get(origEdge, None)

                edgeLabel = ""
                if edgeResources != None:

                    hasPMID = False
                    hasExpDB = False

                    if edgeResources['pmid'] > 0:
                        interactionCountPubmed += 1
                        interactionsWithPubmed.add(oEdge)
                        hasPMID = True
                        edgeLabel += "p"

                    if edgeResources['DIANA'] > 0 or edgeResources['miRTarBase'] > 0 or edgeResources['miRecords'] > 0:

                        if edgeResources['DIANA'] > 0:
                            edgeLabel += "D"

                        if edgeResources['miRTarBase'] > 0:
                            edgeLabel += "T"

                        if edgeResources['miRecords'] > 0:
                            edgeLabel += "R"

                        if not "D" in edgeLabel and not "T" in edgeLabel and not "R" in edgeLabel:
                            print(edgeResources)

                        interactionCountExpDB += 1
                        interactionsWithExpDB.add(oEdge)
                        hasExpDB = True

                    if hasPMID and hasExpDB:
                        interactionCountPubmedExpDB += 1

                objMirna = miRNA(oEdge[1])
                simpleMirna = objMirna.getStringFromParts([miRNAPART.MATURE, miRNAPART.ID, miRNAPART.PRECURSOR])
                simpleEdge = tuple([oEdge[0], simpleMirna])

                simpleMirna2Genes[simpleMirna].add(oEdge[0])
                gene2simpleMirna[oEdge[0]].add(simpleMirna)

                distinctMirna2Genes[oEdge[1]].add(oEdge[0])
                gene2distinctMirna[oEdge[0]].add(oEdge[1])

                networkGraph.add_edge(oEdge[0], oEdge[1], edgeResources=edgeResources, label=edgeLabel, color= 'g' if edgeStatus == "accepted" else "b")

                typeByGene[oEdge[0]][edgeStatus] += 1
                elemsByGene[oEdge[0]][edgeStatus].add( oEdge[1] )

                simpleInteractions.add(simpleEdge)
                distinctInteractions.add(oEdge)

                pmidEvs = set()
                mirtarbaseEvs = set()
                mirecordsEvs = set()
                dianaEvs = set()

                interactionEvidences = set()
                docDiseases = []
                docCells = []
                docGOs = []

                for ev in rel['evidences']:

                    if 'docid' in ev:
                        docid = ev['docid']

                        disEvs = jsonRes['pmidinfo'].get('disease', {}).get(docid, {})

                        for disEv in disEvs:
                            did = disEv['termid']
                            dname = disEv['termname']

                            docDiseases.append((dname, did, docid))

                        cellEvs = jsonRes['pmidinfo'].get('cells', {}).get(docid, {})

                        for cellEv in cellEvs:
                            did = cellEv['termid']
                            dname = cellEv['termname']

                            docCells.append((dname, did, docid))

                            for ct in cellType2AccTerms:

                                ctTerms = cellType2AccTerms[ct]

                                if did in ctTerms:
                                    print(network, oEdge[0], oEdge[1], ct, docid, sep="\t", file=mirna2cellOut)

                        goEvs = jsonRes['pmidinfo'].get('go', {}).get(docid, {})

                        for goEv in goEvs:
                            did = goEv['termid']
                            dname = goEv['termname']

                            docGOs.append((dname, did, docid))

                    if ev['data_source'] == "DIANA":
                        dianaEvs.add( (ev['method'], (ev['direction'])))

                        expDBEdge2Source[oEdge].add(('DIANA', ev['method']))

                        mirnaGeneSourceCounter['DIANA'] += 1
                        interactionEvidences.add("DIANA")

                        intToSources[simpleEdge].add("DIANA")

                    elif ev['data_source'] == "miRTarBase":
                        mirtarbaseEvs.add(
                            (ev['data_id'], ",".join(ev['exp_support']), ev['functional_type'], ev['docid'])
                        )
                        expDBEdge2Source[oEdge].add(('MIRTARBASE', (ev['data_id'], ",".join(ev['exp_support']), ev['functional_type'], ev['docid'])))
                        mirnaGeneSourceCounter['MIRTARBASE'] += 1
                        interactionEvidences.add("MIRTARBASE")
                        intToSources[simpleEdge].add("MIRTARBASE")

                        if ev['docid'] != None:
                            pmidEvidences.add(ev['docid'])

                    elif ev['data_source'] == "pmid":
                        pmidEvs.add((ev['docid'],))

                        expDBEdge2Source[oEdge].add(('PMID', (ev['docid'])))
                        mirnaGeneSourceCounter['PMID'] += 1
                        interactionEvidences.add("PMID")
                        intToSources[simpleEdge].add("PMID")

                        pmidEvidences.add(ev['docid'])

                    elif ev['data_source'] == "mirecords":
                        mirecordsEvs.add((ev['docid']))
                        expDBEdge2Source[oEdge].add(('MIRECORDS', (ev['docid'])))
                        mirnaGeneSourceCounter['mirecords'] += 1
                        interactionEvidences.add("mirecords")
                        intToSources[simpleEdge].add("mirecords")

                        if ev['docid'] != None:
                            pmidEvidences.add(ev['docid'])

                    else:
                        print("Unhandled data source", ev['data_source'])

                interactionEvidences = tuple(sorted(interactionEvidences))
                mirnaGeneIntSourceCounter[interactionEvidences] += 1

                dianaLink = "http://carolina.imis.athena-innovation.gr/diana_tools/web/index.php?r=tarbasev8%2Findex&miRNAs%5B%5D=&genes%5B%5D={geneCap}&genes%5B%5D={geneLow}&sources%5B%5D=1&sources%5B%5D=7&sources%5B%5D=9&publication_year=&prediction_score=&sort_field=&sort_type=&query=1".format(
                    geneCap=oEdge[0].upper(), geneLow=oEdge[1].upper())

                pmidStr = "<br/>".join(
                    [
                        "<a href=\"https://www.ncbi.nlm.nih.gov/pubmed/{pmid}\" target=\"_blank\">{pmid}</a>".format(pmid=elem[0]) for elem in pmidEvs
                    ]
                )

                mirtarbaseStr = "<br/>".join(
                    [
                        "<a href=\"http://mirtarbase.mbc.nctu.edu.tw/php/detail.php?mirtid={mtbid}\">{mtbid}</a>".format(mtbid=elem[0]) for elem in mirtarbaseEvs
                    ]
                )

                mirecordStr = "<br/>".join(
                    [
                        "<a href=\"https://www.ncbi.nlm.nih.gov/pubmed/{pmid}\" target=\"_blank\">{pmid}</a>".format(pmid=elem[0]) for elem in mirecordsEvs
                    ]
                )

                dianaStr = "<br/>".join(
                    [
                        "{method} {direction}".format(method=elem[0], direction=elem[1]) for elem in dianaEvs
                    ]
                )

                goStr = "<br/>".join(
                    [
                        "{method} ({direction}, {docid})".format(method=elem[0], direction=elem[1], docid=elem[2]) for elem in docGOs
                    ]
                )

                cellStr = "<br/>".join(
                    [
                        "{method} ({direction}, {docid})".format(method=elem[0], direction=elem[1], docid=elem[2]) for elem in docCells
                    ]
                )

                diseaseStr = "<br/>".join(
                    [
                        "{method} ({direction}, {docid})".format(method=elem[0], direction=elem[1], docid=elem[2]) for elem in docDiseases
                    ]
                )

                addRow = {
                    'gene rel': oEdge[0] + "<br/>" + oEdge[1],
                    'gene': oEdge[0],
                    'miRNA Group': objMirna.getStringFromParts([miRNAPART.MATURE, miRNAPART.ID, miRNAPART.PRECURSOR]),
                    'miRNA': "<br/>".join(allAcceptedStr),
                    'Original Network': "{edgestate}</br>".format(edgestate=edgeStatus) +
                                        "<a href=\"https://www.ncbi.nlm.nih.gov/pubmed/?term={miRes}+{miShort}\">Search PUBMED</a>".format(miRes=oEdge[1], miShort=objMirna.getStringFromParts([miRNAPART.MATURE, miRNAPART.ID]))+
                                        "</br><a href=\"{dianaLink}\">Search DIANA</a>".format(dianaLink=dianaLink)
                    ,
                    'PubMed': pmidStr,
                    'MIRECORD': mirecordStr,
                    'MIRTARBASE': mirtarbaseStr,
                    'DIANA': dianaStr,
                    "Cells": cellStr,
                    "Disease": diseaseStr,
                    "GO": goStr
                }

                row = DataRow.fromDict(addRow)
                htmlDF.addRow(row)

        for gene in normalizedInteractions:

            for mirna in normalizedInteractions[gene]:

                edgeWasFound = mirna in acceptedInteractions[gene]

                if edgeWasFound:
                    continue

                edgeStatus = "missed"

                interactionCountMissing += 1

                networkGraph.add_edge(gene, mirna, color='r')

                typeByGene[gene][edgeStatus] += 1
                elemsByGene[gene][edgeStatus].add(mirna)

                objMirna = miRNA(mirna)
                dianaLink = "http://carolina.imis.athena-innovation.gr/diana_tools/web/index.php?r=tarbasev8%2Findex&miRNAs%5B%5D=&genes%5B%5D={geneCap}&genes%5B%5D={geneLow}&sources%5B%5D=1&sources%5B%5D=7&sources%5B%5D=9&publication_year=&prediction_score=&sort_field=&sort_type=&query=1".format(
                    geneCap=gene.upper(), geneLow=gene.capitalize())

                cellStr = ""
                diseaseStr = ""
                goStr = ""

                addRow = {
                    'gene rel': gene,
                    'gene': gene,
                    'miRNA Group': objMirna.getStringFromParts([miRNAPART.MATURE, miRNAPART.ID, miRNAPART.PRECURSOR]),
                    'miRNA': mirna,
                    'Original Network': "{edgestate}</br>".format(edgestate=edgeStatus) +
                                        "<a href=\"https://www.ncbi.nlm.nih.gov/pubmed/?term={gene} {miRes}+{miShort}\">Search PUBMED</a>".format(gene=gene,miRes=mirna, miShort=objMirna.getStringFromParts([miRNAPART.MATURE, miRNAPART.ID]))+
                                        "</br><a href=\"{dianaLink}\">Search DIANA</a>".format(dianaLink=dianaLink)
                    ,
                    'PubMed': "",
                    'MIRECORD': "",
                    'MIRTARBASE': "",
                    'DIANA': "",
                    "Cells": cellStr,
                    "Disease": diseaseStr,
                    "GO": goStr
                }

                row = DataRow.fromDict(addRow)
                htmlDF.addRow(row)

        allDegrees = networkGraph.degree()

        singleMirnaGenes = set()
        multiMirnaGenes = set()

        singleGeneMirnas = set()
        multiGeneMirnas = set()

        """
        THIS ANALYSIS ALSO HAS ADDED EDGES FROM MISSING!!!!!!
        """
        for node, nodedegree in allDegrees:

            if node.startswith("miR-") or node.startswith("let-"):

                if nodedegree == 1:
                    singleGeneMirnas.add(node)
                elif nodedegree >= 1:
                    multiGeneMirnas.add(node)

            else:

                if nodedegree == 1:
                    singleMirnaGenes.add(node)
                elif nodedegree >= 1:
                    multiMirnaGenes.add(node)

        print(network)

        print("Simple Interactions", len(simpleInteractions))
        print("Distinct Interactions", len(distinctInteractions))

        print("Number of simple mirnas", len(simpleMirna2Genes))
        print("Number of simple mirnas with 1 gene", sum([1 for x in simpleMirna2Genes if len(simpleMirna2Genes[x]) == 1]))
        print("Number of simple mirnas > 1 gene",
              sum([1 for x in simpleMirna2Genes if len(simpleMirna2Genes[x]) > 1]))

        print("Number of genes", len(gene2simpleMirna))
        print("Number of genes with 1 simple mirna",
              sum([1 for x in gene2simpleMirna if len(gene2simpleMirna[x]) == 1]))
        print("Number of genes with >1 simple mirna",
              sum([1 for x in gene2simpleMirna if len(gene2simpleMirna[x]) > 1]))
        # anzahl gene mit nur einer miRNA
        print("Number of genes with only one miRNA", len(singleMirnaGenes))

        print("Number of distinct mirnas", len(distinctMirna2Genes))
        print("Number of distinct mirnas with 1 gene", sum([1 for x in distinctMirna2Genes if len(distinctMirna2Genes[x]) == 1]))
        print("Number of distinct mirnas > 1 gene",
              sum([1 for x in distinctMirna2Genes if len(distinctMirna2Genes[x]) > 1]))

        print("Number of genes", len(gene2distinctMirna))
        print("Number of genes with 1 distinct mirna",
              sum([1 for x in gene2distinctMirna if len(gene2distinctMirna[x]) == 1]))
        print("Number of genes with >1 distinct mirna",
              sum([1 for x in gene2distinctMirna if len(gene2distinctMirna[x]) > 1]))
        # anzahl gene mit nur einer miRNA
        print("Number of genes with only one miRNA", len(singleMirnaGenes))

        # anzahl gene mit mehr als einer miRNA
        print("Number of genes with more than one miRNA", len(multiMirnaGenes))

        print("Number of miRNAs with only one genes", len(singleGeneMirnas))
        print("Number of miRNAs with more than one genes", len(multiGeneMirnas))
        print("miRNA Gene Interaction source counter", sum([mirnaGeneSourceCounter[x] for x in mirnaGeneSourceCounter]), mirnaGeneSourceCounter)
        print("miRNA Gene Interactions detailed", sum([mirnaGeneIntSourceCounter[x] for x in mirnaGeneIntSourceCounter]), mirnaGeneIntSourceCounter)

        intToSourcesCount = Counter()
        for ed in intToSources:
            uniedge = tuple(sorted(intToSources[ed]))
            intToSourcesCount[uniedge] += 1

        print("miRNA(simple) Gene Interactions detailed", sum([intToSourcesCount[x] for x in intToSourcesCount]), intToSourcesCount)

        totalInteractionCount = interactionCountAccepted+ interactionCountAdditional+ interactionCountMissing
        totalGraphInteractionCount = interactionCountAccepted  + interactionCountAdditional
        # anzahl kanten in graph
        print("Number of interactions (acc, add, graph, mis, total)", interactionCountAccepted, interactionCountAdditional, totalGraphInteractionCount, interactionCountMissing, totalInteractionCount)

        # anzahl kanten nur durch experimental db
        print("Number of interactions supported by experimental db", len(interactionsWithExpDB))

        printSet = interactionsWithExpDB.difference(interactionsWithPubmed)
        if len(printSet) > 50:
            printSet = {}
        print("Number of interactions supported only by experimental db", len(interactionsWithExpDB.difference(interactionsWithPubmed)), printSet)

        for edge in printSet:
            print(edge, expDBEdge2Source.get(edge, None))

        # anzahl kanten nur durch pubmed
        print("Number of interactions supported by pubmed", len(interactionsWithPubmed))
        print("Number of interactions supported only by pubmed", len(interactionsWithPubmed.difference(interactionsWithExpDB)))

        # anzahl kanten mit pubmed UND experimental db
        printSet = interactionsWithExpDB.intersection(interactionsWithPubmed)
        if len(printSet) > 50:
            printSet = {}
        print("Number of interactions supported by pubmed+expdb", len(interactionsWithExpDB.intersection(interactionsWithPubmed)), printSet)
        for edge in printSet:
            print(edge, expDBEdge2Source.get(edge, None))
        # anzahl kanten nur mit exp hinweisen aber zu weak

        # anzahl kanten nur durch pubmed aber zu weak
        print("Number of PMIDs", len(pmidEvidences))

        print()
        print()
        print()
        print()
        for gene in sorted([x for x in typeByGene]):
            if "athero" in network:
                print(gene, typeByGene[gene], elemsByGene[gene]['additional'], elemsByGene[gene]['missing'])
            else:
                print(gene, typeByGene[gene], elemsByGene[gene]['missing'])

        print()
        print()

        networkDF = DataFrame()
        networkDF.addColumns(["gene", "accepted", "additional", "missing", "missing miRNAs"])

        print(network)
        for gene in sorted([x for x in typeByGene]):

            infoDict = {
                "gene": gene,
                "accepted": typeByGene[gene]["accepted"],
                "additional": typeByGene[gene]["additional"],
                "missing": typeByGene[gene]["missing"],
                "missing miRNAs": ", ".join(elemsByGene[gene]['missing'])
            }

            dr = DataRow.fromDict(infoDict)
            networkDF.addRow(dr)

            print("Gene:", gene, "Status: ", ", ".join([": ".join([x, str(typeByGene[gene][x])]) for x in typeByGene[gene]]), "Missing miRNAs: "+",".join(elemsByGene[gene]['missing']))

        print()
        print()

        print(networkDF._makeLatex())

        print()
        print()

        networkGraphs[network] = networkGraph

        htmlDF.export("/mnt/d/yanc_network/" + network.replace(" ", "_") + ".html", ExportTYPE.HTML)
        htmlDF.export("/mnt/d/yanc_network/" + network.replace(" ", "_") + ".tsv", ExportTYPE.TSV)

    print("Network covered genes", network, foundGenes)

    figidx = 0
    for stages in makeStory:

        mergedGraph = networkGraphs[stages[0]]

        for i in range(1, len(stages)):
            mergedGraph = nx.compose(mergedGraph, networkGraphs[stages[i]])

        hasLargeStage = any(['large' in stage for stage in stages])

        pos = nx.spring_layout(mergedGraph)

        for stage in stages:

            networkGraph = networkGraphs[stage]

            edges = networkGraph.edges()
            colors = [networkGraph[u][v]['color'] for u, v in edges]

            d = nx.degree(networkGraph)
            nodes = networkGraph.nodes()
            nodeColors = []
            nodeSizes = []

            allSizes = [x[1] for x in d]
            minSize = min(allSizes)
            maxSize = max(allSizes)
            diffSize = maxSize-minSize

            fontSize = 16
            minNodeSize = 1200
            figSize = (20, 14)
            edgeWidth = 3

            if hasLargeStage:
                fontSize = 8
                minNodeSize = 100
                figSize = (20,30)
                edgeWidth = 0.75

            plt.figure(figidx, figsize=figSize)
            figidx += 1

            maxNodeSize = 3000
            diffNodeSize = maxNodeSize-minNodeSize
            nodeList = []

            for x in nodes:
                if any([x.lower().startswith(y) for y in ['mir', 'let']]):
                    nodeColors.append('blue')
                else:
                    nodeColors.append('green')

                nodeDegree = d[x]

                nodeDegree -= minSize
                nodeDegree = nodeDegree / diffSize

                nodeSize = minNodeSize + diffNodeSize * nodeDegree

                nodeSizes.append( nodeSize )
                nodeList.append(x)

            nx.draw(networkGraph, pos, font_size=fontSize, with_labels=False, node_color=nodeColors, edges=edges, edge_color=colors, nodelist=nodeList, node_size=nodeSizes, width=edgeWidth, font_weight='bold', dpi=1000)

            for p in pos:  # raise text positions

                clist = list(pos[p])

                if p in nodeList:
                    if nodeSizes[nodeList.index(p)] < 1000:
                        clist[1] = clist[1] + 0.005
                    else:
                        clist[1] = clist[1] + 0.02
                pos[p] = tuple(clist)

            nx.draw_networkx_labels(networkGraph, pos, font_weight='bold', font_size=fontSize)

            plt.suptitle(stage)

            plt.savefig("/mnt/d/yanc_network/" + stage.replace(" ", "_") + ".png")
            plt.savefig("/mnt/d/yanc_network/" + stage.replace(" ", "_") + ".pdf")

            CytoscapeGrapher.showGraph(networkGraph, location="/mnt/d/yanc_network/", name="cyjs_" + stage.replace(" ", "_"), title=stage)
            CytoscapeGrapher.exportGraphML(networkGraph, location="/mnt/d/yanc_network/", name="cyjs_" + stage.replace(" ", "_"))

            CytoscapeGrapher.export_d3Json(networkGraph, location="/mnt/c/Users/mjopp/Desktop/d3test/",name= "cyjs_" + stage.replace(" ", "_"), singletons=True )

            if stage == "large_chemokines":

                with open("/mnt/c/Users/mjopp/Desktop/gergely.tsv", "w") as fgout:

                    existingEdges = set()

                    for gene in largeInteractions:

                        for mirna in largeInteractions[gene]:
                            try:
                                miObj = miRNA(mirna)
                                miStr = miObj.getStringFromParts([miRNAPART.MATURE, miRNAPART.ID, miRNAPART.PRECURSOR])
                                selfObj = miRNA(miStr)

                                existingEdges.add((gene, miStr))
                            except:
                                continue

                    for node in exportGraph.nodes():
                        nodename = str(node).lower()

                        if nodename.startswith("mir") or nodename.startswith("let"):
                            fmtNode = "<node y=\"20\" x=\"360\" shape=\"ellipse\" color=\"FFFFFF\" label=\"{elemid}\" id=\"{elemid}\" bgfixed=\"-1\"/>".format(elemid=node)
                            #print(fmtNode)

                    for edge in exportGraph.edges():
                        fmtEdge = "<edge dir=\"false\" bold=\"false\" color=\"000000\" width=\"null\" inhibit=\"false\" rev=\"false\" arrowcolor=\"000000\" to=\"{src}\" from=\"{tgt}\"/>".format(src=edge[0], tgt=edge[1])
                        #print(fmtEdge)

                        if edge[0].lower().startswith("mir") or edge[0].lower().startswith("let"):
                            mirna = edge[0]
                            gene = edge[1]

                        else:

                            mirna = edge[1]
                            gene = edge[0]

                        if (gene, mirna) in existingEdges:
                            continue

                        print(gene, mirna, sep="\t", file=fgout)

    #plt.show()

import numpy,import scipy,import networkx,from networkx


# -*- coding: utf-8 -*-
import pandas as pd
import jieba
import jieba.analyse
import numpy as np
from keras.models import Sequential
from keras.layers import Dense,Dropout,Activation,Embedding
from keras.layers import Conv1D,GlobalMaxPooling1D
from keras.utils import np_utils
from keras.preprocessing import sequence

pos = pd.read_excel("pos.xls", header=None)
neg = pd.read_excel("neg.xls", header=None)
tmp = pd.read_excel("test.xls", header=None)
pos['label']=1
neg['label']=0
all = pos.append(neg, ignore_index=True)

all['word']=list(map(lambda s: list(jieba.cut(s,cut_all=False)), all[0]))
tmp['word']=list(map(lambda s: list(jieba.cut(s,cut_all=False)), tmp[0]))
content = [x for i in all['word'] for x in i]
abc = pd.Series(content).value_counts()
abc[:] = range(1,1+len(abc))
abc[''] = 0
maxlen = 15

def doc2num(s,maxlen):
    s=jieba.cut(s,cut_all=False)
    s=[i for i in s if i in abc]
    s+=max((maxlen-len(s)), 0)*['']
    return list(abc[s])

all['doc2num'] = list(map(lambda s :doc2num(s,maxlen),all[0]))
tmp['doc2num'] = list(map(lambda s:doc2num(s,maxlen),tmp[0]))
idx = list(range(len(all)))
np.random.shuffle(idx)
all = all.loc[idx]
x_train = np.array(list(all['doc2num']))
y_train = np.array(list(all['label']))
y_train = y_train.reshape((-1,1))
x_test = np.array(list(tmp['doc2num']))

#a="流言揭秘：牙齿的颜色，越白越健康？20170210093100摘要你被骗了！在镜头里的那些唇红齿白、巧笑嫣然的美白效果其实都是后期通过电脑修图软件处理生成，并非牙膏的实际使用效果！要知道牙齿并不是越白越健康！根据科学研究表明，最坚固的牙齿是自然健康略带黄色的摘要你被骗了！在镜头里的那些唇红齿白、巧笑嫣然的美白效果其实都是后期通过电脑修图软件处理生成，并非牙膏的实际使用效果！要知道牙齿并不是越白越健康！根据科学研究表明，最坚固的牙齿是自然健康略带黄色的你被骗了！在镜头里的那些唇红齿白、巧笑嫣然的美白效果其实都是后期通过电脑修图软件处理生成，并非牙膏的实际使用效果！要知道牙齿并不是越白越健康！根据科学研究表明，最坚固的牙齿是自然健康略带黄色的牙齿，看上去并不会呢么显白，更不会像纸张或者钢琴键一样白的颜色，而是应该再暗两个色调。当然如果你长期吸烟，那你的牙齿颜色会再深一点。这是因为牙齿外层是牙釉质，具有半透明特性，釉质内层的牙本质颜色并不是白色的，而是偏黄色，它的颜色会部分透过牙釉质，从而使牙齿颜色有一定的黄色。越健康的牙釉质钙化程度越高，其硬度也越高，这样的牙釉质更透明，其深部的牙本质的黄色也更容易透过而呈淡黄色。相反，如果牙釉质发育得不好，钙化程度不高，透明度就会变弱，牙本质颜色不能透过，牙齿就会呈现出一种不正常的白色或乳白色，比如患者患有初期龋齿时，牙面上会因脱矿而形成白垩色斑块。随着年龄的增长，牙釉质在此过程中被磨耗，牙齿会看上去比之前更黄，但这并不影响到牙齿的健康，所以，如果当你的牙齿就连洁牙之后都是黄色，说明这黄色不是来自牙面而是来自釉质本身，换个方向想这也说明你的釉质发育的比别人好，比别人坚固，是健康的证明通常情况下，因为色素沉淀的关系亚洲人的牙齿普遍比西方人更黄，vitapanclassical比色板。东方人种恒牙一般白的话在A2到A3差不多，西方人种会更白。牙齿美白剂最早出现，是为了改善美国人自身牙齿增龄性变化的。白种人随着年龄的增长，到达一定年龄时牙齿就会产生变化，比如过了30岁牙齿就不如从前白了。而这种增龄性在黄种人身上表现得并不那么明显。但并没有研究数据表明深色皮肤的人种拥有更白的牙齿颜色，他们的牙齿看上去更白更亮其实是因为与肤色的对比度比较高而已。所以，牙齿并不是越白越好，看起来自然，与自身肤色、性格协调的牙齿颜色才是最好的选择。在正常颜色范围内的牙齿，并非都需要做美白。如果因为牙齿颜色过深而影响美观者，就需要通过一定的治疗方法加以改善了。另外，除了随着年龄的牙齿变黄，饮食也会导致牙齿泛黄，小编列举了其中八种会导致黄牙的食物与饮料，它们不仅会黏附在牙釉质表面，导致牙齿染色，还会磨损牙釉质，让牙齿收到双重打击。使牙齿提供泛黄双重打击：红茶中富含可以使牙齿染色丹宁酸，因此它被认为是世界上导致牙齿着色问题最棘手的饮料之一。如果你在喝红茶的同时还享用了其他易使牙齿染色的食物和饮料，那伤害会大大增加。与此相反，绿茶被认为影响牙齿着色的因素就小的多，因此用绿茶代替红茶会使你的牙齿更健康。可乐，苏打水和运动饮料中的柠檬酸与磷酸会消磨牙釉质，之前流传过牙齿放在可乐中一夜就完全溶解”已经被证实为谣言，但是据实验证实：牙齿在可乐中浸泡一夜确实会引起严重着色，浸泡一周牙齿会变黑。即使无色的碳酸饮料不会导致牙齿着色，但其中的酸性还是会腐蚀牙釉质，研究表明，许多所谓的运动饮料比可乐具有更强的破坏性。众所周知，由于其深色的多酚和单宁酸，红葡萄酒是牙齿染色的罪魁祸首之一，但你不要以为喝白葡萄酒就能避免这个的问题，因为它一样糟糕，甚至更糟。虽然果汁并不能与可乐和颜色怪异的运动饮料相提并论，但非鲜榨果汁的酸性确实比大多数人所意识到的高。新鲜的果汁并没有经过处理，所以没有太大的问题，但是也同样不要让它们在你的牙齿上停留太久。当你知道你心爱的食物和饮料有可能导致牙齿着色后，第一反应肯定认为是吃完立马刷牙就没事了，但这个想法是错误的。在食用过酸性饮料或食物后，牙釉质会在接下来的一个小时中变得没有呢么坚硬，如果你在这个时间段内刷牙，其实就是在磨损你的牙釉质。你可以选择在牙釉质再次变硬后用软毛牙刷刷牙。当然，更好的选择是在你食用了易使牙齿着色的食物后，做这样两个小步骤：1、在结束后含一大口水在口中并漱口，用来清除酸和染色物质，你可以感受一下喝完可乐后的口腔与喝完可乐再漱口后的口腔，其实是有很大区别的。2、让你的口腔分泌更多的唾液，唾液具有清洁和保护作用，可以冲洗掉口中和牙齿间的食物残渣，抵抗口腔微生物，起到抗菌和杀菌的作用，避免口腔感染，是一道天然屏障。同时，唾液中所含的大量微量元素，可以促进被损伤的牙釉质的再矿化，从而预防蛀牙的进一步发展。所以，小编说了这么多，其实并不是让你完全不去碰你心爱的食物哦，比起牙齿来说，戒掉汽水更有助于缩减你的腰围（笑小编只是为了让你知道随着时间的推移，是什么原因导致牙齿变黄，并提供一个简单的小仪式，让你能够在吃喝之后好好保护你的牙齿哟微科普流言揭秘：牙齿的颜色，越白越健康？"
#a = ''.join(a[:100]+a[-100:])
#b = jieba.analyse.extract_tags(a)
#a = list(jieba.cut(a,cut_all=False))

embedding_vector_length = 32
min_count = 5
batch_size = 32
nb_epoch = 10
nb_filter = 128
filter_length = 3

model = Sequential()
model.add(Embedding(len(abc),
                    embedding_vector_length,
                    input_length=maxlen))
model.add(Conv1D(activation="relu",
                 filters=128,
                 kernel_size=3,
                 padding="valid"))
model.add(GlobalMaxPooling1D())

model.add(Dense(128))
model.add(Dropout(0.2))
model.add(Activation('relu'))

model.add(Dense(1))
model.add(Activation('sigmoid'))
model.compile(loss='binary_crossentropy',
              optimizer='adam',
              metrics=['accuracy'])
model.fit(x_train,y_train,
          batch_size==batch_size,
          epochs=nb_epoch)

pre = model.predict_classes(x_test)

print(pre)


# given list of target seqs (col. headed 'Seqs' in CSV),
# config containing: constant pad positions to remove from both ends, required len. of oligo, number of pools,
# get DF with columns:
# pool_id: name of pool
# seq_id: name of this sequence within the pool
# oligo_f: 'fwd' oligo to synthesize
# oligo_r: 'rev' oligo to synthesize
# pool_primer_f: 'fwd' primer for 'this' pool (binds oligo_r) (+)
# pool_primer_r: 'rev' primer for 'this' pool (binds oligo_f) (+) (!)
# gg_site: golden gate site for this construct
# gg_prod: final output of the construct
# const_primer_f: constant forward primer (binds oligo_f)
# const_primer_r: constant reverse primer (binds oligo_r)

# Process this DF with methods in 'validate_and_get_primers.py'

import sys
import os
import pandas as pd
import numpy as np
import random
from numpy.random import choice
import ConfigParser

DNA = ['A','C','G','T']
DNA_C = {'A':'T', 'C':'G', 'G':'C','T':'A'}

# Represent a set of sequences, divided into two oligos, to be joined by Golden Gate assembly.
# Track the use of Golden Gate sites so that no assembly shares a site.

class sequence_pool(object):
  def __init__(self, num_seqs, site_start, site_end, gg_len):
    # the extra column is so that one col. will always have 0 - covers a potential corner case
    self.gg_mat = np.zeros(shape = (num_seqs, len(DNA)**gg_len + 1), dtype = 'bool')
    self.ss = site_start
    self.se = site_end
    self.gg_len = gg_len
    self.gg_to_bin_id = {} # map 4-base GG sequences to columns in self.gg_mat, and vice versa
    self.curr_col = 0
    self.seq_to_useable_ggs = {} # map a sequence to a (gg_seq -> gg_start) dictionary
    self.seqs = []

  def add(self, seq):
    seq_id = len(self.seqs)
    self.seqs.append(seq)
    useable_ggs = []
    useable_starts = []
    gg_start = self.ss
    gg_end = gg_start + self.gg_len
    # what GG sites could work for this sequence?
    while gg_end <= self.se:
      useable_ggs.append(seq[gg_start:gg_end])
      useable_starts.append(gg_start)
      gg_start += 1; gg_end += 1
    self.seq_to_useable_ggs[seq] = dict(zip(useable_ggs, useable_starts))
    # record this sequence's useable GG sites
    for g in useable_ggs:
      if g not in self.gg_to_bin_id:
        self.gg_to_bin_id[g] = self.curr_col; self.gg_to_bin_id[self.curr_col] = g
        self.curr_col += 1
      self.gg_mat[seq_id, self.gg_to_bin_id[g]] = True

  # assign GGs to sequences; try to assign as many as possible.
  # Heuristic: find the GG used by the fewest sequences, and assign to that one.
  # Take that sequence out of consideration; continue until no remaining bins can be used.
  def assign_seqs(self):
    accepted = np.zeros(shape = (len(self.seqs),), dtype = 'bool')
    ans = []
    while True:
      bin_counts = np.apply_along_axis(np.sum, 0, self.gg_mat)
      u_cts, u_idxes = np.unique(bin_counts, return_index = True)
      if len(u_cts) == 1: # they're all 0
          break
      bin_id = u_idxes[1] # u_cts[0] will be 0
      seq_id = np.min(np.where(self.gg_mat[:,bin_id])[0])
      accepted[seq_id] = True
      self.gg_mat[seq_id,:] = False # this sequence can be removed
      self.gg_mat[:,bin_id] = False # this bin is used
      g = self.gg_to_bin_id[bin_id] # this is a bidirectional dictionary; get the GG sequence
      # get the actual seq and starting position
      seq = self.seqs[seq_id]
      g_s = self.seq_to_useable_ggs[seq][g]
      ans.append((seq, g_s, g))
    rejected = np.array(self.seqs)[np.logical_not(accepted)]
    return(ans, rejected.tolist())

# generate a random pad sequence that won't contain cut sites.
# Context is required to ensure we don't finish a partial site that started before the pad.
def safe_pad(seq, pad_len, forbidden_site_list, is_fwd):
  max_forbid_len = max([len(q) for q in forbidden_site_list])
  for (i,f) in enumerate(forbidden_site_list): # strip N's from binding sites
    f = f.strip('N'); assert('N' not in f)
    forbidden_site_list[i] = f
  if is_fwd:
    test_seq = seq + ''.join( np.random.choice(np.array(DNA), size = pad_len, replace = True).tolist() )
    test_seq_sitecheck = test_seq[-(pad_len + max_forbid_len):]
  else:
    test_seq = ''.join( np.random.choice(np.array(DNA), size = pad_len, replace = True).tolist() ) + seq
    test_seq_sitecheck = test_seq[:(pad_len + max_forbid_len)]
  if any([q in test_seq_sitecheck for q in forbidden_site_list]):
    # recursion approach. Keep trying until a forbidden-site-free sequence is generated.
    return(safe_pad(seq, pad_len, forbidden_site_list, is_fwd))
  return(test_seq)

# Given a list of sequences and a dictionary of params, get the output table as a Pandas DataFrame.
def build_pools_table(seqs, params):
  oligo_len = int(params['oligo_len'])
  gg_site_len = int(params['gg_site_len'])
  fwd_pools = params['fwd_pools']
  rev_pools = params['rev_pools']
  len_fwd_seq = oligo_len - len(params['fwd_toehold']) - gg_site_len - len(fwd_pools[0])
  len_rev_seq = oligo_len - len(params['rev_toehold']) - gg_site_len - len(rev_pools[0])
  # golden gate site can start at len_seqs - len_rev_seq - len_gg_site; must end at/before len_fwd_seq + len_gg_site
  site_start = len(seqs[0]) - len_rev_seq - gg_site_len
  site_end = len_fwd_seq + gg_site_len
  num_pools = len(fwd_pools)
  pools = []
  for i in range(num_pools):
    if len(seqs) == 0:
      break
    this_pool = sequence_pool(len(seqs), site_start, site_end, gg_site_len)
    for s in seqs:
      this_pool.add(s)
    ans, seqs = this_pool.assign_seqs()
    pools.append(ans)
  ### Build the output DataFrame from 'pools' ###
  # 'pools' is a list of lists of tuples; each inner ('pool') list corresponds to a pool.
  # Each pool list contains tuples (seq, g_s); each 'seq' is an original sequence
  # (the thing to be constructed, without any padding), each 'g_s' is the starting position for the GG site.
  # Create a table with columns 'Design', 'gg_start', 'fwd_pool', 'rev_pool'
  dfs = []
  for (i,q) in enumerate(pools):
    pool_dict = {'Design':[], 'gg_start':[], 'fwd_pool':[],'rev_pool':[], 'Experiment':[],'pool_id':[], 'seq_id':[], 'gg_site_seq':[]}
    for j,(seq, g_s, g) in enumerate(q):
      pool_dict['Design'].append(seq)
      pool_dict['gg_start'].append(g_s)
      pool_dict['gg_site_seq'].append(g)
      pool_dict['fwd_pool'].append(fwd_pools[i])
      pool_dict['rev_pool'].append(rev_pools[i])
      pool_dict['Experiment'].append(params['assembly_id'])
      pool_dict['pool_id'].append(i)
      pool_dict['seq_id'].append(j)
    dfs.append(pd.DataFrame(pool_dict))
  table = pd.concat(dfs)

  return(table, seqs) # 'seqs' is now the rejected sequences

def fill_one_oligo(tuple_in, params):
  seq, gg_start, fwd_pool, rev_pool = tuple_in
  oligo_len = int(params['oligo_len'])
  gg_site_len = int(params['gg_site_len'])
  forbidden_site_list = [params['fwd_gg_cut'], params['rev_gg_cut']]
  fwd_prefix = params['fwd_toehold']
  fwd_postfix = fwd_pool # NB: GG site has to be included here!
  rev_prefix = rev_pool # see above
  rev_postfix = params['rev_toehold']
  fwd_seq = fwd_prefix + seq[:(gg_start+gg_site_len)] + fwd_postfix
  rev_seq = rev_prefix + seq[gg_start:] + rev_postfix
  assert(len(fwd_seq) <= oligo_len)
  assert(len(rev_seq) <= oligo_len)
  fwd_seq = safe_pad(fwd_seq, oligo_len - len(fwd_seq), forbidden_site_list, True)
  rev_seq = safe_pad(rev_seq, oligo_len - len(rev_seq), forbidden_site_list, False)
  return(fwd_seq, rev_seq)

# create final oligos; add PCR regions and padding if needed
def fill_pools_table(table, params):
  forbidden_site_list = [params['fwd_gg_cut'], params['rev_gg_cut']]
  oligo_len = int(params['oligo_len'])
  gg_site_len = int(params['gg_site_len'])
  table_zipped = zip(table['Design'], table['gg_start'], table['fwd_pool'], table['rev_pool'])
  fwd_seqs = []; rev_seqs = []
  for q in table_zipped:
    f, r = fill_one_oligo(q, params)
    fwd_seqs.append(f); rev_seqs.append(r)
  table['fwd_oligos'] = fwd_seqs
  table['rev_oligos'] = rev_seqs
  table['fwd_const'] = [params['fwd_homol'] + params['fwd_toehold']]*table.shape[0]
  table['rev_const'] = [params['rev_toehold'] + params['rev_homol']]*table.shape[0]
  # Now that we've added bases to the fwd oligo, we need to update 'gg_start' to compensate.
  table['gg_start'] = [q + len(params['fwd_toehold']) for q in table['gg_start']]
  return(table)

# For modularity: given a list of sequences and a config, get the output table as a Pandas DataFrame.
def seqs_to_df(seqs, cfg):
  assert(all([len(q) == len(seqs[0]) for q in seqs]))
  params = dict(cfg.items('Params'))
  random_seed = int(params['random_seed'])
  random.seed(random_seed); np.random.seed(random_seed)
  # remove padding from sequences
  [pad_len_f, pad_len_r] = [int(q) for q in params['pad_lens'].strip().split(',')]
  seqs = [q[pad_len_f:-pad_len_r] for q in seqs]
  # get distinguishing sequences for pool PCRs
  fwd_pools = params['fwd_pools'].strip().split(',')
  rev_pools = params['rev_pools'].strip().split(',')
  assert(len(fwd_pools) == len(rev_pools))
  assert(all([len(q) == len(fwd_pools[0]) for q in fwd_pools]))
  assert(all([len(q) == len(rev_pools[0]) for q in rev_pools]))
  params['fwd_pools'] = fwd_pools; params['rev_pools'] = rev_pools
  #pools, rejected = build_pools(seqs, params)
  #final_oligos = fill_oligos(pools, params)
  oligo_table, rejected = build_pools_table(seqs, params)
  oligo_table = fill_pools_table(oligo_table, params)
  return(oligo_table, rejected)
  #return(final_oligos, rejected)

def main(cfg):
  seqs = list(pd.read_csv(os.path.expanduser(cfg.get('Files','selected_fn')))['Seqs'])
  oligo_table, rejected = seqs_to_oligos(seqs, cfg)
  fo = os.path.expanduser(cfg.get('Files','oligo_fn'))
  oligo_table.to_csv(fo)

  '''with open(os.path.expanduser(cfg.get('Files','oligo_fn')),'w') as fo:
    for name, oligo in final_oligos:
      fo.write(name + '\n')
      fo.write(oligo + '\n')'''
  with open(os.path.expanduser(cfg.get('Files','rejected_fn')),'w') as fr:
    for seq in rejected:
      fr.write(seq + '\n')

if __name__ == '__main__':
  cfg = ConfigParser.RawConfigParser(allow_no_value=True)
  cfg.read(sys.argv[1])
  main(cfg)


from jax.config import config; config.update("jax_enable_x64", True)

import jax
import numpy as np

from fe import topology

from timemachine.lib import potentials, custom_ops
from timemachine.lib import LangevinIntegrator

from ff.handlers import openmm_deserializer

def get_romol_conf(mol):
    """Coordinates of mol's 0th conformer, in nanometers"""
    conformer = mol.GetConformer(0)
    guest_conf = np.array(conformer.GetPositions(), dtype=np.float64)
    return guest_conf/10 # from angstroms to nm

class BaseFreeEnergy():

    @staticmethod
    def _get_integrator(combined_masses):
        """
        Get a integrator. The resulting impl must be bound to a python handle
        whose lifetime is concurrent with that of the context.
        """
        seed = np.random.randint(np.iinfo(np.int32).max)

        return LangevinIntegrator(
            300.0,
            1.5e-3,
            1.0,
            combined_masses,
            seed
        )

    # this will be eventually gRPC'd out to a worker
    @staticmethod
    def _simulate(lamb, box, x0, v0, final_potentials, integrator, equil_steps, prod_steps):
        all_impls = []
        bonded_impls = []
        nonbonded_impls = []

        # set up observables for du_dps here as well.

        du_dp_obs = []

        for bps in final_potentials:
            obs_list = []

            for bp in bps:
                impl = bp.bound_impl(np.float32)

                if isinstance(bp, potentials.InterpolatedPotential) or isinstance(bp, potentials.LambdaPotential):
                    bp = bp.get_u_fn()

                if isinstance(bp, potentials.Nonbonded):
                    nonbonded_impls.append(impl)
                else:
                    bonded_impls.append(impl)

                all_impls.append(impl)
                obs_list.append(custom_ops.AvgPartialUPartialParam(impl, 5))

            du_dp_obs.append(obs_list)

        intg_impl = integrator.impl()
        # context components: positions, velocities, box, integrator, energy fxns
        ctxt = custom_ops.Context(
            x0,
            v0,
            box,
            intg_impl,
            all_impls
        )

        # equilibration
        for step in range(equil_steps):
            ctxt.step(lamb)

        bonded_du_dl_obs = custom_ops.FullPartialUPartialLambda(bonded_impls, 5)
        nonbonded_du_dl_obs = custom_ops.FullPartialUPartialLambda(nonbonded_impls, 5)

        # add observable
        ctxt.add_observable(bonded_du_dl_obs)
        ctxt.add_observable(nonbonded_du_dl_obs)

        for obs_list in du_dp_obs:
            for obs in obs_list:
                ctxt.add_observable(obs)

        for _ in range(prod_steps):
            ctxt.step(lamb)

        bonded_full_du_dls = bonded_du_dl_obs.full_du_dl()
        nonbonded_full_du_dls = nonbonded_du_dl_obs.full_du_dl()

        bonded_mean, bonded_std = np.mean(bonded_full_du_dls), np.std(bonded_full_du_dls)
        nonbonded_mean, nonbonded_std = np.mean(nonbonded_full_du_dls), np.std(nonbonded_full_du_dls)

        # keep the structure of grads the same as that of final_potentials so we can properly
        # form their vjps.
        grads = []
        for obs_list in du_dp_obs:
            grad_list = []
            for obs in obs_list:
                grad_list.append(obs.avg_du_dp())
            grads.append(grad_list)

        return (bonded_mean, bonded_std), (nonbonded_mean, nonbonded_std), grads

# this class is serializable.
class AbsoluteFreeEnergy(BaseFreeEnergy):

    def __init__(self, mol, ff):
        """
        Compute the absolute free energy of a molecule via 4D decoupling.

        Parameters
        ----------
        mol: rdkit mol
            Ligand to be decoupled

        ff: ff.Forcefield
            Ligand forcefield

        """
        self.mol = mol
        self.ff = ff
        self.top = topology.BaseTopology(mol, ff)

    # this can be used for both the solvent leg and the complex leg
    def host_edge(self, lamb, host_system, host_coords, box, equil_steps=10000, prod_steps=100000):
        """
        Run equilibrium decoupling simulation at a given value of lambda in a host environment.

        Parameters
        ----------
        lamb: float [0, 1]
            0 is the fully interacting system, and 1 is the non-interacting system

        host_system: openmm.System
            OpenMM System object to be deserialized. The host can be simply a box of water, or a fully
            solvated protein

        host_coords: np.array of shape [..., 3]
            Host coordinates, in nanometers. It should be properly minimized and not have clashes
            with the ligand coordinates.

        box: np.array [3,3]
            Periodic boundary conditions, in nanometers.

        equil_steps: float
            Number of steps to run equilibration. Statistics are not gathered.

        prod_steps: float
            Number of steps to run production. Statistics are gathered.

        Returns
        -------
        float, float
            Returns a pair of average du_dl values for bonded and nonbonded terms.

        """

        ligand_masses = [a.GetMass() for a in self.mol.GetAtoms()]
        ligand_coords = get_romol_conf(self.mol)

        host_bps, host_masses = openmm_deserializer.deserialize_system(host_system, cutoff=1.2)
        num_host_atoms = host_coords.shape[0]

        final_potentials = []
        final_vjp_and_handles = []

        for bp in host_bps:
            if isinstance(bp, potentials.Nonbonded):
                host_p = bp
            else:
                final_potentials.append([bp])
                final_vjp_and_handles.append(None)

        hgt = topology.HostGuestTopology(host_p, self.top)

        # setup the parameter handlers for the ligand
        bonded_tuples = [
            [hgt.parameterize_harmonic_bond, self.ff.hb_handle],
            [hgt.parameterize_harmonic_angle, self.ff.ha_handle],
            [hgt.parameterize_proper_torsion, self.ff.pt_handle],
            [hgt.parameterize_improper_torsion, self.ff.it_handle]
        ]

        # instantiate the vjps while parameterizing (forward pass)
        for fn, handle in bonded_tuples:
            params, vjp_fn, potential = jax.vjp(fn, handle.params, has_aux=True)
            final_potentials.append([potential.bind(params)])
            final_vjp_and_handles.append((vjp_fn, handle))

        nb_params, vjp_fn, nb_potential = jax.vjp(hgt.parameterize_nonbonded, self.ff.q_handle.params, self.ff.lj_handle.params, has_aux=True)
        final_potentials.append([nb_potential.bind(nb_params)])
        final_vjp_and_handles.append([vjp_fn])

        combined_masses = np.concatenate([host_masses, ligand_masses])
        combined_coords = np.concatenate([host_coords, ligand_coords])

        return self._simulate(
            lamb,
            box,
            combined_coords,
            np.zeros_like(combined_coords),
            final_potentials,
            self._get_integrator(combined_masses),
            equil_steps,
            prod_steps
        )

# this class is serializable.
class RelativeFreeEnergy(BaseFreeEnergy):

    def __init__(self, mol_a, mol_b, core, ff):
        self.mol_a = mol_a
        self.mol_b = mol_b
        self.core = core
        self.ff = ff
        self.top = topology.SingleTopology(mol_a, mol_b, core, ff)

    def _get_integrator(self, combined_masses):
        """
        Get a integrator. The resulting impl must be bound to a python handle
        whose lifetime is concurrent with that of the context.
        """
        seed = np.random.randint(np.iinfo(np.int32).max)

        return LangevinIntegrator(
            300.0,
            1.5e-3,
            1.0,
            combined_masses,
            seed
        )

    def vacuum_edge(self, lamb, equil_steps=10000, prod_steps=100000):
        """
        Run a vacuum decoupling simulation at a given value of lambda.

        Parameters
        ----------
        lamb: float [0, 1]
            0 is the fully interacting system, and 1 is the non-interacting system

        equil_steps: float
            Number of steps to run equilibration. Statistics are not gathered.

        prod_steps: float
            Number of steps to run production. Statistics are gathered.

        Returns
        -------
        float, float
            Returns a pair of average du_dl values for bonded and nonbonded terms.

        """
        final_potentials = []
        final_vjp_and_handles = []

        ligand_masses_a = [a.GetMass() for a in self.mol_a.GetAtoms()]
        ligand_masses_b = [b.GetMass() for b in self.mol_b.GetAtoms()]

        ligand_coords_a = get_romol_conf(self.mol_a)
        ligand_coords_b = get_romol_conf(self.mol_b)

        bonded_tuples = [
            [self.top.parameterize_harmonic_bond, self.ff.hb_handle],
            [self.top.parameterize_harmonic_angle, self.ff.ha_handle],
            [self.top.parameterize_proper_torsion, self.ff.pt_handle],
            [self.top.parameterize_improper_torsion, self.ff.it_handle]
        ]

        # instantiate the vjps while parameterizing (forward pass)
        for fn, handle in bonded_tuples:
            (src_params, dst_params, uni_params), vjp_fn, (src_potential, dst_potential, uni_potential) = jax.vjp(fn, handle.params, has_aux=True)
            final_potentials.append([src_potential.bind(src_params), dst_potential.bind(dst_params), uni_potential.bind(uni_params)])
            final_vjp_and_handles.append((vjp_fn, handle))

        nb_params, vjp_fn, nb_potential = jax.vjp(self.top.parameterize_nonbonded, self.ff.q_handle.params, self.ff.lj_handle.params, has_aux=True)
        final_potentials.append([nb_potential.bind(nb_params)])
        final_vjp_and_handles.append([vjp_fn])

        combined_masses = np.mean(self.top.interpolate_params(ligand_masses_a, ligand_masses_b), axis=0)

        src_conf, dst_conf = self.top.interpolate_params(ligand_coords_a, ligand_coords_b)
        combined_coords = np.mean(self.top.interpolate_params(ligand_coords_a, ligand_coords_b), axis=0)

        box = np.eye(3) * 100.0

        return self._simulate(
            lamb,
            box,
            combined_coords,
            np.zeros_like(combined_coords),
            final_potentials,
            self._get_integrator(combined_masses),
            equil_steps,
            prod_steps
        )

    def host_edge(self, lamb, host_system, host_coords, box, equil_steps=10000, prod_steps=100000):
        """
        Run equilibrium decoupling simulation at a given value of lambda in a host environment.

        Parameters
        ----------
        lamb: float [0, 1]
            0 is the fully interacting system, and 1 is the non-interacting system

        host_system: openmm.System
            OpenMM System object to be deserialized. The host can be simply a box of water, or a fully
            solvated protein

        host_coords: np.array of shape [..., 3]
            Host coordinates, in nanometers. It should be properly minimized and not have clashes
            with the ligand coordinates.

        box: np.array [3,3]
            Periodic boundary conditions, in nanometers.

        equil_steps: float
            Number of steps to run equilibration. Statistics are not gathered.

        prod_steps: float
            Number of steps to run production. Statistics are gathered.

        Returns
        -------
        float, float
            Returns a pair of average du_dl values for bonded and nonbonded terms.

        """

        ligand_masses_a = [a.GetMass() for a in self.mol_a.GetAtoms()]
        ligand_masses_b = [b.GetMass() for b in self.mol_b.GetAtoms()]

        # extract the 0th conformer
        ligand_coords_a = get_romol_conf(self.mol_a)
        ligand_coords_b = get_romol_conf(self.mol_b)

        host_bps, host_masses = openmm_deserializer.deserialize_system(host_system, cutoff=1.2)
        num_host_atoms = host_coords.shape[0]

        final_potentials = []
        final_vjp_and_handles = []

        # keep the bonded terms in the host the same.
        # but we keep the nonbonded term for a subsequent modification
        for bp in host_bps:
            if isinstance(bp, potentials.Nonbonded):
                host_p = bp
            else:
                final_potentials.append([bp])
                # (ytz): no protein ff support for now, so we skip their vjps
                final_vjp_and_handles.append(None)

        hgt = topology.HostGuestTopology(host_p, self.top)

        # setup the parameter handlers for the ligand
        bonded_tuples = [
            [hgt.parameterize_harmonic_bond, self.ff.hb_handle],
            [hgt.parameterize_harmonic_angle, self.ff.ha_handle],
            [hgt.parameterize_proper_torsion, self.ff.pt_handle],
            [hgt.parameterize_improper_torsion, self.ff.it_handle]
        ]

        # instantiate the vjps while parameterizing (forward pass)
        for fn, handle in bonded_tuples:
            (src_params, dst_params, uni_params), vjp_fn, (src_potential, dst_potential, uni_potential) = jax.vjp(fn, handle.params, has_aux=True)
            final_potentials.append([src_potential.bind(src_params), dst_potential.bind(dst_params), uni_potential.bind(uni_params)])
            final_vjp_and_handles.append((vjp_fn, handle))

        nb_params, vjp_fn, nb_potential = jax.vjp(hgt.parameterize_nonbonded, self.ff.q_handle.params, self.ff.lj_handle.params, has_aux=True)
        final_potentials.append([nb_potential.bind(nb_params)])
        final_vjp_and_handles.append([vjp_fn, (self.ff.q_handle, self.ff.lj_handle)]) # (ytz): note the handlers are a tuple, this is checked later

        combined_masses = np.concatenate([host_masses, np.mean(self.top.interpolate_params(ligand_masses_a, ligand_masses_b), axis=0)])

        src_conf, dst_conf = self.top.interpolate_params(ligand_coords_a, ligand_coords_b)
        combined_coords = np.concatenate([host_coords, np.mean(self.top.interpolate_params(ligand_coords_a, ligand_coords_b), axis=0)])

        # (ytz): us is short form for mean and std dev.
        bonded_us, nonbonded_us, grads = self._simulate(
            lamb,
            box,
            combined_coords,
            np.zeros_like(combined_coords),
            final_potentials,
            self._get_integrator(combined_masses),
            equil_steps,
            prod_steps
        )

        grads_and_handles = []

        for du_dqs, vjps_and_handles in zip(grads, final_vjp_and_handles):
            if vjps_and_handles is not None:
                vjp_fn = vjps_and_handles[0]
                handles = vjps_and_handles[1]

                # we need to get the shapes correct (eg. nonbonded vjp emits an ndarray, not a list.)

                # (ytz): so far nonbonded grads is the only term that map back out to two
                # vjp handlers (charge and lj). the vjp also expects an nd.array, not a list. So we kill
                # two birds with one stone here, but this is quite brittle and should be refactored later on.
                if type(handles) == tuple:
                    # handle nonbonded terms
                    du_dps = vjp_fn(du_dqs[0])
                    for du_dp, handler in zip(du_dps, handles):
                        grads_and_handles.append((du_dp, type(handler)))
                else:
                    du_dp = vjp_fn(du_dqs)
                    # bonded terms return a list, so we need to flatten it here
                    grads_and_handles.append((du_dp[0], type(handles)))

        return bonded_us, nonbonded_us, grads_and_handles