# -*- coding: utf-8 -*-
"""
Created on Sun Apr  5 13:06:45 2020

@author: neil
"""
# Homework in PDFs and simple Monte Carlo

# program to solve the ejecta problem on the moon.  Given PDFs of the ejecta angle and velocity
# we calculate the PDF of the result.  This nicely illustrates either: a MC sensitivity analysis,
# or the MC propagation of errors thru a complex calculation, (a mixture of sines, cosines and powers)
# Of note is the observation that the spread (or error) of the output is much larger than the input!

# This program illustrates sampling from a 'random' distribution to produce
# samples of a random variable.  If you take a lot of samples, and plot the
# histogram of the result, you find a sketch of the underlying distribution.

import matplotlib.pyplot as plt
import numpy as np

n = 100000                       # number of random samples to take from distributions
moong = 1.62                     # mon gravity in m/s^2

fig = plt.figure()               # set up a plot window
fig.canvas.set_window_title ("Moon gravity Ejecta problem")
ax1 = fig.add_subplot(2,2,1)     # a 4 panel plot
ax1.grid(True)

# our input PDFs for the ejecta vertical angle, and for the initial speed
angle = np.random.uniform(30,50,size=n)        # sample the uniform distribution
V  = np.random.triangular(100,200,225,size=n)  # sample a triangular distribution
# first plot the vertical angle PDF
ax1.hist(angle,bins=50,density='True')
ax1.set_ylabel("PDF value")
ax1.set_xlim(0,60)
ax1.set_title("PDF of initial angle (degrees)")
ax1.set_xlabel("Initial Angle of ejecta")

# 2nd plot is the PDF for inital speed
ax2 = fig.add_subplot(2,2,2)
ax2.grid(True)
ax2.hist(V,bins=50,density='True')        # this calls np.histogram to calculate the bins
ax2.set_xlim(0,250)
ax2.set_ylabel("PDF value")
ax2.set_xlabel("Initial Speed of ejecta")
ax2.set_title("PDF of initial Speeds")


# calculate the output vector of length n, of the impact distances for each random input
Xfinal = V*V * np.sin(angle* np.pi/90)/moong   # this line does 'n' calculations!
maxX      = np.max(Xfinal)     # the maximum range achieved, used for plotting

# set up the plot for the output PDFs of how far the ejecta goes
ax3 = fig.add_subplot(2,2,3)
ax3.grid(True)
ax3.hist(Xfinal,bins=50,density='True')        # this calls np.histogram to calculate the bins
ax3.set_xlim(0,1.1*maxX)
ax3.set_ylabel("PDF value")
ax3.set_xlabel("X position of impact")
ax3.set_title("PDF of ejecta impact points")

# the 4th plot is a x,y view of the minimum and max trajectories
max_index = np.argmax(Xfinal)  # this gets the index of the max value, not the max value itself
min_index = np.argmin(Xfinal)
max_angle = angle[max_index]   # max_angle is the angle that gave the greatest range, its not 50degrees!!
max_V     =     V[max_index]   # max_V is the velocity that gave the greatest range
min_angle = angle[min_index]
min_V     =     V[min_index]
# these following are to make nice plots
maxT      = maxX/(max_V*np.cos(max_angle*np.pi/180))  # time of flight for max range
maxY      = max_V*np.sin(max_angle*np.pi/180)*maxT/2 - moong*maxT*maxT/8 # max height reached

t    = np.linspace(0,maxT,100)  # make a list of times to step thru, to find x,y coords of flight
xmax = np.zeros(100)            # vectors to store x,y coords for plotting
xmin = np.zeros(100)
ymax = np.zeros(100)
ymin = np.zeros(100)

for i in range(100):            # calc 100 points for each trajectory
    xmax[i] = max_V*np.cos(max_angle*np.pi/180)*t[i]  # x,y for the max trajectory
    ymax[i] = max_V*np.sin(max_angle*np.pi/180)*t[i] - moong*t[i]*t[i]/2
    xmin[i] = min_V*np.cos(min_angle*np.pi/180)*t[i]  # x,y for the minimum trajectory
    ymin[i] = min_V*np.sin(min_angle*np.pi/180)*t[i] - moong*t[i]*t[i]/2
    
ax4 = fig.add_subplot(2,2,4)
ax4.grid(True)
ax4.plot(xmax,ymax,'bx')  
ax4.plot(xmin,ymin,'rx')
ax4.set_ylabel("height (meters)")
ax4.set_xlabel("X position ")
ax4.set_title("Ejecta trajectory") 
ax4.set_ylim(0,1.2*maxY)   

plt.show()