from org.freehep.math.minuit import FCNBase,MnMigrad,MnUserParameters
from java.lang.Math import *
from jhplot import *
from java.awt import Color

data=P1D("data")      
data.add(5,4,2);      data.add(10,8,3)  # each point has X, Y, standard error on Y
data.add(20,10,3);   data.add(30,12,3)
data.add(40,18,4);   data.add(70,20,4)
data.add(90,26,5);   data.add(150,37,6)
data.add(200,50,6);  data.add(300,60,7)
data.add(500,90,9);  data.add(700,120,10)

class PowerLawFunc(FNon):         # function describing a power law
    def value(self, x):
              return  self.p[0]*pow(x[0],self.p[1])

class PowerLawLogLikeFCN(FCNBase): # define log-likehood. 
  def __init__(self, meas, pos, mvar):
     self.meas=meas;
     self.pos=pos;
     self.mvar=mvar;
  def valueOf(self, par):
     pl = PowerLawFunc("PowerLaw",1,2)
     pl.setParameters(par)
     logsum = 0
     for n in range(len(self.meas)):
                 mu = pl.value([self.pos[n]])
                 k=self.meas[n]
                 logsum = logsum+(k*log(mu) - mu); 
     print "Log likehood=",-1*logsum
     return -1*logsum; 

theFCN = PowerLawLogLikeFCN(data.getArrayY(), data.getArrayX(), data.getArrayErr())
upar = MnUserParameters()
upar.add("p0", 1.0, 0.1)  # initial value and expected error
upar.add("p1", 0.2, 0.1);
migrad =MnMigrad(theFCN, upar)
print "start Migrad "
vmin = migrad.minimize()
if vmin.isValid()==False:
            print "Alternative strategy"
            migrad = MnMigrad(theFCN, upar, 2)
            vmin = migrad.minimize()        
state=vmin.userState()
output=state.params()          # output parameters
print vmin                     # print details of the fit

c1 = HPlot()                   # plot data and the fit function
c1.visible()
min=0;  max=800                # min and max values for plotting
c1.setRangeX(min,max)
c1.setRangeY(0,150)

c1.draw(data)
pl = PowerLawFunc("PowerLaw",1,2)  # plot fit result
pl.setParameters(output)
ff=F1D(pl,min,max)
ff.setPenWidth(2); ff.setColor(Color.blue)
c1.draw(ff)