from java.awt import Color
from jhplot  import *
from jhplot.io  import *
from jhplot.math.num.pdf import *
import urllib

#dir=SystemDir+fSep+"macros"+fSep+"examples"+fSep+"data"+fSep    
#a1=FileAida( dir+"UA1.aida")
#a2=FileAida( dir+"Atlas.aida")

# get from URL
wfile1="UA1.aida"
urllib.urlretrieve ("http://datamelt.org/examples/data/"+wfile1,wfile1)
wfile2="Atlas.aida"
urllib.urlretrieve ("http://datamelt.org/examples/data/"+wfile2,wfile2)
a1=FileAida( wfile1 )
a2=FileAida( wfile2 )

ds_UA1=a1.get("/REF/UA1_1990_S2044935/d01-x01-y03") 
print type(ds_UA1)

p1=P1D(ds_UA1)
p1.setTitle("UA1 data")
p1.setColor(Color.blue)
int=p1.integral()
p1.operScale(1,1.0/int)

c1 = HPlot("Canvas",600,400)
c1.setGTitle("UA1 data")
c1.visible()
c1.setNameX("P_{n}")
c1.setNameY("Probability")
c1.setMarginLeft(90)
c1.setRange(0,100,0.0000001,10)
c1.setLegendPos(0.6,0.9,"NDC")
     
c1.setLogScale(1,1)
c1.draw(p1)


max=50
nbd1=NegativeBinomial(max,0.7)
print "Mean=",nbd1.getMean(),"variance=",nbd1.getVariance()
p1=P1D("NBD (σ =30)")
for i in range(100):
        p1.add(i,nbd1.probability(i))
p1.setStyle("l")   
p1.setSymbolSize(0)     
c1.draw(p1)


max=21
nbd2=NegativeBinomial(max,0.5)
print "Mean=",nbd2.getMean(), "variance=",nbd2.getVariance()
p2=P1D("NBD (σ =42)" )
for i in range(100):
        p2.add(i,nbd2.probability(i))
p2.setStyle("l")   
p2.setSymbolSize(0)    
p2.setColor(Color.blue) 
c1.draw(p2)

max=14
nbd3=NegativeBinomial(max,0.4)
print "Mean=",nbd3.getMean(), "variance=",nbd3.getVariance()
p3=P1D("NBD (σ =52)" )
for i in range(100):
        p3.add(i,nbd3.probability(i))
p3.setStyle("l")   
p3.setSymbolSize(0)    
p3.setColor(Color.green) 
c1.draw(p3)

max=5
nbd4=NegativeBinomial(max,0.2)
print "Mean=",nbd4.getMean(), "variance=",nbd4.getVariance()
p4=P1D("NBD (σ =100)" )
for i in range(100):
    p4.add(i,nbd4.probability(i))
p4.setStyle("l")   
p4.setSymbolSize(0)    
p4.setColor(Color.red) 
c1.draw(p4)