import tardis
import matplotlib.pyplot as plt
import numpy as np

"""
Explanation: How to run TARDIS with a custom ejecta model
This notebook will go through multiple detailed examples of how to properly run TARDIS with a custom ejecta profile specified by a custom density file and a custom abundance file.
End of explanation
"""


model = tardis.run_tardis('./test_config.yml')

"""
Explanation: Your custom density file
First, let's look at an example of a custom density file.

The first line specifies the time in days after the explosion
After a skipped line, each row corresponds to a shell with index specified by the first column.
The second column lists the velocities of the outer boundary of the cell in km/s.
The third column lists the density of the cell.

<font color=red>Important: </font>
The default behavior of TARDIS is to use the first shell as the inner boundary. This means that v_inner_boundary = 9500, and the corresponding density 9e-16 is ignored because it is within the inner boundary. It can be replaced by an arbitrary number.  The outer boundary of the last shell will be used as v_outer_boundary, so the default behavior will set v_outer_boundary = 12000.
Your custom abundance file
Let's look at an example of a custom density file.

The first line indicates which elements (or isotopes) correspond to which columns.
After a skipped line, each row specifies the chemical abundance of one shell. Therefore the numbers in a given row should sum to 1.0

<font color=red>Important: </font>
Note that there are only 2 shells specified in this abundance file (despite the custom density file having 3 lines). This is because the custom density file specifies the boundaries of the shells, while the abundance file specifies the abundances within each shell.
Running TARDIS with the custom files
Now let's run TARDIS using the example custom files.
End of explanation
"""


print('v_inner_boundary = ',model.model.v_boundary_inner)
print('v_outer_boundary = ',model.model.v_boundary_outer)
print('\n')
print('velocities of shell boundaries: ')
print(model.model.velocity)
print('\n')
print('densities loaded by TARDIS: (NOTE that the density in the first line of the file was ignored! Densities are also rescaled.)')
print(model.model.density)

"""
Explanation: You can check to make sure that the model loaded and used by TARDIS during the simulation is consistent with your expectations based on the custom files you provided:
End of explanation
"""


model = tardis.run_tardis('./test_config_ex1.yml')

"""
Explanation: Specifying boundary velocities in the config file
In addition to specifying custom density and abundance files, the user can set the v_inner_boundary and v_outer_boundary velocities in the YAML config file. This can cause some confusion, so we carefully go through some examples.
<font color=red>Important: </font>
Boundary velocities set in the YAML config file must be within the velocity range specified in the custom density file (if one is provided).
Example 1) v_inner_boundary lower than first velocity in density file
In this example, the first velocity in the density file is 9500 km/s. The user can specify in the config file the velocity of the inner boundary to a lower velocity, say v_inner_boundary = 9000 km/s. This will cause TARDIS to raise an error.
End of explanation
"""


model = tardis.run_tardis('./test_config_ex2.yml')

"""
Explanation: Example 2) v_outer_boundary larger than last velocity in density file
In this example, the last velocity in the density file is 12000 km/s. The user can specify in the config file the velocity of the outer boundary to a larger velocity, say v_outer_boundary = 13000 km/s. This will cause TARDIS to raise an error.
End of explanation
"""


model = tardis.run_tardis('./test_config_ex3.yml')

print('v_inner_boundary = ',model.model.v_boundary_inner)
print('v_outer_boundary = ',model.model.v_boundary_outer)
print('\n')
print('velocities of shell boundaries: ')
print(model.model.velocity)
print('\n')
print('densities loaded by TARDIS: (NOTE that the density in the first line of the file was ignored! Densities are also rescaled.)')
print(model.model.density)

"""
Explanation: Example 3) v_boundaries in config file are within density file velocity range
Here the user sets v_inner_boundary = 9700 and v_outer_boundary = 11500 in the config file. Both values fall within the velocity range specified by the custom density file.
End of explanation
"""


