import scipy.optimize
import numpy as np
from defusedxml import ElementTree
from collections import deque
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit

fit_f1 = lambda x, K, alpha: K * np.exp(alpha * x)
fit_f2 = lambda x, a, b: a*x + b



data = {"Germany": deque()
	, "France": deque()
	, "Italy": deque()
#	, "United States": deque()
#	, "Angola": deque()
#	, "China": deque()
}

with open("data/API_NY.GDP.MKTP.KN_DS2_en_xml_v2_10230884.xml") as fin:
	tree = ElementTree.parse(fin)
	for record in tree.getroot().find("data").findall("record"):
		this_data = {field.get("name"): field.text for field in record.findall("field")}
		if(this_data["Country or Area"] in data):
			if(this_data["Value"] != None):
				data[this_data["Country or Area"]].append((this_data["Year"], this_data["Value"]))

			

class Data(object):
	def __init__(self, raw_data):
		self.x = np.array([int(k) for k, v in raw_data])
		self.y = np.array([float(v) for k, v in raw_data])

plots = deque()
for country, values in data.items():
	values = Data(values)

	popt1, pcov = curve_fit(fit_f1, values.x, values.y
			, p0=[values.y[0], 1])
	popt2, pcov = curve_fit(fit_f2, values.x, values.y
			, p0=[values.y[0], (values.y[-1] - values.y[0])/(values.x[-1] - values.x[0])])

	f1 = lambda x: fit_f1(x, popt1[0], popt1[1])
	f2 = lambda x: fit_f2(x, popt2[0], popt2[1])

	p1, = plt.plot(values.x, values.y, label="{}: real".format(country))
	p2, = plt.plot(values.x, f1(values.x), label="%s: exponential fit, K=%.3e, $\\alpha$=%.3e" % (country, *popt1))
	p3, = plt.plot(values.x, f2(values.x), label="%s: linear fit, a=%.3e, b=%.3e" % (country, *popt2))
	plots.extend([p1, p2, p3])

plt.legend(handles=list(plots))
plt.show()