# *****************************************************************************
# This program is free software; you can redistribute it and/or modify it under
# the terms of the GNU General Public License as published by the Free Software
# Foundation; either version 2 of the License, or (at your option) any later
# version.
#
# This program is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more
# details.
#
# You should have received a copy of the GNU General Public License along with
# this program; if not, write to the Free Software Foundation, Inc.,
# 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
#
# Module authors:
#   Damaris Tartarotti Maimone
#   Markus Zolliker <markus.zolliker@psi.ch>
# *****************************************************************************
"""Wrapper for the ADQ data acquisition card for ultrasound"""
import sys
import atexit
import signal
import time
import numpy as np
import ctypes as ct
from scipy.signal import butter, filtfilt


# For different trigger modes
SW_TRIG = 1
# The following external trigger does not work if the level of the trigger is very close to 0.5V.
# Now we have it close to 3V, and it works
EXT_TRIG_1 = 2
EXT_TRIG_2 = 7
EXT_TRIG_3 = 8
LVL_TRIG = 3
INT_TRIG = 4
LVL_FALLING = 0
LVL_RISING = 1

ADQ_CLOCK_INT_INTREF = 0  # internal clock source
ADQ_CLOCK_EXT_REF = 1  # internal clock source, external reference
ADQ_CLOCK_EXT_CLOCK = 2  # External clock source

ADQ_TRANSFER_MODE_NORMAL = 0x00
ADQ_CHANNELS_MASK = 0x3

GHz = 1e9
RMS_TO_VPP = 2 * np.sqrt(2)


class Timer:
    def __init__(self):
        self.data = [(time.time(), 'start')]

    def __call__(self, text=''):
        now = time.time()
        prev = self.data[-1][0]
        self.data.append((now, text))
        return now - prev

    def summary(self):
        return ' '.join(f'{txt} {tim:.3f}' for tim, txt in self.data[1:])

    def show(self):
        first = prev = self.data[0][0]
        print('---', first)
        for tim, txt in self.data[1:]:
            print(f'{(tim - first) * 1000:9.3f} {(tim - prev) * 1000:9.3f} ms {txt}')
            prev = tim


class Adq:
    sample_rate = 2 * GHz
    max_number_of_channels = 2
    ndecimate = 50  # decimation ratio (2GHz / 40 MHz)
    number_of_records = 1
    samples_per_record = 16384
    bw_cutoff = 10E6
    trigger = EXT_TRIG_1
    adq_num = 1
    data = None
    busy = False

    def __init__(self):
        global ADQAPI
        ADQAPI = ct.cdll.LoadLibrary("libadq.so.0")

        ADQAPI.ADQAPI_GetRevision()

        # Manually set return type from some ADQAPI functions
        ADQAPI.CreateADQControlUnit.restype = ct.c_void_p
        ADQAPI.ADQ_GetRevision.restype = ct.c_void_p
        ADQAPI.ADQ_GetPtrStream.restype = ct.POINTER(ct.c_int16)
        ADQAPI.ADQControlUnit_FindDevices.argtypes = [ct.c_void_p]
        # Create ADQControlUnit
        self.adq_cu = ct.c_void_p(ADQAPI.CreateADQControlUnit())
        ADQAPI.ADQControlUnit_EnableErrorTrace(self.adq_cu, 3, '.')

        # Find ADQ devices
        ADQAPI.ADQControlUnit_FindDevices(self.adq_cu)
        n_of_adq = ADQAPI.ADQControlUnit_NofADQ(self.adq_cu)
        if n_of_adq != 1:
            print('number of ADQs must be 1, not %d' % n_of_adq)
            print('it seems the ADQ was not properly closed')
            print('please try again or reboot')
            sys.exit(0)
        atexit.register(self.deletecu)
        signal.signal(signal.SIGTERM, lambda *_: sys.exit(0))

        rev = ADQAPI.ADQ_GetRevision(self.adq_cu, self.adq_num)
        revision = ct.cast(rev, ct.POINTER(ct.c_int))
        out = [f'Connected to ADQ #1, FPGA Revision: {revision[0]}']
        if revision[1]:
            out.append('Local copy')
        else:
            if revision[2]:
                out.append('SVN Managed - Mixed Revision')
            else:
                out.append('SVN Updated')
        print(', '.join(out))
        ADQAPI.ADQ_SetClockSource(self.adq_cu, self.adq_num, ADQ_CLOCK_EXT_REF)

        ##########################
        # Test pattern
        # ADQAPI.ADQ_SetTestPatternMode(self.adq_cu, self.adq_num, 4)
        ##########################
        # Sample skip
        # ADQAPI.ADQ_SetSampleSkip(self.adq_cu, self.adq_num, 1)
        ##########################

        # set trigger mode
        if not ADQAPI.ADQ_SetTriggerMode(self.adq_cu, self.adq_num, self.trigger):
            raise RuntimeError('ADQ_SetTriggerMode failed.')
        if self.trigger == LVL_TRIG:
            if not ADQAPI.ADQ_SetLvlTrigLevel(self.adq_cu, self.adq_num, -100):
                raise RuntimeError('ADQ_SetLvlTrigLevel failed.')
            if not ADQAPI.ADQ_SetTrigLevelResetValue(self.adq_cu, self.adq_num, 1000):
                raise RuntimeError('ADQ_SetTrigLevelResetValue failed.')
            if not ADQAPI.ADQ_SetLvlTrigChannel(self.adq_cu, self.adq_num, 1):
                raise RuntimeError('ADQ_SetLvlTrigChannel failed.')
            if not ADQAPI.ADQ_SetLvlTrigEdge(self.adq_cu, self.adq_num, LVL_RISING):
                raise RuntimeError('ADQ_SetLvlTrigEdge failed.')
        elif self.trigger == EXT_TRIG_1:
            if not ADQAPI.ADQ_SetExternTrigEdge(self.adq_cu, self.adq_num, 2):
                raise RuntimeError('ADQ_SetLvlTrigEdge failed.')
        #    if not ADQAPI.ADQ_SetTriggerThresholdVoltage(self.adq_cu, self.adq_num, trigger, ct.c_double(0.2)):
        #            raise RuntimeError('SetTriggerThresholdVoltage failed.')
        # proabably the folloiwng is wrong.
        # print("CHANNEL:" + str(ct.c_int(ADQAPI.ADQ_GetLvlTrigChannel(self.adq_cu, self.adq_num))))

    def init(self, samples_per_record=None, number_of_records=None):
        """initialize dimensions and store result object"""
        if samples_per_record:
            self.samples_per_record = samples_per_record
        if number_of_records:
            self.number_of_records = number_of_records
        # Setup target buffers for data
        self.target_buffers = (ct.POINTER(ct.c_int16 * self.samples_per_record * self.number_of_records)
                               * self.max_number_of_channels)()
        for bufp in self.target_buffers:
            bufp.contents = (ct.c_int16 * self.samples_per_record * self.number_of_records)()

    def deletecu(self):
        cu = self.__dict__.pop('adq_cu', None)
        if cu is None:
            return
        print('shut down ADQ')
        # Only disarm trigger after data is collected
        ADQAPI.ADQ_DisarmTrigger(cu, self.adq_num)
        ADQAPI.ADQ_MultiRecordClose(cu, self.adq_num)
        # Delete ADQControlunit
        ADQAPI.DeleteADQControlUnit(cu)
        print('ADQ closed')

    def start(self, data):
        # Start acquisition
        ADQAPI.ADQ_MultiRecordSetup(self.adq_cu, self.adq_num,
                                    self.number_of_records,
                                    self.samples_per_record)

        ADQAPI.ADQ_DisarmTrigger(self.adq_cu, self.adq_num)
        ADQAPI.ADQ_ArmTrigger(self.adq_cu, self.adq_num)
        self.data = data

    def get_data(self):
        """get new data if available"""
        ready = False
        data = self.data
        if not data:
            self.busy = False
            return None  # no new data

        if ADQAPI.ADQ_GetAcquiredAll(self.adq_cu, self.adq_num):
            ready = True
            data.timer('ready')
        else:
            if self.trigger == SW_TRIG:
                ADQAPI.ADQ_SWTrig(self.adq_cu, self.adq_num)
        if not ready:
            self.busy = True
            return None
        self.data = None
        t = time.time()
        # Get data from ADQ
        if not ADQAPI.ADQ_GetData(
                    self.adq_cu, self.adq_num, self.target_buffers,
                    self.samples_per_record * self.number_of_records, 2,
                    0, self.number_of_records, ADQ_CHANNELS_MASK,
                    0, self.samples_per_record, ADQ_TRANSFER_MODE_NORMAL):
            raise RuntimeError('no success from ADQ_GetDATA')
        data.retrieve(self)
        return data


class PEdata:
    def __init__(self, adq):
        self.sample_rate = adq.sample_rate
        self.samp_freq = self.sample_rate / GHz
        self.number_of_records = adq.number_of_records
        self.timer = Timer()

    def retrieve(self, adq):
        data = []
        rawsignal = []
        for ch in range(2):
            onedim = np.frombuffer(adq.target_buffers[ch].contents, dtype=np.int16)
            rawsignal.append(onedim[:adq.samples_per_record])
            # convert 16 bit int to a value in the range -1 .. 1
            data.append(onedim.reshape(adq.number_of_records, adq.samples_per_record) / float(2 ** 15))
        # Now this is an array with all records, but the time is artificial
        self.data = data
        self.rawsignal = rawsignal
        self.timer('retrieved')

    def sinW(self, sig, freq, ti, tf):
        # sig: signal array
        # freq
        # ti, tf: initial and end time
        si = int(ti * self.samp_freq)
        nperiods = freq * (tf - ti)
        n = int(round(max(2, int(nperiods)) / nperiods * (tf-ti) * self.samp_freq))
        self.nperiods = n
        t = np.arange(si, len(sig)) / self.samp_freq
        t = t[:n]
        self.pulselen = n / self.samp_freq
        sig = sig[si:si+n]
        a = 2*np.sum(sig*np.cos(2*np.pi*freq*t))/len(sig)
        b = 2*np.sum(sig*np.sin(2*np.pi*freq*t))/len(sig)
        return a, b

    def mix(self, sigin, sigout, freq, ti, tf):
        # sigin, sigout: signal array, incomping, output
        # freq
        # ti, tf: initial and end time of sigin
        a, b = self.sinW(sigin, freq, ti, tf)
        amp = np.sqrt(a**2 + b**2)
        a, b = a/amp, b/amp
        # si = int(ti * self.samp_freq)
        t = np.arange(len(sigout)) / self.samp_freq
        wave1 = sigout * (a * np.cos(2*np.pi*freq*t) + b * np.sin(2*np.pi*freq*t))
        wave2 = sigout * (a * np.sin(2*np.pi*freq*t) - b * np.cos(2*np.pi*freq*t))
        return wave1, wave2

    def averageiq(self, data, freq, ti, tf):
        """Average over records"""
        iorq = np.array([self.mix(data[0][i], data[1][i], freq, ti, tf) for i in range(self.number_of_records)])
        return iorq.sum(axis=0) / self.number_of_records

    def filtro(self, iorq, cutoff):
        # butter lowpass
        nyq = 0.5 * self.sample_rate
        normal_cutoff = cutoff / nyq
        order = 5
        b, a = butter(order, normal_cutoff, btype='low', analog=False)
        iqf = [filtfilt(b, a, iorq[i]) for i in np.arange(len(iorq))]
        return iqf

    def box(self, iorq, ti, tf):
        si = int(self.samp_freq * ti)
        sf = int(self.samp_freq * tf)
        bxa = [sum(iorq[i][si:sf])/(sf-si) for i in np.arange(len(iorq))]
        return bxa

    def gates_and_curves(self, freq, pulse, roi, bw_cutoff):
        """return iq values of rois and prepare plottable curves for iq"""
        self.timer('gates')
        try:
            self.ndecimate = int(round(self.sample_rate / freq))
        except TypeError as e:
            raise TypeError(f'{self.sample_rate}/{freq}  {e}')
        iq = self.averageiq(self.data, freq / GHz, *pulse)
        self.timer('aviq')
        iqf = self.filtro(iq, bw_cutoff)
        self.timer('filtro')
        m = max(1, len(iqf[0]) // self.ndecimate)
        ll = m * self.ndecimate
        iqf = [iqfx[0:ll] for iqfx in iqf]
        self.timer('iqf')
        iqd = np.average(np.resize(iqf, (2, m, self.ndecimate)), axis=2)
        self.timer('avg')
        t_axis = np.arange(m) * self.ndecimate / self.samp_freq
        pulsig = np.abs(self.data[0][0])
        self.timer('pulsig')
        pulsig = np.average(np.resize(pulsig, (m, self.ndecimate)), axis=1)
        result = ([self.box(iqf, *r) for r in roi],  # gates
                  (t_axis, iqd[0], iqd[1], pulsig))  # curves
        self.timer('result')
        # self.timer.show()
        # ns = len(self.rawsignal[0]) * self.number_of_records
        # print(f'{ns} {ns / 2e6} ms')
        return result


class Namespace:
    """holds channel or other data"""
    def __init__(self, **kwds):
        self.__dict__.update(**kwds)


class RUSdata:
    def __init__(self, adq, freq, periods, delay_samples):
        self.sample_rate = adq.sample_rate
        self.freq = freq
        self.periods = periods
        self.delay_samples = delay_samples
        self.samples_per_record = adq.samples_per_record
        self.inp = Namespace(idx=0, name='input')
        self.out = Namespace(idx=1, name='output')
        self.channels = (self.inp, self.out)
        self.timer = Timer()

    def retrieve(self, adq):
        self.timer('start retrieve')
        npts = self.samples_per_record - self.delay_samples
        nbin = max(1, npts // (self.periods * 60))  # for performance reasons, do the binning first
        nreduced = npts // nbin
        ft = 2 * np.pi * self.freq * nbin / self.sample_rate * np.arange(nreduced)
        self.timer('create time axis')
        # complex_sinusoid = np.exp(1j * ft)  # do not use this, below is 33 % faster
        complex_sinusoid = 1j * np.sin(ft) + np.cos(ft)
        self.timer('sinusoid')

        rawsignal = []  # for raw plot
        for chan in self.channels:  # looping over input and output
            # although the ADC is only 14 bit it is represented as unsigend 16 bit numbers,
            # and due to some calculations (calibration) the last 2 bits are not zero
            beg = self.delay_samples
            isignal = np.frombuffer(adq.target_buffers[chan.idx].contents, dtype=np.int16)[beg:beg+nreduced * nbin]
            self.timer('isignal')
            reduced = isignal.reshape((-1, nbin)).mean(axis=1)  # this converts also int16 to float
            self.timer('reduce')
            rawsignal.append(reduced)
            chan.signal = signal = reduced * 2 ** -16  # in V -> peak to peak 1 V ~ +- 0.5 V
            self.timer('divide')
            # calculate RMS * sqrt(2) -> peak sinus amplitude.
            # may be higher than the input range by a factor 1.4 when heavily clipped
            chan.amplitude = np.sqrt((signal ** 2).mean()) * RMS_TO_VPP
            self.timer('amp')
            chan.mixed = signal * complex_sinusoid
            self.timer('mix')
            chan.mean = chan.mixed.mean()
            self.timer('mean')
        self.rawsignal = rawsignal
        if self.inp.mean:
            self.iq = self.out.mean / self.inp.mean
        else:
            self.iq = 0

    def get_quality(self):
        """get signal quality info

        quality info (small values indicate good quality):
        - input_stddev:
          the imaginary part indicates deviations in phase
          the real part indicates deviations in amplitude
        - output_slope:
          the imaginary part indicates a turning phase (rad/sec)
          the real part indicates changes in amplitude (0.01 ~= 1%/sec)
        """
        self.timer('get_quality')
        npts = len(self.channels[0].signal)
        nper = npts // self.periods
        for chan in self.channels:
            mean = chan.mixed.mean()
            chan.reduced = chan.mixed[:self.periods * nper].reshape((-1, nper)).mean(axis=1) / mean

        timeaxis = np.arange(len(self.out.reduced)) * self.sample_rate / self.freq
        result = Namespace(
            input_stddev=self.inp.reduced.std(),
            output_slope=np.polyfit(timeaxis, self.out.reduced, 1)[0])
        self.timer('got_quality')
        self.timer.show()
        ns = len(self.rawsignal[0])
        print(f'{ns} {ns / 2e6} ms')
        return result