from goatools.base import get_godag
godag = get_godag("go-basic.obo", optional_attrs={'relationship'})

# We are going to use a GO DAG subset for this example
godag = get_godag("../tests/data/i126/viral_gene_silence.obo", optional_attrs={'relationship'})

"""
Explanation: How to get GO IDs up both the 'is a' and 'part of' relationships

Download Ontologies, if necessary
Choose a set of GO IDs to print
GO colors for plotting
Traverse up or down the GO DAG with no, all, or some optional relationships(a):
4a) No optional relationships, like part_of and regulates
4b) All relationships
4c) Relationship: part_of
4d) Relationship: regulates
4e) Relationship: positively_regulates
4f) Relationship: negatively_regulates



a) The is_a relationships will always be traversed to find the children or parents of a GO term
Plots with and without relationships
Arrows denote the relationships, is_a (solid black arrows) and part_of (dashed orange arrows), in the plots below: 
| With no optional relationships  | With the optional part_of relationship
|---------------------------------|-------------------------------
|      |   
Notice that the text in the GO term box is different if you are plotting using optional relationships. The bottom GO term text is:

Left plot: GO:0043231 L3 D3 d0
Right plot: GO:0043231 D3 R4 d0

L (Level) is the minimum distance from root term, cellular_component, to the bottom term  
D (Depth), is the maximum distance from root term, cellular_component, to the bottom term using only is_a relationships
R (Depth using optional relationships), is the maximum distance from root term, cellular_component, to the bottom term using both is_a and optional relationships
In this example D (Depth) from the root to the bottom term is 3 (light blue lines).    
And R (Depth using optional relationships) from the root to the bottom is 4 (dark purple lines) because it traverses the part_of relationships
<img src="images/fig1_depth_marked.png" alt="depth explained" width="450">
1. Download Ontologies, if necessary
Load all relationships into a GODag object.     
Specific relationships, like part_of or regulates, will be chosen using GoSubDag.
End of explanation
"""


virus_trigger_silence = 'GO:0060150'

"""
Explanation: 2. Choose a set of GO IDs to print
This example will use all GO IDs above GO:0060150, viral triggering of virus induced gene silencing
End of explanation
"""


from goatools.cli.gos_get import get_go2color

go2hexcolor = get_go2color('../tests/data/i126/viral_gene_silence.txt')

# For printing GO terms and their color names to the print
hex2color = {
    '#e6fad2': 'green',
    '#ffe5b4': 'orange',
    '#d2d2fa': 'purple',
    '#d2fafa': 'cyan',
    '#fad2fa': 'magenta',
    '#d8dcd6': 'grey',
}
GREY = '#d8dcd6'
GO2COLORNAME = {go: hex2color[rgb] for go, rgb in go2hexcolor.items()}

"""
Explanation: 3. GO colors
The GO DAG and go-color files for this example are located in: ./tests/data/i126   
GO colors are stored in viral_gene_silence.txt as 6-digit hax numbers, like #ffe5b4.  
The color file looks like this:
```
Relationship, part_of, shall be orange (#ffe5b4):
ffe5b4 GO:0002376 immune system process
ffe5b4 GO:0050896 response to stimulus
ffe5b4 GO:0051704 multi-organism process
...
```
End of explanation
"""


from goatools.gosubdag.gosubdag import GoSubDag

gosubdag_r0 = GoSubDag({virus_trigger_silence}, godag)

# Decide which GO fields to print. Save in prtfmt

# Print one namedtuple
print(next(iter(gosubdag_r0.go2nt.values())))

# Use some namedtuple fields in print of GO terms
PRTFMT = '{NS} {GO} D-{depth:02} {color:7} {dcnt:2} {GO_name}'

# Report all ancestors of term, GO:001922 (colored grey)
grey_go = 'GO:0019222'
ancestors = gosubdag_r0.rcntobj.go2ancestors[grey_go]

def prt_goids(goids, gosubdag):
    """Print GO IDs"""
    nts = [gosubdag.go2nt[go] for go in goids]
    for nt_go in sorted(nts, key=lambda nt: [nt.depth, nt.dcnt]):
        color = GO2COLORNAME[nt_go.GO]
        print(PRTFMT.format(color=color, **nt_go._asdict()))

prt_goids(ancestors, gosubdag_r0)

"""
Explanation: 4. Traversals up the GO DAG
4a. No optional relationships loaded
Do not traverse any optional relationships, like part_of and regulates
End of explanation
"""


from goatools.gosubdag.plot.gosubdag_plot import GoSubDagPlot

# Color the GO of interest in a copy of the color dictionary
go2col = dict(go2hexcolor)
go2col[grey_go] = GREY

pltobj = GoSubDagPlot(gosubdag_r0, go2color=go2col)
pltobj.plt_dag('viral_r0.png')

"""
Explanation: Plot no optional relationships, just using is_a (runtime script):
Color the term, GO:0019222 regulation of metabolic process, with xkcd's light grey (#d8dcd6)
```
$ go_plot.py -o viral_r0.png GO:0060150 GO:0019222#d8dcd6 --obo=viral_gene_silence.obo --go_color_file=viral_gene_silence.txt
```
<img src="images/viral_r0.png" width="900pt">
Plot no optional relationships, just using is_a (code):
End of explanation
"""


gosubdag_r1 = GoSubDag({virus_trigger_silence}, godag, relationships=True)

# Report all ancestors of grey term
grey_go = 'GO:0010468'  # regulation of gene expression
ancestors = gosubdag_r1.rcntobj.go2ancestors[grey_go]
prt_goids(ancestors, gosubdag_r1)

"""
Explanation: 4b. All relationships
Traverse all optional relationships, like part_of and regulates
End of explanation
"""


# Color the GO of interest in a copy of the color dictionary
go2col = dict(go2hexcolor)
go2col[grey_go] = GREY

pltobj = GoSubDagPlot(gosubdag_r1, go2color=go2col)
pltobj.plt_dag('viral_r1.png')

"""
Explanation: Plot all relationships (runtime script):
$ go_plot.py -o viral_r1.png -r GO:0060150 GO:0010468#d8dcd6 --obo=viral_gene_silence.obo --go_color_file=viral_gene_silence.txt
<img src="images/viral_r1.png" width="2000pt">
Plot all relationships (code):
End of explanation
"""


gosubdag_partof = GoSubDag({virus_trigger_silence}, godag, relationships={'part_of',})

# Report all ancestors of grey term
grey_go = 'GO:0010468'  # regulation of gene expression
ancestors = gosubdag_partof.rcntobj.go2ancestors[grey_go]
prt_goids(ancestors, gosubdag_partof)

"""
Explanation: 4c. Relationship: part_of
End of explanation
"""


# Color the GO of interest in a copy of the color dictionary
go2col = dict(go2hexcolor)
go2col[grey_go] = GREY

pltobj = GoSubDagPlot(gosubdag_partof, go2color=go2col)
pltobj.plt_dag('viral_r_partof.png')

"""
Explanation: Plot part_of relationships (runtime script):
Orange dashed lines represent part_of relationships:
$ go_plot.py -o viral_r_partof.png --relationships=part_of GO:0010468#d8dcd6 GO:0060150 --obo=viral_gene_silence.obo --go_color_file=viral_gene_silence.txt
<img src="images/viral_r_partof.png" width="900pt">
Plot part_of relationships (code)
End of explanation
"""


gosubdag_reg = GoSubDag({virus_trigger_silence}, godag, relationships={'regulates',})

# Report all ancestors of grey term
grey_go = 'GO:0050794'  # regulation of cellular process
ancestors = gosubdag_reg.rcntobj.go2ancestors[grey_go]
prt_goids(ancestors, gosubdag_reg)

"""
Explanation: 4d. Relationship: regulates
End of explanation
"""


# Color the GO of interest in a copy of the color dictionary
go2col = dict(go2hexcolor)
go2col[grey_go] = GREY

pltobj = GoSubDagPlot(gosubdag_reg, go2color=go2col)
pltobj.plt_dag('viral_reg.png')

"""
Explanation: Plot regulates relationships (runtime script):
Purple dashed lines represent regulates relationships:
$ go_plot.py -o viral_reg.png --relationships=regulates GO:0050794#d8dcd6 GO:0060150 --obo=viral_gene_silence.obo --go_color_file=viral_gene_silence.txt
<img src="images/viral_reg.png" width="900pt">
Plot regulates relationships (code):
End of explanation
"""


gosubdag_regp = GoSubDag({virus_trigger_silence}, godag, relationships={'positively_regulates',})

# Report all ancestors of grey term
grey_go = 'GO:0048522'  # positive regulation of cellular process
ancestors = gosubdag_regp.rcntobj.go2ancestors[grey_go]
prt_goids(ancestors, gosubdag_regp)

"""
Explanation: 4e. Relationship: positively_regulates
End of explanation
"""


# Color the GO of interest in a copy of the color dictionary
go2col = dict(go2hexcolor)
go2col[grey_go] = GREY

pltobj = GoSubDagPlot(gosubdag_regp, go2color=go2col)
pltobj.plt_dag('viral_rp.png')

"""
Explanation: Plot positively_regulates relationships (runtime script):
Magenta dashed lines represent positively regulates relationships:    
$ go_plot.py -o viral_rp.png --relationships=positively_regulates GO:0048522#d8dcd6 GO:0060150 --obo=viral_gene_silence.obo --go_color_file=viral_gene_silence.txt
<img src="images/viral_rp.png" width="900pt">
Plot positively_regulates relationships (code):
End of explanation
"""


gosubdag_regn = GoSubDag({virus_trigger_silence}, godag, relationships={'regulates', 'negatively_regulates',})

# Report all ancestors of grey term
grey_go = 'GO:0050794'  # regulation of cellular process
ancestors = gosubdag_regn.rcntobj.go2ancestors[grey_go]
prt_goids(ancestors, gosubdag_regn)

"""
Explanation: 4f. Relationship: negatively_regulates
End of explanation
"""


# Color the GO of interest in a copy of the color dictionary
go2col = dict(go2hexcolor)
go2col[grey_go] = GREY

pltobj = GoSubDagPlot(gosubdag_regn, go2color=go2col)
pltobj.plt_dag('viral_rn.png')

"""
Explanation: Plot negatively_regulates relationships (runtime script):
Cyan dashed lines represent negatively regulates relationships    
NOTE: In this example, the following relationships argument will have no affect because no green GO IDs have the relationship, negatively_regulates:  
--relationships=negatively_regulates
To see negatively_regulates, specify both regulates and negatively_regulates:  
--relationships=regulates,negatively_regulates
Note that specifying both regulates and negatively_regulates will case the magenta GO IDs to be traversed. The magenta GO IDs are accessed by both the regulates and the positively_regulates relationships.
$ go_plot.py -o viral_rn.png --relationships=regulates,negatively_regulates GO:0050794#d8dcd6 GO:0060150 --obo=viral_gene_silence.obo --go_color_file=viral_gene_silence.txt
<img src="images/viral_rn.png" width="900pt">
Plot negatively_regulates relationships (code):
End of explanation
"""


