import numpy as np
import matplotlib.pyplot as plt

class Stock:
    def __init__(self, S0, vol, mean, N, delta):
        self.S0 = S0
        self.vol = vol
        self.mean = mean
        self.N = N
        self.delta = delta

    def graphPrice(self):
        # note: N+1 due to initial S0 value in price list + N steps of sim
        timeSteps = [self.delta * n for n in range(self.N+1)]
        plt.plot(timeSteps, self.priceList())
        plt.show()

    def endPrice(self):
        return self.priceList()[-1]
        
    # return list of lengh N of the points generated by simulation of stock price
    def priceList(self):
        output = [self.S0]
        for i in range(self.N):
            output.append(self.simStep(output[i], self.delta))
        return output

    def simStep(self, St, delta):
        epsilon = np.random.normal(0, 1)
        return St*np.exp((self.mean - (self.vol*self.vol/2))*delta + self.vol*epsilon*np.sqrt(delta))
