"""

This file is unit07_SunJupiterObject.py. It defines the class "SunJupiterVenusSystem"
and initializes various parameters. Venus is now included, but with ten times
its true mass to make its effects on the sun more obvious.

This file also houses various "methods" (functions defined in the class) to allow 
you to generate plots and calculate useful stuff.

I approximate orbits as perfect circles, all lying in the same plane.

The origin of my coordinate system is the center of the sun, but I have 
functions to transform the positions and velocities to a barycentric 
(center of mass) coordinate system.

see the commented out scripts at the end of this file for examples of how to
use it.

distances are in meters and masses in kg.

class methods (functions):

__init__
    class constructor. initialize various constants each time an object of this
    class is instantiated (created).

draw_solar_system(t)
    draw the sun, venus, and jupiter in a heliocentric coordinate system with 
    the planet positions at time t.

planet_theta(t)
    returns the theta angles of the planet positions. the x axis is 0.

get_planet_xy(t)
    returns the planets' x, y positions. sun is at the origin.
    
get_planet_xy_barycentric(t)
    returns the planets' x, y positions in a barycentric (CM) coordinate system
    
get_planet_vxvy(t)
    returns the planets' x, y velocities. sun is at the origin.
    
get_planet_vxvy_barycentric(t)
    returns the planets' x, y velocities in a barycentric (CM) coordinate system
    
    
George Gollin, University of Illinois, January 4, 2018

"""

class SunJupiterVenusSystem:
    
    ###########################################################################
    
    # here is the "class constructor." use it to initialize various parameters.
    # it is called automatically whenever you instantiate (create an instance of)
    # an object of this class.
    
    def __init__(self):
        
        # let's import numpy here 
        import numpy as np
        
        # print a message about instantiating the solar system :-)
        print("\nNow instantiating the star-planet system.")
        
        # Newton's gravitational constant, in SI units.
        self.G_Newton = 6.67408E-11
        
        # exact speed of light, in m/s
        self.clight = 299792458
        
        # define the names of the various objects.
        self.name = ["sun", "jupiter", "BIGvenus"]
        
        # indices for star and its planets, used when referencing some arrays:
        self.planet_index = [0, 1, 2]
                        
        # define the masses, in kg. The first is the sun, the second Jupiter, 
        # the third Venus. (This is the correct mass for Venus.)  
        # See https://en.wikipedia.org/wiki/List_of_Solar_System_objects_by_size
        self.sun_mass = 1988550000.0e21
        self.jupiter_mass = 1898600.00e21
        self.venus_mass = 4868.5e21
        
        ################################################################
        
        # make venus massless to work with a two-body system.
        self.venus_mass = 0.

        # make venus ten times heavier than in reality to work with a 3-body system
        # self.venus_mass = self.venus_mass * 10

        ################################################################

        self.mass = np.array([self.sun_mass, self.jupiter_mass, self.venus_mass])
                
        # take the average of the aphelion and perihelion as the distance 
        # between the sun and Jupiter.
        self.star_jupiter_distance = 778.295e9
        self.orbit_radius_jupiter = self.star_jupiter_distance
        
        # take the average of the aphelion and perihelion as the distance 
        # between the sun and Venus. https://en.wikipedia.org/wiki/Venus 
        self.star_venus_distance = 108.208e9
        self.orbit_radius_venus = self.star_venus_distance
        
        # array of orbit radii, heliocentric (sun-centered) coordinates
        self.orbit_radius = np.array([0., self.star_jupiter_distance, \
        self.star_venus_distance])

        # orbital period for Jupiter, from Wikipedia, is 4332.59 days.
        # orbital period for Venus, from Wikipedia, is 224.701 days.
        self.orbit_period_jupiter = 4332.59 * 24 * 3600
        self.orbit_period_venus = 224.701 * 24 * 3600
        self.orbit_period = \
        np.array([0.001, self.orbit_period_jupiter, self.orbit_period_venus])
                
        # array of orbital speeds for motion in a heliocentric system:
        self.orbit_speed_jupiter = 2 * np.pi * self.orbit_radius_jupiter / \
        self.orbit_period_jupiter
        self.orbit_speed_venus = 2 * np.pi * self.orbit_radius_venus / \
        self.orbit_period_venus

        self.orbit_speed = \
        np.array([0., self.orbit_speed_jupiter, self.orbit_speed_venus])
        
        # put the planets on the positive x axis at t = 0, the star on the negative 
        # x axis. I will take the angle to represent the orientation of the 
        # line from the star to the planet.
        self.theta0 = np.array([-np.pi, 0., 0.])
        
        # color to draw the planets, hex RGB codes
        self.color = ["#cccc00", "#ee3333", "#ffffff"]
        
        # Here are some of the drawing parameters when we make a plot of the 
        # planetary system.
        
        # background colors for plot: dark gray
        self.bgcolor_all_OK = "#333333"
        
        # some parameters relating to possible animations: a frame index
        # and time (in seconds) of the frame.
        self.frame_index = 0
        self.frame_time = 0
        self.last_picture_name = "NothingYet"
        
    ###########################################################################
    # end of class constructor __init__
    ###########################################################################
    
    # draw the solar system at time t (in seconds) in a heliocentric system

    def draw_solar_system(self, t, close_previous_figure = False, \
    full_width = 2000):

        import numpy as np
        import matplotlib.pyplot as plt 
        
        ############# draw the solar system! #############
        
        # full width (in millions of km) for the plot: 2000 will show jupiter's 
        # orbit, while 2 will show the sun's.
    
        # get the positions of the bodies in the planetary system
        xplanet_all, yplanet_all = self.get_planet_xy(t)
        
        # close the previous figure that was created by this routine if we want...
        if close_previous_figure and self.last_picture_name != "NothingYet":
            plt.close(self.last_picture_name)
        
        # give the picture a name, for example SolarSystem0007.
        self.frame_index = self.frame_index + 1
        picture_name = "SolarSystem" + str(self.frame_index).zfill(4)

        # create a (blank) figure so we can set some of its attributes.
        fig = plt.figure(picture_name)
        ax = fig.gca()
    
        # store the name of the picture for later use
        self.last_picture_name = picture_name

        # set labels and title; force square aspect ratio.
        ax.set_xlabel("x (1 = 10^6 km)")
        ax.set_ylabel("y (1 = 10^6 km)")
        ax.set_title("The solar system-- George Gollin") 
        ax.set_aspect("equal")
        
        # set background color
        ax.set_axis_bgcolor(self.bgcolor_all_OK)
           
        # set the plot size (m) and force the axes to be of equal scales...
        plt_halfmax = 0.5 * full_width

        # now define the plot range.
        plt.axis([-plt_halfmax, plt_halfmax, -plt_halfmax, plt_halfmax])
        
        # draw circles for orbits
        for index in self.planet_index:
        
            # get points on the circle
            x, y = self.orbit_radius[index] * np.cos(np.linspace(0, 2*np.pi, 1000, endpoint=False)),  \
            self.orbit_radius[index] * np.sin(np.linspace(0, 2*np.pi, 1000, endpoint=False) )
            
            # now convert to megakm from meters.
            x = x / 1.e9
            y = y / 1.e9

            # plot the circle
            plt.plot(x, y, self.color[index])

            # now put a disk on the plot to represent the planet.
            xplanet = xplanet_all[index]
            yplanet = yplanet_all[index]
            
            # draw the disk
            planet_scale_factor = 100.
            planet_disk = plt.Circle((xplanet / 1.e9, yplanet / 1.e9), \
            plt_halfmax / planet_scale_factor, color = self.color[index])
            
            fig.gca().add_artist(planet_disk)
        
        # now save the figure, using the name we gave it a little while ago as
        # part of the file name. Since we might want to use this in a 
        # multiple-frame sequence, write it into a subdirectory.
        file_name = "animations/" + picture_name + ".png"
        plt.savefig(file_name)
 
    ###########################################################################
    # end of class function draw_solar_system
    ###########################################################################               

    # calculate the theta angles for the planets at a later time than 4/1/16.
    # t is seconds after 4/1/16 00:00:00 z. Recall that one year is ~pi * 10^7.
    
    def planet_theta(self, t):
        
        import numpy as np

        theta_now = self.theta0 + t * 2 * np.pi / self.orbit_period 
        
        return np.array(theta_now)
   
    ###########################################################################
    # end of class function planet_theta
    ###########################################################################

    # return an array with the x,y positions of the various planets and sun
    
    def get_planet_xy(self, t):
        
        # import library
        import numpy as np

        # create array for the current theta.
        theta_now = self.planet_theta(t)

        # create arrays for the x, y positions
        xarray, yarray = self.orbit_radius * np.cos(theta_now), \
        self.orbit_radius * np.sin(theta_now)

        return np.array([xarray, yarray])

    ###########################################################################
    # end of class function get_planet_xy
    ###########################################################################

    # return an array with the x,y positions of the various planets and sun
    # and moon
    
    def get_planet_xy_barycentric(self, t):
        
        # import library
        import numpy as np

        # create arrays for the x, y positions
        xarray, yarray = self.get_planet_xy(t)
        
        # now calculate center of mass position, then subtract that from all 
        # the objects
        x_cm = np.dot(self.mass, xarray) / np.sum(self.mass)
        y_cm = np.dot(self.mass, yarray) / np.sum(self.mass)
        
        xarray = xarray - x_cm
        yarray = yarray - y_cm

        return np.array([xarray, yarray])

    ###########################################################################
    # end of class function get_planet_xy_barycentric
    ###########################################################################

    # return an array with the x,y velocity components of the star and planets
    
    def get_planet_vxvy(self, t):
        
        # import library
        import numpy as np

        # create array for the current theta.
        theta_now = self.planet_theta(t)
        
        # now get the velocity components. i am assuming that orbits are circles.
        vxarray, vyarray = -self.orbit_speed * np.sin(theta_now), \
        self.orbit_speed * np.cos(theta_now)

        return np.array([vxarray, vyarray])
        
    ###########################################################################
    # end of class function get_planet_vxvy
    ###########################################################################

    # return an array with the x,y velocity components of the star and planets
    
    def get_planet_vxvy_barycentric(self, t):
        
        # import library
        import numpy as np

        # create arrays for the vx, vy positions
        vxarray, vyarray = self.get_planet_vxvy(t)
        
        # now calculate center of mass velocities, then subtract that from all 
        # the objects
        vx_cm = np.dot(self.mass, vxarray) / np.sum(self.mass)
        vy_cm = np.dot(self.mass, vyarray) / np.sum(self.mass)

        # now get the velocity components. i am assuming that orbits are circles.
        vxarray = vxarray - vx_cm
        vyarray = vyarray - vy_cm
        
        return np.array([vxarray, vyarray])
        
    ###########################################################################
    # end of class function get_planet_vxvy
    ###########################################################################

###############################################################################
# end of class SolarSystem
###############################################################################

###############################################################################
# here is an example of how to use the class defined in this file. 
# Copy the following to a separate file.
###############################################################################

"""
# This is a demonstrator program to run the code in the SunJupiterObject.py
# class definition.

# George Gollin, Mar 12, 2017

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

# instantiate an object of class SunJupiterSystem and name it SJO:
SJO = unit07_SunJupiterObject.SunJupiterVenusSystem()

# let's draw the sun-Jupiter system. first argument is time (in seconds). 
# second is whether or not to close previous figures when drawing the new orbit
# diagram. third is full width of plot in mega-km.
SJO.draw_solar_system(0., True, 2000)
SJO.draw_solar_system(0., False, 2)

# draw some more, a quarter of a period later.
SJO.draw_solar_system(SJO.orbit_period / 4, False, 2000)
SJO.draw_solar_system(SJO.orbit_period / 4., False, 2)

# get times to plot the sun's x velocity. Jupiter's orbital period is
# a little les than 12 years so let's do 25 years.
tmax = 25 * 365 * 24 * 3600

# create arrays
tarray = tmax * np.linspace(0, 1, 1000, endpoint=False)
vxarray = np.zeros(1000)

index = 0

for t in tarray:
 
    # get vx, vy for the sun + jupiter.   
    vx, vy = SJO.get_planet_vxvy_barycentric(t)
    
    xb, yb = SJO.get_planet_xy_barycentric(t)

    # the sun is in the first array element, so pick that out.
    vxarray[index] = vx[0]
    
    # increment the array index.
    index = index + 1

# now rescale the t array to years from seconds.
tarray = tarray / (3600 * 24 * 365)

# create a (blank) figure so we can set some of its attributes.
fig = plt.figure()
ax = fig.gca()

# set labels and title
ax.set_xlabel("time (years)")
ax.set_ylabel("solar x velocity, m/s")
ax.set_title("solar velocity-- George Gollin") 

# plot everything
plt.plot(tarray, vxarray)

# all done!

"""

###############################################################################
# here is another example of how to use the class defined in this file, in  
# which I modify some of the parameters to describe the Pegasi 51 system.
###############################################################################

"""
# This is a demonstrator program to run the code in the SunJupiterObject.py
# class definition, but reconfigured to describe the 51 Pegasi system

# George Gollin, Mar 12, 2017

# If you are modeling the 51 Pegasi system, see various Wikipedia articles
# for the system's parameters. Note that the values are not consistent 
# across Wikipedia's various articles!
# The mass of 51 Pegasi is estimated to be 1.11 solar masses. 
# Its planet Dimidium has a mass that is at least 0.472 of Jupiter's.
# Dimidium's orbital period is 101.5388 hours, and its orbit radius is 7.89e9 m.

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

# instantiate an object of class SunJupiterSystem and name it PDO:
PDO = unit07_SunJupiterObject.SunJupiterVenusSystem()

##################################################################
# code in here lets us switch from describing the sun, jupiter, venus system
# to the 51 Pegasi system.

# rename the objects and change their masses and so forth.
PDO.name = ["51 Pegasi", "Dimidium", "no second planet"]

# 51 Pegasi mass is 1.11 times that of the sun.
PDO.mass[0] = 1.11 * PDO.mass[0]

# Set the planet mass to the minimum allowed: 0.472 * jupiter's.
PDO.mass[1] = 0.472 * PDO.mass[1]

# Set the nonexistent second planet mass to zero
PDO.mass[2] = 0.

# set the planet's orbital period to be the same as Dimidium:
PDO.orbit_period[1] = 101.5388 * 3600

# Set the nonexistent second planet's period
PDO.orbit_period[2] = 0.001

# set the star-to-planet distance to the Wikipedia orbit value, and use this
# as the planet's orbit radius 
PDO.star_planet_distance = 7.89e9
PDO.orbit_radius[1] = PDO.star_planet_distance

# make the star and non-existent planet orbit radii be negligible
PDO.orbit_radius[0] = 0.001
PDO.orbit_radius[2] = 0.001

# array of orbital speeds for motion around barycenter:
PDO.orbit_speed = 2 * np.pi * PDO.orbit_radius / PDO.orbit_period

# let's draw the star-planet system. first argument is time (in seconds). 
# second is whether or not to close previous figures when drawing the new orbit
# diagram. third is full width of plot in mega-km (billions of meters).
PDO.draw_solar_system(0., True, 20)
PDO.draw_solar_system(0., False, .01)

# draw some more, a quarter of a period later.
PDO.draw_solar_system(PDO.orbit_period / 4, False, 20)
PDO.draw_solar_system(PDO.orbit_period / 4., False, .01)

# get times to plot the star's x velocity. Dimidium's orbital period is
# about 4 days so let's do two weeks.
tmax = 14 * 24 * 3600

##################################################################

# create arrays
tarray = tmax * np.linspace(0, 1, 1000, endpoint=False)
vxarray = np.zeros(1000)
vxarray_planet = np.zeros(1000)

index = 0

for t in tarray:
 
    # get vx, vy for the star + planet(s)   
    vx, vy = PDO.get_planet_vxvy_barycentric(t)

    # the sun is in the first array element, so pick that out.
    vxarray[index] = vx[0]
    vxarray_planet[index] = vx[1]

    # increment the array index.
    index = index + 1

# now rescale the t array to days from seconds.
tarray = tarray / (3600 * 24)

# create a (blank) figure so we can set some of its attributes.
fig = plt.figure()
ax = fig.gca()

# set labels and title
ax.set_xlabel("time (days)")
ax.set_ylabel("51 Pegasi x velocity, m/s")
ax.set_title("51 Pegasi x velocity-- George Gollin") 

# plot everything
plt.plot(tarray, vxarray)

# all done!

"""