#! python3
import pandas as pd
import matplotlib.pyplot as plt

import VisTools.plotting as vp
import ruamel.yaml

yaml = ruamel.yaml.YAML()

data = pd.read_csv('./log_poticalib_ana_02.csv')

print(data.columns)

data['val_poti'] = [int(i, 16) for i in data['val_poti']]

def poll_db(uuid: str):

        db = []
        with open("./pitstop/pitdb.yaml") as f:
            db_raw = yaml.load_all(f)

            for doc in db_raw:
                db.append(doc)

        for doc in db:
            doc = dict(doc)
            if doc['id'] == uuid:
                return doc
        return None

def __calc_vpin(data: pd.DataFrame, poly):
        m = poly

        if m[2] != 0:
            data['x'] = 4 * m[2] * (data['v_pit/v'] - m[0]) + m[1]**2
            data['x'] = np.sqrt(data['x'].tolist())
            data['x'] -= m[1]
            data['x'] /= (2 * m[2])
        else:
            data['x'] = (data['v_pit/v'] - m[0]) / m[1]

        return data

default = poll_db('default')

data = __calc_vpin(data, default['poly18v'])

plt.errorbar(
    data['x'],
    data['v_pit/v'],
    fmt='.'
)

plt.errorbar(
    data['x'],
    data['v_keith/v'],
    fmt='.'
)

res = vp.fit(
    data['x'],
    data['v_keith/v'],
    lambda x,m0,m1,m2: m0+m1*x+m2*x*x,
    default['poly18v']
)

plt.plot(
    data['x'],
    [res[0].n+e*res[1].n+res[2].n*e*e for e in data['x']]
)

print('poly18v:',[round(e.n,4) for e in res])
plt.show()