"""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: Create an EmaFrame instance to define experiment and select input data.
*2: Load a set of collected EMA data into an EmaDataSet instance
*3: Initialize and train an EmaModel instance corresponding to observed data.
*4: Display results and save figures and tables to a directory tree.


*** Version history:

* 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
"""
import numpy as np
from pathlib import Path
import logging
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__

# -------- __main__ check to prevent multiprocessor sub-tasks to re-run this script
if __name__ == '__main__':
    # ------------------------ Set up working directory and result logging:
    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)

    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: Scenarios, Attributes, and Grades

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

    sim_scenarios = {'Stage': ('',),  # only ONE Test Stage with empty label
                     'HA': ('A', 'B'),  # Two Hearing-aid programs
                     'CoSS': [f'{i}.' for i in range(1, 8)],    # Seven CoSS categories
                     }  # nominal variables, same for all (Sub-)Populations
    # NOTE: First scenario dimension is always Stage, even if only ONE category
    # User may set arbitrary stage_key label
    # Dimension 'Stage' may be omitted, if only one category

    emf = EmaFrame(scenarios=sim_scenarios,
                   stage_key='Stage',
                   attribute_grades={'Speech': ['Very Hard',
                                                'Hard',
                                                'Easy',
                                                'Very Easy',
                                                'Perfect'],
                                     'Comfort': ['Bad',
                                                 'Not Good',
                                                 'Not Bad',
                                                 'Good']},  # ordinal rating categories
                   )

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

    # NOTE: if 'grouping' is defined, group directory names must match given labels,
    #   as exemplified in module ema_data doc-string
    # With template example: data_path / 'Age_old' is top of directory tree with data files.
    # If several 'grouping' factors are defined, the directory three structure must match the factors.
    # If no grouping is defined: data_path is top of directory tree with all data files

    ds = EmaDataSet.load(emf, data_path, fmt='xlsx',
                         # grouping=None,  # default: include all subjects as ONE unnamed group
                         grouping={'Age': ('old',),  # analyze only Age=old
                                   },
                         # grouping={'Age': ('young','old')},  # analyze both Age groups separately
                         subject='sheet',   # subject ID defined in xlsx sheet title
                         ema_vars={'HA': 'B',
                                   'CoSS': 'C',
                                   'Speech': 'F',
                                   'Comfort': 'G'}
                         )
    logger.info(f'Using data ds=\n{ds}')

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

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

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

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

    # NOTE: A regression_effects element may include any combination of scenario dimensions, but
    # including ALL interactions -> many model parameters,
    # possibly -> less precise 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,
                              max_n_comp=10,
                              restrict_attribute=False,  # default
                              restrict_threshold=True,  # default
                              # seed=12345  # ONLY if reproducible results are required
                              )
    # max_n_comp = max number of mixture components in population model
    # restrict_attribute=True -> force sensory-variable locations centered at zero
    # restrict_threshold=True -> force middle response threshold(s) centered at zero
    # for each respondent and each sample of each attribute

    ll = emm.learn(max_hours=2., max_minutes=0.)
    logger.info('ll= ' + np.array2string(np.array(ll), precision=5))
    emm.prune()  # keep only active mixture components.
    # NOTE: If the number of components is not reduced,
    # it might be a good idea to re-run with larger initial max_n_comp,
    # BUT max_n_comp always <= half number of subjects

    # -------- Save learned EmaModel (optional):
    with (work_path / model_file).open('wb') as f:
        pickle.dump(emm, f)

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

    emd = EmaDisplaySet.show(emm,
                             scenarios=['CoSS',  # CoSS profile, aggregated across HA
                                        ('CoSS', 'HA'),  # CoSS profiles, conditional on HA
                                        ('HA', 'CoSS'),  # HA profiles, 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
                             population_individual=True,  # random individual in population
                             population_mean=True,  # population mean
                             subjects=False,  # individual results: True -> MANY plots and tables
                             nap=[('Speech', 'HA'),  # NAP(HA) grades aggregated across CoSS
                                  ('Speech', ('HA', 'CoSS')),  # NAP(HA) separated by CoSS
                                  ('Comfort', 'HA')],
                             nap_confidence=0.95,  # level of symmetric confidence interval
                             grade_thresholds=True,  # 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
                             figure_format='pdf',  # or 'eps', 'jpg', 'png'
                             table_format='latex',  # or 'tab' for tab-limited text
                             )
    # NOTE: joint (=interaction) effects are correct only if included in model regression_effects
    # nap results calculated only for Scenario dimension with exactly TWO categories.

    # ------------------------------- save all result displays
    emd.save(result_path)

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

    logging.shutdown()
