"""
Data management.
"""
import sys
import os
import copy
import numpy as np
from numpy.linalg import norm
import zipfile
import numpy.lib.format as npy_format
import pandas as pd
from glob import glob
from os import path
from io import IOBase
import fnmatch
import re
from collections import OrderedDict
from datetime import timedelta
from skimage.util import view_as_windows
# pd.set_option('display.max_colwidth', -1)
pd.set_option('display.max_columns', None)

from .utils import eval_dict, repr_nonscalar
from .grid import Grid

def col_group(x):
    if x[-1] in 'xyz':
        return x[:-1]
    if ',' in x:
        return x.split(',')[0]
    return x

def match_column(pattern, columns):
    """
    match column pattern from available columns
    m -> (mx, my, mz)
    m.region*x
    """
    cols = []
    if pattern in columns:
        cols.append(pattern)
    elif pattern + "x" in columns:
        cols.append(pattern + "x")
        if pattern + "y" in columns:
            cols.append(pattern + "y")
        if pattern + "z" in columns:
            cols.append(pattern + "z")
    else:
        matches = fnmatch.filter(columns, pattern)
        if matches:
            cols.extend(matches)
    return cols

def match_columns(patterns, columns):
    """ match column patterns """
    cols = []
    for p in patterns:
        cols.extend(match_column(p, columns))
    return cols

def csv_columns(file):
    """ Parse column names at the top of a CSV file """
    if isinstance(file, IOBase):
        first_line = file.readline()
    else:
        with open(file) as fp:
            first_line = fp.readline()

    if not first_line.startswith('#'):
        # no header found
        return []

    first_line = first_line.strip('# \n')
    if '(' in first_line:
        # some CSV headers contain extra info such as units per column
        # e.g. "t (s) B_extx (T) ..."
        # remove them from the column names
        first_line = re.sub(r' \([^()]*\)', '', first_line)

    columns = first_line.split()

    return columns

def save_csv(df, file):
    index = df.index.name is not None
    header = not isinstance(df.columns, pd.RangeIndex)

    if not isinstance(file, IOBase):
        file = open(file, 'wt')
    if header:
        file.write('# ')

    # For object columns, we repr non-scalar values before saving to csv
    if (df.dtypes == 'object').any():
        df = df.copy()
        for col, dtype in zip(df.columns, df.dtypes):
            if dtype == 'object':
                df[col] = df[col].map(repr_nonscalar)

    df.to_csv(file, sep=' ', index=index, header=header)

def read_csv(file, index_col=None):
    try:
        skip=0 if isinstance(file, IOBase) else 1
        names = csv_columns(file)
        if names:
            df = pd.read_csv(file, sep=r'\s+', index_col=index_col,
                    names=names, skiprows=skip, engine='c')
        else:
            df = pd.read_csv(file, sep=r'\s+', index_col=index_col,
                    header=None, engine='c')

    except (pd.errors.EmptyDataError, pd.errors.ParserError) as e:
        # empty csv
        df = pd.DataFrame()

    return df

def save_hdf(df, filename):
    filename, key = path.split(filename)
    df.to_hdf(filename, key, mode='a', complevel=9, complib='zlib')

def read_hdf(filename):
    filename, key = archive_key(filename)
    return pd.read_hdf(filename, key)

def hdf_columns(filename):
    filename, key = archive_key(filename)
    store = pd.HDFStore(filename, mode='r')

    if key is None:
        key = store.keys()[0]

    # assuming fixed format
    columns = list(store.get_node('{}/axis0'.format(key)).read().astype(str))
    index_name = store.get_node('{}/axis1'.format(key))._v_attrs.name
    if index_name is not None:
        columns.append(index_name)

    store.close()

    return columns

def read_npy(filename):
    arr = np.load(filename, allow_pickle=True)
    return pd.DataFrame.from_records(arr)

def save_npy(df, filename):
    index = df.index.name is not None
    arr = df.to_records(index=index)
    np.save(filename, arr)

def npy_columns(filename):
    with open(filename, 'rb') as fp:
        version = npy_format.read_magic(fp)
        npy_format._check_version(version)
        shape, fortran_order, dtype = npy_format._read_array_header(fp, version)
        return dtype.names

def read_npz(filename):
    filename, key = archive_key(filename)
    with np.load(filename, allow_pickle=True) as f:
        return pd.DataFrame.from_records(f[key])

def save_npz(df, filename):
    filename, key = archive_key(filename)
    index = df.index.name is not None
    arr = df.to_records(index=index)

    """ np.savez() doesn't support appending to an archive,
    so we need to do it ourselves """
    kw = dict(compression=zipfile.ZIP_DEFLATED, allowZip64=True)
    with zipfile.ZipFile(filename, mode="a", **kw) as zipf:
        fname = key + '.npy'
        # TODO: check for duplicate keys
        # always force zip64, gh-10776
        with zipf.open(fname, 'w', force_zip64=True) as fid:
            npy_format.write_array(fid, arr, allow_pickle=True)

def npz_columns(filename):
    filename, key = archive_key(filename)
    with np.load(filename) as npz:
        fp = npz.zip.open(key + ".npy")
        version = npy_format.read_magic(fp)
        npy_format._check_version(version)
        try:
            # numpy 1.23.5 introduced a safety check for large headers, which
            # needs to be disabled for larger tables by increasing max_header_size
            shape, fortran_order, dtype = npy_format._read_array_header(fp, version,
                    max_header_size=2**64)
        except TypeError:
            # numpy versions older than 1.23.5 do not understand max_header_size
            shape, fortran_order, dtype = npy_format._read_array_header(fp, version)
        return dtype.names

# supported table formats
table_formats = {
    'csv': (read_csv, save_csv, csv_columns),
    'hdf': (read_hdf, save_hdf, hdf_columns),
    'npy': (read_npy, save_npy, npy_columns),
    'npz': (read_npz, save_npz, npz_columns),
}

# map from extension -> format
table_extensions = {
    'csv': 'csv',
    'txt': 'csv',
    'hdf': 'hdf',
    'h5': 'hdf',
    'npy': 'npy',
    'npz': 'npz',
}

archive_formats = ["npz", "hdf"]

table_patterns = ['*.' + ext for ext in table_extensions]

def archive_key(filename):
    base, ext = path.splitext(filename)
    if ext[1:] in archive_formats:
        return filename, None

    head, tail = path.split(filename)
    if not head:
        return tail, None

    return head, tail

def get_format(filename):
    # dir/file.ext
    base, ext = path.splitext(filename)
    if ext:
        return ext[1:]

    # archive.fmt/key
    filename, key = archive_key(filename)
    base, ext = path.splitext(filename)
    if ext:
        return ext[1:]

    # unknown
    return None

def is_archive_format(fmt):
    return fmt in archive_formats

def is_archive(filename):
    fmt = get_format(filename)
    return is_archive_format(fmt)

def is_tablefile(filename):
    return get_format(filename) in table_extensions

def list_dir(filename):
    return os.listdir(filename)

def list_npz(filename):
    with np.load(filename) as npz:
        return npz.files

def list_hdf(filename):
    with pd.HDFStore(filename, mode='r') as store:
        return [k.lstrip('/') for k in store.keys()]

def listfiles(filename):
    files = None

    if not path.exists(filename):
        raise FileNotFoundError(filename)

    if path.isdir(filename):
        files = list_dir(filename)

    elif is_archive(filename):
        fmt = get_format(filename)
        if fmt == 'npz':
            files = list_npz(filename)
        elif fmt == 'hdf':
            files = list_hdf(filename)

    if files:
        files = [path.join(filename, f) for f in sorted(files)]
        return files

    raise ValueError(f"Don't know how to open {filename}, sorry")

def to_table(data):
    if isinstance(data, dict):
        with np.printoptions(linewidth=np.inf, threshold=sys.maxsize):
            data = [[k, repr(v)] for k,v in data.items()]
        df = pd.DataFrame(data)
        return df
        # return pd.DataFrame(data.values(), index=data.keys())

    if not isinstance(data, pd.DataFrame):
        return pd.DataFrame(data)

    return data

def match_any(name, patterns):
    return any([fnmatch.fnmatch(name, pat) for pat in patterns])

def read_table(filename, index_col=None, **kwargs):
    fmt = get_format(filename)
    fmt = table_extensions[fmt]
    reader = table_formats[fmt][0]
    df = reader(filename, **kwargs)
    if index_col is not None and df.index.name != index_col:
        df.set_index(index_col, inplace=True)
    return df

def save_table(data, filename):
    df = to_table(data)
    fmt = get_format(filename)
    fmt = table_extensions[fmt]
    writer = table_formats[fmt][1]
    return writer(df, filename)

def table_columns(filename):
    fmt = get_format(filename)
    fmt = table_extensions[fmt]
    func = table_formats[fmt][2]
    return func(filename)

def read_tables(filenames, index_col=None):
    dfs = [read_table(f, index_col) for f in filenames]

    # merge them
    table = dfs[0]
    for df in dfs[1:]:
        table = table.join(df)

    return table

def read_geometry(filename='geometry.csv'):
    df = read_table(filename)
    pos = np.array(df[['posx', 'posy']])
    angle = np.array(df['angle'])
    return pos, angle

def read_vectors(filenames, quantity='mag', t=None):
    files = []
    columns = []
    for f in filenames:
        cols = table_columns(f)
        # Look for time column
        if 't' not in cols:
            continue

        # Look for quantity
        match = match_column(quantity + '*', cols)
        if not match:
            continue

        # Look for xy data
        match = list(filter(lambda c: c[-1] in 'xy', match))
        if not match:
            continue

        files.append(f)
        columns.extend(match)

    assert files, f"No vector data found for quantity: {quantity}"

    df = read_tables(files, index_col='t')
    df = df[columns]

    if t is not None:
        if np.isscalar(t):
            t = [t]
        df = df.iloc[t]

    assert len(df.columns) % 2 == 0, "Number of columns not divisible by 2, not vector data?"

    UV = np.array(df)
    UV = UV.reshape((UV.shape[0], -1, 2))

    time = np.array(df.index)

    return time, UV

def crop(X, crop, axis=-1):
    """ Crop array X along one or more axes
    crop is a tuple (before, after) for each axis to crop
    axis specifies which axis to start cropping from
    """
    if not isinstance(X, np.ndarray):
        X = np.array(X)

    if np.isscalar(crop):
        crop = (crop, crop)

    crop = np.array(crop)
    if crop.ndim < 2:
        crop = [crop] * len(X.shape)

    if axis < 0:
        axis = X.ndim - len(crop) + axis + 1

    indices = [slice(None) for i in range(X.ndim)]
    for i, (before, after) in enumerate(crop):
        start = before
        stop = -after if after > 0 else None
        indices[axis + i] = slice(start, stop)

    return X[tuple(indices)]

def rolling_window(X, win_shape, step=1, method='sum'):
    """ Apply a function over a rolling window of X """
    assert method in ('sum', 'mean')

    mask = X.mask if np.ma.is_masked(X) else None
    X = np.copy(X)
    dtype = X.dtype
    if mask is not None:
        # view_as_windows doesn't work with masked arrays so set masked values
        # to nan and ignore the nans when computing the mean
        X = X.astype(float)
        X[mask] = np.nan

    if np.isscalar(win_shape):
        win_shape = (win_shape,) * X.ndim

    if np.isscalar(step):
        step = (step,) * X.ndim

    assert len(win_shape) == len(step)

    W = view_as_windows(X, win_shape, step)
    axis = tuple(range(X.ndim, W.ndim))

    if mask is not None:
        if method == 'mean':
            W = np.nanmean(W, axis=axis)
        else:
            W = np.nansum(W, axis=axis)

        # Mask any remaining nan values
        W = np.ma.masked_invalid(W)
        return W.astype(dtype)

    if method == 'mean':
        return np.mean(W, axis=axis)

    return np.sum(W, axis=axis)


def filter_df(df, **kwargs):
    """ Filter dataframe by key=value or range key=slice(start, stop, step) or list key=[key1, key2, key3] """
    for k,v in kwargs.items():
        if callable(v):
            df = df[df[k].apply(v)]

        elif isinstance(v, slice):
            # range (start, stop)
            start = v.start - 1e-8 if v.start is not None else v.start
            stop = v.stop + 1e-8 if v.stop is not None else v.stop
            step = v.step
            df = df.set_index(k).loc[start:stop:step].reset_index()
        else:
            if not isinstance(v, tuple) and not isinstance(v, list):
                v = [v]

            matches = []
            for vi in v:
                # equals
                if isinstance(vi, str) or np.issubdtype(df[k].dtype, np.integer):
                    matches.append(df[k] == vi)
                else:
                    # Determine tolerance based on size of filter value
                    atol = 10**(np.floor(np.log10(np.abs(vi))) - 4)
                    matches.append(np.isclose(df[k], vi, atol=atol))

            df = df[np.logical_or.reduce(matches)]

    return df

class Dataset(object):
    def __init__(self, index=None, params={}, info={}, basepath=None):
        self.index = index
        self.params = params
        self.info = info
        self.basepath = basepath

    @property
    def name(self):
        if self.basepath:
            return path.basename(path.abspath(self.basepath))
        return ''

    def __getitem__(self, i):
        index = self.index.loc[i]
        if isinstance(index, pd.Series):
            # single row returns a Series, but we want a DataFrame
            index = self.index.loc[[i]]
        ds = copy.copy(self)
        ds.index = index
        return ds

    def __repr__(self):
        return 'Dataset({!r}): {} items'.format(self.basepath, len(self))

    def __str__(self):
        s = "Dataset: {}\n\n".format(self.name)
        s += "params:\n"
        for k,v in sorted(self.params.items(), key=lambda i: i[0].lower()):
            s += " {}={}\n".format(k,v)
        s += "\n"
        s += "info:\n"
        for k,v in sorted(self.info.items(), key=lambda i: i[0].lower()):
            s += " {}: {}\n".format(k,v)
        s += "\n"
        s += "index:\n"
        s += str(self.index)
        s += "\n"
        return s

    def keys(self):
        return list(self.index.index)

    def items(self):
        for i in self.index.index:
            yield i, self.__getitem__(i)

    def iterrows(self):
        for i, row in self.index.iterrows():
            yield ((i,) + tuple(row), self.__getitem__(i))

    def __iter__(self):
        for _, ds in self.items():
            yield ds

    def __len__(self):
        return len(self.index)

    def __eq__(self, other):
        return self.basepath == other.basepath and \
                self.params == other.params and \
                (self.index == other.index).all(axis=None)

    def subset(self, i):
        index = self.index.loc[i]
        if isinstance(index, pd.Series):
            # single row returns a Series, but we want a DataFrame
            index = self.index.loc[[i]]
        ds = copy.copy(self)
        ds.index = index
        return ds

    def filter(self, **kwargs):
        ds = copy.copy(self)
        ds.index = filter_df(self.index, **kwargs)

        return ds

    def drop_duplicates(self, **kwargs):
        ds = copy.copy(self)
        ds.index = self.index.drop_duplicates(**kwargs)

        return ds

    def groupby(self, key):
        for k, index in self.index.groupby(key):
            ds = copy.copy(self)
            ds.index = index
            yield k, ds

    def sort_values(self, column):
        ds = copy.copy(self)
        ds.index = ds.index.sort_values(column)
        return ds

    def row(self, row=0):
        return self.index.iloc[row]

    def id(self, row=0):
        return self.row(row).name

    @staticmethod
    def read(basepath):
        index = read_csv(path.join(basepath, 'index.csv'))

        params = read_csv(path.join(basepath, 'params.csv'))
        params = dict(np.array(params))
        params = eval_dict(params)

        info = read_csv(path.join(basepath, 'info.csv'))
        info = dict(np.array(info))
        info = eval_dict(info)

        return Dataset(index, params, info, basepath)

    def save(self, basepath=None):
        if basepath and not self.basepath:
            self.basepath = basepath

        if not basepath:
            basepath = self.basepath

        assert basepath

        save_table(self.index, path.join(basepath, 'index.csv'))
        save_table(self.params, path.join(basepath, 'params.csv'))
        save_table(self.info, path.join(basepath, 'info.csv'))

    def file(self, filename):
        files = []
        for outdir in self.index['outdir']:
            files.append(path.join(self.basepath, outdir, filename))
        if len(files) == 1:
            return files[0]
        return files

    def files(self, patterns=None, squash=True):
        if isinstance(patterns, str):
            patterns = [patterns]

        files = [listfiles(path.join(self.basepath, outdir))
                 for outdir in self.index['outdir']]

        if patterns:
            files = [list(filter(lambda f: match_any(path.basename(f), patterns), fs))
                    for fs in files]

        if len(files) == 1 and squash:
            return files[0]

        return files

    def tablefile(self, tablename, squash=True):
        files = self.files(tablename + '*', squash=False)
        # next() to get the first matching tablefile
        files = [next(filter(is_tablefile, fs)) for fs in files]

        if len(files) == 1 and squash:
            return files[0]

        return files

    def tablefiles(self, patterns=None, squash=True):
        files = self.files(patterns, squash=False)
        files = [list(filter(is_tablefile, fs)) for fs in files]

        if len(files) == 1 and squash:
            return files[0]

        return files

#
# Data conversion
#
def digitize(X, threshold=0):
    return np.where(X > threshold, 1, 0)

def bit_array(x, n_bits):
    """ Convert number x to array of bits """
    return np.array([(x & (1 << bit)) >> bit for bit in range(n_bits)])

def array_bit(a):
    """ Convert bit array a to number """
    return sum([(int(b) << bit) for (bit, b) in enumerate(a)])

def vector_grid(pos, vectors, grid_size=None, crop_width=None,
        win_shape=None, win_step=None, normalize=True, return_grid=False):
    """ Process vectors on a grid """

    # Map the vectors on a grid
    if grid_size:
        # fixed grid width x height
        grid = Grid.fixed_grid(pos, grid_size)
    else:
        # native grid
        grid = Grid(pos)

    XY = grid.center_grid()

    # Determine whether we have a time axis
    assert len(vectors.shape) in (3, 2)
    assert vectors.shape[-1] == 2
    time_axis = 1 if len(vectors.shape) == 3 else 0

    if time_axis:
        # vectors.shape = (time, spin_count, 2)
        # UV.shape = (time, height, width, 2)
        UV = np.array([grid.add_values(UVi) for UVi in vectors])
    else:
        # vectors.shape = (spin_count, 2)
        # UV.shape = (height, width, 2)
        UV = grid.add_values(vectors)

    if crop_width:
        XY = crop(XY, crop_width, axis=0)
        UV = crop(UV, crop_width, axis=time_axis)

    if win_shape:
        if not win_step:
            win_step = win_shape

        # XY.shape = (height, width, 2)
        win_shape += (1,)
        win_step += (1,)
        XY = rolling_window(XY, win_shape, win_step, method='mean')

        if time_axis:
            # UV.shape = (time, height, width, 2)
            win_shape = (1,) + win_shape
            win_step = (1,) + win_step

        UV = rolling_window(UV, win_shape, win_step, method='sum')

    if normalize:
        # Normalize vectors to unit length
        nmax = norm(UV.reshape((-1,2)), axis=-1).max()
        if nmax != 0:
            UV = UV / nmax

    if return_grid:
        return XY, UV, grid

    return XY, UV

#
# Data loading utilities
#
def load_output(dataset, quantity, t=None, grid_size=None, crop_width=None,
        win_shape=None, win_step=None, flatten=True):
    """ Load output vectors from dataset, with optional post-processing """

    assert len(dataset) == 1

    time, UV = read_vectors(dataset.tablefiles(), quantity, t)

    if grid_size or crop_width or win_shape:
        XY, angle = read_geometry(dataset.tablefile('geometry'))
        XY, UV = vector_grid(XY, UV, grid_size, crop_width, win_shape, win_step)

    if flatten:
        UV = UV.reshape((UV.shape[0], -1))

    return UV

def load_input(dataset, t=None, input_column='input'):
    """ Load input data from dataset """

    assert len(dataset) == 1

    df = read_table(dataset.tablefile('input'), index_col='t')
    df = df[input_column]

    if t is not None:
        if np.isscalar(t):
            t = [t]
        df = df.iloc[t]

    return np.array(df)
