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")                         # input 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 PowerLawChi2FCN(FCNBase): # define chi2  function  
  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)
     chi2 = 0
     for n in range(len(self.meas)):
                 delta = pl.value([self.pos[n]]) - self.meas[n]
                 sigma2=self.mvar[n]*self.mvar[n]
                 if (sigma2>0): chi2=chi2+(delta*delta)/sigma2 
     print "Chi2=",chi2 
     return chi2; 

theFCN = PowerLawChi2FCN(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)                 # plot 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)