import ROOT
import array 
from array import *
import sys 

##::: HELPER FUNCTIONS
def ImportLatex( size=0.05, font=42, align=12 ):
 ##::: Import latex for labels on the histograms
  tex = ROOT.TLatex()
  tex.SetNDC()
  tex.SetTextFont( font ) 
  tex.SetTextSize( size )
  tex.SetTextColor( 1 )
  tex.SetTextAlign( align )
  return tex

def DrawAtlas( x, y, align=11 ):
    #print "Size is %.2f"%self.Size
    DrawTLatex( x, y,"#font[72]{ATLAS} %s" % 'Internal', 24, 43, align )

def DrawTLatex( x, y, text, size=18, font=43, align=11 ):
    tex = ImportLatex( size, font, align )
    tex.DrawLatex( x, y, text )

def SetCanvasDefaults( canvas1, tM = 0.015, rM = 0.020, bM = 0.15, lM = 0.12 ):
    canvas1.SetTicks()
    canvas1.SetTopMargin( tM )
    canvas1.SetRightMargin( rM )
    canvas1.SetBottomMargin( bM )
    canvas1.SetLeftMargin( lM ) #original value from Gaetano was .12

def ConfigTGraph( graph ):
  graph.GetXaxis().SetTitle( 'Camera height [mm]' )
  graph.GetYaxis().SetTitle( '#sigma_{-#frac{dI}{dR}}' )
  graph.SetMarkerSize( 1 )
  graph.SetMarkerStyle( 8 )
  graph.GetXaxis().SetRangeUser( 2.5, 4.5 )
  graph.SetMaximum( 3.7 )


##::: Main Analysis 
# Set up canvases for plots
c1 = ROOT.TCanvas( 'c1', 'c1', 600, 600 )
SetCanvasDefaults( c1 )
c2 = ROOT.TCanvas( 'c2', 'c1', 600, 600 )
SetCanvasDefaults( c2 ) 

# Load csv file and read it as a TGraph (Updating to TGraphErrors; current errors are symmetric so no need for TGraphAsymmErrors)
fname3 = '/Users/hannahelizabeth/Desktop/csvdots_146,22/sigmadIdRvZ_imageset.csv' 
gr = ROOT.TGraph( fname3, "%*lg%*lg%lg%lg", ',' ) ## --- https://root.cern.ch/doc/master/classTGraphErrors.html#a5f6ab98471c9e48337cbbedbbfad07e1

#fname3 = '/Users/hannahelizabeth/Desktop/allGaussFits2.csv'
# FileName, x [mm], y [mm], z [mm], Fit Quality (Chi2/NDOF), Amplitude, Amp. Error, Mean, Mean Error, Sigma, Sigma Error
#gr = ROOT.TGraph( fname3, "%*s%*lg%*lg%lg%*lg%*lg%*lg%*lg%*lg%lg%*lg", ',' )
#gr2 = ROOT.TGraph( fname3, "%*s%*lg%*lg%lg%*lg%*lg%*lg%*lg%*lg%lg%*lg", ',' )

# Do the fit
print 'Performing quadratic fit...'
c1.cd()
f2 = ROOT.TF1( 'f2', 'pol2', 2.7, 4.3 )
r_pol2 = gr.Fit( f2, 'RFS' )
r_pol2.Print( 'V' )
#print r_pol2.Chi2()/f2.GetNDF()

#Calculate the error on the minimumX, assuming errors fully correlated 
# p0 + p1*x + p2*x^2
p2 = [ r_pol2.Parameter( 2 ), r_pol2.ParError( 2 ) ] 
p1 = [ r_pol2.Parameter( 1 ), r_pol2.ParError( 1 ) ] 

#quadratic term -- sigma_p2^2 * |p1/2*p2^2|^2
errSigP2 = ( p2[ 1 ]**2 )*( ROOT.TMath.Abs( p1[ 0 ]/( 2*( p2[ 0 ]**2 ) ) )**2 )
#print errSigP2

#linear term -- sigma_p1^2 * |-1/2*p2|^2
errSigP1 = ( p1[ 1 ]**2 )*( ROOT.TMath.Abs( -0.5*p2[ 0 ]**( -1 ) )**2 )
#print errSigP1

#correlation -- 2 * p1/2*p2^2 * -1/2*p2 * sigma_p2 * sigma_p1 * p12_corr 
corr = 2 * ( p1[ 0 ]/( 2*( p2[ 0 ]**2 ) ) ) * ( -0.5*p2[ 0 ]**( -1 ) ) * p2[ 1 ] * p1[ 1 ] * r_pol2.Correlation( 1,2 ) 
#print corr

sigma_pol2 = ROOT.TMath.Sqrt( errSigP2 + errSigP1 + corr )

minStr = 'fitted WD = %.03f #pm %.03f mm' % (f2.GetMinimumX(), sigma_pol2 )
print minStr

# Create plot
ConfigTGraph( gr )
gr.Draw( 'ap' )
gr.SetTitle( '' )
xLat = 0.4
yLat = 0.8
DrawAtlas( xLat, yLat )
DrawTLatex( xLat, yLat-0.05, minStr )
DrawTLatex( xLat, yLat-0.1, '#it{pol2}' )
c1.SaveAs( 'quadratic.pdf' )

# End script execution
sys.exit()

# Code not executed -- not using quartic fit
print 'Performing quartic fit {p3 fixed}...'
c2.cd()
f4 = ROOT.TF1( 'f4', 'pol4',2.7,4.3 )
#f4.FixParameter( 1, 0 )
f4.FixParameter( 3, 0 )
r_pol4 = gr2.Fit( f4, 'RFS' )
r_pol4.Print( 'V' )
print r_pol4.Chi2()/f2.GetNDF()
minStr = str( f4.GetMinimumX() ) + ' +/- CALCULATE'
sigma_pol4 = 0.000000 
minStr = 'fitted WD = %.03f #pm %.03f mm' % ( f4.GetMinimumX(), sigma_pol4 ) 

ConfigTGraph( gr2 )
gr2.Draw( 'ap' )
gr2.SetTitle( '' ) 
DrawAtlas( xLat, yLat )
DrawTLatex( xLat, yLat-0.05, minStr )
DrawTLatex( xLat, yLat-0.1, '#it{pol4, fixed p3}' )
c2.SaveAs( 'quartic.pdf' )