import numpy as np
import matplotlib.pyplot as plt
import matplotlib.mlab as ml
import scipy.optimize as opt
from lmfit import Model

filename='sorted_landscape.dat'
data=np.genfromtxt(filename)

def twoD_Gaussian(x, y, amplitude, xo, yo, sigma1, sigma2, offset):
    xo = float(xo)
    yo = float(yo)
    # a = (np.cos(theta)**2)/(2*sigma_x**2) + (np.sin(theta)**2)/(2*sigma_y**2)
    # b = -(np.sin(2*theta))/(4*sigma_x**2) + (np.sin(2*theta))/(4*sigma_y**2)
    # c = (np.sin(theta)**2)/(2*sigma_x**2) + (np.cos(theta)**2)/(2*sigma_y**2)
    # g = offset + amplitude*np.exp( - (a*((x-xo)**2) + 2*b*(x-xo)*(y-yo)
    #                         + c*((y-yo)**2)))
    g = offset + amplitude*np.exp( - ((x-xo)**2/(2*sigma1**2)+ (y-yo)**2/(2*sigma2**2)))
    return g.ravel()

xscan=data[:,0]*500
yscan=data[:,1]*500 
T=data[:,6]

xi = np.linspace(min(xscan), max(xscan),15)
yi = np.linspace(min(yscan), max(yscan),15)

zi = ml.griddata(xscan, yscan, T, xi, yi, interp='linear')
# print zi

plt.pcolormesh(xi, yi, zi, cmap = plt.get_cmap('rainbow'))
plt.show()

# Fitting
model = Model(twoD_Gaussian, independent_vars=["x", "y"], param_names=["amplitude", "xo", "yo", "sigma1","sigma2","offset"])
params = model.make_params(amplitude=1800, xo=0, yo= 0, sigma1=50,sigma2=50, offset=0)
fitter = model.fit(T, params=params, x=xscan, y=yscan)
print fitter.fit_report()