from __future__ import annotations

import pathlib
import re
from io import BytesIO
from itertools import islice

import numpy as np

from ..data.points import Points
from ..data.nmr import FID, Spectrum
from ..data.bds import BDS
from ..data.dsc import DSC
from ..utils.utils import staggered_range


class AsciiReader:
    # delimiters = ['\t', ' ', ',']

    def __init__(self, fname):
        self.fname = None
        self.header = []
        self.lines = []
        self.delays = None
        self.width = []
        self.line_comment = []
        self._last_read_pos = 0

        if fname:
            self.fname = pathlib.Path(fname)
            self.read_header()
            self.make_preview(10)
            self.look_for_delay()

    def read_header(self, comments='#'):
        with self.fname.open('r') as f:
            for line in f.readlines():
                line = line.rstrip('\n\t\r, ')
                if line.startswith(comments):
                    self.header.append(line.replace(comments, '').strip())
                else:
                    break

    def make_preview(self, num_lines: int):
        if num_lines <= len(self.lines):
            return self.lines[:num_lines], max(self.width[:num_lines]), self.line_comment

        num_lines += len(self.header)
        with self.fname.open('r') as f:
            for i, line in enumerate(islice(f, len(self.header)+len(self.lines), num_lines)):
                line = line.strip('\n\t\r, ')
                line = re.sub(r'[\t, ;]+(?!\w*})', ';', line)
                line = line.split(';')

                try:
                    comment_start = line.index('#')
                    self.line_comment.append(' '.join(line[comment_start:]))
                    line = line[:comment_start]
                except ValueError:
                    self.line_comment.append('')

                is_empty = len(line) == 0

                if not is_empty:
                    self.lines.append(line)
                    self.width.append(len(line))
                else:
                    self.lines.append('')

        return self.lines, max(self.width), self.line_comment

    def look_for_delay(self, fname=None):
        if fname is None:
            fname = self.fname.with_name(f'{self.fname.name}_delay.{self.fname.suffix}')

        try:
            self.delays = np.loadtxt(fname)
        except IOError:
            pass

    def calc_delays(self, start, end, nums, log=True, stagg=False, stag_size=1):
        if log:
            self.delays = np.geomspace(start, end, num=nums)
        else:
            self.delays = np.linspace(start, end, num=nums)

        if stagg:
            self.delays = staggered_range(self.delays, stepsize=stag_size)

    def export(
            self: 'AsciiReader',
            x: int | str = None,
            y: list = None,
            yerr: list = None,
            mode: str = 'points',
            col_names=None,
            num_value: float = None,
    ) -> list:

        mode = mode.lower()
        if mode not in ['points', 'fid', 'spectrum', 'dsc', 'bds']:
            raise ValueError(f'Unknown mode {mode!r} as selected class')

        if yerr is None:
            yerr = []
        elif y is None:
            raise ValueError('y is None and yerr is not None')

        # if (x is None) ^ (y is None):
        #     raise ValueError(f'x is {type(x)} and y is {type(y)}, should be both None or both defined')

        if x is None:
            x = [0]
        elif isinstance(x, int):
            x = [x]
        elif isinstance(x, str) and x == 'index':
            x = []
        else:
            raise ValueError(f'type of x is {type(x)} not `int` or `str`')

        if y is None:
            y = list(range(int(len(x) != 0), max(self.width)))

        if col_names is not None:
            col_names = [col_names[y_i] for y_i in y]

        cols = x + y + yerr
        with self.fname.open('rb') as fh:
            tmp_ = re.sub(b'[;,]', b' ', fh.read())
        raw_data = np.genfromtxt(BytesIO(tmp_), usecols=cols, missing_values='--')
        del tmp_

        if raw_data.ndim == 1:
            # only one row or column
            if len(self.lines) == 1:
                # one row
                raw_data = raw_data.reshape(1, -1)
            else:
                raw_data = raw_data.reshape(-1, 1)

        if self.delays is not None:
            if raw_data.shape[0] % len(self.delays) != 0:
                raise ValueError(f'number of delays {len(self.delays)} does not fit into length of date {raw_data.shape[1]}')
            raw_data = raw_data.reshape((len(self.delays), -1, raw_data.shape[1]))
        else:
            raw_data = raw_data.reshape((1, *raw_data.shape))

        if len(x) == 0 or raw_data.shape[2] == 1:
            raw_data = raw_data.reshape(raw_data.shape[0], raw_data.shape[2], raw_data.shape[1])
            # _temp = np.zeros((raw_data.shape[0], raw_data.shape[2], raw_data.shape[1]))
            # _temp[:, :, 0] = np.arange(raw_data.shape[1])
            # _temp[:, :, 1:] = raw_data
            # raw_data = _temp

            if y:
                y = [i+1 for i in y]
            else:
                y = [1]

        filename = self.fname.stem

        if self.delays:
            delay_names = self.delays
        else:
            delay_names = [num_value]

        imported_sets = []

        for i, value in enumerate(delay_names):
            kwargs = {
                'value': value,
                'filename': self.fname,
                'name': filename,
                'group': filename,
                'y_err': None
            }

            num_y = len(y)
            single_len = 1
            stepsize = 1

            if mode in ['spectrum', 'fid', 'bds']:
                # complex data types
                single_len = 2
                stepsize = 2

            cls = {
                'points': Points,
                'fid': FID,
                'spectrum': Spectrum,
                'bds': BDS,
                'dsc': DSC,
            }[mode]

            for j in range(1, num_y+1, stepsize):
                if col_names is not None:
                    # predefined name
                    kwargs['name'] = col_names[j-1]
                elif num_y > single_len:
                    # more than one axis, append column number
                    kwargs['name'] = f'{filename}_{y[j-1]+1}'

                if j+num_y < raw_data.shape[2]:
                    kwargs['y_err'] = raw_data[i, :, j+num_y]

                imported_sets.append(cls(x=raw_data[i, :, 0], y=raw_data[i, :, j:j+single_len].T, **kwargs))

        return imported_sets
