"""Main script to run Bayesian analysis of data from an Ecological Momentary Assessment (EMA) study.
This script should be used as a template,
to be copied and modified for any desired analysis.

*** Usage: four main steps, see also explicit template example below

*1: Set up an EmaFrame instance to define experiment and to select input data.
*2: Load a set of collected EMA data into an EmaDataSet instance
*3: Initialize and train an EmaModel instance for observed data.
*4: Display results and save figures and tables to a directory tree.

**  Recommended to run a couple of times with the same data set,
    to see random variability in estimated results.


*** Version history:
* Version 1.0.0:
2023-05-17, EmaModel.initialize(..., n_participants_per_comp, ...) defines initial number of GMM components.
            Parameter max_n_comp no longer used.
2023-04-22, adapted to simplified group-key representation in module ema_data.py:
            Group categories specified in EmaFrame.setup(...)
            NOTE: changed signatures EmaFrame.setup(), EmaDataSet.load(), EmaDataSet.save().

* Version 0.9.5:
2023-03-07, ema_display.EmaDisplaySet include observed and model-predicted grade-count histograms
2023-03-01, timestamp_result if needed to avoid over-writing result files

* Version 0.9.3:
2022-08-14, Each ordinal rating scale may be unique or shared by more than one Attribute
2022-07-30, minor cleanup adapting to minor changes in package modules.

* Version 0.9.1:
2022-04-04, all result tables generated and saved as Pandas.DataFrame instances
2022-03-27, NAP and mean-grades results directly from raw data in ema_data.EmaDataSet
2022-03-21, using Pandas DataFrame format in EmaDataSet, allowing many input file formats

* Version 0.7.1:
2022-01-30, allow seed input to EmaModel.initialize for reproducible random results

* Version 0.7: minor update to include new calculation and display functions

* Version 0.6:
2021-12-08, allow user control of model restriction:
            restrict_attribute: sensory-variable location average forced -> 0.
            restrict_threshold: response-threshold median forced -> 0.

* Version 0.5:
2021-11-22, first functional version
"""
# -------- __main__ check to prevent multiprocessor sub-tasks to re-run this script
if __name__ == '__main__':
    import numpy as np
    from pathlib import Path
    import logging
    import datetime as dt
    import pickle

    from EmaCalc.ema_data import EmaFrame, EmaDataSet
    from EmaCalc.ema_model import EmaModel
    from EmaCalc.ema_display import EmaDisplaySet
    from EmaCalc import ema_logging, __version__
    from EmaCalc.ema_display_format import harmonize_ylim

    # ------------------------ Set up working directory and result logging:
    timestamp_result = True  # New result folder for each run, to prevent over-writing
    # timestamp_result = False  # Repeated runs will over-write existing results

    work_path = Path.home() / 'Documents' / 'EMA_sim'  # or whatever...
    data_path = work_path / 'data'  # to use simulation data generated by run_sim.py
    result_path = work_path / 'result'  # or whatever
    model_file = 'test_ema_model.pkl'  # name of saved model file (if saved)

    if timestamp_result:
        t = dt.datetime.now()
        result_path = result_path.with_name(result_path.name +
                                            f'-{t.year}-{t.month:02}-{t.day:02}-{t.hour:02}{t.minute:02}')

    ema_logging.setup(save_path=result_path,
                      log_file='run_ema_log.txt')  # to save the log file
    logger = logging.getLogger(__name__)
    logger.setLevel(logging.INFO)

    logger.info(f'*** Running EmaCalc version {__version__} ***')

    # ------ 1: Define Experimental Framework: Situations, Attributes, and Groups

    # NOTE: This example uses data generated by template script run_sim.py
    # Edit as needed for any other EMA data source

    situations = {'Phase': ('',),  # only One Test phase with empty label
                  'HA': ('A', 'B'),  # Two Hearing-aid programs
                  'CoSS': [f'C{i}' for i in range(1, 8)],  # Seven CoSS categories
                  }  # Categorical variables
    # NOTE: One situation dimension is always phase, even if only ONE phase category.
    # User may set arbitrary phase_key label.
    # Dimension 'Phase' may be omitted, if only one phase category.
    # Category labels may be strings or integers, which can be used as part of file names.
    # If string labels are intended, avoid strings that can be interpreted as numbers,
    # because Pandas will read them as numbers by default.

    emf = EmaFrame.setup(situations=situations,
                         phase_key='Phase',
                         attributes={'Speech': ['Very Hard',
                                                'Hard',
                                                'Easy',
                                                'Very Easy',
                                                'Perfect'],
                                     'Comfort': ['Bad',
                                                 'Not Good',
                                                 'Not Bad',
                                                 'Good']},
                         groups={'Age': ['young', 'old'],
                                 # 'Gender': ('M', 'F'),
                                 }
                         )
    # NOTE: situations, attributes, and groups keys will be parts of result file names,
    # so they can include only characters allowed in file names.

    # NOTE: situations, attributes, and groups are always analyzed as Categorical variables,
    # even if the categories / grades are defined by numeric labels.
    # The rank order of attribute grades is defined by the order given in EmaFrame.setup(...),
    # NOT by, e.g., alphabetical or numerical order of the grade labels.

    # NOTE: if EQUAL grade labels are used for more than one attribute,
    # the model still assumes separate rating scales for each attribute.
    # However, if grade sequences are IDENTICAL objects for more than one attribute,
    # the model uses the SAME rating scale for those attributes.
    #
    # Example:
    # common_ordinal_scale = ['Very Hard',
    #                         'Hard',
    #                         'Easy',
    #                         'Very Easy',
    #                         'Perfect']
    # emf = EmaFrame.setup(situations=situations,
    #                      phase_key='Phase',
    #                      attributes={'Speech': common_ordinal_scale,  # IDENTICAL scale object for both
    #                                  'Comfort': common_ordinal_scale}
    #                      )

    # In either case, response thresholds are always estimated separately for each participant.

    # ------ 2: Load data from previously saved file(s):

    ds = EmaDataSet.load(emf, data_path, fmt='csv', participant='file',
                         path_groups=['Age',  # find 'Age_old' or 'Age_young' in path string
                                      ],
                         # group_join_str='_'  # default in path-string
                         )
    # NOTE: if 'path_groups' elements are specified, group names in path string must match EXACTLY
    # the labels specified in EmaFrame.setup(...,groups={...}), joined by group_join_str,
    # e.g. with directory data_path / 'Age_old'
    # containing data file(s) for group dimension = 'Age', category = 'old'

    logger.info(f'Using data set ds=\n{ds}')

    # ----------------- (Optional) show rating counts for all participants
    for attr in emf.attribute_dtypes.keys():
        a_count = ds.attribute_grade_count(attr, groupby='HA')
        logger.info(str(attr) + '_grade_count:\n' + a_count.to_string())
        logger.info(str(attr) + ': sum grade_count= ' + f'{np.sum(a_count.to_numpy())}')
        # a_count.save(result_path / (str(attr) + '_grades.csv'))

    # ----------------- (Optional) show mean grades and NAP results for all participants
    mean_grades = ds.attribute_grade_mean(groupby=('HA', 'CoSS'))
    mean_grades.save(result_path / 'Attribute_mean_grades.txt', float_format='%.2f')
    # mean_grades.save(result_path / 'Attribute_mean_grades.csv',
    #                  float_format='%.4f')  # if needed for other analysis
    logger.info(f'Attribute mean grades saved in {result_path}'
                + '\nCAUTION: Mean values presume METRIC data, but attribute grades are only ORDINAL!')

    nap = ds.nap_table('HA', nap_cat=['A', 'B'], p=0.95)  # aggregated across other situation dimensions
    # nap = ds.nap_table('HA', nap_cat=['A', 'B'], groupby=('CoSS',), p=0.95)  # grouped results
    nap.save(result_path / 'NAP.txt', float_format='%.2f')  # pretty-formatted by Pandas
    # nap.save(result_path / 'NAP.tex', float_format='%.2f')  # for import to LaTeX doc.
    # nap.save(result_path / 'NAP.csv', sep='\t')  # tab-delimited text, for input to other program
    logger.info(f'NAP results saved in {result_path}')

    # ------ 3: Learn Analysis Model from loaded data set:

    # Model ordinal-regression effects of Situations on each Attribute:

    # regression_effects = ['HA',     # main linear regression effect only
    #                       'CoSS',   # main linear regression effect only
    #                       # 'Phase',  # if there are several phase categories
    #                       ]

    regression_effects = [('HA', 'CoSS')  # joint effects, main AND interaction
                          # 'Phase',  # if there are several phase categories
                          ]

    # NOTE: A regression_effects element may include any combination of situation dimensions, BUT
    # including ALL interactions -> many model parameters,
    # possibly -> less reliable estimation for each parameter.

    # In this example: ['HA', 'CoSS'] -> 2 + (7 - 1) = 8 regression-effect parameters
    #                ['CoSS', 'HA'] -> 7 + (2 - 1) = 8 regression-effect parameters
    #                [('HA', 'CoSS')] -> 2 * 7 = 14 regression-effect parameters

    emm = EmaModel.initialize(ds,
                              effects=regression_effects,
                              # n_participants_per_comp=5,  # default
                              # restrict_attribute=False,  # default
                              # restrict_threshold=True,  # default
                              # seed=12345  # ONLY if reproducible results are required
                              )
    # n_participants_per_comp = initial number of participants per mixture component in population model
    # restrict_attribute=True -> force attribute mean location at zero
    # restrict_threshold=True -> force one mid-scale response threshold at zero
    # for each respondent and each sample of each attribute.

    ll = emm.learn(max_hours=1., max_minutes=0.)
    logger.info(f'*** Data log-likelihood = {ll[-1]:.1f}, indicating model fit to data. ***')
    # *** Recommended to re-run model learning with same data set a couple of times,
    # and then use best-fitting model result.

    emm.prune()  # keep only active mixture components.

    # -------- Save learned EmaModel (optional):
    with (work_path / model_file).open('wb') as f:
        pickle.dump(emm, f)
        logger.info('Model pickle-dumped as ' + f.name)

    # ------ 4: Generate result displays:

    # -------- Re-load learned EmaModel (optional, if saved):
    # with (work_path / model_file).open('rb') as f:
    #     emm = pickle.load(f)
    # -------------------------------------------------
    emd = EmaDisplaySet.show(emm,
                             situations=['CoSS',  # CoSS probabilities, aggregated across HA
                                         ('CoSS', 'HA'),  # CoSS probabilities, conditional on HA
                                         ('HA', 'CoSS'),  # HA probabilities, conditional on CoSS
                                         ],
                             attributes=[('Speech', 'CoSS'),    # Speech, main effect of CoSS
                                         ('Speech', 'HA'),      # Speech, main effect of HA
                                         ('Speech', ('CoSS', 'HA')),    # joint effect of both
                                         ('Comfort', ('CoSS', 'HA'))],  # joint effect of both
                             grade_counts=['Speech',  # total 'Speech' grade-counts, sum across HA and CoSS
                                           ('Speech', 'HA'),    # 'Speech' grade-counts, separated by HA
                                           ('Comfort', 'HA')],
                             random_individual=True,    # random individual in population
                             population_mean=True,      # population mean
                             participants=False,        # individual results: True -> MANY plots and tables
                             grade_thresholds=True,     # True -> median response thresholds in attribute plots
                             percentiles=[2.5, 25, 50, 75, 97.5],   # in profile plots and tables
                             credibility_limit=0.7,     # minimum credibility in difference tables
                             mpl_params={'figure.max_open_warning': 0,
                                         'axes.labelsize': 'x-large'},  # -> matplotlib.rcParam
                             # mpl_style='my_style_sheet',
                             # ... any other ema_display.FMT or ema_display_format.FMT settings
                             )
    # NOTE: joint (=interaction) effects are correct only if included in model regression_effects

    # ------------------------------- (optionally) edit display elements, if desired
    for g_disp in emd.groups.values():
        harmonize_ylim([g_disp.population_mean.attributes[('Speech', ('CoSS', 'HA'))].plot.ax,
                        g_disp.random_individual.attributes[('Speech', ('CoSS', 'HA'))].plot.ax,
                        g_disp.population_mean.attributes[('Comfort', ('CoSS', 'HA'))].plot.ax,
                        g_disp.random_individual.attributes[('Comfort', ('CoSS', 'HA'))].plot.ax
                        ])
        harmonize_ylim([g_disp.population_mean.situations[('CoSS', 'HA')].plot.ax,
                        g_disp.random_individual.situations[('CoSS', 'HA')].plot.ax,
                        ])
    # -> matching y-axis limits: nice for plots to be shown side by side

    # ------------------------------- save all result displays
    emd.save(result_path, figure_format='pdf', table_format='txt', float_format='%.2f')
    # (optionally) save in other format(s), too:
    # emd.save(result_path,
    #          table_format='csv',  # for input to other package
    #          float_format='%.4f',  # any other parameters for Pandas table-writer function
    #          # sep='\t'  # -> tab-delimited
    #          )

    logging.info(f'All results saved in {result_path}')

    logging.shutdown()
