#!/usr/bin/env python
# coding: utf-8

# In[ ]:


get_ipython().run_line_magic('matplotlib', 'inline')

import flavio
import numpy as np
import clusterking as ck
from clusterking.maths.metric import chi2_metric


# In[ ]:


# Define kinematic function using the flavio package
def dBrdq2(w, q):
    return flavio.np_prediction("dBR/dq2(B+->Dtaunu)", w, q)


# In[ ]:


# Set up and configure Scanner
s = ck.scan.WilsonScanner(scale=5, eft="WET", basis="flavio")
# Set kinematic function
s.set_dfunction(
    dBrdq2,
    binning=np.linspace(3.2, 11.6, 10),
    normalize=True
)
# Set sampling points in Wilson space
s.set_spoints_equidist({
    "CVL_bctaunutau": (-0.5, 0.5, 3),
    "CSL_bctaunutau": (-0.5, 0.5, 3),
    "CT_bctaunutau": (-0.1, 0.1, 3)
})


# In[ ]:


# Run scanner and add errors
d = ck.DataWithErrors()    # Create data object to write results to
r = s.run(d)               # Run scanner
r.write()                  # Write results back to data object
d.add_err_poisson(1000)    # statistical uncertainties
d.add_rel_err_uncorr(0.1)  # 0.1% relative system uncertainties, uncorrelated


# In[ ]:


# Clustering
c = ck.cluster.HierarchyCluster()  # Initialize worker class
c.set_metric(chi2_metric)
c.set_max_d(1)   # "Cut off" value for hierarchy
r = c.run(d)     # Run clustering on d
r.write()        # Write results back to data object


# In[ ]:


# Benchmarking
b = ck.Benchmark()  # Initialize worker class
b.set_metric(chi2_metric)
r = b.run(d)        # Run benchmarking
r.write()           # Write results back to data object


# In[ ]:


# Optional: Save data (kinematic distributions, clusters, BPs, ...)
# d.write("btaunu_q2.sql")


# In[ ]:


# Generate plots similar to figures 1, 2, 3, 4, respectively
d.plot_clusters_scatter(["CVL_bctaunutau", "CSL_bctaunutau"])
d.plot_dist_box()
d.plot_clusters_scatter(["CT_bctaunutau", "CVL_bctaunutau", "CSL_bctaunutau"])
d.plot_dist()


# In[ ]:


# Find closest benchmark point to a new point in parameter space
d.find_closest_bpoints({"CT_bctaunutau": 0, "CVL_bctaunutau": -0.2, "CSL_bctaunutau": 0.2}, n=1)