"""Script to display analysis results for ordinal data,
using an OrdinalItemResponseModel object trained and saved earlier,
by script like run_irt.py

This script produces tables and figures
that can easily be imported into a manuscript for publication.

2019-07-21, OK version including standardization result
2019-08-25, adapted for model with automatic selection of traits
2020-06-06, updated for standalone ItemResponseCalc package
"""
import numpy as np
import pickle
from pathlib import Path

import logging
from ItemResponseCalc import ir_logging

from ItemResponseCalc import ir_display as ird


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

TOP_PATH = Path.home() / 'Documents/LeijonKTH/heartech/HA-QualityRegistry'
TOP_PATH = TOP_PATH / 'IRT-IOI-HA'
# change this to top-dir for everything to be read or written here

result_dir = TOP_PATH / 'IOI-HA-results'
# change this to same as in run_irt
# ------------------------------------------------------------------------

run_dir = '2020-05-28-1612'
# change this to the saved result from run_irt

save_dir = result_dir / run_dir

model_file = 'irt-model.pkl'
# = name of saved model generated by run_irt

log_file = 'run_show_log.txt'
# = name of log file saved in same sub-folder
# -------------------------------------------------------------------------

# ------------------------------------- setup logging:
ir_logging.setup(save_dir / log_file)
# -----------------------------------------------------------

logging.info('Showing IRT model results in ' + str(save_dir))

with (save_dir / model_file).open('rb') as f:
    irm = pickle.load(f)

logging.info('Learned Model =\n' + repr(irm))
itm = irm.item_trait_map()

# ------------------------------------- Trait Covariance before standardization
# only informative, may be commented out

c = irm.precision_within.mean_inv
logging.info('Predictive Trait Covariance Within Populations =\n'
             + np.array2string(c, precision=3, suppress_small=True))

c = irm.precision_among.mean_inv
logging.info('Predictive Trait Variance Among Populations= '
             + np.array2string(c, precision=3, suppress_small=True))

# ------------------------------------- Global Predictive Covariance
c = irm.predictive_individual_cov()
logging.info('Predictive Global Trait Cov (before standardization) =\n'
             + np.array2string(c, precision=3, suppress_small=True))

# -------------------------------------- prune and standardize:
irm.prune()  # keep only traits that were really needed to model the data

irm.standardize()  # re-scale all parameters for unity global trait variance
c = irm.predictive_individual_cov()
logging.info('Pruned and Standardized Predictive Global Trait Cov =\n'
             + np.array2string(c, precision=3, suppress_small=True))
s = np.sqrt(np.diag(c))
logging.info('Predictive St.Dev. after Standardization = '
             + np.array2string(s, precision=3))

# ------------------------------------- display with standardized scale
ir_display = ird.ItemResponseDisplaySet.show(irm,  # ***** define trait labels here ? ****
                                             mapping_item=-1,
                                             figure_format='pdf',  # OR eps, png,...
                                             table_format='latex',  # OR tab
                                             )
ir_display.save(save_dir)

# ------------------------------------- Within-group trait correlation
if irm.n_traits > 1:
    cov_within = irm.precision_within.mean_inv
    logging.info('Predictive Trait Covariance Within Populations:\n' +
                 np.array2string(cov_within, precision=3))
    s = np.sqrt(np.diag(cov_within))
    c = cov_within / s
    c /= s[:, np.newaxis]
    # = normalized correlation matrix
    logging.info('Predictive Trait Correlation Within Populations:\n' +
                 np.array2string(c, precision=3))

    # ------------------------------------- PCA of within-group covariance
    (eig_val, eig_vec) = np.linalg.eigh(cov_within)
    logging.debug('Cov eig_val= ' +
                 np.array2string(eig_val, precision=3))
    logging.debug('Cov eig_vec=\n' +
                 np.array2string(eig_vec, precision=3))
    eig_val_cum = np.cumsum(eig_val[::-1]) / np.sum(eig_val)
    logging.info(f'One, Two, etc Principal Trait Factors: ' +
                 ', '.join(f'{ev:.1%}'
                           for ev in eig_val_cum) +
                 ', of Variance within Populations')
    # ***** Check factor analysis with varimax rotation ???

logging.shutdown()
