"""This module is the main template script of the free PyPi package ConfMatrixCalc.

This script can be modified and used to run analysis of phoneme-identification test results.

*** Usage:
Copy this script to a convenient location, edit it as guided by script comments,
and run the edited version.

The data analysis always includes four main steps:
1: Define test stimuli and experimental procedure
2: Load pre-recorded stimulus-response data
3: Learn probabilistic model from observed data
4: Display analysis results.


*** Version History:

* Version 0.8:
2026-02-17, Generalized to handle tests with subsets of stimuli and forced-choice responses.
            Including a limited simulation script to exemplify data format and analysis results.

* Versions <= 0.7.1: As described and evaluated in Leijon et al. (2016).

*** References:

Arne Leijon
A. Leijon, G. E. Henter, and M. Dahlquist. Bayesian analysis of phoneme confusion matrices.
IEEE Trans ASL 24(3):469–482, 2016.
doi: 10.1109/TASLP.2015.2512039

Arne Leijon. ConfMatrixCalc — Bayesian Analysis of Phoneme Confusion Data, 2026.
Unpublished report with all math implementation details for the current package version.
Contact author for a copy.
"""
import numpy as np

from pathlib import Path
import logging
import pickle

from ConfMatrixCalc import cm_logging

from ConfMatrixCalc.cm_data import ConfMatrixFrame, StimRespCountSet

from ConfMatrixCalc.cm_model import ConfMatrixModel
from ConfMatrixCalc.cm_display import ConfMatrixDisplaySet

from ConfMatrixCalc import __version__

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------
# Define data and result paths, either absolute or relative:

home_dir = Path.home()
conf_dir = home_dir / 'Documents/CM_Test'  # or whatever...

data_dir = conf_dir / 'data'
# = top directory where all data files are to be found
result_dir = conf_dir / 'results'
result_dir.mkdir(parents=True, exist_ok=True)
# = top directory for all result files, with subdirectories created as needed.
# NOTE: Existing result files in this directory are OVER-WRITTEN WITHOUT WARNING!!!

model_file = 'cm_model.pkl'  # for saved learned model

cm_logging.setup(log_file='run_cm_log.txt', result_path=result_dir)
logger.info(f'*** Using ConfMatrixCalc v. {__version__} ***')


# 1: ------------ Define test and experimental layout:

# Template example uses data created by run_sim.py:

# Define groups of stimuli with forced-choice responses within each group:
consonants = {'unvoiced_stop': ['p', 't', 'k'],  # -> 3 x 3 (or 3 x 4) stim-response pairs
              'voiced_stop': ['b', 'd', 'g'],
              }
# NOTE: Each stimulus must be defined by a tuple, e.g., ('voiced_stop', 'b'),
# because in some tests the same stimulus item may occur in more than one forced-response group.

cmf = ConfMatrixFrame.setup(stim_groups=consonants,
                            # no_response='-',   # code for allowed missing response, default=None
                            test_factors={'SNR': ['low', 'mid', 'high'],
                                          'HA': ['A', 'B'],
                                          },    # -> Six test conditions
                            group_factors={'Age': ['young', 'old'],
                                           'HL': ['mild', 'moderate']
                                           },  # -> Four possible populations in "dimensions" 'Age' and 'HL'
                            stim_group_head='Consonants',  # header of stim_group column in input files
                            stim_head='Stim',  # header of stimulus column in input file(s)
                            resp_head='Resp',  # header of response column in input file(s)
                            # stim_count=None,  # default: -> counted from recorded data
                            # stim_count=5,  # user-defined n of presentations, for all stimulus items
                            # stim_count={'unvoiced_cons': {'p': 2, 't':2, 'k':2},
                            #             'voiced_cons': {'b': 1, 'd':1, 'g':1}},
                            population_weights={('old', 'mild'): 1.5,  # 2.
                                                ('old', 'moderate'): 2.,
                                                }  # default = 1. for other populations
                            )
# NOTE: String CASE is always distinctive, i.e., 'Old' and 'old' are DIFFERENT categories.

# NOTE: The analysis model estimates CONDITIONAL probabilities for all response alternatives,
#   given each stimulus category.
#   Therefore, to estimate probability of correct responses, or mutual information,
#   the stimulus probabilities must also be used.
# If stim_count is specified in the ConfMatrixFrame instance,
#   it will be used to calculate stimulus probabilities in all results,
# otherwise the actual recorded stim counts are used.

# NOTE: population_weights are needed only if results are to be shown separately
# for only one or a few Group Dimensions, regardless of categories in other Group Dimensions.
# In this example, we defined population Age = old to be twice as large as Age = young,
# both for HL = mild and HL = moderate.
# Each element of population_weights must specify one category in each predefined group factor.


# 2: ------------ Load stimulus-response results:

# Recorded data may be stored in various table-style file formats that Pandas can handle.
# Table "long format" is preferred,
#   using one table row for each recorded stimulus-response pair,
#   and separate columns defining stimulus and response categories,
#   test-conditions, group factors, and participant ID label.

# Participant ID, group-factor categories, and test-factor categories
# may also be extracted from the path string, if not found in a table column.

# Example: Result files from participants in population Age='old', HL='mild',
# tested in test-conditions SNR='low',
# might be stored in directory with Path "...Age_old / HL_mild / SNR_low ..."

# Each input file may include data from one or more participants,
# in one or more test conditions.
# Test results from the same participant in the same test condition
# may occur in more than one file, e.g., for repeated tests.
# See module cm_data.py for more details.

# NOTE: Table column headers and all category labels must agree with values predefined in cmf.
# Must rename table headers and / or convert cell values, if needed for agreement.

# Example: read_excel converter functions, if needed:
age_groups = {'20-29': 'young',
              '30-39': 'young',
              '40-49': 'young',
              '50-59': 'young',
              '60-69': 'old',
              '70-79': 'old',
              '80-89': 'old'}  # re-classification for analysis, if needed

def convert_age(age_decade):  # not used in the template example
    """
    :param age_decade: string like '30-39' found in input file
    :return: recoded age group for analysis
    """
    return age_groups[age_decade]


def convert_hearing(bpta4):  # not used in the template example
        """Categorize hearing
        :param bpta4: best ear pure-tone 4-freq average HL found in input file
        :return: recoded hearing-loss grouping category for analysis
        """
        if bpta4 <= 20:
            return 'NH'
        elif bpta4 <= 40:
            return 'mild'
        else:
            return 'moderate'

# This template example uses long-format Excel data files created by run_sim:
# Directory data_dir contains one file for each simulated participant.

cc_set = StimRespCountSet.load(cmf,
                               data_dir,  # -> top directory, OR just one large file
                               fmt='xlsx',  # Only Excel files are used
                               # participant='ID',  # column header, if stored in file
                               participant='file',  # file name used as participant ID
                               path_groups=['Age', 'HL'],  # group factors in path string
                               # path_test_factors=['SNR'],  # test factors in path string
                               # path_join_str='_',     # find, e.g., "Age_old" in path string
                               # correct_col='Result',  # boolean 0 / 1, with 1 <-> response = stim
                               # rename_cols = {'TestWordList': cmf.stim_group_head,
                               #                'TestWord': cmf.stim_head,
                               #                'ErroneousResponse': cmf.resp_head,
                               #                'AgeGroup': 'Age',
                               #                'BPTA4':'Hearing'},  # new column heads, if needed
                               count_col='Count',  # column head for summed response count
                               # read_fcn= my_file_reader,    # user-defined, see cm_file,
                               #    for any other special file format that Pandas cannot read.
                               # kwargs to pandas.read_excel:
                               sheet_name=0,  # read only first sheet in each xlsx file
                               # usecols=['ID',
                               #          'SNR',
                               #          'HA',
                               #          'Result',
                               #          'ErroneousResponse',
                               #          'BPTA4',
                               #          'AgeGroup',
                               #          ],  # optional, if files include unnecessary data columns
                               # converters={'AgeGroup': convert_age,
                               #             'BPTA4': convert_hearing},
                               # ... any other kwargs allowed by file reader function
                               )
# NOTE: Some study might encode responses as just Correct = boolean 0 or 1,
# with the actual response category in another column, like 'ErroneousResponse',
# only if Correct = 0, i.e., the response was incorrect.
# The load function can interpret this,
#   if columns correct_col and 'ErroneousResponse' -> cmf.resp_head are specified.

logger.info(str(cc_set))  # to check that data was loaded correctly

# 3: ------------ Learn model, given data:
cm_result = ConfMatrixModel.learn(cc_set,
                                  # seed=12345, # for reproducible results
                                  # sampled_conc=True,   # -> sampled population concentration vectors
                                  # n_samples=1000, # of population concentration vectors
                                  )

# (optional) save learned model:
with (result_dir / model_file).open('wb') as f:
    pickle.dump(cm_result, f)

# 4: ------------ Display model results:

cm_display = ConfMatrixDisplaySet.show(cm_result,
                                       sr_factors=['HA',  # Response-prob main effects in selected dimension
                                                   ],
                                       mi_factors=['SNR', 'HA', # MI main effects in each test dimension
                                                   ('SNR', 'HA'),  # joint (interaction) effects of both
                                                   ],
                                       pc_factors=['SNR', 'HA', # PC main effects in each test dimension
                                                   ('SNR', 'HA'),  # joint (interaction) effects of both
                                                   ],
                                       group_factors=['Age',  # main effects of 'Age' factor
                                                      ('HL', 'Age'),  # interaction effects of both factors
                                                      ],
                                       random_individual=True,  # unseen random individual in population
                                       population_mean=False,  # needs sampled population model
                                       participants=False,  # True -> display for each participant, MANY files
                                       participant_group=True,  # -> overview of participant group
                                       percentiles=[2.5, 25, 50, 75, 97.5], # in predictive distributions
                                       credibility_limit=0.7,   # min joint probability for differences
                                       mpl_params={'figure.max_open_warning': 0,  # suppress warning
                                                   'figure.constrained_layout.use': True,  # -> compact figs
                                                   'axes.labelsize': 'x-large',
                                                   # 'axes.facecolor': '0.9',  # -> light gray
                                                   # 'legend.facecolor': 'w',
                                                   }, # any other Matplotlib.rcParams settings
                                       # plot_labels=dict(mi='MI (bits / phoneme)',  # plot y_labels
                                       #                  pc='Prob. Correct (%)',
                                       #                  sr='Resp.Prob.'),
                                       # mpl_style='my_style_sheet',  # use predefined matplotlib style sheet
                                       # colors='rbgk',   # default, used cyclically
                                       # markers='oxs*_'  # default, used cyclically
                                       # ... any other parameters in cm_display.FMT or cm_display_format.FMT
                                       )

cm_display.save(result_dir,
                figure_format='pdf',    # or eps, png; any format that matplotlib can handle
                table_format='txt',     # or tex, csv, xlsx, odt; any format that Pandas can handle
                float_format='%.2f')    # for numerical results in tables

# cm_display.save(result_dir,             # save only tables again in other format
#                 table_format='csv',     # to allow further result processing
#                 float_format='%.5f')    # with more precise results

logger.info('*** All results saved to ' + str(result_dir))
logging.shutdown()

