# timeseries.py
#
# Read the time_series.dat and return a TimeSeries object of 1D numpy
# arrrays
"""
Contains the classes and methods to read the time series file.
"""


def ts(*args, **kwargs):
    """
    ts(file_name='time_series.dat', datadir='data',
       quiet=False, comment_char='#', sim=None, unique_clean=False)

    Read Pencil Code time series data.

    Parameters
    ----------
    file_name : string
        Name of the time series file.
        For supernova data change file_name to 'sn_series.dat'.

    datadir : string
        Directory where the data is stored.

    quiet : bool
        Flag for switching off output.

    comment_char : string
        Comment character in the time series file.

    sim : obj
      Simulation object from which to take the datadir.

    unique_clean : bool
      Set True, np.unique is used to clean up the ts, e.g. remove errors
      at the end of crashed runs.
    """

    ts_tmp = TimeSeries()
    ts_tmp.read(*args, **kwargs)
    return ts_tmp


class TimeSeries(object):
    """
    TimeSeries -- holds Pencil Code time series data.
    """

    def __init__(self):
        """
        Fill members with default values.
        """

        self.t = []
        self.keys = []

    def keys(self):
        for i in self.__dict__.keys():
            print(i)

    def read(
        self,
        file_name="time_series.dat",
        datadir="data",
        quiet=False,
        comment_char="#",
        sim=None,
        unique_clean=False,
        time_range=None,
    ):
        """
        read(file_name='time_series.dat', datadir='data',
             quiet=False, comment_char='#', sim=None, unique_clean=False)

        Read Pencil Code time series data.

        Parameters
        ----------
        file_name : string
            Name of the time series file.
            For supernova data change file_name to 'sn_series.dat'.

        datadir : string
            Directory where the data is stored.

        quiet : bool
            Flag for switching off output.

        comment_char : string
            Comment character in the time series file.

        sim : obj
          Simulation object from which to take the datadir.

        unique_clean : bool
          Set True, np.unique is used to clean up the ts, e.g. remove errors
          at the end of crashed runs.
        """

        import numpy as np
        import os.path
        import re

        if sim:
            from pencil.sim import __Simulation__

            if isinstance(sim, __Simulation__):
                datadir = sim.datadir

        datadir = os.path.expanduser(datadir)
        with open(os.path.join(datadir, file_name), "r") as infile:
            lines = infile.readlines()

        nlines_init = len(lines)
        data = np.zeros((nlines_init, len(self.keys)))
        nlines = 0
        for i, line in enumerate(lines):
            if re.search("^%s--" % comment_char, line):
                # Read header and create keys for dictionary.
                line = line.strip("{0}-\n".format(comment_char))
                keys_new = re.split("-+", line)
                if keys_new != self.keys:
                    n_newrows = abs(len(keys_new) - len(self.keys))
                    data = np.append(data, np.zeros((nlines_init, n_newrows)), axis=1)
                    self.keys = keys_new
            else:
                try:
                    row = np.array(list(map(float, re.split(" +", line.strip(" \n")))))
                    data[nlines, :] = row
                    nlines += 1
                except ValueError:
                    print(f"Invalid data on line {i}. Skipping.")

        # Clean up data.
        data = np.resize(data, (nlines, len(self.keys)))

        if not quiet:
            print("Read {0} lines.".format(nlines))

        # Assemble into a TimeSeries class.
        for i in range(0, len(self.keys)):
            setattr(self, self.keys[i], data[:, i])

        # Do unique clean up.
        if unique_clean:
            clean_t, unique_indices = np.unique(self.t, return_index=True)

            if np.size(clean_t) != np.size(self.t):
                for key in self.keys:
                    setattr(self, key, getattr(self, key)[unique_indices])
        if time_range:
            if isinstance(time_range, list):
                time_range = time_range
            else:
                time_range = [time_range]
            if len(time_range) == 1:
                start_time = 0.
                end_time = time_range[0]
            elif len(time_range) == 2:
                start_time = time_range[0]
                end_time = time_range[1]
            ilist = list()
            for i, time in zip(range(self.t.size),self.t):
                if time >= start_time:
                    if time <= end_time:
                        ilist.append(i)
            for key in self.keys:
                tmp = self.__getattribute__(key)[ilist]
                self.__delattr__(key)
                setattr(self, key, tmp)