"""Interface to source(s) of individual item response data,
e.g., a flat database or Excel sheet(s) or a csv file,
accessible by some pandas.read_xxx function, or similar user-defined function.

Item Responses are accessed via an iterator,
yielding a list of integer-encoded ordinal responses for each table row.
Each table row must include response data from ONE unique respondent.

*** Main Access Function
item_response_table --- creates interface to ONE source of item response data
default_file_reader --- selects pandas reader depending on file suffix

*** Classes

ItemResponseFile --- interface to single file containing item response data.
ItemResponseSQL --- interface to SQL database accessed by pd.read_sql using SQLAlchemy

ItemResponseTable --- superclass interface to all concrete sources of tabulated data
    acting as an iterator over data from individual respondents.

Tables --- iterable chaining several sub-iterables with tabulated data. *****

*** Version history:
* Version 0.6.0:
2023-07-xx, allow reading SQLAlchemy or sqlite3 database
2022-09-15, using pandas read_xxx functions for input data sources
"""
# **** test reading by SQLAlchemy.engine
# **** allow reading by chunks, as complete pd.DataFrame ? sqlite chunks?

import numpy as np
import pandas as pd
from pathlib import Path
from itertools import chain

import logging

logger = logging.getLogger(__name__)
# logger.setLevel(logging.DEBUG)


# ------------------------------------------------ Reader exceptions
class ReadError(RuntimeError):
    """Any type of exception while attempting to read data from a source
    """


# ---------------------------------------------------------------
def item_response_table(items,
                        source,
                        sql=None,
                        accept_fcn=None,
                        sample_factor=1,
                        read_fcn=None,
                        **kwargs):
    """Factory creating an access object to a source of item response data.
        :param items: dict with (column_name, list(ordinal categories))
            for items and allowed item responses, with one column for each item.
        :param source: source for reading input data,
            e.g., a file Path, path string, or URL,
            OR an SQLAlchemy engine object, if reading from a general SQL database,
            OR a sqlite3.connection object, if using sqlite3 database
        :param sql: (optional) string with SQL SELECT statement, for pd.read_sql
            OR an SQL table name (if source is an SQLAlchemy engine)
            OR a SQLAlchemy Select object, generated by its select(...) function
        :param read_fcn: (optional) any user-defined file-reader function
            with signature similar to pd.read_excel or pd.read_csv.
            If None, a default pandas function is determined by file_path.suffix.
        :param accept_fcn: (optional) test function for each row to be valid input
        :param sample_factor: (optional) scalar integer for down-sampling among records
            Only one record is returned for every sample_factor records in the file
        :param kwargs: (optional) additional arguments for the reader function,
            depending on the selected reader function, e.g., for read_excel:
            index_col = column name or index for respondent ID,
                used to eliminate duplicate records from same respondent
            usecols = list of column names, if there are many redundant items
            dtype = dict with column datatype(s),
            converters = dict with elements (item_column, recoding_fcn)
            sheet_name = sheet name or list of sheet names, in case excel-type file
        :return: generator yielding item response arrays r, with
            r[i] = Response index for i-th Item, with origin=1,
            coded as integer in range 1,..., L_i for valid response, and
            r[i] = 0 for a missing response.
    """
    if sql is None:
        return ItemResponseFile(path=Path(source),
                                items=items, accept_fcn=accept_fcn, sample_factor=sample_factor,
                                read_fcn=read_fcn, **kwargs)
    else:  # con must be SQLAlchemy engine or sqlite3.connection object
        return ItemResponseSQL(con=source, sql=sql,
                               items=items, accept_fcn=accept_fcn, sample_factor=sample_factor,
                               **kwargs)


def default_file_reader(file_path):
    """Find a function that can read the given file path
    :param file_path: a Path instance
    :return: a pandas reader function, e.g., pd.read_excel
    """
    suffix = file_path.suffix
    if suffix in ['.xlsx', '.xls', '.odf', '.ods', '.odt']:
        return pd.read_excel
    elif suffix in ['.csv', '.txt']:
        return pd.read_csv
    elif suffix in ['.sav', '.zsav']:
        return pd.read_spss
    elif suffix in ['.dta']:
        return pd.read_stata
    else:
        raise ReadError(f'Must assign explicit read_fcn to read file format {suffix}')


# ---------------------------------------------- interface classes
class ItemResponseTable:
    """Superclass interface to a source of ordinal item responses,
    that can be accessed by some pandas.read_xxx or similar functions.
    Item Response data must be stored in rows and named items,
    e.g., a csv or Excel file.
    """
    def __init__(self, items, accept_fcn=None, sample_factor=1):
        """
        :param items: dict with (column_name, list(ordinal categories))
            for item responses to be returned to caller
        :param accept_fcn: (optional) test function for each row to be valid input
        :param sample_factor: (optional) scalar integer for down-sampling among records
            Only one record is returned for every sample_factor records in the file
        :return: generator yielding item response arrays r, with
            r[i] = Response index for i-th Item, with origin=1,
            coded as integer in range 1,..., L_i for valid response, and
            r[i] = 0 for a missing response.
        """
        self.items = items
        if accept_fcn is None:
            accept_fcn = _default_accept_all  # lambda row: True
        self.accept_fcn = accept_fcn
        self.sample_factor = sample_factor

    def __repr__(self):
        return (self.__class__.__name__ + '(\n\t'
                + ',\n\t'.join((f'{k}={repr(v)}' for (k, v) in self.__dict__.items()))
                + ')')

    @property
    def column_dtypes(self):
        return {c: pd.CategoricalDtype(categories=list(cats),
                                       ordered=True)
                for (c, cats) in self.items.items()}

    def __len__(self):
        """Number of (down-sampled) data records in file,
        not including header,
        exact or crudely estimated.
        This default method simply reads and counts all rows.
        Subclasses may implement faster method.
        :return: scalar integer
        """
        df = self.load()
        try:
            if type(df) is pd.DataFrame:
                return df.shape[0] // self.sample_factor
            elif type(df) is dict:  # dict of sheets from pd.read_excel()
                return sum(df_i.shape[0]
                           for df_i in df.values()) // self.sample_factor
            else:  # iterable of DataFrame objects, if reading by chunks
                return sum(df_i.shape[0]
                           for df_i in df) // self.sample_factor
        except RuntimeError as e:
            raise ReadError(f'Cannot count input data; Error: {e}')

    def __iter__(self):
        """Iterator over all ordinal item-response vectors from source defined by self.
        :return: generator yielding item response vectors r, with
            r[i] = Ordinal Response Code for i-th Item,
            coded as integer in range 0,..., L_i-1 for valid response, and
            r[i] = -1 for a missing response.
        """
        def gen_codes(df):
            """Generate code vectors from rows a DataFrame instance
            :param df: a pd.DataFrame instance
            :return: iterator
            """
            df = _check_unique_index(df)
            try:
                df_codes = [df[c].astype(dt).array.codes
                            for (c, dt) in self.column_dtypes.items()]
                df_codes = np.array(df_codes).T
                return (r for (i, r) in enumerate(df_codes)
                        if self.accept_fcn(r) and 0 == i % self.sample_factor)
            except RuntimeError as e:
                raise ReadError(f'Cannot extract code vectors from {df}; Error: {e}')
        # ---------------------------------------------------------

        df = self.load()
        if type(df) is pd.DataFrame:
            return gen_codes(df)
        elif type(df) is dict:  # dict of sheets from pd.read_excel()
            return chain.from_iterable((gen_codes(df_s)
                                        for (s, df_s) in df.items()))
        else:  # iterable of DataFrame objects, if reading by chunks
            try:
                return chain.from_iterable((gen_codes(df_i) for df_i in df))
            except RuntimeError as e:
                raise ReadError(f'Cannot use input data; Error: {e}')

    def load(self):
        """Abstract method to load all data from a file or other source
        :return: a pd.DataFrame or dict of pd.DataFrame objects
        """
        raise NotImplementedError


# -----------------------------------------------------------
class ItemResponseFile(ItemResponseTable):
    """Access to file storage of item response data,
    e.g., xlsx, csv, or similar
    """
    def __init__(self, items, path, accept_fcn=None, sample_factor=1, read_fcn=None, **file_kwargs):
        """
        :param items: dict with (column_name, list(ordinal categories))
            for item responses to be returned to caller
        :param path: Path instance defining input file
        :param accept_fcn: (optional) test function for each row to be valid input
        :param sample_factor: (optional) scalar integer for down-sampling among records
            Only one record is returned for every sample_factor records in the file
        :param read_fcn: (optional) reader function, e.g., pd.read_excel,
            or any user-supplied function with similar signature.
            If None, a default pandas function is determined by file_path.suffix.
        :param file_kwargs: (optional) additional arguments for the reader function, e.g.,
            usecols = list of column names, if there are many redundant items
            dtype = dict with column datatype(s),
            converters = dict with elements (item_colum, recoding_fcn)
            sheet_name = single sheet name to access, in case excel-type file
        """
        super().__init__(items=items, accept_fcn=accept_fcn, sample_factor=sample_factor)
        self.path = path
        if read_fcn is None:
            read_fcn = default_file_reader(path)
        self.read_fcn = read_fcn
        self.file_kwargs = file_kwargs

    def load(self):
        """Load data from file into local memory
        :return: a pd.DataFrame object or dict with such objects (Excel file sheets).
        """
        logger.debug(f'Loading from {self.path} with {self.read_fcn}')
        try:
            return self.read_fcn(self.path,
                                 **self.file_kwargs)
        except Exception as e:
            raise ReadError(f'{self.read_fcn.__name__}({self.path}) error: {e}')


# -------------------------------------------------------
class ItemResponseSQL(ItemResponseTable):
    """Access class using SQLAlchemy specification of database storage.
    """
    def __init__(self, con, sql, items, accept_fcn=None, sample_factor=1, **kwargs):
        """
        :param con: SQLAlchemy engine instance (preferably),
            OR SQLAlchemy or sqlite3 connection instance, input for pd.read_sql
        :param sql: SQL select string or SQLAlchemy selectable, input for pd.read_sql
        :param items: dict with (column_name, list(ordinal categories))
            for item responses to be returned to caller
        :param accept_fcn: (optional) test function for each row to be valid input
        :param sample_factor: (optional) scalar integer for down-sampling among records
            Only one record is returned for every sample_factor records in the file
        :param kwargs: (optional) additional arguments for pd.read_sql, e.g.,
            columns = list of column names, if there are many redundant items
            index_col = column for respondent ID
            chunksize = number of rows for each input chunk
            params = any other database-dependent parameters
        """
        # ******* should CLOSE con instance if it is sqlite3.connection instance
        super().__init__(items=items, accept_fcn=accept_fcn, sample_factor=sample_factor)
        self.con = con
        self.sql = sql
        self.kwargs = kwargs

    def load(self):
        """Load data using pd.read_sql with SQLAlchemy instances as input
         :return: pd.DataFrame or iterator of DataFrame objects
         """
        logger.debug(f'Loading read_sql({self.sql}, {self.con}')
        try:
            return pd.read_sql(sql=self.sql, con=self.con,
                               **self.kwargs)
        except Exception as e:
            raise ReadError(f'read_sql({self.sql}, {self.con}) error: {e}')


class Tables:
    """Iterable chaining sub-iterables
    """
    def __init__(self, *table_list):
        """
        :param table_list: list of iterables, e.g., lists or ItemResponseTable objects.
        """
        self.table_list = table_list

    def __len__(self):
        """Total number of data records in all included sub-tables,
        exact or crudely estimated.
        :return: scalar integer
        """
        return sum(len(t) for t in self.table_list)

    def __iter__(self):
        """iterator of records in all included sub-tables
        :return: iterator over all response records in all sub-tables
        """
        return chain.from_iterable(self.table_list)


# ----------------------------------------------------- Module help functions:
def _default_accept_all(r):
    """Default function for ItemResponseTable property accept_fcn.
    :param r: list with one record from input
    :return: always True, without any checking
    """
    return True


def _check_unique_index(df):
    """Check if input dataframe has unique index labels,
    and remove duplicates
    :param df: a pd.DataFrame instance
    :return: df_unique = corresponding pd.DataFrame with duplicates removed
    """
    if df.index.is_unique:
        return df
    else:
        n_unique = len(df.index.unique())
        n_df = df.shape[0]
        logger.debug(f'DataFrame has {n_unique} unique index labels, less than all {n_df} rows. ' +
                       'Keeping last record from each respondent.')
        df_u = df.groupby(level=[0]).last()
        return df_u


# -------------------------------------------------------- Module TEST:
if __name__ == '__main__':
    # from pathlib import Path
    import sqlite3
    import ir_logging


    HAQ_PATH = Path.home() / 'Documents/LeijonKTH/heartech/HA-QualityRegistry/IRT-IOI-HA/IOI-HA-Data'
    # ioiha_data_path = HAQ_PATH / 'IRT-IOI-HA' / 'IOI-HA-data'
    # to ioi-ha data sets OTHER THAN the Swedish Quality Registry
    # ------------------------------------------------------------------

    print('\n*** Testing item_response_table with Hickson data set:\n')
    work_path = HAQ_PATH
    test_file = HAQ_PATH / 'Hickson' / 'Short_Eartrak AUS.xlsx'

    # ------------------------------------- setup logging for test:
    ir_logging.setup(work_path / 'test_ir_source.txt')
    logger = logging.getLogger(__name__)

    column_names = [f'Q0{i}' for i in range(1, 8)]
    ir_hickson = item_response_table(items={c: [1, 2, 3, 4, 5]
                                          for c in column_names},
                                   source=test_file, index_col='EarRowId')  #0)
    print('ir_hickson = ', ir_hickson)

    ir_df = ir_hickson.load()
    print('ir_df.head =\n', ir_df.head(10))

    print('printing a few rows:')
    n = 0
    for r in ir_hickson:
        print(r)
        n += 1
        if n > 10:
            break

    print(f'len(ir_hickson) = ', len(ir_hickson))

    sql_path = work_path / 'test_sql'
    sql_path.mkdir(parents=True, exist_ok=True)
    # *** create sqlite3 database with copy of ir_df:
    # con = sqlite3.connect(sql_path / 'testHickson.db' )
    # cur = con.cursor()
    # # cur.execute("CREATE TABLE ioiha(" + ','.join(ir_df.columns) + ")")
    # res = cur.execute("SELECT name FROM sqlite_master")
    # print(res.fetchone())

    # cur.execute("""
    # INSERT INTO ioiha VALUES
    #     ('tes1', 'arne', 1, 2, 3, 4, 5, 6 ,7),
    #     ('test2', 'eddie', 1, 7.5, 3, 4, 5, 6, 9) """)
    # con.commit()
    # res = cur.execute("SELECT * FROM ioiha")
    # r = res.fetchall()
    # print(r)
    # con.close()

    # print('*** test read_sql')
    # con = sqlite3.connect(sql_path / 'testHickson.db' )
    # ir_test = pd.read_sql_query("SELECT * FROM ioiha", con, index_col='EarRowId')
    # print(ir_test)
    # # con.close()

    print('*** test Hickson to_sql')

    con = sqlite3.connect(sql_path / 'testHickson.db' )
    n_rows = ir_df.to_sql('ioiha', con, index=True, if_exists='replace')
    print(f'{n_rows} saved by Pandas to_sql')
    # NOTE: Hickson's xlsx has integer codes and '.' for missing data, so the result is mixed: dtype= object
    # When saving to sqlite by to_sql, dtype object -> string in all items, except Q00 which is still integer
    # SQL requires all elements in one column to have SAME dtype

    # print('*** test read_sql_query')
    # ir_test = pd.read_sql_query("SELECT * FROM ioiha", con, index_col='EarRowId')
    # print(ir_test)
    # con.close()

    print('*** test read_sql')
    con = sqlite3.connect(sql_path / 'testHickson.db' )
    ir_test = pd.read_sql("SELECT * FROM ioiha", con, index_col='EarRowId')
    print(ir_test)
    con.close()

    print('\n*** Test item_response_table using sqlite3 connection and query input ***')
    con = sqlite3.connect(sql_path / 'testHickson.db' )
    ir_sql = item_response_table(source=con,
                                   sql="SELECT * FROM ioiha",
                                   index_col='EarRowId',
                                   items={c: [str(i) for i in range(1, 6)]
                                          for c in column_names}
                                   )  # NOTE: string-encoded integers here
    ir_df = ir_sql.load()
    print(ir_df.head(20))

    print('printing a few rows:')
    n = 0
    for r in ir_sql:
        print(r)
        n += 1
        if n > 10:
            break

    con.close()

    print('\n*** Test SQLAlchemy input directly ***')
    from sqlalchemy import create_engine
    engine = create_engine("sqlite+pysqlite:///" + str(sql_path / 'testHickson.db'))
    from sqlalchemy import text
    with engine.connect() as conn:
        result = conn.execute(text("select * from ioiha")).fetchmany(20)
        print(np.array(result))

    print('\n*** Test item_response_table using SQLAlchemy engine ***')
    ir_sql = item_response_table(source=engine,
                                   sql='ioiha',  # *** "SELECT * FROM ioiha",
                                   index_col='EarRowId',
                                   items={c: [str(i) for i in range(1, 6)]
                                          for c in column_names}
                                   )  # NOTE: string-encoded integers here
    ir_df = ir_sql.load()
    print(ir_df.head(20))

    print('finished')


