import warnings
warnings.filterwarnings('ignore')

from tardis import run_tardis
import tardis

tardis.logger.setLevel(0)
tardis.logging.captureWarnings(False)

"""
Explanation: Using Custom Callbacks When Running TARDIS
End of explanation
"""


def display_table(sim):
    '''Display a table of velocities and 
    radiative temperatures at each iteration
    '''
    
    # We have direct access to the attributes of the simulation
    columns = zip(sim.model.v_inner[::5].to('km/s'), 
                  sim.model.t_rad[::5].to('K'))
    print("Iteration:", sim.iterations_executed)
    print("  {:<15}    {:<15}".format('v_inner', 't_rad'))
    format_string = "  {0.value:<8.2f} {0.unit:<6s}\
    {1.value:<8.2f} {1.unit:<6s}"
    for velocity, temperature in columns:
        print(format_string.format(velocity, temperature))


"""
Explanation: The command run_tardis allows users to provide a set of callbacks to the simulation.  These callbacks are called at the end of each iteration.  This example will show you how to create a custom callback and run a model with TARDIS.  As an example, we create a custom callback that prints out some basic information about our model at every iteration.  Specifically, we'll print out a table of the inner velocities of each shell as well as the radiative temperature of each shell.  The first thing to note is that the callback function must have the simulation object as the first argument.  This grants the user access to the state of the simulation at each iteration.
End of explanation
"""


sim = run_tardis('tardis_example.yml', 
                 simulation_callbacks=[[display_table]])

"""
Explanation: Now we give the callback to run_tardis.  run_tardis offers the keyword argument simulation_callbacks which takes a list of lists containing the callback as well as any optional arguments you wish to include with your callback.  For this example our function requires no extra arguments and we only have a single callback, so we give run_tardis a 2D list containing the callback as its only element.
End of explanation
"""


def append_t_rad_to_table(sim, table):
    '''append the array for the radiative temperature
    at each iteration to a given table'''
    
    table.append(sim.model.t_rad.copy())

"""
Explanation: Running Callbacks with Extra Arguments
The callbacks provided to run_tardis can also take extra arguments.  As an example, we'll make a callback that appends the value of the radiative temperature at each iteration to a list so we can watch how the model converges.  The callback will take a list we want to append to as an argument.  We'll send both this new callback and our original display_table callback to run_tardis as an example of using multiple callbacks at once.
End of explanation
"""


t_rad_table = [] # list to store t_rad at each iteration

callbacks = [[display_table], 
             [append_t_rad_to_table, t_rad_table]]
sim = run_tardis('tardis_example.yml', 
                 simulation_callbacks=callbacks)

"""
Explanation: In order to add our new callback, we just create another entry in our list of callbacks.  Since append_t_rad_to_table takes an extra argument, we will provide that argument in the inner list containing the callback.
End of explanation
"""


%pylab notebook
for t_rad in t_rad_table:
    plot(t_rad)
ylabel(r'$T_{rad}\ [K]$')
xlabel('Iteration Number')

"""
Explanation: Now we can look at the way the radiative temperature changes in each shell every iteration.
End of explanation
"""


