import numpy as np
import scipy.optimize as opt
import matplotlib.pyplot as plt


# functions for fitting
def fx1(x, a, b):
    """
    returns f(x) = a+b*x 
    """
    f = a + b*x
    return f

def fx2(x,a,b,c,d,e):
    """
    returns f(x) = (a*x)/d*ln(b*x +c) + e
    """
    f = (a*x)/d*np.log(b*x +c) + e
    return f

def fx3(x,a,b,c,d):
    """
    returns f(x) = a*ln(b*x +c) + d
    """
    f = a*np.log(b*x +c)+d
    return f

infile = open("curvefit_happy.csv","r")
infile.readline()
# inflie.readline() gets rid of the column headers
# read in file generated by happy_numbers.py

lines=[] # empty list for lines from file
for line in infile:
    lines.append(line)

infile.close()

x = np.array([])
y = np.array([])

for line in lines:
    data = line.strip().split(",")
    x = np.append(x, float(data[0])) #x array for happy numbers
    y = np.append(y, float(data[1])) #y array for the happy number count. Total happy numbers is len(x) or the last index in the y array
                                     #for 10,000 happy numbers this is 1442

# generating fit parameters by using curve_fit. Maxfev is the amount of function calls, we found the defualt of 200 per variable was to low
p1, cov = opt.curve_fit(fx1,x,y,maxfev=1000000)
p2, cov2 = opt.curve_fit(fx2,x,y,maxfev=1000000)
p3, cov3 = opt.curve_fit(fx3,x,y,maxfev=1000000)

# fit parameters, you will need them
print ("Fit-parameters = ", p1,'\n',p2,'\n',p3)

# this creates the fitted coordinate arrays with the fit parameters
yp1 = fx1(x,p1[0],p1[1],) #y=a*x +b
yp2 = fx2(x,p2[0],p2[1],p2[2],p2[3],p2[4]) #y= (a*x)/d*ln(b*x +c) + e
yp3 = fx3(x,p3[0],p3[1],p3[2],p3[3]) #y = a*ln(b*x +c) + d

# Create plots of fitted curves
axes = plt.gca()

plt.figure(1)
plt.plot(x,y,'b--')
plt.plot(x,yp1,'r-')
plt.ylabel('Happy number value')
plt.xlabel('Cumulative amount of happy numbers')
plt.grid()
plt.title('Happy number curve fit y=x')

plt.figure(2)
plt.plot(x,y,'b--')
plt.plot(x,yp2,'r-')
plt.ylabel('Happy number value')
plt.xlabel('Cumulative amount of happy numbers')
plt.grid()
plt.title('Happy number curve fit y = x/ln(x)')

plt.figure(3)
plt.plot(x,y,'b--')
plt.plot(x,yp3,'r-')
plt.ylabel('Happy number value')
plt.xlabel('Cumulative amount of happy numbers')
plt.grid()
plt.title('Happy number curve fit y = ln(x)')
#we find that y=ln(x) provides the best fit, however x/ln(x) is also very close, y=x is not that great.
#it is also recommended to look at the images in full screen for the best judgement of the fit by eye (chi-by-eye)
plt.show()
