from solarmach import SolarMACH, print_body_list

# optional: get list of available bodies/spacecraft
print(print_body_list().index)

# necessary:
body_list = ['STEREO-A', 'Earth', 'BepiColombo', 'PSP', 'Solar Orbiter', 'Mars']
vsw_list = [400, 400, 400, 400, 400, 400, 400]
date = '2021-10-28 15:15:00'

# optional:
reference_long = 273                             # Carrington longitude of reference (None to omit)
reference_lat = 0                                # Carrington latitude of reference (None to omit)
plot_spirals = True                              # plot Parker spirals for each body
plot_sun_body_line = True                        # plot straight line between Sun and body
show_earth_centered_coord = False                # display Earth-aligned coordinate system
reference_vsw = 400                              # define solar wind speed at reference
transparent = False                              # make output figure background transparent
numbered_markers = True                          # plot each body with a numbered marker
filename = 'Solar-MACH_'+date.replace(' ', '_')  # define filename of output figure

# optional
# if input coordinates for reference are Stonyhurst, convert them to Carrington for further use
import astropy.units as u
from astropy.coordinates import SkyCoord
from sunpy.coordinates import frames
reference_long = 2                               # Stonyhurst longitude of reference (None to omit)
reference_lat = 26                               # Stonyhurst latitude of reference (None to omit)
coord = SkyCoord(reference_long*u.deg, reference_lat*u.deg, frame=frames.HeliographicStonyhurst, obstime=date)
coord = coord.transform_to(frames.HeliographicCarrington(observer='Sun'))
reference_long = coord.lon.value                 # Carrington longitude of reference
reference_lat = coord.lat.value                  # Carrington latitude of reference

# initialize
sm = SolarMACH(date, body_list, vsw_list, reference_long, reference_lat)

# make plot
sm.plot(
   plot_spirals=plot_spirals,
   plot_sun_body_line=plot_sun_body_line,
   show_earth_centered_coord=show_earth_centered_coord,
   reference_vsw=reference_vsw,
   transparent=transparent,
   numbered_markers=numbered_markers,
   outfile=filename+'.png'
)

# obtain data as Pandas DataFrame
display(sm.coord_table)