"""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__


# ------------------------ 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()
