"""
Calibration constants and model type can be found in the table TimsCalibration
For very old timsTOF data you might find ModelType 1, for quite some time we just use ModelType 2

c0 = ndelay
c1 = ncycles
c2 = dv_start
c3 = dv_end
c4 = ttrans = transit time in cycles
c5 = polynomgrad (im moment immer 1)
c6 = C0
c7 = C1
c8 = vmin
c9 = vmax

given a scan number scan_num
voltage V = dv_start + m * (scan_num - ttrans - ndelay)
where m = (dv_end - dv_start) / ncycles

you should check that the voltage is in the range [vmin, vmax]. If not the calibration is not valid

Using the voltage you get the inverse mobility
1/K0 = 1 / (C0 + C1 / V) = 1 / (c6 + c7/V)

"""

from bdal.io.timsdata import TimsData
import os, sys
import numpy as np

if len(sys.argv) < 2:
    raise RuntimeError("need arguments: tdf_directory")

analysis_dir = sys.argv[1]


tims_file = TimsData(analysis_dir)
frameId=4 # just some frame number

c = tims_file.conn
num_scans = c.execute("select distinct(NumScans) from Frames").fetchone()[0]
scan_numbers = np.arange(0, num_scans)

# temperature compensation
def get_voltage(conn, frame_id):
    tims_calib_id = c.execute("select TimsCalibration from Frames where Id={}".format(frame_id)).fetchone()[0]
    ndelay = c.execute("select C0 from TimsCalibration where Id={}".format(tims_calib_id)).fetchone()[0]
    ncycles = c.execute("select C1 from TimsCalibration where Id={}".format(tims_calib_id)).fetchone()[0]
    dv_start = c.execute("select C2 from TimsCalibration where Id={}".format(tims_calib_id)).fetchone()[0]
    dv_end = c.execute("select C3 from TimsCalibration where Id={}".format(tims_calib_id)).fetchone()[0]
    # transit time in cycles
    ttrans = c.execute("select C4 from TimsCalibration where Id={}".format(tims_calib_id)).fetchone()[0]
    # bounds for voltage
    vmin = c.execute("select C8 from TimsCalibration where Id={}".format(tims_calib_id)).fetchone()[0]
    vmax = c.execute("select C9 from TimsCalibration where Id={}".format(tims_calib_id)).fetchone()[0]

    slope = (dv_end - dv_start) / ncycles
    def scan_num_voltage_transformation(scan_num):
        v = dv_start + slope*(scan_num - ttrans - ndelay)
        vmin_indices = np.where(v < vmin)
        vmax_indices = np.where(v > vmax)
        if (vmax_indices[0].size != 0 or vmax_indices[0] != 0):
            raise RuntimeError("invalid tims calibration")
        return v

    return scan_num_voltage_transformation

def get_oneoverK0_transformation(conn, frame_id):
    # type one calibration, values taken from tdf-sdk example data
    mz_calib_id = c.execute("select TimsCalibration from Frames where Id={}".format(frame_id)).fetchone()[0]
    c6 = c.execute("select C6 from TimsCalibration where Id={}".format(mz_calib_id)).fetchone()[0]
    c7 = c.execute("select C7 from TimsCalibration where Id={}".format(mz_calib_id)).fetchone()[0]

    def oneoverK0_transformation(scan_num):
        voltage = get_voltage(conn, frame_id)
        return 1 / (c6 + c7 / voltage(scan_num))

    return oneoverK0_transformation

voltage = get_voltage(c, frameId)
v_sdk1 = tims_file.scanNumToVoltage(frameId, scan_numbers)

# pick one voltage example to check
print ("voltage for {}: {}".format(scan_numbers[3], voltage(scan_numbers[3])))
print ("voltage for from sdk {}: {}".format(scan_numbers[3], v_sdk1[3]))


oneoverK0 = get_oneoverK0_transformation(c, frameId)
mob_sdk1 = tims_file.scanNumToOneOverK0(frameId, scan_numbers)

# pick an example for 1/K0 check
print ("1/K0 for {}: {}".format(scan_numbers[14], oneoverK0(scan_numbers[14])))
print ("1/K0 from sdk for {}: {}".format(scan_numbers[14], mob_sdk1[14]))


# check differences of 1/K0
mob1 = oneoverK0(scan_numbers)
diff_mob = np.abs(mob1 - mob_sdk1);
# get max value of diffs
print ("max difference between mobility values: {}".format(diff_mob.max()))

