"""Main template script to analyze ordinal response data using Item Response Theory.
This template should be copied and edited by the user for the desired application.

The model estimates the distribution of one or more traits underlying responses to a questionnaire,
for one or several populations, represented by groups of respondents.

The results show trait distributions within and between populations,
and reveals jointly credible differences between populations.

Populations should preferably be represented by large groups of respondents,
recruited at random.

Analysis results show the inter-individual variance within and between populations,
and the mapping between traits and item ratings,
but individual trait values are not displayed.

The analysis script always includes three main steps:
1: Define and collect item response data from one or more sources,
    and specify a top directory for all results.
2: Learn an OrdinalItemResponseModel instance using the collected response data
3: Display the main characteristics of the learned model.

*** Version History

* Version 0.6.0
2023-07-11, cleanup this script template with examples using new ir_source module
2023-07-06, result file format choice moved from ItemResponseDisplaySet.show() -> .save()
2023-07-01, new module ir_source for pandas.read_xxx for input files
2022-09-15, use pandas.read_xxx functions to access input data files.
2020-06-06, first standalone version for ItemResponseCalc package
2020-05-30, template example based on IOI-HA data analysis
"""
# -------- __main__ check to prevent multiprocessor sub-tasks to re-run this script
if __name__ == '__main__':

    import numpy as np
    import pickle
    from pathlib import Path
    import datetime as dt
    import logging
    # from sqlalchemy import create_engine  # to read from SQL database

    from ItemResponseCalc import __version__
    from ItemResponseCalc import ir_logging
    from ItemResponseCalc.item_response_data import Questionnaire
    from ItemResponseCalc.item_response_data import ItemResponseDataSet
    from ItemResponseCalc.ir_source import item_response_table, Tables
    # Tables used to join data from two or more sources as one group
    from ItemResponseCalc.item_response_model import OrdinalItemResponseModel
    from ItemResponseCalc.ir_display import ItemResponseDisplaySet

    # -------------------------------- ItemResponseModel parameters:
    max_n_traits = 4
    # = max number of traits to explain total set of item responses
    # = 1, if the test instrument is known to measure a uni-dimensional trait

    # n_scale_samples = 50
    # -> sampled set of scale threshold parameters
    n_scale_samples = 1
    # -> MAP-estimated single scale threshold parameters
    n_subject_samples = 50
    # = number of sampled trait(s) for each respondent
    # NOTE: computation time is proportional to
    # n_subjects * n_scale_samples * n_subject_samples
    # n_scale_samples = 1 may give similar result as sampled scale model

    TEST_DOWNSAMPLE = 1
    # -> no down-sampling -> all records included
    # TEST_DOWNSAMPLE = 10
    # -> smaller data subset only for faster initial test run

    # ----------------------------- STEP 1: Define input sources and result path:

    timestamp_result = True  # Prevent over-writing result files
    # timestamp_result = False  # Repeated runs will save results in same directory

    TOP_PATH = Path.home() / 'Documents/LeijonKTH/heartech/HA-QualityRegistry'
    TOP_PATH = TOP_PATH / 'IRT-IOI-HA'  # or whatever...
    # Change this example to top-dir for everything to be read or written

    ioiha_data_path = TOP_PATH / 'IOI-HA-data'  # or whatever...
    # path to directory containing item response data files

    result_dir = TOP_PATH / 'IOI-HA-results'
    # path where results will be saved, for several runs if needed

    ioiha_file = ioiha_data_path / 'IOI-HA-English.txt'
    # text file containing the questions and response alternatives
    # If not available, a Questionnaire object must be created manually
    # as suggested below.

    if timestamp_result:
        timefmt = '{0.year}-{0.month:02}-{0.day:02}-{0.hour:02}{0.minute:02}'
        result_run = timefmt.format(dt.datetime.now())
        # -> name of sub-folder in result_dir
        save_dir = result_dir / result_run
    else:
        save_dir = result_dir
    save_dir.mkdir(parents=True, exist_ok=True)
    # sub-path for results of this particular run

    model_file = 'irt-model.pkl'
    # = name of saved model in save_dir

    log_file = 'run_irt_log.txt'
    # = name of log file saved in same sub-folder
    ir_logging.setup(save_dir / log_file)
    logger = logging.getLogger(__name__)

    logger.info('*** Using ItemResponseCalc v.' + __version__)
    if TEST_DOWNSAMPLE > 1:
        logger.info(f'*** Test with down-sampling {TEST_DOWNSAMPLE}')

    def accept_record(r):
        """User-defined inclusion criterion for accepting a response record,
        applied to responses recoded as integer with origin 0,
        and missing responses encoded as -1.
        :param r: list with integer-encoded responses in ONE record from a source
        :return: boolean == True if the record is acceptable
        """
        n_missing = sum(r_i < 0 for r_i in r)
        return n_missing <= 3
    # ----------------------------------------------------------

    # ------------------ 1A: Define questionnaire instrument:
    # create a Questionnaire object to specify
    # number of items and response alternatives

    # Example: load from text file, edit as needed:
    q_items = Questionnaire.load(ioiha_file)

    # if the text file is not available, create q_items manually, e.g.:
    # Example: seven items, all with 5 response levels, like the IOI-HA:
    # r_levels = [5, 5, 5, 5, 5, 5, 5]
    # = list with number of response levels, one element for each item
    # q_items = Questionnaire(item_response_levels=r_levels)

    # ------------------ 1B: Define sources of data to be included in analysis:
    data_groups = dict()  # space for all included groups

    # Example: Include xlsx file(s) from Hickson, Australia:
    # Edit as needed:
    au1_file = ioiha_data_path / 'Hickson' / 'Short_Eartrak AUS.xlsx'
    logger.info(f'AU-H-10: Using {au1_file}')

    au1 = item_response_table(source=au1_file,
                              items={f'Q0{item}': list(range(1, 6))
                                     for item in range(1, 8)},  # as encoded in the xlsx file
                              accept_fcn=accept_record,  # None -> no restriction
                              sample_factor=TEST_DOWNSAMPLE,
                              index_col=0    # with respondent ID, to avoid duplicate records
                              )
    # NOTE: items = a dict with column header and a sequence of allowed ordinal responses,
    #   as encoded in this input source, not necessarily == Questionnaire.responses.
    # If index_col is defined, only the LAST record from each respondent is included.

    data_groups['AU-H-10'] = au1

    # Example: second input source:
    # au2_file = ioiha_data_path / 'Hickson' / 'Eartrakdata_10_May_08.xlsx'
    # logger.info(f'AU-H-10e: Using {au2_file}')
    #
    # au2 = item_response_table(source=au2_file,
    #                           items={f'Q0{item}': list(range(1, 6))
    #                                  for item in range(1, 8)},
    #                           accept_fcn=accept_record,
    #                           sample_factor=TEST_DOWNSAMPLE,
    #                           index_col=0)
    # data_groups['AU-H-10e'] = au2

    # More examples: Combine two csv files into one group:
    # nle1_file = ioiha_data_path / 'Kramer' / 'moeder_final - kopieforArne.csv'
    # logger.info(f'NL-02: Using {nle1_file}')
    #
    # nle1 = item_response_table(source=nle1_file,
    #                            items={f'ioi_v{item}': list(range(1, 6))
    #                                   for item in range(1, 8)},     # as encoded in the csv file
    #                            accept_fcn=accept_record,
    #                            sample_factor=TEST_DOWNSAMPLE)   # no respondent ID given
    #
    # nle2_file = ioiha_data_path / 'Kramer' / 'NLSH-Kramer-T2IOI-HA.csv'
    # logger.info(f'NL-16: Using {nle2_file}')
    #
    # kramer2_columns = {'T2IOIHA1': range(5),  # responses coded as 0,.., 4
    #                    'T2IOIHA2': range(1, 6),  # coded as 1,.., 5
    #                    'T2IOIHA3': range(1, 6),
    #                    'T2IOIHA4': range(5),  # coded as 0,.., 4
    #                    'T2IOIHA5': range(1, 6),
    #                    'T2IOIHA6': range(1, 6),
    #                    'T2IOIHA7': range(1, 6)}
    #
    # nle2 = item_response_table(source=nle2_file,
    #                            items=kramer2_columns,
    #                            accept_fcn=accept_record,
    #                            sample_factor=TEST_DOWNSAMPLE)
    # data_groups['NL-02-16'] = Tables(nle1, nle2)

    # SQL example: access data from a database:
    # engine = create_engine("sqlite+pysqlite:///" + str(sql_path / 'testDataBase.db'))
    # ir_sql = item_response_table(source=engine,
    #                              sql='ioiha',  # table name in database
    #                              index_col='EarRowId',
    #                              items={c: [str(i) for i in range(1, 6)]
    #                                     for c in column_names}
    #                              )  # NOTE: string-encoded responses in this DB
    # data_groups['test'] = ir_sql

    # ------------------ 1C: Collect all into ONE DataSet object:

    ids = ItemResponseDataSet(questionnaire=q_items,
                              groups=data_groups)
    # NOT YET reading any response data
    # In case of any input error, it will show here:
    logger.info('Item Response Data = \n' + repr(ids))

    # (Optional) Show response counts, to verify that input was OK
    r_counts = ids.response_counts()
    # = list with one pandas.DataFrame object for each questionnaire item
    logger.info('Response Counts =\n'
                 + '\n\n'.join(item.question + '\n' + c_df.to_string()
                               for (item, c_df) in zip(ids.questionnaire.items,
                                                       r_counts)
                               )
                 )

    # ------------------ STEP 2: Learn ItemResponseModel from data:

    irm = OrdinalItemResponseModel.initialize(ids,
                                              n_traits=max_n_traits,
                                              n_scale_samples=n_scale_samples,
                                              n_subject_samples=n_subject_samples,
                                              trait_scale=3.  # *** = default
                                              )

    logger.info(f'Learning model with max {irm.n_traits} traits'
                 + f' for {irm.n_groups} groups'
                 + f' with {irm.n_subjects} included respondent records in total')
    logger.info(f'The model uses {n_scale_samples} scale samples'
                 + f' and {n_subject_samples} samples for each subject')

    LL = irm.learn(max_hours=5)
    # *** Might take MANY hours, try setting shorter time limit first
    # When learning is complete, it finishes before the time limit.

    logger.info(f'Learned model with {irm.n_groups} groups'
                 + f' with {irm.n_subjects} subjects in total')
    # LL = np.array(LL)
    # logger.info('Log Likelihood: ' + np.array2string(LL, precision=1))

    # ------------------ Optionally, save model for later display generation
    with (save_dir / model_file).open('wb') as f:
        pickle.dump(irm, f)
    logger.info(f'Model saved in {save_dir}')
    logger.info('Learned model = \n' + repr(irm))

    irm.prune()  # keep only traits that were really needed to model the data
    irm.standardize()  # re-scale all parameters for unity trait variance

    # ------------------ STEP 3: Display Model Results:
    # mapping_item used to display trait scales by reference to ONE item rating scale
    ir_display = ItemResponseDisplaySet.show(irm,
                                             mapping_item=-1,  # last questionnaire item
                                             percentiles=[2.5, 50., 97.5],
                                             mpl_params={'figure.max_open_warning': 0,  # suppress warning
                                                         'figure.autolayout': True,  # -> tight layout
                                                         'axes.labelsize': 'x-large'},  # -> matplotlib.rcParam
                                             # mpl_style='my_style_sheet',
                                             # ... any other ema_display.FMT settings
                                             # ... any ema_display_format.FMT settings
                                             )
    ir_display.save(save_dir,
                    figure_format = 'pdf',  # any format that Matplotlib can handle
                    table_format = 'txt',   # any format that Pandas can handle
                    float_format = '%.2f')  # any other parameter(s) for pandas.write_xxx

    # --------------------------- (optionally) save in other format(s), too:
    # ir_display.save(save_dir,
    #          table_format='csv',  # only tables, for input to other package
    #          float_format='%.4f',
    #          )
    # ir_display.save(save_dir,
    #          figure_format='eps',  # only figures, for input to other package
    #          )

    logger.info(f'All results saved in {save_dir}')

    # ------------------ (Optional) PCA of within-population covariance
    if irm.n_traits > 1:
        cov_within = irm.precision_within.mean_inv
        (eig_val, eig_vec) = np.linalg.eigh(cov_within)
        logger.debug('Cov eig_val= ' +
                      np.array2string(eig_val, precision=3))
        logger.debug('Cov eig_vec=\n' +
                      np.array2string(eig_vec, precision=3))
        eig_val_cum = np.cumsum(eig_val[::-1]) / np.sum(eig_val)
        logger.info(f'One, Two, etc Principal Trait Factors: ' +
                     ', '.join(f'{ev:.1%}'
                               for ev in eig_val_cum) +
                     ', of Variance within Populations')

    logging.shutdown()
