# *****************************************************************************
# 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:
#   Markus Zolliker <markus.zolliker@psi.ch>
# *****************************************************************************
"""frappy support for ultrasound"""

import math
import os
import time

import numpy as np

from frappy_psi.adq_mr import Adq, PEdata, RUSdata
from frappy.core import Attached, BoolType, Done, FloatRange, HasIO, \
    IntRange, Module, Parameter, Readable, Writable, Drivable, StringIO, StringType, \
    IDLE, BUSY, DISABLED, ERROR, TupleOf, ArrayOf, Command
from frappy.properties import Property
#from frappy.modules import Collector

Collector = Readable


def fname_from_time(t, extension):
    tm = time.localtime(t)
    dirname = os.path.join('..', 'data', time.strftime("%Y-%m-%d_%H", tm))
    filename = time.strftime("%Y-%m-%d_%H-%M-%S_", tm)
    filename = filename + ("%.1f" % t)[-1]
    if not os.path.isdir(dirname):
        os.makedirs(dirname)
    return os.path.join(dirname, filename)


class Roi(Readable):
    main = Attached()

    value = Parameter('amplitude', FloatRange(), default=0)
    phase = Parameter('phase', FloatRange(unit='deg'), default=0)
    i = Parameter('in phase', FloatRange(), default=0)
    q = Parameter('out of phase', FloatRange(), default=0)
    time = Parameter('start time', FloatRange(unit='nsec'), readonly=False)
    size = Parameter('interval (symmetric around time)', FloatRange(unit='nsec'), readonly=False)
    enable = Parameter('calculate this roi', BoolType(), readonly=False, default=True)
    pollinterval = Parameter(export=False)

    interval = (0, 0)

    def initModule(self):
        super().initModule()
        self.main.register_roi(self)
        self.calc_interval()

    def calc_interval(self):
        self.interval = (self.time - 0.5 * self.size, self.time + 0.5 * self.size)

    def read_status(self):
        return (IDLE, '') if self.enable else (DISABLED, 'disabled')

    def write_time(self, value):
        self.time = value
        self.calc_interval()
        return Done

    def write_size(self, value):
        self.size = value
        self.calc_interval()
        return Done


class Pars(Module):
    description = 'relevant parameters from SEA'

    timestamp = Parameter('unix timestamp', StringType(), default='0', readonly=False)
    temperature = Parameter('T', FloatRange(unit='K'), default=0, readonly=False)
    mf = Parameter('field', FloatRange(unit='T'), default=0, readonly=False)
    sr = Parameter('rotation angle', FloatRange(unit='deg'), default=0, readonly=False)


class FreqStringIO(StringIO):
    end_of_line = '\r'


class Frequency(HasIO, Writable):
    value = Parameter('frequency', unit='Hz')
    amp = Parameter('amplitude', FloatRange(unit='dBm'), readonly=False)

    last_change = 0
    ioClass = FreqStringIO
    dif = None

    def register_dif(self, dif):
        self.dif = dif

    def write_target(self, value):
        self.communicate('FREQ %.15g;FREQ?' % value)
        self.last_change = time.time()
        if self.dif:
            self.dif.read_value()

    def write_amp(self, amp):
        reply = self.communicate('AMPR %g;AMPR?' % amp)
        return float(reply)

    def read_amp(self):
        reply = self.communicate('AMPR?')
        return float(reply)


class FrequencyDif(Readable):
    freq = Attached(Frequency)
    base = Parameter('base frequency', FloatRange(unit='Hz'), default=0)
    value = Parameter('difference to base frequency', FloatRange(unit='Hz'), default=0)

    def initModule(self):
        super().initModule()
        self.freq.register_dif(self)

    def read_value(self):
        return self.freq - self.base


class Base(Collector):
    freq = Attached()
    adq = Attached(Adq)
    sr = Parameter('samples per record', datatype=IntRange(1, 1E9), default=16384)
    pollinterval = Parameter(datatype=FloatRange(0, 120))  # allow pollinterval = 0
    _data = None
    _data_args = None

    def read_status(self):
        adqstate = self.adq.get_status()
        if adqstate == Adq.BUSY:
            return BUSY, 'acquiring'
        if adqstate == Adq.UNDEFINED:
            return ERROR, 'no data yet'
        if adqstate == Adq.READY:
            return IDLE, 'new data available'
        return IDLE, ''

    def get_data(self):
        data = self.adq.get_data(*self._data_args)
        if id(data) != id(self._data):
            self._data = data
            return data
        return None


class PulseEcho(Base):
    value = Parameter("t, i, q, pulse curves",
                      TupleOf(*[ArrayOf(FloatRange(), 0, 16283) for _ in range(4)]), default=[[]] * 4)
    nr = Parameter('number of records', datatype=IntRange(1, 9999), default=500)
    bw = Parameter('bandwidth lowpassfilter', datatype=FloatRange(unit='Hz'), default=10E6)
    control = Parameter('control loop on?', BoolType(), readonly=False, default=True)
    time = Parameter('pulse start time', FloatRange(unit='nsec'),
                     readonly=False)
    size = Parameter('pulse length (starting from time)', FloatRange(unit='nsec'),
                     readonly=False)
    pulselen = Parameter('adjusted pulse length (integer number of periods)', FloatRange(unit='nsec'), default=1)

    starttime = None

    def initModule(self):
        super().initModule()
        self.adq = Adq()
        self.adq.init(self.sr, self.nr)
        self.roilist = []

    def write_nr(self, value):
        self.adq.init(self.sr, value)

    def write_sr(self, value):
        self.adq.init(value, self.nr)

    def write_bw(self, value):
        self.adq.bw_cutoff = value

    def register_roi(self, roi):
        self.roilist.append(roi)

    def go(self):
        self.starttime = time.time()
        self.adq.start()

    def read_value(self):
        if self.get_rawdata():  # new data available
            roilist = [r for r in self.roilist if r.enable]
            freq = self.freq.value
            gates = self.adq.gates_and_curves(self._data, freq,
                                              (self.time, self.time + self.size),
                                              [r.interval for r in roilist])
            for i, roi in enumerate(roilist):
                roi.i = a = gates[i][0]
                roi.q = b = gates[i][1]
                roi.value = math.sqrt(a ** 2 + b ** 2)
                roi.phase = math.atan2(a, b) * 180 / math.pi
        return self.adq.curves

        # TODO: CONTROL
        # inphase = self.roilist[0].i
        # if self.control:
        #     newfreq = freq + inphase * self.slope - self.basefreq
        #     # step = sorted((-self.maxstep, inphase * self.slope, self.maxstep))[1]
        # if self.old:
        #     fdif = freq - self.old[0]
        #     idif = inphase - self.old[1]
        #     if abs(fdif) >= self.minstep:
        #         self.slope = - fdif / idif
        # else:
        #     fdif = 0
        #     idif = 0
        #     newfreq = freq + self.minstep
        # self.old = (freq, inphase)
        # if self.skipctrl > 0:  # do no control for some time after changing frequency
        #     self.skipctrl -= 1
        # elif self.control:
        #     self.freq = sorted((self.freq - self.maxstep, newfreq, self.freq + self.maxstep))[1]


class RUS(Base):
    value = Parameter('averaged (I, Q) tuple', TupleOf(FloatRange(), FloatRange()))
    periods = Parameter('number of periods', IntRange(1, 9999), default=12)
    scale = Parameter('scale,taking into account input attenuation', FloatRange(), default=0.1)
    input_phase_stddev = Parameter('input signal quality', FloatRange(unit='rad'))
    output_phase_slope = Parameter('output signal phase slope', FloatRange(unit='rad/sec'))
    output_amp_slope = Parameter('output signal amplitude change', FloatRange(unit='1/sec'))
    phase = Parameter('phase', FloatRange(unit='deg'))
    amp = Parameter('amplitude', FloatRange())

    starttime = None
    _data_args = None

    def initModule(self):
        super().initModule()
        self.adq = Adq()
        # self.write_periods(self.periods)

    def read_value(self):
        if self._data_args is None:
            return self.value  # or may we raise as no value is defined yet?
        data = self.get_data(RUSdata, *self._data_args)
        if data:
            # data available
            data.calc_quality()
            self.input_phase_stddev = data.input_stddev.imag
            self.output_phase_slope = data.output_slope.imag
            self.output_amp_slope = data.output_slope.real

            iq = data.iq * self.scale
            self.phase = np.arctan2(iq.imag, iq.real) * 180 / np.pi
            self.amp = np.abs(iq.imag, iq.real)
            return iq.real, iq.imag
        return self.value

    def go(self):
        self.starttime = time.time()
        freq = self.freq.value
        self._data_args = (RUSdata, freq, self.periods)
        self.sr = round(self.periods * self.adq.sample_rate / freq)
        self.adq.init(self.sr, 1)
        self.adq.start()
        self.read_status()


class ControlLoop:
    maxstep = Parameter('max frequency step', FloatRange(unit='Hz'), readonly=False,
                      default=10000)
    minstep = Parameter('min frequency step for slope calculation', FloatRange(unit='Hz'),
                        readonly=False, default=4000)
    slope = Parameter('inphase/frequency slope', FloatRange(), readonly=False,
                      default=1e6)


# class Frequency(HasIO, Readable):
#     pars = Attached()
#     curves = Attached(mandatory=False)
#     maxy = Property('plot y scale', datatype=FloatRange(), default=0.5)
#
#     value = Parameter('frequency@I,q', datatype=FloatRange(unit='Hz'), default=0)
#     basefreq = Parameter('base frequency', FloatRange(unit='Hz'), readonly=False)
#     nr = Parameter('number of records', datatype=IntRange(1,10000), default=500)
#     sr = Parameter('samples per record', datatype=IntRange(1,1E9), default=16384)
#     freq = Parameter('target frequency', FloatRange(unit='Hz'), readonly=False)
#     bw = Parameter('bandwidth lowpassfilter', datatype=FloatRange(unit='Hz'),default=10E6)
#     amp = Parameter('amplitude', FloatRange(unit='dBm'), readonly=False)
#     control = Parameter('control loop on?', BoolType(), readonly=False, default=True)
#     rusmode = Parameter('RUS mode on?', BoolType(), readonly=False, default=False)
#     time = Parameter('pulse start time', FloatRange(unit='nsec'),
#                       readonly=False)
#     size = Parameter('pulse length (starting from time)', FloatRange(unit='nsec'),
#                     readonly=False)
#     pulselen = Parameter('adjusted pulse length (integer number of periods)', FloatRange(unit='nsec'), default=1)
#     maxstep = Parameter('max frequency step', FloatRange(unit='Hz'), readonly=False,
#                       default=10000)
#     minstep = Parameter('min frequency step for slope calculation', FloatRange(unit='Hz'),
#                         readonly=False, default=4000)
#     slope = Parameter('inphase/frequency slope', FloatRange(), readonly=False,
#                       default=1e6)
#     plot = Parameter('create plot images', BoolType(), readonly=False, default=True)
#     save = Parameter('save data', BoolType(), readonly=False, default=True)
#     pollinterval = Parameter(datatype=FloatRange(0,120))