import io
import re
from typing import Sequence, Union

import numpy as np
import pandas as pd
import pytz

from . import utils

logger = utils.make_logger('pytrthree')

TRTHFile = Union[str, io.TextIOWrapper]


class TRTHIterator:
    """
    Helper class to parse a set of TRTH .csv.gz files
    and yield DataFrame grouped by RIC.
    """

    def __init__(self, files, chunksize=10 ** 6):
        """
        Validates input files and initializes iterator.
        :param files: Compressed CSV files downloaded from the TRTH API
        :param chunksize: Number of rows to be parsed per iteration.
                          Higher number causes higher memory usage.
        """
        self.files = self._validate_input(files)
        self.chunksize = chunksize
        self.iter = self.make_next()

    def __iter__(self):
        return self

    def __next__(self):
        return next(self.iter)

    @staticmethod
    def _validate_input(files: Union[TRTHFile, Sequence[TRTHFile]]) -> Sequence[TRTHFile]:
        if isinstance(files, (str, io.TextIOWrapper)):
            files = [files]
        for f in files:
            if not isinstance(f, (str, io.TextIOWrapper)):
                raise ValueError(f'Invalid input: {f}')

        output = []
        for file in files:
            fname = file.name if isinstance(file, io.TextIOWrapper) else file
            try:
                _, ftype = utils.parse_rid_type(fname)
            except (IndexError, ValueError):
                logger.debug(f'Ignoring {fname}')
                continue
            if ftype in {'confirmation', 'report'}:
                logger.debug(f'Ignoring {fname}')
                continue
            else:
                output.append(file)

        return sorted(output)

    def make_next(self):
        """Iterates over input files and generates single-RIC DataFrames"""
        for file in self.files:
            logger.info(file)
            lastrow = None
            chunks = pd.read_csv(file, iterator=True, chunksize=self.chunksize)
            for chunk in chunks:
                for _, df in chunk.groupby('#RIC'):
                    processed_df = self.pre_process(df.copy(), lastrow)
                    yield processed_df
                    lastrow = None
                lastrow = processed_df.iloc[-1]

    @staticmethod
    def pre_process(df, lastrow=None) -> pd.DataFrame:
        """Generates a unique DateTimeIndex and drops datetime-related columns"""
        def find_columns(df, pattern):
            try:
                return [i for i in df.columns if re.search(pattern, i)][0]
            except IndexError:
                return None

        # Remove characters that cause problems in MongoDB/Pandas (itertuples)
        df.columns = [re.sub('\.|-|#', '', col) for col in df.columns]

        # Make DateTimexIndex
        date_col = find_columns(df, 'Date')
        time_col = find_columns(df, 'Time')
        gmt_col = find_columns(df, 'GMT')
        if time_col:
            df.index = pd.to_datetime(df[date_col].astype(str) + ' ' + df[time_col])
            df.drop([date_col, time_col], axis=1, inplace=True)
        else:
            df.index = pd.to_datetime(df[date_col].astype(str))
            df.drop(date_col, axis=1)

        # Add 10 ns to repeated timestamps to make timeseries index unique.
        offset = pd.DataFrame(df.index).groupby(0).cumcount() * np.timedelta64(10, 'ns')
        df.index += offset.values

        # Make DateTimeIndex timezone-aware
        if gmt_col:
            assert len(df[gmt_col].value_counts()) == 1
            df.index = df.index + pd.Timedelta(hours=df.ix[0, gmt_col])
            df.index = df.index.tz_localize(pytz.FixedOffset(9 * 60))
            df.drop(gmt_col, axis=1, inplace=True)
        else:
            df.index = df.index.tz_localize(pytz.timezone('utc'))

        # Make sure rows separated by chunks have different timestamps
        if lastrow is not None:
            if lastrow['RIC'] == df.ix[0, 'RIC'] and lastrow.name == df.index[0]:
                logger.info('Adjusting first row timestamp')
                df.index.values[0] += np.timedelta64(5, 'ns')

        return df

@@ -1897,38 +1897,8 @@
 es:%0A
-            logger.info(file)%0A
     
@@ -1912,32 +1912,32 @@
  lastrow = None%0A
+
             chun
@@ -2015,16 +2015,19 @@
      for
+ i,
  chunk i
@@ -2028,22 +2028,33 @@
 hunk in 
+enumerate(
 chunks
+)
 :%0A      
@@ -2068,12 +2068,174 @@
    f
-or _
+name = file.name if isinstance(file, io.TextIOWrapper) else file%0A                logger.info('%7B%7D chunk #%7B%7D'.format(fname.split('/')%5B-1%5D, i+1))%0A                for ric
 , df
@@ -2358,16 +2358,22 @@
    yield
+ (ric,
  process
@@ -2377,16 +2377,17 @@
 essed_df
+)
 %0A       
@@ -4209,21 +4209,23 @@
  logger.
-info(
+debug(f
 'Adjusti
@@ -4246,16 +4246,35 @@
 imestamp
+: %7Bdf.ix%5B0, %22RIC%22%5D%7D
 ')%0A     

Improve iteration