# ================================================================
# /home/jetsu/opetus/both/programs/ExerciseSineFit.py
# ================================================================
# - Keep computations and plots apart.
# ================================================================
import os
import numpy as np
import pylab as pl
from scipy import optimize
os.system('clear')
# ================================================================
def Data(n):
    T=np.random.uniform(0,10,n)
    k=IndexSortOrder(T)
    T=T[k]
    BETAsim=-1.+np.random.uniform(0.,2.,3)
    EY=np.random.normal(0,0.1,n)
    G=Model(T,BETAsim)
    Y=G+EY
    Write1(T,Y,EY)
    return T,Y,EY,BETAsim 
# ================================================================
def IndexSortOrder(y):
    k=sorted(range(len(y)), key=lambda i:y[i])
    return k
# ================================================================
def Write1(T,Y,EY):
    file1=open('SineModel.dat',"w")
    for i in range(np.size(T)):
        file1.write("%10.5f %10.5f %10.5f\n" %(T[i],Y[i],EY[i]))
    file1.close()            
    return
# ================================================================
def Model(T,BETA):
    P=2.4
    F=1./P
    ny=2.*np.pi*F*T
    G=BETA[0]+BETA[1]*np.cos(ny)+BETA[2]*np.sin(ny)
    return G
# ================================================================
def Funct(BETA,T,Y,EY):
    G=Model(T,BETA)
    return (Y-G)/EY
# ================================================================
def LSF(T,Y,EY):
    BETAtrial=np.ones(3)
    ny=optimize.leastsq(Funct,BETAtrial,args=(T,Y,EY))
    BETA=ny[0]
    return BETA
# ================================================================
def Plot1(T,Y,EY,TT,GG):
    pl.axes([0.1,0.1,0.8,0.8])
    pl.errorbar(T,Y,EY,fmt='ok',ms=6)
    pl.plot(TT,GG,'r')
    pl.savefig('SineFit.eps')
# ================================================================
#                       Main program
#  - Computations ================================================
n=100
T,Y,EY,BETAsim=Data(n)
BETAfinal=LSF(T,Y,EY)
print('Simulated data ...... BETAsim =',BETAsim)
print('Least Squares Fit ... BETAfinal=',BETAfinal)
DT=np.max(T)-np.min(T)
TT=np.min(T)+(np.arange(101)/100.)*DT
GG=Model(TT,BETAfinal)
#  - Plots =======================================================
Plot1(T,Y,EY,TT,GG)