<gh_stars>1-10
## import numpy as np
## import pyhsmm
##
## from tqdm import tqdm
##
## import ethologger.utils.auxiliary as aux
##
##
## class _MixinModel:
##     default_cfg = {
##         "name": "WeakLimitHDPHSMM",
##         "kwargs": {
##             "alpha_a_0": 1.0,
##             "alpha_b_0": 1.0 / 4,
##             "gamma_a_0": 1.0,
##             "gamma_b_0": 1.0 / 4,
##             "init_state_concentration": 6.0,
##         },
##     }
##
##     def add_data(self, obs_seq):
##         ## assert isinstance(obs_seq, np.ndarray) and len(obs_seq.shape) <= 2
##         mv = True if len(obs_seq.shape) == 2 else False
##
##         if mv:
##             print("Adding multivariate observations...")
##         else:
##             print("Adding univariate observations...")
##
##         self.model.add_data(obs_seq)
##
##     def add_multiple_data(self, obs):
##         ## assert isinstance(obs, np.ndarray) and len(obs.shape) <= 3
##         for i in range(obs.shape[0]):
##             self.add_data(obs[i])
##
##     def resample_model(self, num_resample, save_iter=None):
##         for idx in tqdm(range(num_resample)):
##             if save_iter is None or (idx + 1) % save_iter != 0:
##                 self.model.resample_model()
##             else:
##                 self.models_tosave.append(self.model.resample_and_copy())
##
##     def save_models(self, path):
##         import pickle
##
##         with open(path, "wb") as outfile:
##             pickle.dump(self.models_tosave, outfile, protocol=-1)
##
##     def run_viterbi(self, obs_seq):
##         return self.model.heldout_viterbi(obs_seq)
##
##     def name_check(self):
##         if self.name != "HDPHSMM" and self.name != "WeakLimitHDPHSMM":
##             raise ValueError(
##                 f"Given ssm['name'] {self.name} is not defined. Allowed values are 'HDPHSMM' and 'WeakLimitHDPHSMM'."
##             )
##
##
## class _MixinDistn:
##     @staticmethod
##     def get_distn(name, opt="observation"):
##         if opt == "observation":
##             distn_name_dict = {
##                 "Categorical": pyhsmm.distributions.Categorical,
##                 "CategoricalAndConcentration": pyhsmm.distributions.CategoricalAndConcentration,
##                 "Multinomial": pyhsmm.distributions.Multinomial,
##                 "Poisson": pyhsmm.distributions.Poisson,
##                 "NegativeBinomial": pyhsmm.distributions.NegativeBinomial,
##                 "NegativeBinomialIntegerR": pyhsmm.distributions.NegativeBinomialIntegerR,
##                 "NegativeBinomialFixedR": pyhsmm.distributions.NegativeBinomialFixedR,
##                 "NegativeBinomialIntegerR2": pyhsmm.distributions.NegativeBinomialIntegerR2,
##                 "Geometric": pyhsmm.distributions.Geometric,
##                 ## "Gaussian": pyhsmm.distributions.Gaussian,
##                 ## "ScalarGaussianNonconjNIX": pyhsmm.distributions.ScalarGaussianNonconjNIX,
##                 ## "ScalarGaussianNonconjNIG": pyhsmm.distributions.ScalarGaussianNonconjNIG,
##             }
##         elif opt == "duration":
##             distn_name_dict = {
##                 "Poisson": pyhsmm.distributions.PoissonDuration,
##                 "NegativeBinomial": pyhsmm.distributions.NegativeBinomialDuration,
##                 "NegativeBinomialIntegerR": pyhsmm.distributions.NegativeBinomialIntegerRDuration,
##                 "NegativeBinomialFixedR": pyhsmm.distributions.NegativeBinomialFixedRDuration,
##                 "NegativeBinomialIntegerR2": pyhsmm.distributions.NegativeBinomialIntegerR2Duration,
##                 "Geometric": pyhsmm.distributions.GeometricDuration,
##             }
##         else:
##             raise ValueError(
##                 f"Given opt {opt} is not defined, allowed values are 'observation' and 'duration'."
##             )
##
##         distn = distn_name_dict.get(name, None)
##         if distn is None:
##             raise ValueError(f"Given prob. distn. {name} is not defined for opt {opt}.")
##
##         return distn
##
##     # This will be expanded, better to include.
##     # We can use expressions for cov. matrix etc. in the .yaml config.
##     # Then they can be evaluated here.
##     @staticmethod
##     def get_keyw_args(keyw_args):
##         for key, val in keyw_args.items():
##             if isinstance(val, list):
##                 keyw_args[key] = np.array(val)
##             elif isinstance(val, str):
##                 keyw_args[key] = eval(val)
##             else:
##                 continue
##         return keyw_args
##
##     @classmethod
##     def expand_distn(cls, distns_dict, opt):
##         distns = []
##         for distn_name, distn_kwargs in distns_dict.items():
##             distn = cls.get_distn(distn_name, opt=opt)
##             len_kwargs = map(len, distn_kwargs.values())
##             if aux.is_all_equal(len_kwargs):
##                 kwarg_keys, kwarg_values = zip(*distn_kwargs.items())
##                 for val in zip(*kwarg_values):
##                     kwarg_dict = {}
##                     for i in range(len(kwarg_keys)):
##                         kwarg_dict[kwarg_keys[i]] = val[i]
##                     distns.append(distn(**cls.get_keyw_args(kwarg_dict)))
##             else:
##                 raise ValueError(
##                     f"Inconsistent number of kwargs are given for distribution {distn_name}."
##                 )
##         return distns
##
##
## class _MixinObservation(_MixinDistn):
##     default_obs_cfg = {
##         "Multinomial": {"K": [4, 4, 4, 4, 4, 4], "alpha_0": [1, 1, 1, 1, 1, 1]}
##     }
##
##     def construct_observation_distn(self):
##         self.observation_distns = self.expand_distn(
##             self.observation_dicts, opt="observation"
##         )
##
##
## class _MixinDuration(_MixinDistn):
##     default_dur_cfg = {
##         "Poisson": {
##             "alpha_0": [2 * 30, 2 * 30, 2 * 30, 2 * 30, 2 * 30, 2 * 30],
##             "beta_0": [2, 2, 2, 2, 2, 2],
##         }
##     }
##
##     def construct_duration_distn(self):
##         self.duration_distns = self.expand_distn(self.duration_dicts, opt="duration")
##
##
## class HDPStickyHMM(_MixinModel, _MixinObservation):
##     def __init__(self, main_cfg, distn_cfg):
##         self.cfg = main_cfg.get("ssm", _MixinModel.default_cfg)
##         self.name = self.cfg["name"]
##         self.name_check()
##
##         self.observation_dicts = distn_cfg.get(
##             "observations", _MixinObservation.default_obs_cfg
##         )
##         self.construct_observation_distn()
##
##         self.model_kwargs = self.get_keyw_args(
##             self.cfg.get("kwargs", _MixinModel.default_cfg["kwargs"])
##         )
##
##         self.model = pyhsmm.models.WeakLimitStickyHDPHMM(
##             obs_distns=self.obs_distns, **self.model_kwargs
##         )
##         self.models_tosave = []
##
##
## class HDPHSMM(_MixinModel, _MixinObservation, _MixinDuration):
##     def __init__(self, main_cfg, distn_cfg):
##         self.cfg = main_cfg.get("ssm", _MixinModel.default_cfg)
##         self.name = self.cfg["name"]
##         self.name_check()
##
##         self.observation_dicts = distn_cfg.get(
##             "observations", _MixinObservation.default_obs_cfg
##         )
##         self.duration_dicts = distn_cfg.get("durations", _MixinDuration.default_dur_cfg)
##         self.construct_observation_distn()
##         self.construct_duration_distn()
##
##         if len(self.observation_distns) != len(self.duration_distns):
##             raise ValueError(
##                 "#(observation distributions) != #(duration distributions)"
##             )
##
##         self.model_kwargs = self.get_keyw_args(
##             self.cfg.get("kwargs", _MixinModel.default_cfg["kwargs"])
##         )
##
##         self.model = pyhsmm.models.WeakLimitHDPHSMM(
##             obs_distns=self.observation_distns,
##             dur_distns=self.duration_distns,
##             **self.model_kwargs,
##         )
##         self.models_tosave = []
