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

h1 = H2D("Test",16,-5.0, 5.0, 16, -5.0, 5.0)
rand = Random();
for i in range(10000):
               h1.fill(2*rand.nextGaussian(),rand.nextGaussian())

class Gauss2D(FNon):  # 2D Gaussian 
   def value(self, x):
        if self.p[2]>0 and self.p[4]>0:
          d1=(x[0]-self.p[1])/self.p[2]
          d2=(x[1]-self.p[3])/self.p[4]
          return  self.p[0]*exp( -d1*d1-d2*d2 )
        else:
                  return 1.0e10 

class GaussChi2(FCNBase): # define chi2  function  
  def __init__(self,  meas, posX, posY, errors):
     self.meas=meas;
     self.posX=posX;
     self.posY=posY;
     self.errX=0.5*(self.posX[1]-self.posX[0])
     self.errY=0.5*(self.posY[1]-self.posY[0])
     self.errors=errors;
  def valueOf(self, par):
     pl = Gauss2D("2D",2,5)  # 2 varibles, 5 parameters
     pl.setParameters(par)
     chi2 = 0
     for n1 in range(len(self.posX)):
            for n2 in range(len(self.posY)):
                 delta = pl.value([self.posX[n1],self.posY[n2]])-self.meas[n1][n2]
                 sigma2=self.errors[n1][n2]*self.errors[n1][n2]
                 if (sigma2>0): chi2=chi2+(delta*delta)/sigma2; 
     print "Chi2=",chi2 
     return chi2; 

theFCN = GaussChi2(h1.binHeights(), h1.getLowerEdgesX(), h1.getLowerEdgesY(), h1.binErrors())
upar = MnUserParameters()
upar.add("p0", 1000.0, 1)  # initial value and expected error
upar.add("p1", 0.0, 0.1);
upar.add("p2", 1.0, 0.1)  # initial value and expected error
upar.add("p3", 0.0, 0.1);
upar.add("p4", 1.0, 0.1);
migrad =MnMigrad(theFCN, upar)
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 minimization

c1 = HPlot3D("Plot",600,400,2,1)
c1.visible(1)
c1.setRange(-5,5,-5,5,0,500)
c1.draw(h1)
c1.cd(2,1)                                # plot fit function
c1.setRange(-5,5,-5,5,0,500)
pl =Gauss2D("Gauss2D",2,5)  # 2 varibles, 5 parameters
pl.setParameters(output)        # set fitted parameters
ff=F2D(pl,-5,5,-5,5)
c1.draw(ff)