# -*- coding: utf-8 -*-
"""
Created on Sun Sep 27 15:16:37 2020

@author: jbobowsk
"""

# This tutorial shows how to create a function in Python in one .py file and
# then import it as a module and execute it in another .py Python script.

# The first example function will be used to calculate the factorial of n.
# We will call this function 'fact' and it will be saved in the script
# mod1.py'.

# The format for entering this function is shown in the lines below.
# They are commented out in this script, but these lines of code appear in
# the file mod1.m.

# def fact(n):
#     z = 1
#     for i in range(n):
#         z = z*(i+1)
#     return z

# To call an execute this function, the file mod1.py must me in the same
# directory as the current script that you're working in.  Then, simply
# import mod1.py and then call the function using `mod1.fact(n)' where n is 
# the integer that you want to take the factorial of.
import mod1

# Here's the factorial of 5.
x = mod1.fact(5)
print(x)

# Here's the factorial of 6.
x = mod1.fact(6)
print(x)

# mod1.py also has a function to calculate the number of combinations for
# selecting m objects for n.  c = n!/((n - m)!m!).  Here's the code that is
# contained in mod1.py.  Note that the 'choose' function relies on the 'fact' 
# function that is also contained in mod1.py.

# def choose(n, m):
#     c = fact(n)/(fact(n - m)*fact(m))
#     return int(c)

# Here's 6 choose 3.
x = mod1.choose(6, 3)
print(x)

# Of course, a factorial function already exists and we didn't need to write 
# our own.
import math
print(math.factorial(5), math.factorial(6))

# There is also a module/function for calculating combinations.
from scipy.special import comb
print(comb(6, 3, exact = True))

# Here's one more basic example.  We will calculate the quadrature sum of an
# array of numbers.  The array will be passed to the function which will return
# the scaler quadrature sum.  This type of calculation is used often in 
# propagation of errors.  It is also used to find the magnitude of a vector.

# The following code is included in mod1.py:
# def quad(xx):
#     total = 0
#     for i in range(len(xx)):
#         total = total + xx[i]**2
#     return total**0.5

# Here's the quadrature sum of 1, 2, 3: sqrt(1^2 + 2^2 + 3^2) = sqrt(14)
xx = [1, 2, 3]
mag = mod1.quad(xx)
print(mag)

# As a final example, we'll create a slightly more complicated function.
# This function, called 'primes(a, b)' in the module mod2.py, will find all 
# of the prime numbers between input integers a and b.  The function code 
# is shown below without any explanations.  If you open the file mod2.py you 
# will find some comments that help explain the purpose of the
# various lines of code.  The function returns to values.  The first is an
# array of the prime numbers found is sequential order.  The second item returned
# is the number of primes that were found.  The code also checks that the user
# supplied appropriate inputs.  The requirements are that a and b are positive
# integers and that b > a.

#def primes(a, b):
#     import numpy as np
#     primes = []
#     if a % 1 == 0 and b % 1 == 0 and a > 0 and b > 0 and b > a:
#         x = 0
#         for i in range (a, b + 1):
#             cnt = 0
#             if i != 1 and i % 2 == 1 and sum(map(int, str(i))) % 3 != 0\
#             and i % 10 != 5 or i == 2 or i == 3 or i == 5:
#                 j = 7
#                 while j < i/(j - 2) and cnt == 0:
#                     if i % j == 0:
#                         cnt = 1
#                     j = j + 2
#                 if cnt == 0:
#                     x = x + 1
#                     primes = primes + [i]
#         f = np.array(primes)
#         n = len(f)
#     else:
#         f = 'ERROR: a and b in primes(a, b) must be positive integers with b > a.'
#         n = -1
#     return f, n 

# Here's what happens when mod2.primes(a, b) is called with bad inputs.
import mod2
xx, x = mod2.primes(6.2, -4)
print(xx)
print(x)

# Here's a couple of simple tests with sensible inputs.
a = 1
b = 100
xx, x = mod2.primes(a, b)
print('There are', x, 'prime numbers between', a, 'and', b)
print(xx)

a = 500
b = int(10e3)
xx, x = mod2.primes(a, b)
print('There are', x, 'prime numbers between', a, 'and', b)
print(xx)

# Let's try calling the function with some more interesting inputs.
a = 92939223
b = 92939600
xx, x = mod2.primes(a, b)
print('There are', x, 'prime numbers between', a, 'and', b)
print(xx)

# We can use our prime number function to find the number of prime numbers
# between, say, 1 and 1e6 and write those numbers to a file.
from datetime import datetime
max_n = int(1e6)
print(datetime.now())
xx, x = mod2.primes(1, max_n)
print(datetime.now())
print(x)
import numpy as np
np.savetxt("primes_1e6.txt", xx, fmt="%s")
# This took my laptop about 23 seconds.

# Here's a plot of the number of prime numbers less than x versus x.
import matplotlib.pyplot as plt
plt.plot(xx, np.arange(1, len(xx) + 1, 1), linewidth = 2)
plt.xlabel('x')
plt.ylabel('Number of primes less than x')
plt.axis((0, 1e6, 0, 8e4))