import pyhf
import json
import copy
import jsonpatch
import numpy as np
import matplotlib.pyplot as plt

%matplotlib inline

def make_model(channel_list):
    spec = json.load(open("lhood.json"))
    patch = json.load(open("patch.json"))
    spec = jsonpatch.apply_patch(spec, patch)
    spec["channels"] = [c for c in spec["channels"] if c["name"] in channel_list]

    ws = pyhf.Workspace(spec)
    model = ws.model(
        measurement_name="NormalMeasurement",
        modifier_settings={
            "normsys": {"interpcode": "code4"},
            "histosys": {"interpcode": "code4p"},
        },
    )
    data = ws.data(model)
    return ws, model, data

def fitresults(constraints=None):
    _, model, data = make_model(["CRtt_meff"])

    pyhf.set_backend("numpy", pyhf.optimize.minuit_optimizer(verbose=True))

    constraints = constraints or []
    init_pars = model.config.suggested_init()
    fixed_params = model.config.suggested_fixed()

    for idx, fixed_val in constraints:
        init_pars[idx] = fixed_val
        fixed_params[idx] = True

    result = pyhf.infer.mle.fit(
        data,
        model,
        init_pars=init_pars,
        fixed_params=fixed_params,
        return_uncertainties=True,
    )
    bestfit = result[:, 0]
    errors = result[:, 1]
    return model, data, bestfit, errors

def calc_impact(idx, b, e, i, width, poi_index):
    _, _, bb, ee = fitresults([(idx, b + e)])
    poi_up_post = bb[poi_index]

    _, _, bb, ee = fitresults([(idx, b - e)])
    poi_dn_post = bb[poi_index]

    _, _, bb, ee = fitresults([(idx, b + width)])
    poi_up_pre = bb[poi_index]

    _, _, bb, ee = fitresults([(idx, b - width)])
    poi_dn_pre = bb[poi_index]
    return np.asarray([poi_dn_post, poi_up_post, poi_dn_pre, poi_up_pre])

def get_impact_data():
    model, _, b, e = fitresults()
    widths = pyhf.tensorlib.concatenate(
        [
            model.config.param_set(k).width()
            for k, v in model.config.par_map.items()
            if model.config.param_set(k).constrained
        ]
    )
    initv = pyhf.tensorlib.concatenate(
        [
            model.config.param_set(k).suggested_init
            for k, v in model.config.par_map.items()
            if model.config.param_set(k).constrained
        ]
    )
    labels = np.asarray(
        [
            f"{k}[{i:02}]" if model.config.param_set(k).n_parameters > 1 else k
            for k in model.config.par_order
            if model.config.param_set(k).constrained
            for i in range(model.config.param_set(k).n_parameters)
        ]
    )
    poi_free = b[model.config.poi_index]
    impacts = []

    for i, width in enumerate(widths):
        if i == model.config.poi_index:
            continue
        if i % 5 == 0:
            print(i)
        impct = calc_impact(i, b[i], e[i], width, initv[i], model.config.poi_index)
        impacts.append(impct - poi_free)
    return np.asarray(impacts), labels

impcord = np.argsort(np.max(np.abs(impacts[:, :2]), axis=1))
simpacts = impacts[impcord]
slabels = labels[impcord]