#!/usr/bin/python
import os,sys

units = 'z (cm); u or v (cm); fit residual (ns)'

## __________
def read_in_file(fn):
    coords = []
    for l in file(fn,'r').readlines():
        ls = l.split()
        z,uv,res = float(ls[0]),float(ls[1]),float(ls[4])
        coords.append((z,uv,res))
    return coords

## __________
def make_tg2d(coords,ti): 
    gr = TGraph2D()
    i = 0
    for (z,uv,res) in coords:
        if res > -999.:
            i+=1
            gr.SetPoint(i,z,uv,-1.*res)
#            gr.SetPoint(i,z,uv,res)
    ## style-time
    gr.SetTitle(ti+';'+units)
    return gr.Clone() #important!

## ____________________
if __name__=='__main__':

    data_coords = read_in_file('data_res_of_res.txt')
    sim_coords  = read_in_file('sim_res_of_res.txt')

    diff_coords = []
    for (z_d,uv_d,res_d),(z_s,uv_s,res_s) in zip(data_coords,sim_coords):
        if z_d!=z_s and uv_d!=uv_s:
            print 'serious error'
            sys.exit(-1)
        if res_d>-999. and res_s>-999.:
            diff_coords.append((z_d,uv_d,res_d-res_s))

    from ROOT import *
    from clean_canvas import *
    from nice_colours import *
    gStyle = makegstyle(gStyle)
    gStyle = nice_colours(gStyle)

    g_data = make_tg2d(data_coords,'data')
    g_sim  = make_tg2d(sim_coords,'simulation')
    g_diff = make_tg2d(diff_coords,'difference')

    c1 = TCanvas()
    c1.Divide(2,2)
    c1.cd(1)
    g_data.Clone().Draw('surf2')
    c1.cd(2)
    g_sim.Clone().Draw('surf2')
    c1.cd(3)
    g_diff.Clone().Draw('surf2')
    c1.cd(4)
    g_diff.SetTitle('difference on #pm 4ns scale;'+units)
    g_diff.SetMaximum(4.)
    g_diff.SetMinimum(-4.)
    g_diff.Clone().Draw('surf2')

    # now alex's fit function
    def ff(x,a):
        return -1.*(a[1]*x[1]**2+a[0]*x[0]**2+a[3]*x[1]+a[2]*x[0])
#        return a[1]*x[1]**2+a[0]*x[0]**2+a[3]*x[1]+a[2]*x[0]
        
    alex_ff = TF2('alexs fit function; '+units,ff,-1.5,1.7,0.0,1.8,4)
#    alex_ff.SetParameter(0,-0.31257219)
#    alex_ff.SetParameter(1,-0.96603435)
#    alex_ff.SetParameter(2,-0.0034988)
#    alex_ff.SetParameter(3,1.08264566)
#    alex_ff.SetParNames('a0','a1','a2','a3')
    alex_ff.SetParameters(-0.31257219, -0.96603435, -0.0034988 , 1.08264566)

    c2 = TCanvas()
    alex_ff.Draw('surf2')
    
    au = raw_input('>')