from . import nn, rl, util, RaggedArray, ContinuousSpace, FiniteSpace, optim, thutil, gmmil
import numpy as np
from contextlib import contextmanager
import theano; from theano import tensor

from scipy.optimize import fmin_l_bfgs_b

class BehavioralCloningOptimizer(object):
    def __init__(self, mdp, policy, lr, batch_size, obsfeat_fn, ex_obs, ex_a, eval_sim_cfg, eval_freq, train_frac):
        self.mdp, self.policy, self.lr, self.batch_size, self.obsfeat_fn = mdp, policy, lr, batch_size, obsfeat_fn

        # Randomly split data into train/val
        assert ex_obs.shape[0] == ex_a.shape[0]
        num_examples = ex_obs.shape[0]
        num_train = int(train_frac * num_examples)
        shuffled_inds = np.random.permutation(num_examples)
        train_inds, val_inds = shuffled_inds[:num_train], shuffled_inds[num_train:]
        assert len(train_inds) >= 1 and len(val_inds) >= 1
        print '{} training examples and {} validation examples'.format(len(train_inds), len(val_inds))
        self.train_ex_obsfeat, self.train_ex_a = self.obsfeat_fn(ex_obs[train_inds]), ex_a[train_inds]
        self.val_ex_obsfeat, self.val_ex_a = self.obsfeat_fn(ex_obs[val_inds]), ex_a[val_inds]

        self.eval_sim_cfg = eval_sim_cfg
        self.eval_freq = eval_freq

        self.total_time = 0.
        self.curr_iter = 0

    def step(self):
        with util.Timer() as t_all:
            # Subsample expert transitions for SGD
            inds = np.random.choice(self.train_ex_obsfeat.shape[0], size=self.batch_size)
            batch_obsfeat_B_Do = self.train_ex_obsfeat[inds,:]
            batch_a_B_Da = self.train_ex_a[inds,:]
            # Take step
            loss = self.policy.step_bclone(batch_obsfeat_B_Do, batch_a_B_Da, self.lr)

            # Roll out trajectories when it's time to evaluate our policy
            val_loss = val_acc = trueret = avgr = ent = np.nan
            avglen = -1
            if self.eval_freq != 0 and self.curr_iter % self.eval_freq == 0:
                val_loss = self.policy.compute_bclone_loss(self.val_ex_obsfeat, self.val_ex_a)
                # Evaluate validation accuracy (independent of standard deviation)
                if isinstance(self.mdp.action_space, ContinuousSpace):
                    val_acc = -np.square(self.policy.compute_actiondist_mean(self.val_ex_obsfeat) - self.val_ex_a).sum(axis=1).mean()
                else:
                    assert self.val_ex_a.shape[1] == 1
                    # val_acc = (self.policy.sample_actions(self.val_ex_obsfeat)[1].argmax(axis=1) == self.val_ex_a[1]).mean()
                    val_acc = -val_loss # val accuracy doesn't seem too meaningful so just use this

        # Log
        self.total_time += t_all.dt
        fields = [
            ('iter', self.curr_iter, int),
            ('bcloss', loss, float), # supervised learning loss
            ('valloss', val_loss, float), # loss on validation set
            ('valacc', val_acc, float), # loss on validation set
            ('trueret', trueret, float), # true average return for this batch of trajectories
            ('avgr', avgr, float), # average reward encountered
            ('avglen', avglen, int), # average traj length
            ('ent', ent, float), # entropy of action distributions
            ('ttotal', self.total_time, float), # total time
        ]
        self.curr_iter += 1
        return fields

class TransitionClassifier(nn.Model):
    '''Reward/adversary for generative-adversarial training'''

    def __init__(self, obsfeat_space, action_space, hidden_spec, max_kl, adam_lr, adam_steps, ent_reg_weight, enable_inputnorm, include_time, time_scale, favor_zero_expert_reward, varscope_name):
        self.obsfeat_space, self.action_space = obsfeat_space, action_space
        self.hidden_spec = hidden_spec
        self.max_kl = max_kl
        self.adam_steps = adam_steps
        self.ent_reg_weight = ent_reg_weight; assert ent_reg_weight >= 0
        self.include_time = include_time
        self.time_scale = time_scale
        self.favor_zero_expert_reward = favor_zero_expert_reward

        with nn.variable_scope(varscope_name) as self.__varscope:
            # Map (s,a) pairs to classifier scores (log probabilities of classes)
            obsfeat_B_Df = tensor.matrix(name='obsfeat_B_Df')
            a_B_Da = tensor.matrix(name='a_B_Da', dtype=theano.config.floatX if self.action_space.storage_type == float else 'int64')
            t_B = tensor.vector(name='t_B')

            scaled_t_B = self.time_scale * t_B

            if isinstance(self.action_space, ContinuousSpace):
                # For a continuous action space, map observation-action pairs to a real number (reward)
                trans_B_Doa = tensor.concatenate([obsfeat_B_Df, a_B_Da], axis=1)
                trans_dim = self.obsfeat_space.dim + self.action_space.dim
                # Normalize
                with nn.variable_scope('inputnorm'):
                    self.inputnorm = (nn.Standardizer if enable_inputnorm else nn.NoOpStandardizer)(self.obsfeat_space.dim + self.action_space.dim)
                normedtrans_B_Doa = self.inputnorm.standardize_expr(trans_B_Doa)
                if self.include_time:
                    net_input = tensor.concatenate([normedtrans_B_Doa, scaled_t_B[:,None]], axis=1)
                    net_input_dim = trans_dim + 1
                else:
                    net_input = normedtrans_B_Doa
                    net_input_dim = trans_dim
                # Compute scores
                with nn.variable_scope('hidden'):
                    net = nn.FeedforwardNet(net_input, (net_input_dim,), self.hidden_spec)
                with nn.variable_scope('out'):
                    out_layer = nn.AffineLayer(net.output, net.output_shape, (1,), initializer=np.zeros((net.output_shape[0], 1)))
                scores_B = out_layer.output[:,0]

            else:
                # For a finite action space, map observation observations to a vector of rewards

                # Normalize observations
                with nn.variable_scope('inputnorm'):
                    self.inputnorm = (nn.Standardizer if enable_inputnorm else nn.NoOpStandardizer)(self.obsfeat_space.dim)
                normedobs_B_Df = self.inputnorm.standardize_expr(obsfeat_B_Df)
                if self.include_time:
                    net_input = tensor.concatenate([normedobs_B_Df, scaled_t_B[:,None]], axis=1)
                    net_input_dim = self.obsfeat_space.dim + 1
                else:
                    net_input = normedobs_B_Df
                    net_input_dim = self.obsfeat_space.dim
                # Compute scores
                with nn.variable_scope('hidden'):
                    net = nn.FeedforwardNet(net_input, (net_input_dim,), self.hidden_spec)
                with nn.variable_scope('out'):
                    out_layer = nn.AffineLayer(
                        net.output, net.output_shape, (self.action_space.size,),
                        initializer=np.zeros((net.output_shape[0], self.action_space.size)))
                scores_B = out_layer.output[tensor.arange(normedobs_B_Df.shape[0]), a_B_Da[:,0]]

        if self.include_time:
            self._compute_scores = thutil.function([obsfeat_B_Df, a_B_Da, t_B], scores_B) # scores define the conditional distribution p(label | (state,action))
        else:
            compute_scores_without_time = thutil.function([obsfeat_B_Df, a_B_Da], scores_B)
            self._compute_scores = lambda _obsfeat_B_Df, _a_B_Da, _t_B: compute_scores_without_time(_obsfeat_B_Df, _a_B_Da)

        if self.favor_zero_expert_reward:
            # 0 for expert-like states, goes to -inf for non-expert-like states
            # compatible with envs with traj cutoffs for good (expert-like) behavior
            # e.g. mountain car, which gets cut off when the car reaches the destination
            rewards_B = thutil.logsigmoid(scores_B)
        else:
            # 0 for non-expert-like states, goes to +inf for expert-like states
            # compatible with envs with traj cutoffs for bad (non-expert-like) behavior
            # e.g. walking simulations that get cut off when the robot falls over
            rewards_B = -tensor.log(1.-tensor.nnet.sigmoid(scores_B))
        if self.include_time:
            self._compute_reward = thutil.function([obsfeat_B_Df, a_B_Da, t_B], rewards_B)
        else:
            compute_reward_without_time = thutil.function([obsfeat_B_Df, a_B_Da], rewards_B)
            self._compute_reward = lambda _obsfeat_B_Df, _a_B_Da, _t_B: compute_reward_without_time(_obsfeat_B_Df, _a_B_Da)

        param_vars = self.get_trainable_variables()

        # Logistic regression loss, regularized by negative entropy
        labels_B = tensor.vector(name='labels_B')
        weights_B = tensor.vector(name='weights_B')
        losses_B = thutil.sigmoid_cross_entropy_with_logits(scores_B, labels_B)
        ent_B = thutil.logit_bernoulli_entropy(scores_B)
        loss = ((losses_B - self.ent_reg_weight*ent_B)*weights_B).sum(axis=0)
        lossgrad_P = thutil.flatgrad(loss, param_vars)

        if self.include_time:
            self._adamstep = thutil.function(
                [obsfeat_B_Df, a_B_Da, t_B, labels_B, weights_B], loss,
                updates=thutil.adam(loss, param_vars, lr=adam_lr))
        else:
            adamstep_without_time = thutil.function(
                [obsfeat_B_Df, a_B_Da, labels_B, weights_B], loss,
                updates=thutil.adam(loss, param_vars, lr=adam_lr))
            self._adamstep = lambda _obsfeat_B_Df, _a_B_Da, _t_B, _labels_B, _weights_B: adamstep_without_time(_obsfeat_B_Df, _a_B_Da, _labels_B, _weights_B)

    @property
    def varscope(self): return self.__varscope

    def compute_reward(self, obsfeat_B_Df, a_B_Da, t_B):
        return self._compute_reward(obsfeat_B_Df, a_B_Da, t_B)

    def fit(self, obsfeat_B_Df, a_B_Da, t_B, exobs_Bex_Do, exa_Bex_Da, ext_Bex):
        # Transitions from the current policy go first, then transitions from the expert
        obsfeat_Ball_Df = np.concatenate([obsfeat_B_Df, exobs_Bex_Do])
        a_Ball_Da = np.concatenate([a_B_Da, exa_Bex_Da])
        t_Ball = np.concatenate([t_B, ext_Bex])

        # Update normalization
        self.update_inputnorm(obsfeat_Ball_Df, a_Ball_Da)

        B = obsfeat_B_Df.shape[0] # number of examples from the current policy
        Ball = obsfeat_Ball_Df.shape[0] # Ball - b = num examples from expert

        # Label expert as 1, current policy as 0
        labels_Ball = np.zeros(Ball)
        labels_Ball[B:] = 1.

        # Evenly weight the loss terms for the expert and the current policy
        weights_Ball = np.zeros(Ball)
        weights_Ball[:B] = 1./B
        weights_Ball[B:] = 1./(Ball - B); assert len(weights_Ball[B:]) == Ball-B

        # Optimize
        for _ in range(self.adam_steps):
            loss, kl, num_bt_steps = self._adamstep(obsfeat_Ball_Df, a_Ball_Da, t_Ball, labels_Ball, weights_Ball), None, 0

        # Evaluate
        scores_Ball = self._compute_scores(obsfeat_Ball_Df, a_Ball_Da, t_Ball); assert scores_Ball.shape == (Ball,)
        accuracy = .5 * (weights_Ball * ((scores_Ball < 0) == (labels_Ball == 0))).sum()
        accuracy_for_currpolicy = (scores_Ball[:B] <= 0).mean()
        accuracy_for_expert = (scores_Ball[B:] > 0).mean()
        assert np.allclose(accuracy, .5*(accuracy_for_currpolicy + accuracy_for_expert))

        return [
            ('rloss', loss, float), # reward function fitting loss
            ('racc', accuracy, float), # reward function accuracy
            ('raccpi', accuracy_for_currpolicy, float), # reward function accuracy
            ('raccex', accuracy_for_expert, float), # reward function accuracy
            ('rkl', kl, float),
            ('rbt', num_bt_steps, int),
            # ('rpnorm', util.maxnorm(self.get_params()), float),
            # ('snorm', util.maxnorm(scores_Ball), float),
        ]

    def update_inputnorm(self, obs_B_Do, a_B_Da):
        if isinstance(self.action_space, ContinuousSpace):
            self.inputnorm.update(np.concatenate([obs_B_Do, a_B_Da], axis=1))
        else:
            self.inputnorm.update(obs_B_Do)

    def plot(self, ax, idx1, idx2, range1, range2, n=100):
        assert len(range1) == len(range2) == 2 and idx1 != idx2
        x, y = np.mgrid[range1[0]:range1[1]:(n+0j), range2[0]:range2[1]:(n+0j)]

        if isinstance(self.action_space, ContinuousSpace):
            points_B_Doa = np.zeros((n*n, self.obsfeat_space.storage_size + self.action_space.storage_size))
            points_B_Doa[:,idx1] = x.ravel()
            points_B_Doa[:,idx2] = y.ravel()
            obsfeat_B_Df, a_B_Da = points_B_Doa[:,:self.obsfeat_space.storage_size], points_B_Doa[:,self.obsfeat_space.storage_size:]
            assert a_B_Da.shape[1] == self.action_space.storage_size
            t_B = np.zeros(a_B_Da.shape[0]) # XXX make customizable
            z = self.compute_reward(obsfeat_B_Df, a_B_Da, t_B).reshape(x.shape)
        else:
            obsfeat_B_Df = np.zeros((n*n, self.obsfeat_space.storage_size))
            obsfeat_B_Df[:,idx1] = x.ravel()
            obsfeat_B_Df[:,idx2] = y.ravel()
            a_B_Da = np.zeros((obsfeat_B_Df.shape[0], 1), dtype=np.int32) # XXX make customizable
            t_B = np.zeros(a_B_Da.shape[0]) # XXX make customizable
            z = self.compute_reward(obsfeat_B_Df, a_B_Da, t_B).reshape(x.shape)

        ax.pcolormesh(x, y, z, cmap='viridis')
        ax.contour(x, y, z, levels=np.log(np.linspace(2., 3., 10)))
        # ax.contourf(x, y, z, levels=[np.log(2.), np.log(2.)+.5], alpha=.5) # high-reward region is highlighted

class LinearReward(object):
    # things to keep in mind
    # - continuous vs discrete actions
    # - simplex or l2 ball
    # - input norm
    # - shifting so that 0 == expert or 0 == non-expert

    def __init__(self,
            obsfeat_space, action_space,
            mode, enable_inputnorm, favor_zero_expert_reward,
            include_time,
            time_scale,
            exobs_Bex_Do, exa_Bex_Da, ext_Bex,
            sqscale=.01,
            quadratic_features=False):

        self.obsfeat_space, self.action_space = obsfeat_space, action_space
        assert mode in ['l2ball', 'simplex']
        print 'Linear reward function type: {}'.format(mode)
        self.simplex = mode == 'simplex'
        self.favor_zero_expert_reward = favor_zero_expert_reward
        self.include_time = include_time
        self.time_scale = time_scale
        self.sqscale = sqscale
        self.quadratic_features = quadratic_features
        self.exobs_Bex_Do, self.exa_Bex_Da, self.ext_Bex = exobs_Bex_Do, exa_Bex_Da, ext_Bex
        with nn.variable_scope('inputnorm'):
            # Standardize both observations and actions if actions are continuous
            # otherwise standardize observations only.
            self.inputnorm = (nn.Standardizer if enable_inputnorm else nn.NoOpStandardizer)(
                (obsfeat_space.dim + action_space.dim) if isinstance(action_space, ContinuousSpace)
                    else obsfeat_space.dim)
            self.inputnorm_updated = False
        self.update_inputnorm(self.exobs_Bex_Do, self.exa_Bex_Da) # pre-standardize with expert data

        # Expert feature expectations
        self.expert_feat_Df = self._compute_featexp(self.exobs_Bex_Do, self.exa_Bex_Da, self.ext_Bex)
        # The current reward function
        feat_dim = self.expert_feat_Df.shape[0]
        print 'Linear reward: {} features'.format(feat_dim)
        if self.simplex:
            # widx is the index of the most discriminative reward function
            self.widx = np.random.randint(feat_dim)
        else:
            # w is a weight vector
            self.w = np.random.randn(feat_dim)
            self.w /= np.linalg.norm(self.w) + 1e-8

        self.reward_bound = 0.

    def _featurize(self, obsfeat_B_Do, a_B_Da, t_B):
        assert self.inputnorm_updated
        assert obsfeat_B_Do.shape[0] == a_B_Da.shape[0] == t_B.shape[0]
        B = obsfeat_B_Do.shape[0]

        # Standardize observations and actions
        if isinstance(self.action_space, ContinuousSpace):
            trans_B_Doa = self.inputnorm.standardize(np.concatenate([obsfeat_B_Do, a_B_Da], axis=1))
            obsfeat_B_Do, a_B_Da = trans_B_Doa[:,:obsfeat_B_Do.shape[1]], trans_B_Doa[:,obsfeat_B_Do.shape[1]:]
            assert obsfeat_B_Do.shape[1] == self.obsfeat_space.dim and a_B_Da.shape[1] == self.action_space.dim
        else:
            assert a_B_Da.shape[1] == 1 and np.allclose(a_B_Da, a_B_Da.astype(int)), 'actions must all be ints'
            obsfeat_B_Do = self.inputnorm.standardize(obsfeat_B_Do)

        # Concatenate with other stuff to get final features
        scaledt_B_1 = t_B[:,None]*self.time_scale
        if isinstance(self.action_space, ContinuousSpace):
            if self.quadratic_features:
                feat_cols = [obsfeat_B_Do, a_B_Da]
                if self.include_time:
                    feat_cols.extend([scaledt_B_1])
                feat = np.concatenate(feat_cols, axis=1)
                quadfeat = (feat[:,:,None] * feat[:,None,:]).reshape((B,-1))
                feat_B_Df = np.concatenate([feat,quadfeat,np.ones((B,1))], axis=1)
            else:
                feat_cols = [obsfeat_B_Do, a_B_Da, (self.sqscale*obsfeat_B_Do)**2, (self.sqscale*a_B_Da)**2]
                if self.include_time:
                    feat_cols.extend([scaledt_B_1, scaledt_B_1**2, scaledt_B_1**3])
                feat_cols.append(np.ones((B,1)))
                feat_B_Df = np.concatenate(feat_cols, axis=1)

        else:
            assert not self.quadratic_features
            # Observation-only features
            obsonly_feat_cols = [obsfeat_B_Do, (.01*obsfeat_B_Do)**2]
            if self.include_time:
                obsonly_feat_cols.extend([scaledt_B_1, scaledt_B_1**2, scaledt_B_1**3])
            obsonly_feat_B_f = np.concatenate(obsonly_feat_cols, axis=1)

            # To get features that include actions, we'll have blocks of obs-only features,
            # one block for each action.
            assert a_B_Da.shape[1] == 1
            action_inds = [np.flatnonzero(a_B_Da[:,0] == a) for a in xrange(self.action_space.size)]
            assert sum(len(inds) for inds in action_inds) == B
            action_block_size = obsonly_feat_B_f.shape[1]
            # Place obs features into their appropriate blocks
            blocked_feat_B_Dfm1 = np.zeros((obsonly_feat_B_f.shape[0], action_block_size*self.action_space.size))
            for a in range(self.action_space.size):
                blocked_feat_B_Dfm1[action_inds[a],a*action_block_size:(a+1)*action_block_size] = obsonly_feat_B_f[action_inds[a],:]
            assert np.isfinite(blocked_feat_B_Dfm1).all()
            feat_B_Df = np.concatenate([blocked_feat_B_Dfm1, np.ones((B,1))], axis=1)

        if self.simplex:
            feat_B_Df = np.concatenate([feat_B_Df, -feat_B_Df], axis=1)

        assert feat_B_Df.ndim == 2 and feat_B_Df.shape[0] == B
        return feat_B_Df

    def _compute_featexp(self, obsfeat_B_Do, a_B_Da, t_B):
        return self._featurize(obsfeat_B_Do, a_B_Da, t_B).mean(axis=0)

    def fit(self, obsfeat_B_Do, a_B_Da, t_B, _unused_exobs_Bex_Do, _unused_exa_Bex_Da, _unused_ext_Bex):
        # Ignore expert data inputs here, we'll use the one provided in the constructor.

        # Current feature expectations
        curr_feat_Df = self._compute_featexp(obsfeat_B_Do, a_B_Da, t_B)

        # Compute adversary reward
        if self.simplex:
            v = curr_feat_Df - self.expert_feat_Df
            self.widx = np.argmin(v)
            return [('vmin', v.min(), float)]
        else:
            self.w = self.expert_feat_Df - curr_feat_Df
            l2 = np.linalg.norm(self.w)
            self.w /= l2 + 1e-8
            return [('l2', l2, float)]

    def compute_reward(self, obsfeat_B_Do, a_B_Da, t_B):
        feat_B_Df = self._featurize(obsfeat_B_Do, a_B_Da, t_B)
        r_B = (feat_B_Df[:,self.widx] if self.simplex else feat_B_Df.dot(self.w)) / float(feat_B_Df.shape[1])
        assert r_B.shape == (obsfeat_B_Do.shape[0],)

        if self.favor_zero_expert_reward:
            self.reward_bound = max(self.reward_bound, r_B.max())
        else:
            self.reward_bound = min(self.reward_bound, r_B.min())
        shifted_r_B = r_B - self.reward_bound
        if self.favor_zero_expert_reward:
            assert (shifted_r_B <= 0).all()
        else:
            assert (shifted_r_B >= 0).all()

        return shifted_r_B

    def update_inputnorm(self, obs_B_Do, a_B_Da):
        if isinstance(self.action_space, ContinuousSpace):
            self.inputnorm.update(np.concatenate([obs_B_Do, a_B_Da], axis=1))
        else:
            self.inputnorm.update(obs_B_Do)
        self.inputnorm_updated = True

import pickle

class ImitationOptimizer(object):
    def __init__(self, mdp, discount, lam, policy, sim_cfg, step_func, reward_func, value_func,
                 policy_obsfeat_fn, reward_obsfeat_fn, policy_ent_reg, ex_obs, ex_a, ex_t):
        self.mdp, self.discount, self.lam, self.policy = mdp, discount, lam, policy
        self.sim_cfg = sim_cfg
        self.step_func = step_func
        self.reward_func = reward_func
        self.value_func = value_func
        # assert value_func is not None, 'not tested'
        self.policy_obsfeat_fn = policy_obsfeat_fn
        self.reward_obsfeat_fn = reward_obsfeat_fn
        self.policy_ent_reg = policy_ent_reg
        util.header('Policy entropy regularization: {}'.format(self.policy_ent_reg))

        assert ex_obs.ndim == ex_a.ndim == 2 and ex_t.ndim == 1 and ex_obs.shape[0] == ex_a.shape[0] == ex_t.shape[0]
        self.ex_pobsfeat, self.ex_robsfeat, self.ex_a, self.ex_t = policy_obsfeat_fn(ex_obs), reward_obsfeat_fn(ex_obs), ex_a, ex_t

        self.total_num_trajs = 0
        self.total_num_sa = 0
        self.total_time = 0.
        self.curr_iter = 0
        self.last_sampbatch = None # for outside access for debugging

    def step(self):
        with util.Timer() as t_all:
            # Sample trajectories using current policy
            # print 'Sampling'
            with util.Timer() as t_sample:
                sampbatch = self.mdp.sim_mp(
                    policy_fn=lambda obsfeat_B_Df: self.policy.sample_actions(obsfeat_B_Df),
                    obsfeat_fn=self.policy_obsfeat_fn,
                    cfg=self.sim_cfg)
                samp_pobsfeat = sampbatch.obsfeat
                self.last_sampbatch = sampbatch
                # if self.curr_iter % 50 == 0:
                #     print('Save batch (s,a) pair information into pickle file sa_info.pk ...')
                #     with open('sa_info.pk', 'wb') as saf:
                #         pickle.dump(samp_pobsfeat, saf)
                #print len(sampbatch.trajs)
                #print sampbatch.trajs[0].obs_T_Do

            # Compute baseline / advantages
            # print 'Computing advantages'

            with util.Timer() as t_adv:
                # Compute observation features for reward input
                samp_robsfeat_stacked = self.reward_obsfeat_fn(sampbatch.obs.stacked)
                # Reward is computed wrt current reward function
                # TODO: normalize rewards
                rcurr_stacked = self.reward_func.compute_reward(samp_robsfeat_stacked, sampbatch.a.stacked, sampbatch.time.stacked)
                assert rcurr_stacked.shape == (samp_robsfeat_stacked.shape[0],)

                # If we're regularizing the policy, add negative log probabilities to the rewards
                # Intuitively, the policy gets a bonus for being less certain of its actions
                orig_rcurr_stacked = rcurr_stacked.copy()
                if self.policy_ent_reg is not None and self.policy_ent_reg != 0:
                    assert self.policy_ent_reg > 0
                    # XXX probably faster to compute this from sampbatch.adist instead
                    actionlogprobs_B = self.policy.compute_action_logprobs(samp_pobsfeat.stacked, sampbatch.a.stacked)
                    policyentbonus_B = -self.policy_ent_reg * actionlogprobs_B
                    rcurr_stacked += policyentbonus_B
                else:
                    policyentbonus_B = np.zeros_like(rcurr_stacked)

                rcurr = RaggedArray(rcurr_stacked, lengths=sampbatch.r.lengths)

                # Compute advantages using these rewards
                advantages, qvals, vfunc_r2, simplev_r2 = rl.compute_advantage(
                    rcurr, samp_pobsfeat, sampbatch.time, self.value_func, self.discount, self.lam)

            # Take a step
            # print 'Fitting policy'
            with util.Timer() as t_step:

                params0_P = self.policy.get_params()
                step_print = self.step_func(
                    self.policy, params0_P,
                    samp_pobsfeat.stacked, sampbatch.a.stacked, sampbatch.adist.stacked,
                    advantages.stacked)

                self.policy.update_obsnorm(samp_pobsfeat.stacked)

            # Fit reward function
            # print 'Fitting reward'
            with util.Timer() as t_r_fit:
                if True:#self.curr_iter % 20 == 0:
                    # Subsample expert transitions to the same sample count for the policy
                    inds = np.random.choice(self.ex_robsfeat.shape[0], size=samp_pobsfeat.stacked.shape[0])
                    exbatch_robsfeat = self.ex_robsfeat[inds,:]
                    exbatch_pobsfeat = self.ex_pobsfeat[inds,:] # only used for logging
                    exbatch_a = self.ex_a[inds,:]
                    exbatch_t = self.ex_t[inds]
                    rfit_print = self.reward_func.fit(samp_robsfeat_stacked, sampbatch.a.stacked, sampbatch.time.stacked, exbatch_robsfeat, exbatch_a, exbatch_t)
                else:
                    rfit_print = []

            # Fit value function for next iteration
            # print 'Fitting value function'
            with util.Timer() as t_vf_fit:
                if self.value_func is not None:
                    # Recompute q vals # XXX: this is only necessary if fitting reward after policy
                    # qnew = qvals

                    # TODO: this should be a byproduct of reward fitting

                    if isinstance(self.reward_func, gmmil.MMDReward):
                        rnew = RaggedArray(self.reward_func.current_reward, lengths=sampbatch.r.lengths)
                    else:
                        rnew = RaggedArray(
                            self.reward_func.compute_reward(samp_robsfeat_stacked, sampbatch.a.stacked, sampbatch.time.stacked),
                            lengths=sampbatch.r.lengths)

                    qnew, _ = rl.compute_qvals(rnew, self.discount)
                    vfit_print = self.value_func.fit(samp_pobsfeat.stacked, sampbatch.time.stacked, qnew.stacked)
                else:
                    vfit_print = []

        # Log
        self.total_num_trajs += len(sampbatch)
        self.total_num_sa += sum(len(traj) for traj in sampbatch)
        self.total_time += t_all.dt
        fields = [
            ('iter', self.curr_iter, int),
            ('trueret', sampbatch.r.padded(fill=0.).sum(axis=1).mean(), float), # average return for this batch of trajectories
            ('iret', rcurr.padded(fill=0.).sum(axis=1).mean(), float), # average return on imitation reward
            ('avglen', int(np.mean([len(traj) for traj in sampbatch])), int), # average traj length
            ('ntrajs', self.total_num_trajs, int), # total number of trajs sampled over the course of training
            ('nsa', self.total_num_sa, int), # total number of state-action pairs sampled over the course of training
            ('ent', self.policy._compute_actiondist_entropy(sampbatch.adist.stacked).mean(), float), # entropy of action distributions
            ('vf_r2', vfunc_r2, float),
            ('tdvf_r2', simplev_r2, float),
            ('dx', util.maxnorm(params0_P - self.policy.get_params()), float), # max parameter difference from last iteration
        ] + step_print + vfit_print + rfit_print + [
            ('avgr', rcurr_stacked.mean(), float), # average regularized reward encountered
            ('avgunregr', orig_rcurr_stacked.mean(), float), # average unregularized reward
            ('avgpreg', policyentbonus_B.mean(), float), # average policy regularization
            # ('bcloss', -self.policy.compute_action_logprobs(exbatch_pobsfeat, exbatch_a).mean(), float), # negative log likelihood of expert actions
            # ('bcloss', np.square(self.policy.compute_actiondist_mean(exbatch_pobsfeat) - exbatch_a).sum(axis=1).mean(axis=0), float),
            ('tsamp', t_sample.dt, float), # time for sampling
            ('tadv', t_adv.dt + t_vf_fit.dt, float), # time for advantage computation
            ('tstep', t_step.dt, float), # time for step computation
            ('ttotal', self.total_time, float), # total time
        ]
        self.curr_iter += 1
        return fields

import numpy,from scipy,import theano,from theano


#! /usr/bin/python3
# -*- coding: utf-8 -*-

"""
    ********
    Crosslet
    ********

    :class:`~nenupy.crosslet.crosslet.Crosslet` is the main class
    for both :class:`~nenupy.crosslet.xstdata.XST_Data` and
    :class:`~nenupy.crosslet.tvdata.TV_Data`, which inherit from
    it.

    This enables *beamforming* with the
    :meth:`~nenupy.crosslet.crosslet.Crosslet.beamform` method
    (see also :ref:`tuto_beamforming` for a detailed tutorial)
    and *imaging* with the
    :meth:`~nenupy.crosslet.crosslet.Crosslet.image` method
    (see also :ref:`tuto_tv` for a detailed tutorial) from
    cross-correlation statistics data.
"""

__author__ = 'Alan Loh'
__copyright__ = 'Copyright 2020, nenupy'
__credits__ = ['Alan Loh']
__maintainer__ = 'Alan'
__email__ = 'alan.loh@obspm.fr'
__status__ = 'Production'
__all__ = [
    'Crosslet'
]

import numpy as np
import astropy.units as un
try:
    from tqdm import tqdm
except ModuleNotFoundError:
    def tqdm(func):
        """ Overdefine tqdm to return `range` or `enumerate`
        """
        return func

from nenupy.astro import wavelength, HpxSky, eq_zenith
from nenupy.instru import nenufar_loc, read_cal_table, ma_pos
from nenupy.crosslet import UVW
from nenupy.beamlet.sdata import SData

import logging
log = logging.getLogger(__name__)

# ============================================================= #
# ------------------------- Functions ------------------------- #
# ============================================================= #
try:
    from numba import jit,prange
except ModuleNotFoundError:
    from functools import wraps
    def jit(**kwargs):
        """ Override numba.jit function by defining a decorator
            which does nothing but returning the decorated
            function.
        """
        def inner_function(function):
            @wraps(function)
            def wrapper(*args, **kwargs):
                function(*args, **kwargs)
            return wrapper
        return inner_function
    log.warning(
        'numba module not found, operations will take longer.'
    )

@jit(nopython=True, parallel=True, fastmath=True)
def ft_mul(x, y):
    return x * y

@jit(nopython=True, parallel=True, fastmath=True)
def ft_phase(ul, vm):
    return np.exp(
        - 2.j * np.pi * (ul + vm)
    )

@jit(nopython=True, parallel=True, fastmath=True)
def ft_sum(vis, exptf):
    return np.mean(
        vis * exptf,
    )
# ============================================================= #

# ============================================================= #
# ------------------------- Crosslet -------------------------- #
# ============================================================= #
class Crosslet(object):
    """ :class:`~nenupy.crosslet.crosslet.Crosslet` class is not
        designed to be called directly but rather as a base class
        for both :class:`~nenupy.crosslet.xstdata.XST_Data` and
        :class:`~nenupy.crosslet.tvdata.TV_Data` (see their
        related documentation for further instructions on how
        to load these data sets).
    """

    def __init__(self):
        self.freqs = None
        self.mas = None
        self.dt = None
        self.times = None
        self.vis = None

    # --------------------------------------------------------- #
    # --------------------- Getter/Setter --------------------- #
    @property
    def mas(self):
        """ Mini-Arrays used to get the cross-correlation
            statistics data.

            :setter: Mini-Array list

            :getter: Mini-Array list

            :type: :class:`~numpy.ndarray`
        """
        return self._mas
    @mas.setter
    def mas(self, m):
        if m is not None:
            self._mas_idx = np.arange(m.size)
            self._ant1, self._ant2 = np.tril_indices(m.size, 0)
        self._mas = m
        return

    @property
    def xx(self):
        """ Extracts the XX polarization from the cross-correlation
            statistics data. Also refered to as ``'NW'``.
            Array is shaped like (time, frequencies, baselines).

            :getter: XX polarization

            :type: :class:`~numpy.ndarray`
        """
        xx = self.vis[:, :, self._get_cross_idx('X', 'X')]
        log.info(
            'XX loaded.'
        )
        return xx

    @property
    def yy(self):
        """ Extracts the YY polarization from the cross-correlation
            statistics data. Also refered to as ``'NE'``.
            Array is shaped like (time, frequencies, baselines).

            :getter: YY polarization

            :type: :class:`~numpy.ndarray`
        """
        yy = self.vis[:, :, self._get_cross_idx('Y', 'Y')]
        log.info(
            'YY loaded.'
        )
        return yy

    @property
    def yx(self):
        """ Extracts the YX polarization from the cross-correlation
            statistics data.
            Array is shaped like (time, frequencies, baselines).

            :getter: YX polarization

            :type: :class:`~numpy.ndarray`
        """
        yx = self.vis[:, :, self._get_cross_idx('Y', 'X')]
        log.info(
            'YX loaded.'
        )
        return yx

    @property
    def xy(self):
        """ Extracts the XY polarization from the cross-correlation
            statistics data. This polarization is not recorded by
            default for the NenuFAR auto-correlations. However
            this can be computed as it is the complex conjugate
            of :attr:`~nenupy.crosslet.crosslet.Crosslet.yx`
            auto-correlations.
            Array is shaped like (time, frequencies, baselines).

            :getter: XY polarization

            :type: :class:`~numpy.ndarray`
        """
        # Deal with lack of auto XY cross in XST-like data
        ma1, ma2 = np.tril_indices(self.mas.size, 0)
        auto = ma1 == ma2
        cross = ~auto
        _xy = np.zeros(
            (list(self.yx.shape[:-1]) + [ma1.size]),
            dtype=np.complex
        )
        _xy[:, :, auto] = self.yx[:, :, auto].conj()
        # Get XY correlations
        indices = self._get_cross_idx('X', 'Y')
        _xy[:, :, cross] = self.vis[:, :, indices]
        log.info(
            'XY loaded.'
        )
        return _xy

    @property
    def stokes_i(self):
        r""" Computes the stokes parameter I from the
            cross-correlation statistics data using
            :attr:`~nenupy.crosslet.crosslet.Crosslet.xx` and
            :attr:`~nenupy.crosslet.crosslet.Crosslet.yy`.

            .. math::
                \mathbf{I}(t, \nu) = \frac{1}{2} \left[
                    \mathbf{XX}(t, \nu) + \mathbf{YY}(t, \nu)
                \right]

            Array is shaped like (time, frequencies, baselines).

            :getter: Stokes I data

            :type: :class:`~numpy.ndarray`
        """
        return 0.5*(self.xx + self.yy)

    @property
    def mask_auto(self):
        """ Masks auto-correlations. In particular, this is used
            to compute the image (see
            :meth:`~nenupy.crosslet.crosslet.Crosslet.image`).

            :getter: Auto-correlation mask

            :type: :class:`~numpy.ndarray`
        """
        ma1, ma2 = np.tril_indices(self.mas.size, 0)
        auto = ma1 == ma2
        return ~auto

    # --------------------------------------------------------- #
    # ------------------------ Methods ------------------------ #
    # def plot(self):
    #     """
    #     """
    #     mat = np.zeros(
    #         (self.mas.size, self.mas.size),
    #         dtype=np.float
    #     )
    #     mat[np.tril_indices(self.mas.size, 0)] = np.abs(self.xy[0, 0, :])
    #     fig = plt.figure(figsize=(10, 10))
    #     plt.imshow(mat, origin='lower')
    #     return

    def beamform(self, az, el, pol='NW', ma=None, calibration='default'):
        r""" Converts cross correlation statistics data XST,
            :math:`\mathbf{X}(t, \nu)`, in beamformed data BST,
            :math:`B(t, \nu)`, where :math:`t` and :math:`\nu` are
            the time and the frequency respectively.
            :math:`\mathbf{X}(t, \nu)` is a subset of XST data at
            the required polarization ``pol``.

            This is done for a given phasing direction in local
            sky coordinates :math:`\varphi` (azimuth, ``az``) and
            :math:`\theta` (elevation, ``el``), with a selection
            of Mini-Arrays ``ma`` (numbered :math:`a`).

            .. math::
                B (t, \nu) = \operatorname{Re} \left\{
                \sum
                    \left[
                        \underset{\scriptscriptstyle a \times 1}{\mathbf{C}}
                        \cdot
                        \underset{\scriptscriptstyle 1 \times a}{\mathbf{C}^{H}}
                    \right](\nu)
                    \cdot
                    \underset{\scriptscriptstyle a \times a}{\mathbf{X}} (t, \nu)
                    \cdot
                    \left[
                        \underset{\scriptscriptstyle a \times 1}{\mathbf{P}}
                        \cdot
                        \underset{\scriptscriptstyle 1 \times a}{\mathbf{P}^{H}}
                    \right](\nu)
                \right\}

            .. math::
                \rm{with} \quad
                \cases{
                    \mathbf{C}(\nu ) = e^{2 \pi i \nu \mathbf{ t }} \quad \rm{the~calibration~ file}\\
                    \mathbf{P} (\nu) = e^{-2 \pi i \frac{\nu}{c} (\mathbf{b} \cdot \mathbf{u})} \quad \rm{phasing}\\
                    \underset{\scriptscriptstyle a \times 3}{\mathbf{b}} = \mathbf{a}_{1} - \mathbf{a}_{2} \quad \rm{baseline~positions}\\
                    \underset{\scriptscriptstyle 3 \times 1}{\mathbf{u}} = \left[
                        \cos(\theta)\cos(\varphi),
                        \cos(\theta)\sin(\varphi),
                        \sin(\theta) \right]
                }

            :param az:
                Azimuth coordinate used for beamforming (default
                unit is degrees in `float` input).
            :type az: `float` or :class:`~astropy.units.Quantity`
            :param el:
                Elevation coordinate used for beamforming (default
                unit is degrees in `float` input).
            :type el: `float` or :class:`~astropy.units.Quantity`
            :param pol:
                Polarization (either ``'NW'`` or ``'NE'``.
            :type pol: `str`
            :param ma:
                Subset of Mini-Arrays (minimum 2) used for
                beamforming.
            :type ma: `list` or :class:`~numpy.ndarray`
            :param calibration:
                Antenna delay calibration file (i.e, :math:`\mathbf{C}`).
                If ``'none'``, no calibration is applied.
                If ``'default'``, the standard calibration file is
                used, otherwise the calibration file name should
                be given (see also :func:`~nenupy.instru.instru.read_cal_table`).
            :type calibration: `str`

            :returns: Beamformed data.
            :rtype: :class:`~nenupy.beamlet.sdata.SData`

            :Example:
                >>> from nenupy.crosslet import XST_Data
                >>> xst = XST_Data('20191129_141900_XST.fits')
                >>> bf = xst.beamform(
                        az=180,
                        el=90,
                        pol='NW',
                        ma=[17, 44],
                        calibration='default'
                    )
        """
        log.info(
            'Beamforming towards az={}, el={}, pol={}'.format(
                az,
                el,
                pol
            )
        )
        # Mini-Array selection
        if ma is None:
            ma = self.mas.copy()
        mas = self._mas_idx[np.isin(self.mas, ma)]
        # Calibration table
        if calibration.lower() == 'none':
            # No calibration
            cal = np.ones(
                (self.sb_idx.size, mas.size)
            )
        else:
            pol_idx = {'NW': [0], 'NE': [1]}
            cal = read_cal_table(
                calfile=calibration
            )
            cal = cal[np.ix_(
                self.sb_idx,
                mas,
                pol_idx[pol]
            )].squeeze()
        # Matrix of BSTs
        c = np.zeros(
            (
                self.times.size,
                self.freqs.size,
                mas.size,
                mas.size
            ),
            dtype=np.complex
        )
        # Matrix of phasings
        p = np.ones(
            (
                self.freqs.size,
                mas.size,
                mas.size
            ),
            dtype=np.complex
        )
        # Pointing direction
        if isinstance(az, un.Quantity):
            az = az.to(un.deg).value
        if isinstance(el, un.Quantity):
            el = el.to(un.deg).value
        az = np.radians(az)
        el = np.radians(el)
        u = np.array([
            np.cos(el) * np.cos(az),
            np.cos(el) * np.sin(az),
            np.sin(el)
        ])
        # Polarization selection
        mask = np.isin(self._ant1, mas) & np.isin(self._ant2, mas)
        log.info(
            'Loading data...'
        )
        if pol.upper() == 'NW':
            cpol = self.xx[:, :, mask]
        else:
            cpol = self.yy[:, :, mask]
        log.info(
            'Data of shape {} loaded for beamforming'.format(
                cpol.shape
            )
        )
        # Put the Xcorr in a matrix
        trix, triy = np.tril_indices(mas.size, 0)
        c[:, :, trix, triy] = cpol
        c[:, :, triy, trix] = c[:, :, trix, triy].conj()
        # Calibrate the Xcorr with the caltable
        for fi in prange(c.shape[1]):
            cal_i = np.expand_dims(cal[fi], axis=1)
            cal_i_h = np.expand_dims(cal[fi].T.conj(), axis=0)
            mul = np.dot(cal_i, cal_i_h)
            c[:, fi, :, :] *= mul[np.newaxis, :, :]
        # Phase the Xcorr
        dphi = np.dot(
            ma_pos[self._ant1[mask]] - ma_pos[self._ant2[mask]],
            u
        )
        wavel = wavelength(self.freqs).value
        p[:, trix, triy] = np.exp(
            -2.j*np.pi/wavel[:, None] * dphi
        )
        p[:, triy, trix] = p[:, trix, triy].conj()
        data = np.sum((c * p).real, axis=(2, 3))
        log.info(
            'Beamforming complete.'
        )
        return SData(
            data=np.expand_dims(data, axis=2),
            time=self.times,
            freq=self.freqs,
            polar=np.array([pol.upper()])
        )

    def image(self, resolution=1, fov=50):
        r""" Converts NenuFAR-TV-like data sets containing
            visibilities (:math:`V(u,v,\nu , t)`) into images
            :math:`I(l, m, \nu)` phase-centered at the local
            zenith while time averaging the visibilities.
            The Field of View ``fov`` argument defines the
            diameter angular size (zenith-centered) above which
            the image is not computed.

            .. math::
                I(l, m, \nu) = \int
                    \langle V(u, v, \nu, t) \rangle_t e^{
                        2 \pi i \frac{\nu}{c} \left(
                            \langle u(t) \rangle_t l + \langle v(t) \rangle_t m
                        \right)
                    }
                    \, du \, dv

            :param resolution:
                Resoltion (in degrees if a `float` is given) of
                the HEALPix grid (passed to initialize the
                :class:`~nenupy.astro.hpxsky.HpxSky` object).
            :type resolution: `float` or :class:`~astropy.units.Quantity`
            :param fov:
                Field of view diameter of the image (in degrees
                if a `float` is given).
            :type fov: `float` or :class:`~astropy.units.Quantity`

            :returns: HEALPix sky object embedding the computed
                image.
            :rtype: :class:`~nenupy.astro.hpxsky.HpxSky`

            :Example:
                >>> from nenupy.crosslet import TV_Data
                >>> import astropy.units as u
                >>> tv = TV_Data('20191204_132113_nenufarTV.dat')
                >>> im = tv.image(
                        resolution=0.2*u.deg,
                        fov=60*u.deg
                    )

            .. seealso::
                :class:`~nenupy.astro.hpxsky.HpxSky`,
                :meth:`~nenupy.astro.hpxsky.HpxSky.lmn`,
                :meth:`~nenupy.crosslet.uvw.UVW.from_tvdata`

            .. warning::
                This method is intended to be used for NenuFAR-TV
                data and relatively small XST datasets. It is not
                suited to long observations for which a MS
                conversion is required before using imaging
                dedicated softwares.
        """
        if not isinstance(fov, un.Quantity):
            fov *= un.deg
        f_idx = 0 # Frequency index
        # Sky preparation
        sky = HpxSky(resolution=resolution)
        exposure = self.times[-1] - self.times[0]
        sky.time = self.times[0] + exposure/2.
        sky._is_visible = sky._ho_coords.alt >= 90*un.deg - fov/2.
        phase_center = eq_zenith(sky.time)
        l, m, n = sky.lmn(phase_center=phase_center)
        # UVW coordinates
        uvw = UVW.from_tvdata(self)
        u = np.mean( # Mean in time
            uvw.uvw[:, :, 0],
            axis=0
        )[self.mask_auto]/wavelength(self.freqs[f_idx]).value
        v = np.mean(
            uvw.uvw[:, :, 1],
            axis=0
        )[self.mask_auto]/wavelength(self.freqs[f_idx]).value
        w = np.mean( # Mean in time
            uvw.uvw[:, :, 2],
            axis=0
        )[self.mask_auto]/wavelength(self.freqs[f_idx]).value
        # Mulitply (u, v) by (l, m) and compute FT exp
        ul = ft_mul(
            x=np.tile(u, (l.size, 1)).T,
            y= np.tile(l, (u.size, 1))
        )
        vm = ft_mul(
            x=np.tile(v, (m.size, 1)).T,
            y=np.tile(m, (v.size, 1))
        )
        phase = ft_phase(ul, vm)
        # Phase visibilities
        vis = np.mean( # Mean in time
            self.stokes_i,
            axis=0
        )[f_idx, :][self.mask_auto]
        im = np.zeros(l.size)
        for i in tqdm(prange(l.size)):
            im[i] = np.real(
                ft_sum(vis, phase[:, i])
            )
        sky.skymap[sky._is_visible] = im
        return sky

    def nearfield(self, resolution):
        """
        """
        return

    # --------------------------------------------------------- #
    # ----------------------- Internal ------------------------ #
    def _get_cross_idx(self, c1='X', c2='X'):
        """
        """
        corr = np.array(['X', 'Y']*self.mas.size)
        i_ant1, i_ant2 = np.tril_indices(self.mas.size*2, 0)
        corr_mask = (corr[i_ant1] == c1) & (corr[i_ant2] == c2)
        indices = np.arange(i_ant1.size)[corr_mask]
        return indices
# ============================================================= #


"""Contains the classes that deal with constant temperature dynamics.

Copyright (C) 2013, Joshua More and Michele Ceriotti

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program. If not, see <http.//www.gnu.org/licenses/>.

Contains the algorithms which propagate the thermostatting steps in the constant
temperature ensembles. Includes the new GLE thermostat, which can be used to
run PI+GLE dynamics, reducing the number of path integral beads required.

Classes:
   Thermostat: Base thermostat class with the generic methods and attributes.
   ThermoLangevin: Holds the algorithms for a langevin thermostat.
   ThermoPILE_L: Holds the algorithms for a path-integral langevin equation
      thermostat, with a thermostat coupled directly to the
      centroid coordinate of each bead.
   ThermoPILE_G: Holds the algorithms for a path-integral langevin equation
      thermostat, with a thermostat coupled to the kinetic energy for
      the entire system.
   ThermoSVR: Holds the algorithms for a stochastic velocity rescaling
      thermostat.
   ThermoGLE: Holds the algorithms for a generalised langevin equation
      thermostat.
   ThermoNMGLE: Holds the algorithms for a generalised langevin equation
      thermostat in the normal mode representation.
   ThermoNMGLEG: Holds the algorithms for a generalised langevin equation
      thermostat in the normal mode representation, with kinetic energy as
      well as potential energy sampling optimization.
"""

__all__ = ['Thermostat', 'ThermoLangevin', 'ThermoPILE_L', 'ThermoPILE_G',
           'ThermoSVR', 'ThermoGLE', 'ThermoNMGLE', 'ThermoNMGLEG']

import numpy as np
from ipi.utils.depend   import *
from ipi.utils.units    import *
from ipi.utils.mathtools import matrix_exp, stab_cholesky, root_herm
from ipi.utils.prng import Random
from ipi.utils.messages import verbosity, warning, info
from ipi.engine.beads import Beads
from ipi.engine.normalmodes import NormalModes

class Thermostat(dobject):
   """Base thermostat class.

   Gives the standard methods and attributes needed in all the thermostat
   classes.

   Attributes:
      prng: A pseudo random number generator object.
      ndof: The number of degrees of freedom that the thermostat will be
         attached to.

   Depend objects:
      dt: The time step used in the algorithms. Depends on the simulation dt.
      temp: The simulation temperature. Higher than the system temperature by
         a factor of the number of beads. Depends on the simulation temp.
      ethermo: The total energy exchanged with the bath due to the thermostat.
      p: The momentum vector that the thermostat is coupled to. Depends on the
         beads p object.
      m: The mass vector associated with p. Depends on the beads m object.
      sm: The square root of the mass vector.
   """

   def __init__(self, temp = 1.0, dt = 1.0, ethermo=0.0):
      """Initialises Thermostat.

      Args:
         temp: The simulation temperature. Defaults to 1.0.
         dt: The simulation time step. Defaults to 1.0.
         ethermo: The initial heat energy transferred to the bath.
            Defaults to 0.0. Will be non-zero if the thermostat is
            initialised from a checkpoint file.
      """

      dset(self,"temp",   depend_value(name='temp', value=temp))
      dset(self,"dt",     depend_value(name='dt', value=dt))
      dset(self,"ethermo",depend_value(name='ethermo',value=ethermo))

   def bind(self, beads=None, atoms=None, pm=None, prng=None, fixdof=None):
      """Binds the appropriate degrees of freedom to the thermostat.

      This takes an object with degrees of freedom, and makes their momentum
      and mass vectors members of the thermostat. It also then creates the
      objects that will hold the data needed in the thermostat algorithms
      and the dependency network.

      Args:
         beads: An optional beads object to take the mass and momentum vectors
            from.
         atoms: An optional atoms object to take the mass and momentum vectors
            from.
         pm: An optional tuple containing a single momentum value and its
            conjugate mass.
         prng: An optional pseudo random number generator object. Defaults to
            Random().
         fixdof: An optional integer which can specify the number of constraints
            applied to the system. Defaults to zero.

      Raises:
         TypeError: Raised if no appropriate degree of freedom or object
            containing a momentum vector is specified for
            the thermostat to couple to.
      """

      if prng is None:
         warning("Initializing thermostat from standard random PRNG", verbosity.medium)
         self.prng = Random()
      else:
         self.prng = prng

      if not beads is None:
         dset(self,"p",beads.p.flatten())
         dset(self,"m",beads.m3.flatten())
      elif not atoms is None:
         dset(self,"p",dget(atoms, "p"))
         dset(self,"m",dget(atoms, "m3"))
      elif not pm is None:
         dset(self,"p",pm[0])
         dset(self,"m",pm[1])
      else:
         raise TypeError("Thermostat.bind expects either Beads, Atoms, NormalModes, or a (p,m) tuple to bind to")

      if fixdof is None:
         self.ndof = len(self.p)
      else:
         self.ndof = float(len(self.p) - fixdof)

      dset(self, "sm",
         depend_array(name="sm", value=np.zeros(len(dget(self,"m"))),
            func=self.get_sm, dependencies=[dget(self,"m")]))

   def get_sm(self):
      """Retrieves the square root of the mass matrix.

      Returns:
         A vector of the square root of the mass matrix with one value for
         each degree of freedom.
      """

      return np.sqrt(self.m)

   def step(self):
      """Dummy thermostat step."""

      pass

class ThermoLangevin(Thermostat):
   """Represents a langevin thermostat.

   Depend objects:
      tau: Thermostat damping time scale. Larger values give a less strongly
         coupled thermostat.
      T: Coefficient of the diffusive contribution of the thermostat, i.e. the
         drift back towards equilibrium. Depends on tau and the time step.
      S: Coefficient of the stochastic contribution of the thermostat, i.e.
         the uncorrelated Gaussian noise. Depends on T and the temperature.
   """

   def get_T(self):
      """Calculates the coefficient of the overall drift of the velocities."""

      return np.exp(-0.5*self.dt/self.tau)

   def get_S(self):
      """Calculates the coefficient of the white noise."""

      return np.sqrt(Constants.kb*self.temp*(1 - self.T**2))

   def __init__(self, temp = 1.0, dt = 1.0, tau = 1.0, ethermo=0.0):
      """Initialises ThermoLangevin.

      Args:
         temp: The simulation temperature. Defaults to 1.0.
         dt: The simulation time step. Defaults to 1.0.
         tau: The thermostat damping timescale. Defaults to 1.0.
         ethermo: The initial heat energy transferred to the bath.
            Defaults to 0.0. Will be non-zero if the thermostat is
            initialised from a checkpoint file.
      """

      super(ThermoLangevin,self).__init__(temp, dt, ethermo)

      dset(self,"tau",depend_value(value=tau,name='tau'))
      dset(self,"T",
         depend_value(name="T",func=self.get_T,
            dependencies=[dget(self,"tau"), dget(self,"dt")]))
      dset(self,"S",
         depend_value(name="S",func=self.get_S,
            dependencies=[dget(self,"temp"), dget(self,"T")]))

   def step(self):
      """Updates the bound momentum vector with a langevin thermostat."""

      p = depstrip(self.p).copy()
      sm = depstrip(self.sm)

      p /= sm

      self.ethermo += np.dot(p,p)*0.5
      p *= self.T
      p += self.S*self.prng.gvec(len(p))
      self.ethermo -= np.dot(p,p)*0.5

      p *= sm

      self.p = p

class ThermoPILE_L(Thermostat):
   """Represents a PILE thermostat with a local centroid thermostat.

   Attributes:
      _thermos: The list of the different thermostats for all the ring polymer
         normal modes.
      nm: A normal modes object to attach the thermostat to.
      prng: Random number generator used in the stochastic integration
         algorithms.

   Depend objects:
      tau: Centroid thermostat damping time scale. Larger values give a
         less strongly coupled centroid thermostat.
      tauk: Thermostat damping time scale for the non-centroid normal modes.
         Depends on the ring polymer spring constant, and thus the simulation
         temperature.
      pilescale: A float used to reduce the intensity of the PILE thermostat if
         required.
   """

   def __init__(self, temp = 1.0, dt = 1.0, tau = 1.0, ethermo=0.0, scale=1.0):
      """Initialises ThermoPILE_L.

      Args:
         temp: The simulation temperature. Defaults to 1.0.
         dt: The simulation time step. Defaults to 1.0.
         tau: The centroid thermostat damping timescale. Defaults to 1.0.
         ethermo: The initial conserved energy quantity. Defaults to 0.0. Will
            be non-zero if the thermostat is initialised from a checkpoint file.
         scale: A float used to reduce the intensity of the PILE thermostat if
            required.

      Raises:
         TypeError: Raised if the thermostat is used with any object other than
            a beads object, so that we make sure that the objects needed for the
            normal mode transformation exist.
      """

      super(ThermoPILE_L,self).__init__(temp,dt,ethermo)
      dset(self,"tau",depend_value(value=tau,name='tau'))
      dset(self,"pilescale",depend_value(value=scale,name='pilescale'))

   def bind(self, nm=None, prng=None, bindcentroid=True, fixdof=None):
      """Binds the appropriate degrees of freedom to the thermostat.

      This takes a beads object with degrees of freedom, and makes its momentum
      and mass vectors members of the thermostat. It also then creates the
      objects that will hold the data needed in the thermostat algorithms
      and the dependency network.

      Gives the interface for both the PILE_L and PILE_G thermostats, which
      only differ in their treatment of the centroid coordinate momenta.

      Args:
         nm: An optional normal mode object to take the mass and momentum
            vectors from.
         prng: An optional pseudo random number generator object. Defaults to
            Random().
         bindcentroid: An optional boolean which decides whether a Langevin
            thermostat is attached to the centroid mode of each atom
            separately, or the total kinetic energy. Defaults to True, which
            gives a thermostat bound to each centroid momentum.
         fixdof: An optional integer which can specify the number of constraints
            applied to the system. Defaults to zero.

      Raises:
         TypeError: Raised if no appropriate degree of freedom or object
            containing a momentum vector is specified for
            the thermostat to couple to.
      """

      if nm is None or not type(nm) is NormalModes:
         raise TypeError("ThermoPILE_L.bind expects a NormalModes argument to bind to")
      if prng is None:
         self.prng = Random()
      else:
         self.prng = prng

      prev_ethermo = self.ethermo

      # creates a set of thermostats to be applied to individual normal modes
      self._thermos = [ ThermoLangevin(temp=1, dt=1, tau=1) for b in range(nm.nbeads) ]
      # optionally does not bind the centroid, so we can re-use all of this
      # in the PILE_G case
      if not bindcentroid:
         self._thermos[0] = None

      self.nm = nm

      dset(self,"tauk",
         depend_array(name="tauk", value=np.zeros(nm.nbeads-1,float),
            func=self.get_tauk, dependencies=[dget(self,"pilescale"), dget(nm,"dynomegak")] ) )

      # must pipe all the dependencies in such a way that values for the nm thermostats
      # are automatically updated based on the "master" thermostat
      def make_taugetter(k):
         return lambda: self.tauk[k-1]
      it = 0
      for t in self._thermos:
         if t is None:
            it += 1
            continue
         if it > 0:
            fixdof = None # only the centroid thermostat may have constraints

         # bind thermostat t to the it-th bead
         t.bind(pm=(nm.pnm[it,:],nm.dynm3[it,:]),prng=self.prng, fixdof=fixdof)
         # pipes temp and dt
         deppipe(self,"temp", t, "temp")
         deppipe(self,"dt", t, "dt")

         # for tau it is slightly more complex
         if it == 0:
            deppipe(self,"tau", t, "tau")
         else:
            # Here we manually connect _thermos[i].tau to tauk[i].
            # Simple and clear.
            dget(t,"tau").add_dependency(dget(self,"tauk"))
            dget(t,"tau")._func = make_taugetter(it)
         dget(self,"ethermo").add_dependency(dget(t,"ethermo"))
         it += 1

      # since the ethermo will be "delegated" to the normal modes thermostats,
      # one has to split
      # any previously-stored value between the sub-thermostats
      if bindcentroid:
         for t in self._thermos:
            t.ethermo = prev_ethermo/nm.nbeads
         dget(self,"ethermo")._func = self.get_ethermo;
         # if we are not binding the centroid just yet, this bit of the piping
         # is delegated to the function which is actually calling this

   def get_tauk(self):
      """Computes the thermostat damping time scale for the non-centroid
      normal modes.

      Returns:
         An array with the damping time scales for the non-centroid modes.
      """

      # Also include an optional scaling factor to reduce the intensity of NM thermostats
      return  np.array([ self.pilescale/(2*self.nm.dynomegak[k])  for k in range(1,len(self._thermos)) ])

   def get_ethermo(self):
      """Computes the total energy transferred to the heat bath for all the
      thermostats.
      """

      et = 0.0;
      for t in self._thermos:
         et += t.ethermo
      return et

   def step(self):
      """Updates the bound momentum vector with a PILE thermostat."""

      # super-cool! just loop over the thermostats! it's as easy as that!
      for t in self._thermos:
         t.step()

class ThermoSVR(Thermostat):
   """Represents a stochastic velocity rescaling thermostat.

   Depend objects:
      tau: Centroid thermostat damping time scale. Larger values give a
         less strongly coupled centroid thermostat.
      K: Scaling factor for the total kinetic energy. Depends on the
         temperature.
      et: Parameter determining the strength of the thermostat coupling.
         Depends on tau and the time step.
   """

   def get_et(self):
      """Calculates the damping term in the propagator."""

      return np.exp(-0.5*self.dt/self.tau)

   def get_K(self):
      """Calculates the average kinetic energy per degree of freedom."""

      return Constants.kb*self.temp*0.5

   def __init__(self, temp = 1.0, dt = 1.0, tau = 1.0, ethermo=0.0):
      """Initialises ThermoSVR.

      Args:
         temp: The simulation temperature. Defaults to 1.0.
         dt: The simulation time step. Defaults to 1.0.
         tau: The thermostat damping timescale. Defaults to 1.0.
         ethermo: The initial conserved energy quantity. Defaults to 0.0. Will
            be non-zero if the thermostat is initialised from a checkpoint file.
      """

      super(ThermoSVR,self).__init__(temp,dt,ethermo)

      dset(self,"tau",depend_value(value=tau,name='tau'))
      dset(self,"et",
         depend_value(name="et",func=self.get_et,
            dependencies=[dget(self,"tau"), dget(self,"dt")]))
      dset(self,"K",
         depend_value(name="K",func=self.get_K, dependencies=[dget(self,"temp")]))

   def step(self):
      """Updates the bound momentum vector with a stochastic velocity rescaling
      thermostat. See G Bussi, D Donadio, M Parrinello,
      Journal of Chemical Physics 126, 014101 (2007)
      """

      K = np.dot(depstrip(self.p),depstrip(self.p)/depstrip(self.m))*0.5

      # rescaling is un-defined if the KE is zero
      if K == 0.0:
         return

      # gets the stochastic term (basically a Gamma distribution for the kinetic energy)
      r1 = self.prng.g
      if (self.ndof-1)%2 == 0:
         rg = 2.0*self.prng.gamma((self.ndof-1)/2)
      else:
         rg = 2.0*self.prng.gamma((self.ndof-2)/2) + self.prng.g**2

      alpha2 = self.et + self.K/K*(1 - self.et)*(r1**2 + rg) + 2.0*r1*np.sqrt(self.K/K*self.et*(1 - self.et))
      alpha = np.sqrt(alpha2)
      if (r1 + np.sqrt(2*K/self.K*self.et/(1 - self.et))) < 0:
         alpha *= -1

      self.ethermo += K*(1 - alpha2)
      self.p *= alpha

class ThermoPILE_G(ThermoPILE_L):
   """Represents a PILE thermostat with a global centroid thermostat.

   Simply replaces the Langevin thermostat for the centroid normal mode with
   a global velocity rescaling thermostat.
   """

   def __init__(self, temp = 1.0, dt = 1.0, tau = 1.0, ethermo=0.0, scale = 1.0):
      """Initialises ThermoPILE_G.

      Args:
         temp: The simulation temperature. Defaults to 1.0.
         dt: The simulation time step. Defaults to 1.0.
         tau: The centroid thermostat damping timescale. Defaults to 1.0.
         ethermo: The initial conserved energy quantity. Defaults to 0.0. Will
            be non-zero if the thermostat is initialised from a checkpoint file.
         scale: A float used to reduce the intensity of the PILE thermostat if
            required.
      """

      super(ThermoPILE_G,self).__init__(temp,dt,tau,ethermo)
      dset(self,"pilescale",depend_value(value=scale,name='pilescale'))

   def bind(self, nm=None, prng=None, fixdof=None):
      """Binds the appropriate degrees of freedom to the thermostat.

      This takes a beads object with degrees of freedom, and makes its momentum
      and mass vectors members of the thermostat. It also then creates the
      objects that will hold the data needed in the thermostat algorithms
      and the dependency network.

      Uses the PILE_L bind interface, with bindcentroid set to false so we can
      specify that thermostat separately, by binding a global
      thermostat to the centroid mode.

      Args:
         beads: An optional beads object to take the mass and momentum vectors
            from.
         prng: An optional pseudo random number generator object. Defaults to
            Random().
         fixdof: An optional integer which can specify the number of constraints
            applied to the system. Defaults to zero.

      """

      # first binds as a local PILE, then substitutes the thermostat on the centroid
      prev_ethermo = self.ethermo
      super(ThermoPILE_G,self).bind(nm=nm,prng=prng,bindcentroid=False, fixdof=fixdof)

      #centroid thermostat
      self._thermos[0] = ThermoSVR(temp=1, dt=1, tau=1)

      t = self._thermos[0]
      t.bind(pm=(nm.pnm[0,:],nm.dynm3[0,:]),prng=self.prng, fixdof=fixdof)
      deppipe(self,"temp", t, "temp")
      deppipe(self,"dt", t, "dt")
      deppipe(self,"tau", t, "tau")
      dget(self,"ethermo").add_dependency(dget(t,"ethermo"))

      # splits any previous ethermo between the thermostats, and finishes to bind ethermo to the sum function
      for t in self._thermos:
         t.ethermo = prev_ethermo/nm.nbeads
      dget(self,"ethermo")._func = self.get_ethermo;

class ThermoGLE(Thermostat):
   """Represents a GLE thermostat.

   This is similar to a langevin thermostat, in that it uses Gaussian random
   numbers to simulate a heat bath acting on the system, but simulates a
   non-Markovian system by using a Markovian formulation in an extended phase
   space. This allows for a much greater degree of flexibility, and this
   thermostat, properly fitted, can give the an approximation to the correct
   quantum ensemble even for a classical, 1-bead simulation. More reasonably,
   using this thermostat allows for a far smaller number of replicas of the
   system to be used, as the convergence of the properties
   of the system is accelerated with respect to number of beads when PI+GLE
   are used in combination. (See M. Ceriotti, D. E. Manolopoulos, M. Parinello,
   J. Chem. Phys. 134, 084104 (2011)).

   Attributes:
      ns: The number of auxilliary degrees of freedom.
      s: An array holding all the momenta, including the ones for the
         auxilliary degrees of freedom.

   Depend objects:
      A: Drift matrix giving the damping time scales for all the different
         degrees of freedom.
      C: Static covariance matrix.
         Satisfies A.C + C.transpose(A) = B.transpose(B), where B is the
         diffusion matrix, giving the strength of the coupling of the system
         with the heat bath, and thus the size of the stochastic
         contribution of the thermostat.
      T: Matrix for the diffusive contribution of the thermostat, i.e. the
         drift back towards equilibrium. Depends on A and the time step.
      S: Matrix for the stochastic contribution of the thermostat, i.e.
         the uncorrelated Gaussian noise. Depends on C and T.
   """

   def get_T(self):
      """Calculates the matrix for the overall drift of the velocities."""

      return matrix_exp(-0.5*self.dt*self.A)

   def get_S(self):
      """Calculates the matrix for the coloured noise."""

      SST = Constants.kb*(self.C - np.dot(self.T,np.dot(self.C,self.T.T)))

      # Uses a symetric decomposition rather than Cholesky, since it is more stable
      return root_herm(SST)

   def get_C(self):
      """Calculates C from temp (if C is not set explicitly)"""

      rC = np.identity(self.ns + 1,float)*self.temp
      return rC[:]

   def __init__(self, temp = 1.0, dt = 1.0, A = None, C = None, ethermo=0.0):
      """Initialises ThermoGLE.

      Args:
         temp: The simulation temperature. Defaults to 1.0.
         dt: The simulation time step. Defaults to 1.0.
         A: An optional matrix giving the drift matrix. Defaults to a single
            value of 1.0.
         C: An optional matrix giving the covariance matrix. Defaults to an
            identity matrix times temperature with the same dimensions as the
            total number of degrees of freedom in the system.
         ethermo: The initial heat energy transferred to the bath.
            Defaults to 0.0. Will be non-zero if the thermostat is
            initialised from a checkpoint file.
      """

      super(ThermoGLE,self).__init__(temp,dt,ethermo)

      if A is None:
         A = np.identity(1,float)
      dset(self,"A",depend_value(value=A.copy(),name='A'))

      self.ns = len(self.A) - 1;

      # now, this is tricky. if C is taken from temp, then we want it to be updated
      # as a depend of temp. Otherwise, we want it to be an independent beast.
      if C is None:
         C = np.identity(self.ns+1,float)*self.temp
         dset(self,"C",
            depend_value(name='C', func=self.get_C,
               dependencies=[dget(self,"temp")]))
      else:
         dset(self,"C",depend_value(value=C.copy(),name='C'))

      dset(self,"T",
         depend_value(name="T",func=self.get_T,
            dependencies=[dget(self,"A"), dget(self,"dt")]))
      dset(self,"S",
         depend_value(name="S",func=self.get_S,
            dependencies=[dget(self,"C"), dget(self,"T")]))

      self.s = np.zeros(0)

   def bind(self, beads=None, atoms=None, pm=None, prng=None, fixdof=None):
      """Binds the appropriate degrees of freedom to the thermostat.

      This takes an object with degrees of freedom, and makes their momentum
      and mass vectors members of the thermostat. It also then creates the
      objects that will hold the data needed in the thermostat algorithms
      and the dependency network.

      Args:
         beads: An optional beads object to take the mass and momentum vectors
            from.
         atoms: An optional atoms object to take the mass and momentum vectors
            from.
         pm: An optional tuple containing a single momentum value and its
            conjugate mass.
         prng: An optional pseudo random number generator object. Defaults to
            Random().
         fixdof: An optional integer which can specify the number of constraints
            applied to the system. Defaults to zero.

      Raises:
         TypeError: Raised if no appropriate degree of freedom or object
            containing a momentum vector is specified for
            the thermostat to couple to.
      """

      super(ThermoGLE,self).bind(beads,atoms,pm,prng,fixdof)

      # allocates, initializes or restarts an array of s's
      if self.s.shape != (self.ns + 1, len(dget(self,"m"))):
         if len(self.s) > 0:
            warning("Mismatch in GLE s array size on restart, will reinitialise to free particle.", verbosity.low)
         self.s = np.zeros((self.ns + 1, len(dget(self,"m"))))

         # Initializes the s vector in the free-particle limit
         info(" GLE additional DOFs initialised to the free-particle limit.", verbosity.low)
         SC = stab_cholesky(self.C*Constants.kb)
         self.s[:] = np.dot(SC, self.prng.gvec(self.s.shape))
      else:
         info("GLE additional DOFs initialised from input.", verbosity.medium)

   def step(self):
      """Updates the bound momentum vector with a GLE thermostat"""

      p = depstrip(self.p).copy()

      self.s[0,:] = self.p/self.sm

      self.ethermo += np.dot(self.s[0],self.s[0])*0.5
      self.s[:] = np.dot(self.T,self.s) + np.dot(self.S,self.prng.gvec(self.s.shape))
      self.ethermo -= np.dot(self.s[0],self.s[0])*0.5

      self.p = self.s[0]*self.sm

class ThermoNMGLE(Thermostat):
   """Represents a 'normal-modes' GLE thermostat.

   An extension to the GLE thermostat which is applied in the
   normal modes representation, and which allows to use a different
   GLE for each normal mode

   Attributes:
      ns: The number of auxilliary degrees of freedom.
      nb: The number of beads.
      s: An array holding all the momenta, including the ones for the
         auxilliary degrees of freedom.

   Depend objects:
      A: Drift matrix giving the damping time scales for all the different
         degrees of freedom (must contain nb terms).
      C: Static covariance matrix.
         Satisfies A.C + C.transpose(A) = B.transpose(B), where B is the
         diffusion matrix, giving the strength of the coupling of the system
         with the heat bath, and thus the size of the stochastic
         contribution of the thermostat.
   """

   def get_C(self):
      """Calculates C from temp (if C is not set explicitly)."""

      rv = np.ndarray((self.nb, self.ns+1, self.ns+1), float)
      for b in range(0,self.nb):
         rv[b] = np.identity(self.ns + 1,float)*self.temp
      return rv[:]

   def __init__(self, temp = 1.0, dt = 1.0, A = None, C = None, ethermo=0.0):
      """Initialises ThermoGLE.

      Args:
         temp: The simulation temperature. Defaults to 1.0.
         dt: The simulation time step. Defaults to 1.0.
         A: An optional matrix giving the drift matrix. Defaults to a single
            value of 1.0.
         C: An optional matrix giving the covariance matrix. Defaults to an
            identity matrix times temperature with the same dimensions as the
            total number of degrees of freedom in the system.
         ethermo: The initial heat energy transferred to the bath.
            Defaults to 0.0. Will be non-zero if the thermostat is
            initialised from a checkpoint file.
      """

      super(ThermoNMGLE,self).__init__(temp,dt,ethermo)

      if A is None:
         A = np.identity(1,float)
      dset(self,"A",depend_value(value=A.copy(),name='A'))

      self.nb = len(self.A)
      self.ns = len(self.A[0]) - 1;

      # now, this is tricky. if C is taken from temp, then we want it to be
      # updated as a depend of temp.
      # Otherwise, we want it to be an independent beast.
      if C is None:
         dset(self,"C",depend_value(name='C', func=self.get_C, dependencies=[dget(self,"temp")]))
      else:
         dset(self,"C",depend_value(value=C.copy(),name='C'))

   def bind(self, nm=None, prng=None, fixdof=None):
      """Binds the appropriate degrees of freedom to the thermostat.

      This takes an object with degrees of freedom, and makes their momentum
      and mass vectors members of the thermostat. It also then creates the
      objects that will hold the data needed in the thermostat algorithms
      and the dependency network. Actually, this specific thermostat requires
      being called on a beads object.

      Args:
         nm: An optional normal modes object to take the mass and momentum
            vectors from.
         prng: An optional pseudo random number generator object. Defaults to
            Random().
         fixdof: An optional integer which can specify the number of constraints
            applied to the system. Defaults to zero.

      Raises:
         TypeError: Raised if no beads object is specified for
            the thermostat to couple to.
      """

      if nm is None or not type(nm) is NormalModes:
         raise TypeError("ThermoNMGLE.bind expects a NormalModes argument to bind to")

      if prng is None:
         self.prng = Random()
      else:
         self.prng = prng

      if (nm.nbeads != self.nb):
         raise IndexError("The parameters in nm_gle options correspond to a bead number "+str(self.nb)+ " which does not match the number of beads in the path" + str(nm.nbeads) )

      # allocates, initializes or restarts an array of s's
      if self.s.shape != (self.nb, self.ns + 1, nm.natoms *3) :
         if len(self.s) > 0:
            warning("Mismatch in GLE s array size on restart, will reinitialise to free particle.", verbosity.low)
         self.s = np.zeros((self.nb, self.ns + 1, nm.natoms*3))

         # Initializes the s vector in the free-particle limit
         info(" GLE additional DOFs initialised to the free-particle limit.", verbosity.low)
         for b in range(self.nb):
            SC = stab_cholesky(self.C[b]*Constants.kb)
            self.s[b] = np.dot(SC, self.prng.gvec(self.s[b].shape))
      else:
         info("GLE additional DOFs initialised from input.", verbosity.medium)

      prev_ethermo = self.ethermo

      # creates a set of thermostats to be applied to individual normal modes
      self._thermos = [ThermoGLE(temp=1, dt=1, A=self.A[b], C=self.C[b]) for b in range(self.nb)]

      # must pipe all the dependencies in such a way that values for the nm
      # thermostats are automatically updated based on the "master" thermostat
      def make_Agetter(k):
         return lambda: self.A[k]
      def make_Cgetter(k):
         return lambda: self.C[k]

      it = 0
      for t in self._thermos:
         t.s = self.s[it]  # gets the s's as a slice of self.s
         t.bind(pm=(nm.pnm[it,:],nm.dynm3[it,:]), prng=self.prng) # bind thermostat t to the it-th normal mode

         # pipes temp and dt
         deppipe(self,"temp", t, "temp")
         deppipe(self,"dt", t, "dt")

         # here we pipe the A and C of individual NM to the "master" arrays
         dget(t,"A").add_dependency(dget(self,"A"))
         dget(t,"A")._func = make_Agetter(it)
         dget(t,"C").add_dependency(dget(self,"C"))
         dget(t,"C")._func = make_Cgetter(it)
         dget(self,"ethermo").add_dependency(dget(t,"ethermo"))
         it += 1

      # since the ethermo will be "delegated" to the normal modes thermostats,
      # one has to split
      # any previously-stored value between the sub-thermostats
      for t in self._thermos:
         t.ethermo = prev_ethermo/self.nb

      dget(self,"ethermo")._func = self.get_ethermo;

   def step(self):
      """Updates the thermostat in NM representation by looping over the
      individual DOFs.
      """

      for t in self._thermos:
         t.step()

   def get_ethermo(self):
      """Computes the total energy transferred to the heat bath for all the nm
      thermostats.
      """

      et = 0.0;
      for t in self._thermos:
         et += t.ethermo
      return et

class ThermoNMGLEG(ThermoNMGLE):
   """Represents a 'normal-modes' GLE thermostat + SVR.

   An extension to the above NMGLE thermostat which also adds a stochastic velocity
   rescaling to the centroid.

   Depend objects:
      tau: Thermostat damping time scale. Larger values give a less strongly
         coupled thermostat.
   """

   def __init__(self, temp = 1.0, dt = 1.0, A = None, C = None, tau=1.0, ethermo=0.0):

      super(ThermoNMGLEG,self).__init__(temp, dt, A, C, ethermo)
      dset(self,"tau",depend_value(value=tau,name='tau'))

   def bind(self, nm=None, prng=None, fixdof=None):
      """Binds the appropriate degrees of freedom to the thermostat.

      This takes an object with degrees of freedom, and makes their momentum
      and mass vectors members of the thermostat. It also then creates the
      objects that will hold the data needed in the thermostat algorithms
      and the dependency network. Actually, this specific thermostat requires
      being called on a beads object.

      Args:
         nm: An optional normal modes object to take the mass and momentum
            vectors from.
         prng: An optional pseudo random number generator object. Defaults to
            Random().
         fixdof: An optional integer which can specify the number of constraints
            applied to the system. Defaults to zero.
      """

      super(ThermoNMGLEG,self).bind(nm, prng, fixdof)

      t = ThermoSVR(self.temp, self.dt, self.tau)

      t.bind(pm=(nm.pnm[0,:],nm.dynm3[0,:]), prng=self.prng) # bind global thermostat to centroid

      # pipes temp and dt
      deppipe(self,"temp", t, "temp")
      deppipe(self,"dt", t, "dt")
      deppipe(self,"tau", t, "tau")

      dget(self,"ethermo").add_dependency(dget(t,"ethermo"))
      self._thermos.append(t)


"""
Common RGB Colour Models Utilities
==================================

Defines various RGB colour models common utilities.
"""

from __future__ import annotations

from colour.colorimetry import CCS_ILLUMINANTS
from colour.hints import ArrayLike, Boolean, Literal, NDArray, Optional, Union
from colour.models.rgb import RGB_COLOURSPACES, RGB_to_XYZ, XYZ_to_RGB

__author__ = "Colour Developers"
__copyright__ = "Copyright 2013 Colour Developers"
__license__ = "New BSD License - https://opensource.org/licenses/BSD-3-Clause"
__maintainer__ = "Colour Developers"
__email__ = "colour-developers@colour-science.org"
__status__ = "Production"

__all__ = [
    "XYZ_to_sRGB",
    "sRGB_to_XYZ",
]

def XYZ_to_sRGB(
    XYZ: ArrayLike,
    illuminant: ArrayLike = CCS_ILLUMINANTS[
        "CIE 1931 2 Degree Standard Observer"
    ]["D65"],
    chromatic_adaptation_transform: Optional[
        Union[
            Literal[
                "Bianco 2010",
                "Bianco PC 2010",
                "Bradford",
                "CAT02 Brill 2008",
                "CAT02",
                "CAT16",
                "CMCCAT2000",
                "CMCCAT97",
                "Fairchild",
                "Sharp",
                "Von Kries",
                "XYZ Scaling",
            ],
            str,
        ]
    ] = "CAT02",
    apply_cctf_encoding: Boolean = True,
) -> NDArray:
    """
    Convert from *CIE XYZ* tristimulus values to *sRGB* colourspace.

    Parameters
    ----------
    XYZ
        *CIE XYZ* tristimulus values.
    illuminant
        Source illuminant chromaticity coordinates.
    chromatic_adaptation_transform
        *Chromatic adaptation* transform.
    apply_cctf_encoding
        Whether to apply the *sRGB* encoding colour component transfer function
        / inverse electro-optical transfer function.

    Returns
    -------
    :class:`numpy.ndarray`
        *sRGB* colour array.

    Notes
    -----
    +------------+-----------------------+---------------+
    | **Domain** | **Scale - Reference** | **Scale - 1** |
    +============+=======================+===============+
    | ``XYZ``    | [0, 1]                | [0, 1]        |
    +------------+-----------------------+---------------+

    +------------+-----------------------+---------------+
    | **Range**  | **Scale - Reference** | **Scale - 1** |
    +============+=======================+===============+
    | ``RGB``    | [0, 1]                | [0, 1]        |
    +------------+-----------------------+---------------+

    Examples
    --------
    >>> import numpy as np
    >>> XYZ = np.array([0.20654008, 0.12197225, 0.05136952])
    >>> XYZ_to_sRGB(XYZ)  # doctest: +ELLIPSIS
    array([ 0.7057393...,  0.1924826...,  0.2235416...])
    """

    sRGB = RGB_COLOURSPACES["sRGB"]

    return XYZ_to_RGB(
        XYZ,
        illuminant,
        sRGB.whitepoint,
        sRGB.matrix_XYZ_to_RGB,
        chromatic_adaptation_transform,
        sRGB.cctf_encoding if apply_cctf_encoding else None,
    )

def sRGB_to_XYZ(
    RGB: ArrayLike,
    illuminant: ArrayLike = CCS_ILLUMINANTS[
        "CIE 1931 2 Degree Standard Observer"
    ]["D65"],
    chromatic_adaptation_transform: Optional[
        Union[
            Literal[
                "Bianco 2010",
                "Bianco PC 2010",
                "Bradford",
                "CAT02 Brill 2008",
                "CAT02",
                "CAT16",
                "CMCCAT2000",
                "CMCCAT97",
                "Fairchild",
                "Sharp",
                "Von Kries",
                "XYZ Scaling",
            ],
            str,
        ]
    ] = "CAT02",
    apply_cctf_decoding: Boolean = True,
) -> NDArray:
    """
    Convert from *sRGB* colourspace to *CIE XYZ* tristimulus values.

    Parameters
    ----------
    RGB
        *sRGB* colourspace array.
    illuminant
        Source illuminant chromaticity coordinates.
    chromatic_adaptation_transform
        *Chromatic adaptation* transform.
    apply_cctf_decoding
        Whether to apply the *sRGB* decoding colour component transfer function
        / electro-optical transfer function.

    Returns
    -------
    :class:`numpy.ndarray`
        *CIE XYZ* tristimulus values.

    Notes
    -----
    +------------+-----------------------+---------------+
    | **Domain** | **Scale - Reference** | **Scale - 1** |
    +============+=======================+===============+
    | ``RGB``    | [0, 1]                | [0, 1]        |
    +------------+-----------------------+---------------+

    +------------+-----------------------+---------------+
    | **Range**  | **Scale - Reference** | **Scale - 1** |
    +============+=======================+===============+
    | ``XYZ``    | [0, 1]                | [0, 1]        |
    +------------+-----------------------+---------------+

    Examples
    --------
    >>> import numpy as np
    >>> RGB = np.array([0.70573936, 0.19248266, 0.22354169])
    >>> sRGB_to_XYZ(RGB)  # doctest: +ELLIPSIS
    array([ 0.2065429...,  0.1219794...,  0.0513714...])
    """

    sRGB = RGB_COLOURSPACES["sRGB"]

    return RGB_to_XYZ(
        RGB,
        sRGB.whitepoint,
        illuminant,
        sRGB.matrix_RGB_to_XYZ,
        chromatic_adaptation_transform,
        sRGB.cctf_decoding if apply_cctf_decoding else None,
    )


from os.path import exists
from pathlib import Path
import numpy as np
from . import util

# comparison of the correlations between CD4/CD8 expression levels and
# either the old cd8 score or this new one on two 10x datasets
# Looks like this new one is a better predictor (which is should be!)
#
# based on this I am switching this to be the default
#
#                   R    Pearson P  Spearman  Kendall  dataset
# old corr: CD4   -0.171  1.10e-19  5.25e-20  1.07e-19 hs_pbmc3
# new corr: CD4   -0.184  1.60e-22  1.03e-23  2.27e-23 hs_pbmc3
# old corr: CD8A   0.360  3.56e-86  1.42e-91  4.98e-85 hs_pbmc3
# new corr: CD8A   0.384  1.57e-98 7.91e-109 6.17e-100 hs_pbmc3
# old corr: CD8B   0.312  5.94e-64  3.23e-74  1.25e-69 hs_pbmc3
# new corr: CD8B   0.329  3.76e-71  1.55e-91  1.28e-84 hs_pbmc3
# old corr: CD4   -0.136  8.38e-08  8.22e-08  1.02e-07 hs_pbmc
# new corr: CD4   -0.154  1.37e-09  4.40e-09  5.52e-09 hs_pbmc
# old corr: CD8A   0.348  3.83e-45  2.75e-46  2.76e-43 hs_pbmc
# new corr: CD8A   0.400  2.49e-60  1.46e-63  1.40e-57 hs_pbmc
# old corr: CD8B   0.258  8.65e-25  2.64e-30  1.75e-28 hs_pbmc
# new corr: CD8B   0.303  5.09e-34  5.89e-45  1.38e-41 hs_pbmc

amino_acids = ['A', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'K', 'L', \
               'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'V', 'W', 'Y']

# read the model parameters
all_models = {}

for ab in 'AB':
    model_params = {
        'window_size':None,
        'min_lenbin':None,
        'max_lenbin':None,
        'NV':None,
        'NJ':None,
        'NL':None,
        'NC':None,
        }
    model_file = Path.joinpath( Path(util.path_to_data),f'cd8_logreg_params_{ab}.txt')
    mtags, weights = [], []
    #print(f'loading cd8 model for chain {ab} {model_file}')
    with open(model_file, 'r') as data:
        for line in data:
            l = line.split()
            tag, value = l
            if tag in model_params:
                model_params[tag] = int(value)
            else:
                mtags.append(tag)
                weights.append(float(value))
    NV, NJ, NL, NC = [model_params[x] for x in 'NV NJ NL NC'.split()]
    NTOT = NV + NJ + NL + NC
    assert mtags[-1] == 'BIAS'
    bias = float(weights[-1])
    assert len(weights) == NTOT+1
    weights = np.array(weights)
    vgenes = mtags[:NV-1] # last is UNK gene
    jgenes = mtags[NV:NV+NJ-1]
    assert all(x.startswith(f'TR{ab}V') for x in vgenes)
    assert all(x.startswith(f'TR{ab}J') for x in jgenes)
    vgene_indexer = {x:i for i,x in enumerate(vgenes)}
    jgene_indexer = {x:i for i,x in enumerate(jgenes)}

    model_params['weights'] = weights
    model_params['vgene_indexer'] = vgene_indexer
    model_params['jgene_indexer'] = jgene_indexer

    all_models[ab] = model_params

def get_lenbin(L, min_lenbin, max_lenbin):
    if L <= min_lenbin:
        return min_lenbin
    elif L >= max_lenbin:
        return max_lenbin
    else:
        return L

def get_allele(g):
    assert g.count('*')==1
    return g[:g.index('*')]

def encode_single_chain_tcr(vgene, jgene, cdr3, model_params):
    ''' vgene and jgene still include '*01' (e.g.)
    '''
    window_size = model_params['window_size']
    min_lenbin = model_params['min_lenbin']
    max_lenbin = model_params['max_lenbin']
    vgene_indexer = model_params['vgene_indexer']
    jgene_indexer = model_params['jgene_indexer']
    NV = len(vgene_indexer)+1
    NJ = len(jgene_indexer)+1
    NL = max_lenbin-min_lenbin+1
    NC = 20 * (2*window_size+1)
    NTOT = NV + NJ + NL + NC

    # 1-hot encode the V/J genes and length
    x = np.zeros((NTOT+1,))
    iv = vgene_indexer.get(get_allele(vgene), NV-1)
    x[iv] = 1.0

    ij = jgene_indexer.get(get_allele(jgene), NJ-1)
    x[NV+ij] = 1.0

    il = get_lenbin(len(cdr3), min_lenbin, max_lenbin) - min_lenbin
    x[NV+NJ+il] = 1.0

    # 1-hot encode aas in the nterminal and cterminal cdr3 windows
    # k-hot encode the middle of the tcr
    nterm = min(window_size, len(cdr3)//2)
    cterm = min(window_size, len(cdr3)-nterm)
    #middle = len(cdr3)-nterm-cterm
    for i in range(window_size):
        if i<nterm:
            istart = NV+NJ+NL+20*i
            x[istart+amino_acids.index(cdr3[i])] += 1.0
        if i<cterm:
            istart = NV+NJ+NL+20*window_size+20*i
            x[istart+amino_acids.index(cdr3[-1-i])] += 1.0
    for aa in cdr3[nterm:-cterm]:
        assert len(cdr3)>2*window_size # sanity check
        istart = NV+NJ+NL+40*window_size
        x[istart+amino_acids.index(aa)] += 1.0

    x[NTOT] = 1. # for the bias term
    return x

def make_cd8_score_table_column(tcrs, use_sigmoid=False):
    ''' We fit a logistic regression model on single-chain bulk TCR data
    from sorted CD4/CD8 subsets. Here we just average the alpha-chain and beta-chain
    scores.
    '''
    global all_models
    score_totals = np.zeros((len(tcrs),))
    for iab, ab in enumerate('AB'):
        model = all_models[ab]
        tcr_vectors = [ encode_single_chain_tcr(x[iab][0], x[iab][1], x[iab][2], model)
                        for x in tcrs ]
        tcr_vectors = np.vstack(tcr_vectors)
        weights = model['weights']
        assert tcr_vectors.shape == (len(tcrs), weights.shape[0])
        ab_scores = np.dot( tcr_vectors, weights)
        if use_sigmoid:
            ab_scores = 1.0/(1.0+np.exp(-1*ab_scores))
        score_totals += 0.5 * ab_scores
    return score_totals


import model
import utils
import numpy as np

class ThomasModel(model.Model):

    def __init__(self, dog_day_prob, message_prob, n_messengers = 1):
        """
        >>> m = ThomasModel(0.1, 0.25)
        >>> m.probs
        {(1, 1, 1, 0, 0): 0.0015625000000000003, (1, 0, 1, 1, 0): 0.0046874999999999981, (1, 1, 0, 1, 1): 0.0046874999999999981, (1, 0, 0, 0, 0): 0.014062500000000011, (0, 0, 0, 0, 1): 0.12656249999999999, (1, 0, 1, 0, 0): 0.0046874999999999981, (0, 0, 1, 1, 1): 0.042187499999999989, (1, 1, 1, 1, 0): 0.0015625000000000003, (0, 1, 0, 1, 1): 0.042187499999999989, (1, 0, 0, 1, 0): 0.014062500000000011, (0, 0, 1, 0, 1): 0.042187499999999989, (0, 1, 1, 1, 0): 0.014062500000000011, (0, 0, 1, 0, 0): 0.042187499999999989, (1, 1, 1, 0, 1): 0.0015625000000000003, (0, 1, 0, 0, 0): 0.042187499999999989, (1, 0, 1, 1, 1): 0.0046874999999999981, (1, 1, 0, 1, 0): 0.0046874999999999981, (1, 1, 0, 0, 0): 0.0046874999999999981, (1, 0, 1, 0, 1): 0.0046874999999999981, (0, 0, 1, 1, 0): 0.042187499999999989, (0, 1, 1, 0, 0): 0.014062500000000011, (0, 0, 0, 1, 0): 0.12656249999999999, (1, 0, 0, 0, 1): 0.014062500000000011, (0, 0, 0, 0, 0): 0.12656249999999999, (1, 1, 1, 1, 1): 0.0015625000000000003, (0, 0, 0, 1, 1): 0.12656249999999999, (1, 0, 0, 1, 1): 0.014062500000000011, (0, 1, 0, 0, 1): 0.042187499999999989, (0, 1, 1, 0, 1): 0.014062500000000011, (0, 1, 0, 1, 0): 0.042187499999999989, (1, 1, 0, 0, 1): 0.0046874999999999981, (0, 1, 1, 1, 1): 0.014062500000000011}
        >>> utils.print_partition(m.partitions)
        player 0
        (1, 1, None, None)   [(1, 1, 0, 0, 0), (1, 1, 0, 0, 1), (1, 1, 1, 0, 0), (1, 1, 1, 0, 1)]
        (1, 1, 0, None)   [(1, 1, 0, 1, 0), (1, 1, 0, 1, 1)]
        (0, 1, 0, None)   [(0, 1, 0, 1, 0), (0, 1, 0, 1, 1)]
        (None, None, None, None)   [(0, 0, 0, 0, 0), (0, 0, 0, 0, 1), (0, 0, 0, 1, 0), (0, 0, 0, 1, 1), (0, 0, 1, 0, 0), (0, 0, 1, 0, 1), (0, 0, 1, 1, 0), (0, 0, 1, 1, 1), (1, 0, 0, 0, 0), (1, 0, 0, 0, 1), (1, 0, 0, 1, 0), (1, 0, 0, 1, 1), (1, 0, 1, 0, 0), (1, 0, 1, 0, 1), (1, 0, 1, 1, 0), (1, 0, 1, 1, 1)]
        (0, 1, 1, 1)   [(0, 1, 1, 1, 1)]
        (1, 1, 1, 1)   [(1, 1, 1, 1, 1)]
        (0, 1, 1, 0)   [(0, 1, 1, 1, 0)]
        (0, 1, None, None)   [(0, 1, 0, 0, 0), (0, 1, 0, 0, 1), (0, 1, 1, 0, 0), (0, 1, 1, 0, 1)]
        (1, 1, 1, 0)   [(1, 1, 1, 1, 0)]
        player 1
        (1, 0, 1, None)   [(1, 0, 1, 0, 0), (1, 0, 1, 0, 1), (1, 0, 1, 1, 0), (1, 0, 1, 1, 1)]
        (0, 0, 1, None)   [(0, 0, 1, 0, 0), (0, 0, 1, 0, 1), (0, 0, 1, 1, 0), (0, 0, 1, 1, 1)]
        (None, None, None, None)   [(0, 0, 0, 0, 0), (0, 0, 0, 0, 1), (0, 0, 0, 1, 0), (0, 0, 0, 1, 1), (0, 1, 0, 0, 0), (0, 1, 0, 0, 1), (0, 1, 0, 1, 0), (0, 1, 0, 1, 1), (1, 0, 0, 0, 0), (1, 0, 0, 0, 1), (1, 0, 0, 1, 0), (1, 0, 0, 1, 1), (1, 1, 0, 0, 0), (1, 1, 0, 0, 1), (1, 1, 0, 1, 0), (1, 1, 0, 1, 1)]
        (0, 1, 1, 1)   [(0, 1, 1, 1, 1)]
        (1, 1, 1, 1)   [(1, 1, 1, 1, 1)]
        (1, 1, 1, None)   [(1, 1, 1, 0, 0), (1, 1, 1, 1, 0)]
        (1, 1, 1, 0)   [(1, 1, 1, 0, 1)]
        (0, 1, 1, 0)   [(0, 1, 1, 0, 1)]
        (0, 1, 1, None)   [(0, 1, 1, 0, 0), (0, 1, 1, 1, 0)]
        """
        self.dog_day_prob = dog_day_prob
        self.message_prob = message_prob
        self.n_messengers = n_messengers
        self.message_length = 4
        super(ThomasModel, self).__init__([(0,1)]*(1 + self.message_length*self.n_messengers), 2)

    def run(self, variables):
        """
        >>> m = ThomasModel(0.1, 0.25)
        >>> m.run((1,0,1,0,0))
        [(None, None, None, None), (1, 0, 1, None)]
        >>> m.run((1,0,1,1,0))
        [(None, None, None, None), (1, 0, 1, None)]
        >>> m.run((1,1,0,0,0))
        [(1, 1, None, None), (None, None, None, None)]
        >>> m.run((1,1,0,1,0))
        [(1, 1, 0, None), (None, None, None, None)]
        >>> m.run((1,1,1,1,0))
        [(1, 1, 1, 0), (1, 1, 1, None)]
        >>> m = ThomasModel(0.1, 0.25, 2)
        >>> m.run((1,1,1,1,0,0,1,0,0))
        [(1, 1, 1, 0, None, None, None, None), (1, 1, 1, None, 1, 0, 1, None)]
        >>> m.run((1,0,0,0,0,1,1,1,1))
        [(None, None, None, None, 1, 1, 1, 1), (None, None, None, None, 1, 1, 1, 1)]
        """

        out = [[] for i in range(self.n_players)]

        dog_day = variables[0]
        for i in range(self.n_messengers):
            this_out = self.messenger(dog_day, *variables[(1 + i*self.message_length):(1 + (i+1)*self.message_length)])
            for j in range(self.n_players):
                out[j] += this_out[j]

        out = [tuple(out[i]) for i in range(self.n_players)]

        return out

    # encodes the generative model of the messenger
    # d indicates if it is a "dog day"
    def messenger(self, d, m0, m1, tp, wtp):

        out = []

        if m0 == 1: # visiting player 0
            if tp == 1: # tell plan
                if m1 == 1: # visiting player 1
                    out += [[d, m0, m1, wtp]]
                else:
                    out += [[d, m0, m1, None]]
            else:
                out += [[d, m0, None, None]]
        else:
            out += [[None, None, None, None]]

        if m1 == 1: # visiting player 1
            if m0 == 1 and wtp == 1: # will tell player 0 plan to player 1, if it exists
                out += [[d, m0, m1, tp]]
            else:
                out += [[d, m0, m1, None]]
        else:
            out += [[None, None, None, None]]

        return out

    def evaluate(self, variables, log = False):
        """
        >>> m = ThomasModel(0.1, 0.25)
        >>> np.abs(m.evaluate((1,0,1,0,0)) - 0.1 * 0.75 * 0.25 * 0.5 * 0.5) < 1e-8
        True
        >>> np.abs(np.exp(m.evaluate((1,0,1,0,0), True)) - 0.1 * 0.75 * 0.25 * 0.5 * 0.5) < 1e-8
        True
        >>> m = ThomasModel(0.1, 0.25, 2)
        >>> np.abs(np.exp(m.evaluate((1,0,1,0,0,1,1,0,0), True)) - 0.1 * 0.75 * 0.25 * 0.5 * 0.5 * 0.25 * 0.25 * 0.5 * 0.5) < 1e-8
        True
        """

        log_p = 0

        log_p += utils.log_bernoulli(variables[0], self.dog_day_prob)
        for i in range(self.n_messengers):
            log_p += self.messenger_log_prob(*variables[(1 + i*self.message_length):(1 + (i+1)*self.message_length)])

        if log:
            return log_p
        else:
            return np.exp(log_p)

    def messenger_log_prob(self, m0, m1, tp, wtp):

        log_p = 0
        log_p += utils.log_bernoulli(m0, self.message_prob)
        log_p += utils.log_bernoulli(m1, self.message_prob)
        log_p += utils.log_bernoulli(tp, 0.5)
        log_p += utils.log_bernoulli(wtp, 0.5)

        return log_p

if __name__ == "__main__":
    import doctest
    doctest.testmod()