import numpy as np
from scipy.integrate import trapz
from pygolem_lite import Shot

def integrate_hxr(shot, remove_offset=True):
    if not isinstance(shot, Shot):
        shot = Shot(shot)
    try:
        t, hxr = shot['hxr']
    except:
        return np.nan
    if hxr.plasma_start < 0:
        return np.nan
    if remove_offset:
        offset = np.mean(hxr[t < hxr.plasma_start - 2e-3])
        hxr = hxr - offset
    return float(trapz(hxr, hxr.tvec))