#!/usr/bin/python
import os,sys
from michel_spectrum import *
from approx_michel_spectrum import *
from scipy.integrate import quad
from ROOT import *
from clean_canvas import *
gStyle = makegstyle(gStyle)

## __________
def git(vx,vy,ti,col=2,mark=20):
    gr = TGraph()
    gr.SetMarkerColor(col)
    gr.SetMarkerStyle(mark)
    gr.SetTitle(ti)
    for i,(px,py) in enumerate(zip(vx,vy)):
        gr.SetPoint(i,px,py)
    return gr

## __________
def make_Nu_Nd_graphs(xs,mp,a,b):
    ups, downs, sums, diffs = [], [], [], []        
    for x in xs:
        up, down = Nu(x,a,b,mp), Nd(x,a,b,mp)
        ups.append(up)
        downs.append(down)
        sums.append(up+down)
        diffs.append(up-down)
    grs = [git(xs,ups,'Nf',1,1),git(xs,downs,'Nb',1,1),git(xs,sums,'sum',1,1),git(xs,diffs,'diff',1,1),
           git(sums,diffs,'correlation sum diff',1,1)]
    return grs

## __________
def make_asy_graphs(xs,mp,a,b,mpname):
    asys = []
    for x in xs:
        asys.append(asy(x,a,b,mp))
    grs = [git(xs,asys,'asymmetry_for_scan_'+mpname)]
    return grs

## __________
def gen_mps(mpname,pm):
    mp_sm = {'rho':0.75,'eta':0.00,'xi':1.00,'delta':0.75}
    sm_value = mp_sm[mpname]
    mps = []
    for v in [sm_value,sm_value-pm,sm_value+pm]:
        mp = mp_sm
        mp[mpname] = v
        mps.append(mp.copy())
    return mps

## __________
def play_nu_nd(mpname,pm,a,b):

    mark = [20,22,23]
    mps = gen_mps(mpname,pm)

    c = TCanvas(mpname+'_nund_scan',mpname+'_nund_scan')
    c.Divide(2,3)
    for i,(imark,mp) in enumerate(zip(mark,mps)):
        grs = make_Nu_Nd_graphs(xs,mp,a,b)
        for g in grs:
            g.SetMarkerStyle(imark)
            g.SetMarkerSize(0.3)
        for j in range(5):
            c.cd(j+1)
            if i==0:
                grs[j].Clone().Draw('AC')
            else:
                grs[j].Clone().Draw('C')
    return c

## __________
def play_asy(mpname,pm,a,b):

    mark = [20,22,23]
    mps = gen_mps(mpname,pm)

    c = TCanvas(mpname+'_asy_scan',mpname+'_asy_scan')
#    c.Divide(1)
    for i,(imark,mp) in enumerate(zip(mark,mps)):
        grs = make_asy_graphs(xs,mp,a,b,mpname)

        #integrate using Gaussian quadrature
        I1 = quad(asy, 0.0, 1.0, args=(a,b,mp))[0]
        print 'mp=',mp,'I=',I1
        
        for g in grs:
            g.SetMarkerStyle(imark)
            g.SetMarkerSize(0.3)
        for j in range(1):
#            c.cd(j+1)
            if i==0:
                grs[j].Clone().Draw('AC')
#                grs[j].Clone().Draw('APL')
            else:
                grs[j].Clone().Draw('C')
#                grs[j].Clone().Draw('PL')
##        c100 = TCanvas('c100','c100')
##        grs[0].Clone().Draw('AC')
##        au = raw_input('>')
    return c



## ____________________
if __name__=='__main__':

    xs = [0.+ix*(1./1000.) for ix in range(0,1001)]

    ## __________ verify can generate michel spectrum __________

    grs_full_michel, grs_base, grs_simple_funcs = [], [], []

    for c in [-1.,0.,1.]:

        full_michel, base, simple_funcs = [], [], []
        for x in xs:
            full_michel.append( michel_spectrum(x,c) )
            base.append( michel_spectrum(x,c,{'rho':0.75,'eta':0.00,'xi':1.00,'delta':0.75},RCs='base') )
            simple_funcs.append( approx_michel_spectrum(x,c) )

        grs_full_michel.append(git(xs,full_michel,'full michel spectrum inc all RCs',2,20))
        grs_base.append(git(xs,base,'base michel spectrum from full code',4,21))
        grs_simple_funcs.append(git(xs,simple_funcs,'base, no RC, simple funcs',1,22))

    c1 = TCanvas('verifying_can_generate_spectrum','verifying_can_generate_spectrum')

    grs_full_michel[0].Draw('AP')
    grs_base[0].Draw('P')
    grs_simple_funcs[0].Draw('P')

    grs_full_michel[1].Draw('P')
    grs_base[1].Draw('P')
    grs_simple_funcs[1].Draw('P')

    grs_full_michel[2].Draw('P')
    grs_base[2].Draw('P')
    grs_simple_funcs[2].Draw('P')

    ## __________ play with Nu and Nd __________
    a = 0.0
    b = 0.5
    c_rho = play_nu_nd('rho',0.25,a,b)
    c_eta = play_nu_nd('eta',0.5,a,b)
    c_xi = play_nu_nd('xi',1.00,a,b)
    c_delta = play_nu_nd('delta',0.25,a,b)

    ## _________ play with asymmetry __________
    a = 0.0
    b = 1.0
    ca_rho = play_asy('rho',0.1,a,b)
    ca_eta = play_asy('eta',0.1,a,b)
    ca_xi = play_asy('xi',0.1,a,b)
    ca_delta = play_asy('delta',0.1,a,b)
    
    au = raw_input('>')