"""This module defines functions that read / write speech stim-response data
and ConfMatrixCalc analysis result tables,
from/to a table file of any kind that package Pandas can handle.

A single data file for READING may include speech stim-response records from ONE or SEVERAL respondents,
in ONE or SEVERAL respondent groups, in ONE or SEVERAL test conditions.

The participant id may be stored in a designated column of the table,
otherwise the file name can be used as participant id,
or (in Excel-type file) the participant may be identified by the sheet name.

The table must always include columns with Stimulus and Response categories,
and may or may not include Count results for each (stim, resp) combination.

The table may include group categories, test-condition categories,
and also other data not used for analysis.

Regardless of the file storage format, even if an input file represents only ONE participant,
the data is always delivered by a generator function
stim_resp_gen, which is an iterable yielding pandas DataFrame instances.

Some pandas file readers can also do some
data conversions and/or type checking.
Column names may be modified before the data are delivered to caller.

For writing CM data or any other table to a file,
simply cast the DataFrame instance to subclass Table,
and use Table(df).save(...) method.

*** Version History:

* Version 0.8:
2025-12-01, changed _stim_resp_df: yields only data found in file.
            Caller must decode path string, and rename columns if needed.
2025-11-22, tested seems OK
2025-11-20, copied from EmaCalc 1.0.2, adapted for ConfMatrixCalc
"""
# *** simplify to get ONE df chunk from each file, possibly with several participants or test-conds
# *** NO, some chunks might include forbidden data types, while other should be accepted
import pandas as pd
import logging

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


# ------------------------------------------------ File exceptions
class FileReadError(RuntimeError):
    """Any type of exception while attempting to read stim-response data from a file
    """


class FileWriteError(RuntimeError):
    """Any type of exception while attempting to write a pd.DataFrame table to a file
    """


# --------------------------------------------------
class Table(pd.DataFrame):
    """Subclass adding a general save method,
    automatically switching to desired file format.
    """
    def save(self, file_path,
             allow_over_write=True,
             write_fcn=None,
             **file_kwargs):
        """Save self to a table-style file.
        :param file_path: Path instance defining file location and full name incl. suffix.
        :param allow_over_write: (optional) boolean, if False, find new unused path stem.
        :param write_fcn: (optional) user-supplied function with signature
            write_fcn(table, path, **kwargs).
            If None, default pd.DataFrame method is used, determined by file_path.suffix
        :param file_kwargs: (optional) any additional arguments for the
            write function for the specific file format.
        :return: None
        """
        if not allow_over_write:
            file_path = safe_file_path(file_path)
        suffix = file_path.suffix
        try:
            # ******** file_path.parent.mkdir here? NO, done by container class
            if write_fcn is None:
                if suffix in ['.xlsx', '.xls', '.odf', '.ods', '.odt']:
                    # self.to_excel(file_path, sheet_name=file_path.stem, **file_kwargs) # use default
                    self.to_excel(file_path, **file_kwargs)
                elif suffix in ['.csv']:
                    self.to_csv(file_path, **file_kwargs)
                elif suffix in ['.txt']:
                    self.to_string(file_path, **file_kwargs)
                elif suffix in ['.tex']:
                    # with warnings.catch_warnings():
                    #     # suppress Pandas FutureWarning about to_latex method
                    #     warnings.simplefilter('ignore')
                    #     self.to_latex(file_path, **file_kwargs)
                    self.to_latex(file_path, **file_kwargs)
                else:
                    raise FileWriteError(f'No default write method for file type {suffix}')
            else:
                write_fcn(self, file_path, **file_kwargs)
        except Exception as e:
            raise FileWriteError(f'Could not write to {file_path}. Error: {e}')


# ----------------------------------------- General File Reader:
def stim_resp_gen(file_path,
                  read_fcn=None,
                  **file_kwargs
                  ):
    """Reader of data stored in a table-style file that can be read by Pandas.read_xxx()
    or user-supplied function with similar signature.
    Each file row is assumed to include a (stim, resp) pair, but may also include other data fields.
    :param file_path: Path to existing file for reading
    :param read_fcn: (optional) basic read 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: any additional arguments for the
        read function for the specific file format
    :return: generator object yielding pd.DataFrame instance(s),
        with one row for each stim-response pair,
        and one column specifying participant id,
        possibly copied from file name or sheet name if requested by argument participant.
        Any other column data, e.g., group and/or test-condition categories,
        are also included, to be checked by caller.
    """
    # ******* simplify: return only data found in file, let caller decode or rename if needed ? *******
    if read_fcn is None:
        read_fcn = _default_reader(file_path)
    logger.debug(f'Reading from {file_path} with {read_fcn}')
    # if participant == 'sheet':  # ******* not needed, ALWAYS return a sheet column if
    if read_fcn is pd.read_excel:
        if 'sheet_name' not in file_kwargs.keys():
            file_kwargs['sheet_name'] = None  # -> include ALL sheets by default
            # because pd.read_excel uses sheet_name=0 by default
    try:
        df = read_fcn(file_path, **file_kwargs)
    except Exception as e:
        raise FileReadError(f'{read_fcn.__name__}({file_path}) error:\n {e}')
    if type(df) is pd.DataFrame:
        yield df
    elif type(df) is dict:  # one item per sheet, like pd.read_excel
        for (sheet_name, df_s) in df.items():
            try:
                if 'sheet' not in df_s.columns:
                    df_s['sheet'] = sheet_name
                yield df_s
            except Exception as e:
                raise FileReadError(f'Cannot use {df_s.__class__} object: {e}')
    elif type(df) is list:  # several results from user-defined reader
        for df_s in df:
            yield df_s
    else:
        raise TypeError(f'Cannot use {df.__class__} result from {read_fcn}.')


# ------------------------------------------ help functions:
def safe_file_path(p):
    """Ensure previously non-existing file path, to avoid over-writing,
    by adding a sequence number to the path stem
    :param p: file path
    :return: p_safe = similar file path with modified name
    """
    f_stem = p.stem
    f_suffix = p.suffix
    i = 0
    while p.exists():
        i += 1
        p = p.with_name(f_stem + f'-{i}' + f_suffix)
    return p


def _default_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 FileReadError(f'Cannot (yet) read file format {suffix}')


# ----------------------------------------- TEST:
if __name__ == '__main__':
    from pathlib import Path
    import cm_logging

    work_path = Path.home() / 'Documents' / 'CM_sim'  # or whatever...
    data_path = work_path / 'data_xlsx'  # / 'Age_old'  # to use simulation data generated by run_sim.py
    # data_path = work_path / 'data' / 'Age_old'  # to use simulation data generated by run_sim.py

    cm_logging.setup()  # to save the log file

    cm_file = stim_resp_gen(data_path / 'test.xlsx',
                            # participant='sheet',
                            )
    # for (s, sr) in cm_file:
    for sr in cm_file:
        print(sr.head(100))
        print('dtypes=\n', sr.dtypes)
        # count = sr.value_counts(subset=['HA', 'CoSS'], sort=False)
        # print(count)

    path = data_path / 'test_write' / 'test_participant.csv'
    path.parent.mkdir(parents=True, exist_ok=True)
    # ********* check CoSS data type ? should be string ****'
    Table(sr).save(path)
