import pynucastro as pyna

"""
Explanation: Plotting Examples
There are a variety of methods that can be used to visualize a network.
End of explanation
"""


library = pyna.ReacLibLibrary()

sub = library.linking_nuclei(["p", "n", "he4", "c12", "n13", "n14", "o16",
                              "ne20", "ne21", "na23", "mg23", "mg24", "al27",
                              "si27", "si28"])
rc = pyna.RateCollection(libraries=sub)

"""
Explanation: Rotated View
Let's build a quick network that represents He, C, and O burning
End of explanation
"""


rc.plot()

"""
Explanation: The standard network plot of this is:
End of explanation
"""


comp = pyna.Composition(rc.get_nuclei())
comp.set_solar_like()

rc.plot(rho=1e4, T=1e8, comp=comp, rotated=True)

"""
Explanation: Since everything lies close to the diagonal, it is hard to see the structure.  We can instead look at a rotated view of the network.
While we're at it, we'll also color the rates by reaction rate strength
End of explanation
"""


rc.plot(rho=1e4, T=1e8, comp=comp, rotated=True, hide_xalpha=True)

"""
Explanation: This looks a bit crowded, since it is showing all of the links to ${}^4\mathrm{He}$ from reactions of the form $A(x, \alpha)B$.  We can hide those links with hide_xalpha=True
End of explanation
"""


rc.plot(rho=1e4, T=1e8, comp=comp, rotated=True, hide_xalpha=True,
        rate_filter_function=lambda r: pyna.Nucleus("p") in r.reactants)

"""
Explanation: We can also use a rate filter to show only certain links.  For example, to show only links involving proton capture, we can do:
End of explanation
"""


fig = rc.plot(rho=1e4, T=1e8, comp=comp, rotated=True, hide_xalpha=True,
              curved_edges=True)

"""
Explanation: We can also plot the edges curved, to help distinguish the forward and reverse rates
End of explanation
"""


files = ["c12-pg-n13-ls09", 
         "c13-pg-n14-nacr",
         "n13--c13-wc12",
         "n13-pg-o14-lg06",
         "n14-pg-o15-im05",
         "n15-pa-c12-nacr",
         "o14--n14-wc12",
         "o15--n15-wc12"]
rc = pyna.RateCollection(files)

comp = pyna.Composition(rc.get_nuclei())
comp.set_solar_like()

rc.gridplot(comp=comp, color_field="X", scale="log", area=36)

"""
Explanation: Plot nuclides on a grid
Nuclides in a network may also be visualized as cells on a grid of Z vs. N, colored by some quantity. This can be more interpretable for large networks. Calling gridplot without any arguments will just plot the grid - to see anything interesting we need to supply some conditions. Here is a plot of nuclide mass fraction on a log scale, with a 36 square inch figure:
End of explanation
"""


rc.gridplot(comp=comp, rho=1e4, T=1e8, color_field="ydot", area=36,
        cmap="RdGy", cbar_bounds=(-0.1, 0.1))

"""
Explanation: The plot is configurable through a large number of keyword arguments. Here we want to look at the rates at which nuclides are being created or destroyed, so we color by $\dot{Y}$, the rate of change of molar abundance. Density and temperature need to be supplied to evaluate the rates. A full list of valid keyword arguments can be found in the API documentation.
End of explanation
"""


ff = lambda nuc: nuc.Z > 2
rc.gridplot(comp=comp, rho=1e4, T=1e8, color_field="activity", scale="log",
            filter_function=ff, area=20, cmap="Blues")

"""
Explanation: Unlike the network plot, this won't omit hydrogen and helium by default. To just look at the heavier nuclides, we can define a function to filter by proton number.  Here we also plot
activity, which is the sum of all rates contributions to Ydot ignoring the sign.
End of explanation
"""


rc.plot_network_chart(rho=1.e4, T=1.e8, comp=comp)

"""
Explanation: Network Chart
A network chart shows how each rate affects each nucleus
End of explanation
"""


