"""This script simulates a phoneme-discrimination experiment
mainly to illustrate the preferred format of recorded data
from a similar real experiment.

This simulation generates data files that can be used as input to the main script
run_cm.py, in its template version,
to illustrate the kind of results produced by the analysis model.

*** Demo usage:
Copy this script to a convenient location, edit if needed, and run it.

*** Version history:

* Version 0.8:
2026-02-15, first version of generalized simulator for extended types of forced-choice speech tests.

* Version <= 0.7.1:
Could simulate only some specific examples used for tests
like those presented in the IEEE Trans ASLP paper 2016
"""
import numpy as np

from pathlib import Path
import logging

from ConfMatrixCalc import cm_logging

from ConfMatrixCalc.cm_data import ConfMatrixFrame
from ConfMatrixCalc.cm_simulation import SimPopulation, SimExperiment

from ConfMatrixCalc import __version__

logger = logging.getLogger(__name__)
# logger.setLevel(logging.INFO)


# ------------------------------------------------------------------------
# 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 simulated data files are stored for later analysis
data_dir.mkdir(parents=True, exist_ok=True)
# NOTE: Existing result files in this directory are OVER-WRITTEN WITHOUT WARNING!!!

cm_logging.setup(log_file='run_sim.txt', result_path=conf_dir)

# ---------------------------------------------------------

logger.info(f'*** Simulation using ConfMatrixCalc v. {__version__} ***')

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

# Define groups of stimuli for forced choice response categories 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 missing response
                            test_factors={'SNR': ['low', 'mid', 'high'],
                                          'HA': ['A', 'B'],  # two hearing aids
                                          },
                            group_factors={'Age': ['young', 'old'],
                                           'HL': ['mild', 'moderate']
                                           },
                            stim_group_head='Consonants',  # header of stim_group column in data files
                            stim_head='Stim',  # header of stimulus column in data file(s)
                            resp_head='Resp',  # header of response column in data file(s)
                            )
# print('rvs_index = ', cmf.rvs_index(10))
# print(f'stim_count= ', cmf.stim_count)
# print(f'stim_prob= ', cmf.stim_prob())
# # print(f'correct_ind= ', cmf.correct_response_indices())
# print(f'sr_slices= ', list(cmf.gen_sr_slices()))
# print('')

# 1: ------------ Define populations with different properties:
# Each population is defined by a mixture of Dirichlet distributions,
# for the allowed array of response categories in each stim group.
# Each Dirichlet component is defined by an array of concentration vectors,
# one for each test condition and each (stim_group, stim) category.

# Define effects of test conditions for the simulated populations:
tc_effects = {('low', 'A'): -0.4,
              ('low', 'B'): -0.2,  # B slightly better than A in SNR=low
              # ('mid', 'A'): 0.,   # default: no effect in SNR=mid
              # ('mid', 'B'): 0.,   # default: no effect in SNR=mid
              ('high', 'A'): +0.4,
              ('high', 'B'): +0.2   # A slightly better than B in SNR=high
              }  # used for all populations, in this example
# NOTE: Each test condition MUST be defined as a tuple,
#   with one category for each test "dimension" defined in cmf.test_factors

# Example: three populations specified by given mean prob.correct for each stimulus group:
pop_old_mild = SimPopulation.from_pc({'unvoiced_stop': 0.6,
                                      'voiced_stop': 0.7},  # mean prob.correct
                                     conc_scale=10.,    # sum concentration parameters
                                     conc_rel_sd=0.1,   # random variations of conc. params.
                                     test_cond_effects=tc_effects,
                                     cmf=cmf,
                                     n_comp=10, # in Dirichlet mixture model
                                     name='old_mild'    # prefix for participants
                                     )
# See cm_simulation.SimPopulation.from_pc() for detailed doc of arguments.

logger.info('pop_old_mild=\n' + str(pop_old_mild))

pop_old_moderate = SimPopulation.from_pc({'unvoiced_stop': 0.5,
                                      'voiced_stop': 0.5},
                                conc_scale=10.,
                                conc_rel_sd=0.1,
                                test_cond_effects=tc_effects,
                                cmf=cmf,
                                n_comp=10,
                                name='old_mod')
logger.info('pop_old_moderate=\n' + str(pop_old_moderate))

pop_young = SimPopulation.from_pc({'unvoiced_stop': 0.8,
                                      'voiced_stop': 0.8},
                                conc_scale=10.,
                                conc_rel_sd=0.1,
                                test_cond_effects=tc_effects,
                                cmf=cmf,
                                n_comp=10,
                                name='young')  # only ONE population, with mild HL
logger.info('pop_young=\n' + str(pop_young))

# Show true population properties in each test conditions
#   to allow later comparison with results estimated by run_cm.py:
quantiles = [0.025, 0.25, 0.5, 0.75, 0.975]
for pop in [pop_old_mild, pop_old_moderate, pop_young]:
    (pc, mi) = pop.pc_mi_quantiles(quantiles)  # two pandas.DataFrame instances
    logger.info(f'Population {repr(pop.name)}; true pc=\n' + pc.to_string(float_format='%.1f'))
    logger.info(f'Population {repr(pop.name)}; true mi=\n' + mi.to_string(float_format='%.3f'))
    pc.save(conf_dir / f'true_pc_{pop.name}.csv', float_format='%.2f')
    mi.save(conf_dir / f'true_mi_{pop.name}.csv', float_format='%.4f')

# 2: ------------ Generate groups of simulated "participants" from each population:

groups = {('old','mild'): pop_old_mild.random_group(n_participants=10),
          ('old','moderate'): pop_old_moderate.random_group(n_participants=20),
          ('young', 'mild'): pop_young.random_group(n_participants=10)}
#   no group from possible population ('young', 'moderate')


# 3: ------------ Define experiment with the participant groups:
cm_exp = SimExperiment(cmf, groups)

# 4: ------------ Generate a data set from the complete experiment:
sr_set = cm_exp.gen_dataset(min_pres=5, # n of presentations per stim category
                            # max_pres=3,  # > min_pres, if random stim items
                            # random_stim=False
                            )

# *** (optional) show stim-response counts, for input data check in run_cm.pu
logger.info('Generated ' + str(sr_set))
logger.info('sr_set.counts:\n'
            + '\n'.join('\n'.join(f'group {str(g)}, {repr(sg)}:\n' + c_gsg.to_string()
                                  for (sg, c_gsg) in count_g)
                        for (g, count_g) in sr_set.sum_stim_response_counts()))

# 5: ------------ Export all data to files, as if recorded in a real experiment
sr_set.save(data_dir,   # one sub-dir for each group, one file for each participant
            fmt='xlsx',  # as Excel files, OR
            # fmt='csv',  # default
            # fmt='txt',    # or 'tex' or any other that Pandas can handle
            )
# NOTE: some Pandas file-writer functions may need other modules, manually installed.

logger.info(f'Simulated stim-response data saved to {data_dir}')
# This data set can now be used as input to script run_cm.py
