# -*- coding: utf-8 -*-
from math import *
import pyfits
#import matplotlib.pyplot as plt
from pylab import *
from scipy import *
from scipy import optimize
from scipy import special

def graafik(pilt, x_kesk, y_kesk, alpha, pool_laius, samm):
    fail=open("data.txt", "w")
    im=pyfits.open(pilt)
    im.info()
    pix=im[0].data
    kaugus=[]
    kiirus=[]
    a=tan(alpha)
    b=y_kesk-a*x_kesk
    for x in range (0, pix.shape[0]-pool_laius, samm):
        y=a*x+b
        summa=0
        arvestatud_punkte=0
        for l in range(-pool_laius, pool_laius+1):
            x_p=x-l*sin(alpha)
            y_p=y+l*cos(alpha)
            if (pix[x_p,y_p]!=pix[1,1]):
                summa+=pix[x_p,y_p]
                arvestatud_punkte+=1
        if summa!=0:
            kaugus.append(x)
            kiirus.append(summa/(arvestatud_punkte*1.0))
            fail.write("%f %f\n\r"%(kaugus[len(kaugus)-1], kiirus[len(kiirus)-1]))
    fail.close()
    return kaugus, kiirus
    #plt.axis('equal')
    #plot(kaugus, kiirus, color='blue', linestyle='solid', marker='s',markerfacecolor='cyan', markersize=5)
    #xlabel("kaugus")
    #ylabel("kiirus")
    #show()

def fitting(x_katse,y_katse):
    fitfunc=lambda p,x: p[0]*special.erf(p[3]*(x+p[2]))+p[1]
    errfunc=lambda p,x,y:fitfunc(p,x)-y
    p0=[-144,1650,-250,0.05]
    p1, succes=optimize.leastsq(errfunc, p0[:], args=(x_katse,y_katse))
    r=linspace(x_katse[0], x_katse[len(x_katse)-1], 100)
    plot(x_katse, y_katse, "co", r, fitfunc(p1, r), "b-", r, fitfunc(p0, r), "r--")
    xlabel("kaugus")
    ylabel("kiirus")
    print "%f %f %f"%(p1[0], p1[1], p1[2])
    show()

def main ():
    a, b=graafik("N6384_GHAFAS_vlos.fits", 248.7, 239.2, 122, 50, 2)
    fitting(a,b)

main()