import os
import re
import itertools

from collections import OrderedDict

import numpy as np
from numpy import dtype

from ..data.nmr import FID, Spectrum
from ..utils.utils import staggered_range

__all__ = ['TNTReader']

"""
Read Tecmag files. Based on <https://www.github.com/chatcannon/pytnt>_
"""

TNTMAGIC_re = re.compile(r'^TNT1\.\d\d\d$')
TNTMAGIC = np.dtype('a8')
TLV = np.dtype([('tag', 'a4'), ('bool', '<u4'), ('length', '<u4')])

TMAG = np.dtype([
    ('npts', '<i4', 4),
    ('actual_npts', '<i4', 4),
    ('acq_points', '<i4'),
    ('npts_start', '<i4', 4),
    ('scans', '<i4'),
    ('actual_scans', '<i4'),
    ('dummy_scans', '<i4'),
    ('repeat_times', '<i4'),
    ('sadimension', '<i4'),
    ('samode', '<i4'),
    ('magnet_field', '<f8'),
    ('ob_freq', '<f8', 4),
    ('base_freq', '<f8', 4),
    ('offset_freq', '<f8', 4),
    ('ref_freq', '<f8'),
    ('NMR_frequency', '<f8'),
    ('obs_channel', '<i2'),
    ('space2', 'a42'),

    ('sw', '<f8', 4),
    ('dwell', '<f8', 4),
    ('filter', '<f8'),
    ('experiment_time', '<f8'),
    ('acq_time', '<f8'),
    ('last_delay', '<f8'),
    ('spectrum_direction', '<i2'),
    ('hardware_sideband', '<i2'),
    ('Taps', '<i2'),
    ('Type', '<i2'),
    ('bDigRec', '<u4'),
    ('nDigitalCenter', '<i4'),
    ('space3', 'a16'),

    ('transmitter_gain', '<i2'),
    ('receiver_gain', '<i2'),
    ('NumberOfReceivers', '<i2'),
    ('RG2', '<i2'),
    ('receiver_phase', '<f8'),
    ('space4', 'a4'),

    ('set_spin_rate', '<u2'),
    ('actual_spin_rate', '<u2'),

    ('lock_field', '<i2'),
    ('lock_power', '<i2'),
    ('lock_gain', '<i2'),
    ('lock_phase', '<i2'),
    ('lock_freq_mhz', '<f8'),
    ('lock_ppm', '<f8'),
    ('H2O_freq_ref', '<f8'),
    ('space5', 'a16'),

    ('set_temperature', '<f8'),
    ('actual_temperature', '<f8'),

    ('shim_units', '<f8'),
    ('shims', '<i2', 36),
    ('shim_FWHM', '<f8'),

    ('HH_dcpl_attn', '<i2'),
    ('DF_DN', '<i2'),
    ('F1_tran_mode', '<i2', 7),
    ('dec_BW', '<i2'),
    ('grd_orientation', 'a4'),
    ('LatchLP', '<i4'),
    ('grd_Theta', '<f8'),
    ('grd_Phi', '<f8'),
    ('space6', 'a264'),

    ('start_time', '<u4'),
    ('finish_time', '<u4'),
    ('elapsed_time', '<i4'),

    ('date', 'a32'),
    ('nuclei', 'a16', 4),
    ('sequence', 'a32'),
    ('lock_solvent', 'a16'),
    ('lock_nucleus', 'a16')
])

GridAndAxis = dtype([
    ('majorTickInc', '<f8', 12),
    ('minorIntNum', '<i2', 12),
    ('labelPrecision', '<i2', 12),
    ('gaussPerCentimeter', '<f8'),
    ('gridLines', '<i2'),
    ('axisUnits', '<i2'),
    ('showGrid', '<u4'),
    ('showGridLabels', '<u4'),
    ('adjustOnZoom', '<u4'),
    ('showDistanceUnits', '<u4'),
    ('axisName', 'a32'),
    ('space', 'a52'),
])

TMG2 = dtype([
    ('real_flag', '<u4'),
    ('imag_flag', '<u4'),
    ('magn_flag', '<u4'),
    ('axis_visible', '<u4'),
    ('auto_scale', '<u4'),
    ('line_display', '<u4'),
    ('show_shim_units', '<u4'),

    ('integral_display', '<u4'),
    ('fit_display', '<u4'),
    ('show_pivot', '<u4'),
    ('label_peaks', '<u4'),
    ('keep_manual_peaks', '<u4'),
    ('label_peaks_in_units', '<u4'),
    ('integral_dc_average', '<u4'),
    ('integral_show_multiplier', '<u4'),
    ('Boolean_space', '<u4', 9),

    ('all_ffts_done', '<u4', 4),
    ('all_phase_done', '<u4', 4),

    ('amp', '<f8'),
    ('ampbits', '<f8'),
    ('ampCtl', '<f8'),
    ('offset', '<i4'),

    ('axis_set', GridAndAxis),

    ('display_units', '<i2', 4),
    ('ref_point', '<i4', 4),
    ('ref_value', '<f8', 4),
    ('z_start', '<i4'),
    ('z_end', '<i4'),
    ('z_select_start', '<i4'),
    ('z_select_end', '<i4'),
    ('last_zoom_start', '<i4'),
    ('last_zoom_end', '<i4'),
    ('index_2D', '<i4'),
    ('index_3D', '<i4'),
    ('index_4D', '<i4'),

    ('apodization_done', '<i4', 4),
    ('linebrd', '<f8', 4),
    ('gaussbrd', '<f8', 4),
    ('dmbrd', '<f8', 4),
    ('sine_bell_shift', '<f8', 4),
    ('sine_bell_width', '<f8', 4),
    ('sine_bell_skew', '<f8', 4),
    ('Trapz_point_1', '<i4', 4),
    ('Trapz_point_2', '<i4', 4),
    ('Trapz_point_3', '<i4', 4),
    ('Trapz_point_4', '<i4', 4),
    ('trafbrd', '<f8', 4),
    ('echo_center', '<i4', 4),

    ('data_shift_points', '<i4'),
    ('fft_flag', '<i2', 4),
    ('unused', '<f8', 8),
    ('pivot_point', '<i4', 4),
    ('cumm_0_phase', '<f8', 4),
    ('cumm_1_phase', '<f8', 4),
    ('manual_0_phase', '<f8'),
    ('manual_1_phase', '<f8'),
    ('phase_0_value', '<f8'),
    ('phase_1_value', '<f8'),
    ('session_phase_0', '<f8'),
    ('session_phase_1', '<f8'),

    ('max_index', '<i4'),
    ('min_index', '<i4'),
    ('peak_threshold', '<f4'),
    ('peak_noise', '<f4'),
    ('integral_dc_points', '<i2'),
    ('integral_label_type', '<i2'),
    ('integral_scale_factor', '<f4'),
    ('auto_integrate_shoulder', '<i4'),
    ('auto_integrate_noise', '<f8'),
    ('auto_integrate_threshold', '<f8'),
    ('s_n_peak', '<i4'),
    ('s_n_noise_start', '<i4'),
    ('s_n_noise_end', '<i4'),
    ('s_n_calculated', '<f4'),

    ('Spline_point', '<i4', 14),
    ('Spline_point_avr', '<i2'),
    ('Poly_point', '<i4', 8),
    ('Poly_point_avr', '<i2'),
    ('Poly_order', '<i2'),

    ('space', 'a610'),

    ('line_simulation_name', 'a32'),
    ('integral_template_name', 'a32'),
    ('baseline_template_name', 'a32'),
    ('layout_name', 'a32'),
    ('relax_information_name', 'a32'),
    ('username', 'a32'),
    ('user_string_1', 'a16'),
    ('user_string_2', 'a16'),
    ('user_string_3', 'a16'),
    ('user_string_4', 'a16')
])

delay_re = re.compile(rb'de\d+:\d')
rf_re = re.compile(rb'rf\d+:\d')


class TNTReader:
    def __init__(self, fname=None):

        self.fname = fname
        self.shape = [1, 1, 1, 1]
        self.delays = {}
        self.data = None
        self.dw = [1, 1, 1, 1]
        self.fft_flag = [0]
        self.freq_shift = 0
        self.onedimensional = False

        if self.fname is not None:
            self.readtntfile()

            if all([s == 1 for s in self.shape[1:]]):
                self.onedimensional = True
            else:
                self.onedimensional = False

    def __call__(self, fname):
        self.fname = fname
        self.shape = [1, 1, 1, 1]
        self.freq_shift = 0
        self.data = None
        self.dw = [1, 1, 1, 1]
        self.delays = {}

        self.fft_flag = [0]
        self.readtntfile()

        if all([s == 1 for s in self.shape[1:]]):
            self.onedimensional = True
        else:
            self.onedimensional = False

        return self

    def readtntfile(self):
        tnt_sections = {}
        with open(self.fname, 'rb') as f:
            tntmagic = np.fromstring(f.read(TNTMAGIC.itemsize), TNTMAGIC, count=1)[0]
            if not TNTMAGIC_re.match(str(tntmagic, encoding='utf-8')):
                raise ValueError("Invalid magic number (is '{}' really a TNMR file?): {}".format(self.fname, tntmagic))

            tnthdrbytes = f.read(TLV.itemsize)
            # loop to find sections and store them
            while TLV.itemsize == len(tnthdrbytes):
                tlv = np.fromstring(tnthdrbytes, TLV)[0]
                data_length = tlv['length']
                hdrdict = {'offset': f.tell(),
                           'length': data_length,
                           'bool': bool(tlv['bool'])}
                if data_length <= 4096:
                    hdrdict['data'] = f.read(data_length)
                else:
                    f.seek(data_length, os.SEEK_CUR)
                tnt_sections[tlv['tag'].decode()] = hdrdict
                tnthdrbytes = f.read(TLV.itemsize)

            # irgendwas aus dem Pulssequenzbereich sinnvoll herauszubekommen, ist zum Kotzen
            f.seek(tnt_sections["PSEQ"]["offset"])
            search_region = f.read()
            for tables in [delay_re, rf_re]:
                for match in tables.finditer(search_region):
                    offset = match.start() - 4
                    delay_name = self.read_pascal_string(search_region[offset:])

                    offset = match.start() + len(delay_name)
                    delay = self.read_pascal_string(search_region[offset:])
                    if len(delay) > 1:
                        delay = delay.split()
                        delay = np.array([self.si_conversion(d) for d in delay])
                        self.delays[delay_name] = delay

            tmag = np.fromstring(tnt_sections['TMAG']['data'], TMAG, count=1)[0]
            data = np.memmap(self.fname, np.dtype('<c8'), mode='c',
                             offset=tnt_sections['DATA']['offset'],
                             shape=tmag['actual_npts'].prod())
            self.data = np.reshape(data, tmag['actual_npts'], order='F')

            tmag2 = np.fromstring(tnt_sections['TMG2']['data'], TMG2, count=1)[0]
            self.dw = tmag['dwell']
            self.shape = tmag['actual_npts']
            self.fft_flag = tmag2['fft_flag']
            self.freq_shift = tmag['ref_freq']

    def add_delay(self, start, stop, dims, key, staggered=False, stag_steps=1, log=False):
        if log:
            times = np.logspace(np.log10(start), np.log10(stop),
                                num=self.shape[dims])
        else:
            times = np.linspace(start, stop, num=self.shape[dims])

        if staggered:
            times = staggered_range(times, stepsize=stag_steps)

        self.delays[key] = times

    def export(self, delay_keys=[]):
        ret_val = OrderedDict()
        fname = os.path.splitext(os.path.split(self.fname)[-1])[0]
        if self.shape[1] == 1:
            if all(ffts == 0 for ffts in self.fft_flag):
                x = self.dw[0] * np.arange(self.shape[0])
                ret_val[self.fname] = FID(x, self.data[:, 0, 0, 0], accus=1, name=fname,
                                          filename=self.fname)
            else:
                x = -np.fft.fftshift(np.fft.fftfreq(self.shape[0], self.dw[0]))
                x -= self.freq_shift
                ret_val[self.fname] = Spectrum(x, self.data[:, 0, 0, 0], accus=1, name=fname, filename=self.fname)
        else:
            d = []
            for k in delay_keys:
                if k is not None:
                    d.append(list(map(str, self.delays[k])))
                else:
                    d.append([''])

            ijk = list(itertools.product(*d))
            all_fft = all([ffts == 0 for ffts in self.fft_flag])
            for v, (i, j, k) in enumerate(itertools.product(*map(range, self.shape[1:]))):
                d2, d3, d4 = ijk[v]
                x = self.dw[0] * np.arange(self.shape[0])
                if all_fft:
                    ret_val['/'.join(ijk[v])] = FID(x, self.data[:, i, j, k], accus=1, value=d2,
                                                    name='{}'.format(d2), group=fname, filename=self.fname)
                else:
                    x = -np.fft.fftshift(np.fft.fftfreq(self.shape[0], self.dw[0]))
                    x -= self.freq_shift
                    ret_val['/'.join(ijk[v])] = Spectrum(x, self.data[:, i, j, k], accus=1, value=d2,
                                                         name='{}'.format(d2), group=fname, filename=self.fname)

        return ret_val

    @staticmethod
    def si_conversion(sistring):
        si_dict = {'n': 1e-9, 'u': 1e-6, 'm': 1e-3, 's': 1e0}
        try:
            ret_val = float(sistring)
        except ValueError:
            ret_val = float(sistring[:-1]) * si_dict[sistring[-1]]
        return ret_val

    @staticmethod
    def read_pascal_string(data, number_type='<i4'):
        number_type = np.dtype(number_type)
        length = int(np.fromstring(data, dtype=number_type, count=1))
        number_size = int(number_type.itemsize)
        text = data[number_size:number_size + length]
        if not isinstance(text, str):
            text = str(text, encoding='utf-8')
        return text
