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


from pandas import read_json
from numpy import nan
import pandapower as pp
try:
    import pplog as logging
except ImportError:
    import logging
logger = logging.getLogger(__name__)

def IEEE123Bus():
    """
    Create the IEEE 34 bus from IEEE PES Test Feeders:
    "https://site.ieee.org/pes-testfeeders/resources/”.

    OUTPUT:
        **net** - The pandapower format network.
    """
    net = pp.create_empty_network()

    # Linedata
    # CF-1
    line_data = {'c_nf_per_km': 4.2976310, 'r_ohm_per_km': 0.1901810,
                 'x_ohm_per_km': 0.3896204,
                 'c0_nf_per_km': 2.0019129, 'r0_ohm_per_km': 0.4806720,
                 'x0_ohm_per_km': 1.2037409,
                 'max_i_ka': 0.53, 'type': 'ol'}
    pp.create_std_type(net, line_data, name='CF-1', element='line')

    # CF-2
    line_data = {'c_nf_per_km': 4.2976310, 'r_ohm_per_km': 0.1901810,
                 'x_ohm_per_km': 0.3896204,
                 'c0_nf_per_km': 2.0019129, 'r0_ohm_per_km': 0.4806720,
                 'x0_ohm_per_km': 1.2037409,
                 'max_i_ka': 0.53, 'type': 'ol'}
    pp.create_std_type(net, line_data, name='CF-2', element='line')

    # CF-3
    line_data = {'c_nf_per_km': 4.2976310, 'r_ohm_per_km': 0.1901810,
                 'x_ohm_per_km': 0.3896204,
                 'c0_nf_per_km': 2.0019129, 'r0_ohm_per_km': 0.4806720,
                 'x0_ohm_per_km': 1.2037409,
                 'max_i_ka': 0.53, 'type': 'ol'}
    pp.create_std_type(net, line_data, name='CF-3', element='line')

    # CF-4
    line_data = {'c_nf_per_km': 4.2976310, 'r_ohm_per_km': 0.1901810,
                 'x_ohm_per_km': 0.3896204,
                 'c0_nf_per_km': 2.0019129, 'r0_ohm_per_km': 0.4806720,
                 'x0_ohm_per_km': 1.2037409,
                 'max_i_ka': 0.53, 'type': 'ol'}
    pp.create_std_type(net, line_data, name='CF-4', element='line')

    # CF-5
    line_data = {'c_nf_per_km': 4.2976310, 'r_ohm_per_km': 0.1901810,
                 'x_ohm_per_km': 0.3896204,
                 'c0_nf_per_km': 2.0019129, 'r0_ohm_per_km': 0.4806720,
                 'x0_ohm_per_km': 1.2037409,
                 'max_i_ka': 0.53, 'type': 'ol'}
    pp.create_std_type(net, line_data, name='CF-5', element='line')

    # CF-6
    line_data = {'c_nf_per_km': 4.2976310, 'r_ohm_per_km': 0.1901810,
                 'x_ohm_per_km': 0.3896204,
                 'c0_nf_per_km': 2.0019129, 'r0_ohm_per_km': 0.4806720,
                 'x0_ohm_per_km': 1.2037409,
                 'max_i_ka': 0.53, 'type': 'ol'}
    pp.create_std_type(net, line_data, name='CF-6', element='line')

    # CF-7
    line_data = {'c_nf_per_km': 2.3489280, 'r_ohm_per_km': 0.15857392,
                 'x_ohm_per_km': 0.3641649,
                 'c0_nf_per_km': 1.6934436, 'r0_ohm_per_km': 0.2539544,
                 'x0_ohm_per_km': 0.6033307,
                 'max_i_ka': 0.53, 'type': 'ol'}
    pp.create_std_type(net, line_data, name='CF-7', element='line')

    # CF-8
    line_data = {'c_nf_per_km': 2.3489280, 'r_ohm_per_km': 0.15857392,
                 'x_ohm_per_km': 0.3641649,
                 'c0_nf_per_km': 1.6934436, 'r0_ohm_per_km': 0.2539544,
                 'x0_ohm_per_km': 0.6033307,
                 'max_i_ka': 0.53, 'type': 'ol'}
    pp.create_std_type(net, line_data, name='CF-8', element='line')

    # CF-9
    line_data = {'c_nf_per_km': 0.9360542, 'r_ohm_per_km': 0.2753088,
                 'x_ohm_per_km': 0.2790992,
                 'c0_nf_per_km': 0.9360542, 'r0_ohm_per_km': 0.2753088,
                 'x0_ohm_per_km': 0.2790992,
                 'max_i_ka': 0.23, 'type': 'ol'}
    pp.create_std_type(net, line_data, name='CF-9', element='line')

    # CF-10
    line_data = {'c_nf_per_km': 0.9360542, 'r_ohm_per_km': 0.2753088,
                 'x_ohm_per_km': 0.2790992,
                 'c0_nf_per_km': 0.9360542, 'r0_ohm_per_km': 0.2753088,
                 'x0_ohm_per_km': 0.2790992,
                 'max_i_ka': 0.23, 'type': 'ol'}
    pp.create_std_type(net, line_data, name='CF-10', element='line')

    # CF-11
    line_data = {'c_nf_per_km': 0.9360542, 'r_ohm_per_km': 0.2753088,
                 'x_ohm_per_km': 0.2790992,
                 'c0_nf_per_km': 0.9360542, 'r0_ohm_per_km': 0.2753088,
                 'x0_ohm_per_km': 0.2790992,
                 'max_i_ka': 0.23, 'type': 'ol'}
    pp.create_std_type(net, line_data, name='CF-11', element='line')

    # CF-12
    line_data = {'c_nf_per_km': 41.7711813, 'r_ohm_per_km': 0.6302153,
                 'x_ohm_per_km': 0.3002672,
                 'c0_nf_per_km': 41.7711813, 'r0_ohm_per_km': 1.5821560,
                 'x0_ohm_per_km': 0.7791580,
                 'max_i_ka': 0.31, 'type': 'cs'}
    pp.create_std_type(net, line_data, name='CF-12', element='line')

    # Busses
    bus_1   = pp.create_bus(net, name='Bus 1',   vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_2   = pp.create_bus(net, name='Bus 2',   vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_3   = pp.create_bus(net, name='Bus 3',   vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_4   = pp.create_bus(net, name='Bus 4',   vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_5   = pp.create_bus(net, name='Bus 5',   vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_6   = pp.create_bus(net, name='Bus 6',   vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_7   = pp.create_bus(net, name='Bus 7',   vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_8   = pp.create_bus(net, name='Bus 8',   vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_9   = pp.create_bus(net, name='Bus 9',   vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_10  = pp.create_bus(net, name='Bus 10',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_11  = pp.create_bus(net, name='Bus 11',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_12  = pp.create_bus(net, name='Bus 12',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_13  = pp.create_bus(net, name='Bus 13',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_14  = pp.create_bus(net, name='Bus 14',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_15  = pp.create_bus(net, name='Bus 15',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_16  = pp.create_bus(net, name='Bus 16',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_17  = pp.create_bus(net, name='Bus 17',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_18  = pp.create_bus(net, name='Bus 18',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_19  = pp.create_bus(net, name='Bus 19',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_20  = pp.create_bus(net, name='Bus 20',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_21  = pp.create_bus(net, name='Bus 21',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_22  = pp.create_bus(net, name='Bus 22',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_23  = pp.create_bus(net, name='Bus 23',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_24  = pp.create_bus(net, name='Bus 24',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_25  = pp.create_bus(net, name='Bus 25',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_26  = pp.create_bus(net, name='Bus 26',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_27  = pp.create_bus(net, name='Bus 27',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_28  = pp.create_bus(net, name='Bus 28',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_29  = pp.create_bus(net, name='Bus 29',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_30  = pp.create_bus(net, name='Bus 30',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_31  = pp.create_bus(net, name='Bus 31',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_32  = pp.create_bus(net, name='Bus 32',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_33  = pp.create_bus(net, name='Bus 33',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_34  = pp.create_bus(net, name='Bus 34',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_35  = pp.create_bus(net, name='Bus 35',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_36  = pp.create_bus(net, name='Bus 36',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_37  = pp.create_bus(net, name='Bus 37',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_38  = pp.create_bus(net, name='Bus 38',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_39  = pp.create_bus(net, name='Bus 39',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_40  = pp.create_bus(net, name='Bus 40',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_41  = pp.create_bus(net, name='Bus 41',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_42  = pp.create_bus(net, name='Bus 42',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_43  = pp.create_bus(net, name='Bus 43',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_44  = pp.create_bus(net, name='Bus 44',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_45  = pp.create_bus(net, name='Bus 45',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_46  = pp.create_bus(net, name='Bus 46',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_47  = pp.create_bus(net, name='Bus 47',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_48  = pp.create_bus(net, name='Bus 48',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_49  = pp.create_bus(net, name='Bus 49',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_50  = pp.create_bus(net, name='Bus 50',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_51  = pp.create_bus(net, name='Bus 51',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_52  = pp.create_bus(net, name='Bus 52',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_53  = pp.create_bus(net, name='Bus 53',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_54  = pp.create_bus(net, name='Bus 54',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_55  = pp.create_bus(net, name='Bus 55',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_56  = pp.create_bus(net, name='Bus 56',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_57  = pp.create_bus(net, name='Bus 57',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_58  = pp.create_bus(net, name='Bus 58',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_59  = pp.create_bus(net, name='Bus 59',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_60  = pp.create_bus(net, name='Bus 60',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_61  = pp.create_bus(net, name='Bus 61',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_62  = pp.create_bus(net, name='Bus 62',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_63  = pp.create_bus(net, name='Bus 63',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_64  = pp.create_bus(net, name='Bus 64',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_65  = pp.create_bus(net, name='Bus 65',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_66  = pp.create_bus(net, name='Bus 66',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_67  = pp.create_bus(net, name='Bus 67',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_68  = pp.create_bus(net, name='Bus 68',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_69  = pp.create_bus(net, name='Bus 69',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_70  = pp.create_bus(net, name='Bus 70',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_71  = pp.create_bus(net, name='Bus 71',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_72  = pp.create_bus(net, name='Bus 72',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_73  = pp.create_bus(net, name='Bus 73',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_74  = pp.create_bus(net, name='Bus 74',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_75  = pp.create_bus(net, name='Bus 75',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_76  = pp.create_bus(net, name='Bus 76',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_77  = pp.create_bus(net, name='Bus 77',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_78  = pp.create_bus(net, name='Bus 78',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_79  = pp.create_bus(net, name='Bus 79',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_80  = pp.create_bus(net, name='Bus 80',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_81  = pp.create_bus(net, name='Bus 81',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_82  = pp.create_bus(net, name='Bus 82',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_83  = pp.create_bus(net, name='Bus 83',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_84  = pp.create_bus(net, name='Bus 84',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_85  = pp.create_bus(net, name='Bus 85',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_86  = pp.create_bus(net, name='Bus 86',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_87  = pp.create_bus(net, name='Bus 87',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_88  = pp.create_bus(net, name='Bus 88',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_89  = pp.create_bus(net, name='Bus 89',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_90  = pp.create_bus(net, name='Bus 90',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_91  = pp.create_bus(net, name='Bus 91',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_92  = pp.create_bus(net, name='Bus 92',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_93  = pp.create_bus(net, name='Bus 93',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_94  = pp.create_bus(net, name='Bus 94',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_95  = pp.create_bus(net, name='Bus 95',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_96  = pp.create_bus(net, name='Bus 96',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_97  = pp.create_bus(net, name='Bus 97',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_98  = pp.create_bus(net, name='Bus 98',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_99  = pp.create_bus(net, name='Bus 99',  vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_100 = pp.create_bus(net, name='Bus 100', vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_101 = pp.create_bus(net, name='Bus 101', vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_102 = pp.create_bus(net, name='Bus 102', vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_103 = pp.create_bus(net, name='Bus 103', vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_104 = pp.create_bus(net, name='Bus 104', vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_105 = pp.create_bus(net, name='Bus 105', vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_106 = pp.create_bus(net, name='Bus 106', vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_107 = pp.create_bus(net, name='Bus 107', vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_108 = pp.create_bus(net, name='Bus 108', vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_109 = pp.create_bus(net, name='Bus 109', vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_110 = pp.create_bus(net, name='Bus 110', vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_111 = pp.create_bus(net, name='Bus 111', vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_112 = pp.create_bus(net, name='Bus 112', vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_113 = pp.create_bus(net, name='Bus 113', vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_114 = pp.create_bus(net, name='Bus 114', vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_135 = pp.create_bus(net, name='Bus 135', vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_149 = pp.create_bus(net, name='Bus 149', vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_152 = pp.create_bus(net, name='Bus 152', vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_160 = pp.create_bus(net, name='Bus 160', vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_197 = pp.create_bus(net, name='Bus 197', vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_250 = pp.create_bus(net, name='Bus 250', vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_300 = pp.create_bus(net, name='Bus 300', vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_450 = pp.create_bus(net, name='Bus 450', vn_kv=4.16, type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)
    bus_150 = pp.create_bus(net, name='Bus 150', vn_kv=115., type='b', zone='123_BUS', max_vm_pu=1.05, min_vm_pu=0.95)

    # Lines
    pp.create_line(net, bus_149, bus_1,   length_km=0.12192, std_type='CF-1',  name='Line 1')
    pp.create_line(net, bus_1,   bus_2,   length_km=0.05334, std_type='CF-10', name='Line 2')
    pp.create_line(net, bus_1,   bus_3,   length_km=0.07620, std_type='CF-11', name='Line 3')
    pp.create_line(net, bus_1,   bus_7,   length_km=0.09144, std_type='CF-1',  name='Line 4')
    pp.create_line(net, bus_3,   bus_4,   length_km=0.06096, std_type='CF-11', name='Line 5')
    pp.create_line(net, bus_3,   bus_5,   length_km=0.09906, std_type='CF-11', name='Line 6')
    pp.create_line(net, bus_5,   bus_6,   length_km=0.07620, std_type='CF-11', name='Line 7')
    pp.create_line(net, bus_7,   bus_8,   length_km=0.06096, std_type='CF-1',  name='Line 8')
    pp.create_line(net, bus_8,   bus_12,  length_km=0.06858, std_type='CF-10', name='Line 9')
    pp.create_line(net, bus_8,   bus_9,   length_km=0.06858, std_type='CF-9',  name='Line 10')
    pp.create_line(net, bus_8,   bus_13,  length_km=0.09144, std_type='CF-1',  name='Line 11')
    # pp.create_line(net, bus_9,   bus_14,  length_km=0.12954, std_type='CF-9',  name='Line 12')
    pp.create_line(net, bus_13,  bus_34,  length_km=0.04572, std_type='CF-11', name='Line 13')
    pp.create_line(net, bus_13,  bus_18,  length_km=0.25146, std_type='CF-2',  name='Line 14')
    pp.create_line(net, bus_14,  bus_11,  length_km=0.07620, std_type='CF-9',  name='Line 15')
    pp.create_line(net, bus_14,  bus_10,  length_km=0.07620, std_type='CF-9',  name='Line 16')
    pp.create_line(net, bus_15,  bus_16,  length_km=0.11430, std_type='CF-11', name='Line 17')
    pp.create_line(net, bus_15,  bus_17,  length_km=0.10668, std_type='CF-11', name='Line 18')
    pp.create_line(net, bus_18,  bus_19,  length_km=0.07620, std_type='CF-9',  name='Line 19')
    pp.create_line(net, bus_18,  bus_21,  length_km=0.09144, std_type='CF-2',  name='Line 20')
    pp.create_line(net, bus_19,  bus_20,  length_km=0.09906, std_type='CF-9',  name='Line 21')
    pp.create_line(net, bus_21,  bus_22,  length_km=0.16002, std_type='CF-10', name='Line 22')
    pp.create_line(net, bus_21,  bus_23,  length_km=0.07620, std_type='CF-2',  name='Line 23')
    pp.create_line(net, bus_23,  bus_24,  length_km=0.15240, std_type='CF-11', name='Line 24')
    pp.create_line(net, bus_23,  bus_25,  length_km=0.08382, std_type='CF-2',  name='Line 25')
    # pp.create_line(net, bus_25,  bus_26,  length_km=0.10668, std_type='CF-7',  name='Line 26')
    pp.create_line(net, bus_25,  bus_28,  length_km=0.06096, std_type='CF-2',  name='Line 27')
    pp.create_line(net, bus_26,  bus_27,  length_km=0.08382, std_type='CF-7',  name='Line 28')
    pp.create_line(net, bus_26,  bus_31,  length_km=0.06858, std_type='CF-11', name='Line 29')
    pp.create_line(net, bus_27,  bus_33,  length_km=0.15240, std_type='CF-9',  name='Line 30')
    pp.create_line(net, bus_28,  bus_29,  length_km=0.09144, std_type='CF-2',  name='Line 31')
    pp.create_line(net, bus_29,  bus_30,  length_km=0.10668, std_type='CF-2',  name='Line 32')
    pp.create_line(net, bus_30,  bus_250, length_km=0.06096, std_type='CF-2',  name='Line 33')
    pp.create_line(net, bus_31,  bus_32,  length_km=0.15240, std_type='CF-11', name='Line 34')
    pp.create_line(net, bus_34,  bus_15,  length_km=0.08382, std_type='CF-11', name='Line 35')
    pp.create_line(net, bus_35,  bus_36,  length_km=0.10668, std_type='CF-8',  name='Line 36')
    pp.create_line(net, bus_35,  bus_40,  length_km=0.06096, std_type='CF-1',  name='Line 37')
    pp.create_line(net, bus_36,  bus_37,  length_km=0.08382, std_type='CF-9',  name='Line 38')
    pp.create_line(net, bus_36,  bus_38,  length_km=0.06858, std_type='CF-10', name='Line 39')
    pp.create_line(net, bus_38,  bus_39,  length_km=0.15240, std_type='CF-10', name='Line 40')
    pp.create_line(net, bus_40,  bus_41,  length_km=0.09144, std_type='CF-11', name='Line 41')
    pp.create_line(net, bus_40,  bus_42,  length_km=0.10668, std_type='CF-1',  name='Line 42')
    pp.create_line(net, bus_42,  bus_43,  length_km=0.06096, std_type='CF-10', name='Line 43')
    pp.create_line(net, bus_42,  bus_44,  length_km=0.06096, std_type='CF-1',  name='Line 44')
    pp.create_line(net, bus_44,  bus_45,  length_km=0.06096, std_type='CF-9',  name='Line 45')
    pp.create_line(net, bus_44,  bus_47,  length_km=0.07620, std_type='CF-1',  name='Line 46')
    pp.create_line(net, bus_45,  bus_46,  length_km=0.09144, std_type='CF-9',  name='Line 47')
    pp.create_line(net, bus_47,  bus_48,  length_km=0.04572, std_type='CF-4',  name='Line 48')
    pp.create_line(net, bus_47,  bus_49,  length_km=0.07620, std_type='CF-4',  name='Line 49')
    pp.create_line(net, bus_49,  bus_50,  length_km=0.07620, std_type='CF-4',  name='Line 50')
    pp.create_line(net, bus_50,  bus_51,  length_km=0.07620, std_type='CF-4',  name='Line 51')
    pp.create_line(net, bus_52,  bus_53,  length_km=0.06096, std_type='CF-1',  name='Line 52')
    pp.create_line(net, bus_53,  bus_54,  length_km=0.03810, std_type='CF-1',  name='Line 53')
    pp.create_line(net, bus_54,  bus_55,  length_km=0.08382, std_type='CF-1',  name='Line 54')
    pp.create_line(net, bus_54,  bus_57,  length_km=0.10668, std_type='CF-3',  name='Line 55')
    pp.create_line(net, bus_55,  bus_56,  length_km=0.08382, std_type='CF-1',  name='Line 56')
    pp.create_line(net, bus_57,  bus_58,  length_km=0.07620, std_type='CF-10', name='Line 57')
    pp.create_line(net, bus_57,  bus_60,  length_km=0.22860, std_type='CF-3',  name='Line 58')
    pp.create_line(net, bus_58,  bus_59,  length_km=0.07620, std_type='CF-10', name='Line 59')
    pp.create_line(net, bus_60,  bus_61,  length_km=0.16764, std_type='CF-5',  name='Line 60')
    pp.create_line(net, bus_60,  bus_62,  length_km=0.07620, std_type='CF-12', name='Line 61')
    pp.create_line(net, bus_62,  bus_63,  length_km=0.05334, std_type='CF-12', name='Line 62')
    pp.create_line(net, bus_63,  bus_64,  length_km=0.10668, std_type='CF-12', name='Line 63')
    pp.create_line(net, bus_64,  bus_65,  length_km=0.12954, std_type='CF-12', name='Line 64')
    pp.create_line(net, bus_65,  bus_66,  length_km=0.09906, std_type='CF-12', name='Line 65')
    pp.create_line(net, bus_67,  bus_68,  length_km=0.06096, std_type='CF-9',  name='Line 66')
    pp.create_line(net, bus_67,  bus_72,  length_km=0.08382, std_type='CF-3',  name='Line 67')
    pp.create_line(net, bus_67,  bus_97,  length_km=0.07620, std_type='CF-3',  name='Line 68')
    pp.create_line(net, bus_68,  bus_69,  length_km=0.08382, std_type='CF-9',  name='Line 69')
    pp.create_line(net, bus_69,  bus_70,  length_km=0.09906, std_type='CF-9',  name='Line 69')
    pp.create_line(net, bus_70,  bus_71,  length_km=0.08382, std_type='CF-9',  name='Line 70')
    pp.create_line(net, bus_72,  bus_73,  length_km=0.08382, std_type='CF-11', name='Line 71')
    pp.create_line(net, bus_72,  bus_76,  length_km=0.06096, std_type='CF-3',  name='Line 72')
    pp.create_line(net, bus_73,  bus_74,  length_km=0.10668, std_type='CF-11', name='Line 73')
    pp.create_line(net, bus_74,  bus_75,  length_km=0.12192, std_type='CF-11', name='Line 74')
    pp.create_line(net, bus_76,  bus_77,  length_km=0.12192, std_type='CF-6',  name='Line 75')
    pp.create_line(net, bus_76,  bus_86,  length_km=0.21336, std_type='CF-3',  name='Line 76')
    pp.create_line(net, bus_77,  bus_78,  length_km=0.03048, std_type='CF-6',  name='Line 77')
    pp.create_line(net, bus_78,  bus_79,  length_km=0.06858, std_type='CF-6',  name='Line 78')
    pp.create_line(net, bus_78,  bus_80,  length_km=0.14478, std_type='CF-6',  name='Line 79')
    pp.create_line(net, bus_80,  bus_81,  length_km=0.14478, std_type='CF-6',  name='Line 80')
    pp.create_line(net, bus_81,  bus_82,  length_km=0.07620, std_type='CF-6',  name='Line 81')
    pp.create_line(net, bus_81,  bus_84,  length_km=0.20574, std_type='CF-11', name='Line 82')
    pp.create_line(net, bus_82,  bus_83,  length_km=0.07620, std_type='CF-6',  name='Line 83')
    pp.create_line(net, bus_84,  bus_85,  length_km=0.14478, std_type='CF-11', name='Line 84')
    pp.create_line(net, bus_86,  bus_87,  length_km=0.13716, std_type='CF-6',  name='Line 85')
    pp.create_line(net, bus_87,  bus_88,  length_km=0.05334, std_type='CF-9',  name='Line 86')
    pp.create_line(net, bus_87,  bus_89,  length_km=0.08382, std_type='CF-6',  name='Line 87')
    pp.create_line(net, bus_89,  bus_90,  length_km=0.06858, std_type='CF-10', name='Line 88')
    pp.create_line(net, bus_89,  bus_91,  length_km=0.06858, std_type='CF-6',  name='Line 89')
    pp.create_line(net, bus_91,  bus_92,  length_km=0.09144, std_type='CF-11', name='Line 90')
    pp.create_line(net, bus_91,  bus_93,  length_km=0.06858, std_type='CF-6',  name='Line 91')
    pp.create_line(net, bus_93,  bus_94,  length_km=0.08382, std_type='CF-9',  name='Line 92')
    pp.create_line(net, bus_93,  bus_95,  length_km=0.09144, std_type='CF-6',  name='Line 93')
    pp.create_line(net, bus_95,  bus_96,  length_km=0.06096, std_type='CF-10', name='Line 94')
    pp.create_line(net, bus_97,  bus_98,  length_km=0.08382, std_type='CF-3',  name='Line 95')
    pp.create_line(net, bus_98,  bus_99,  length_km=0.16764, std_type='CF-3',  name='Line 96')
    pp.create_line(net, bus_99,  bus_100, length_km=0.09144, std_type='CF-3',  name='Line 97')
    pp.create_line(net, bus_100, bus_450, length_km=0.24384, std_type='CF-3',  name='Line 98')
    pp.create_line(net, bus_101, bus_102, length_km=0.06858, std_type='CF-11', name='Line 99')
    pp.create_line(net, bus_101, bus_105, length_km=0.08382, std_type='CF-3',  name='Line 100')
    pp.create_line(net, bus_102, bus_103, length_km=0.09906, std_type='CF-11', name='Line 101')
    pp.create_line(net, bus_103, bus_104, length_km=0.21336, std_type='CF-11', name='Line 102')
    pp.create_line(net, bus_105, bus_106, length_km=0.06858, std_type='CF-10', name='Line 103')
    pp.create_line(net, bus_105, bus_108, length_km=0.09906, std_type='CF-3',  name='Line 104')
    pp.create_line(net, bus_106, bus_107, length_km=0.17526, std_type='CF-10', name='Line 105')
    pp.create_line(net, bus_108, bus_109, length_km=0.13716, std_type='CF-9',  name='Line 106')
    pp.create_line(net, bus_108, bus_300, length_km=0.30480, std_type='CF-3',  name='Line 107')
    pp.create_line(net, bus_109, bus_110, length_km=0.09144, std_type='CF-9',  name='Line 108')
    pp.create_line(net, bus_110, bus_111, length_km=0.17526, std_type='CF-9',  name='Line 109')
    pp.create_line(net, bus_110, bus_112, length_km=0.03810, std_type='CF-9',  name='Line 110')
    pp.create_line(net, bus_112, bus_113, length_km=0.16002, std_type='CF-9',  name='Line 111')
    pp.create_line(net, bus_113, bus_114, length_km=0.09906, std_type='CF-9',  name='Line 112')
    pp.create_line(net, bus_135, bus_35,  length_km=0.11430, std_type='CF-4',  name='Line 113')
    pp.create_line(net, bus_149, bus_1,   length_km=0.12192, std_type='CF-1',  name='Line 114')
    pp.create_line(net, bus_152, bus_52,  length_km=0.12192, std_type='CF-1',  name='Line 115')
    # pp.create_line(net, bus_160, bus_67,  length_km=0.10668, std_type='CF-6',  name='Line 116')
    pp.create_line(net, bus_197, bus_101, length_km=0.07620, std_type='CF-3',  name='Line 117')
    pp.create_line(net, bus_51, bus_300, length_km=0.19812, std_type='CF-4',  name='Line 118')

    # Substation / Regulator 1
    pp.create_transformer_from_parameters(net, bus_150, bus_149, sn_mva=5., vn_hv_kv=115.,
                                            vn_lv_kv=4.16, vkr_percent=1.0, vk_percent=8.062257,
                                            pfe_kw=0.0, i0_percent=0.0, shift_degree=30.0,
                                            tap_side='lv', tap_neutral=0, tap_max=2, tap_min=-2,
                                            tap_step_percent=2.5, tap_pos=0, tap_phase_shifter=False,
                                            name='Regulator 1')

    # Regulator 2
    pp.create_transformer_from_parameters(net, bus_9, bus_14, sn_mva=1., vn_hv_kv=4.16,
                                          vn_lv_kv=4.16, vkr_percent=3.566, vk_percent=5.0784,
                                          pfe_kw=0.0, i0_percent=0.0, shift_degree=0.0,
                                          tap_side='lv', tap_neutral=0, tap_max=16, tap_min=-16,
                                          tap_step_percent=0.625, tap_pos=0, tap_phase_shifter=False,
                                          name='Regulator 2')

    # Regulator 3
    pp.create_transformer_from_parameters(net, bus_25, bus_26, sn_mva=1., vn_hv_kv=4.16,
                                          vn_lv_kv=4.16, vkr_percent=1.6916, vk_percent=4.2372,
                                          pfe_kw=0.0, i0_percent=0.0, shift_degree=0.0,
                                          tap_side='lv', tap_neutral=0, tap_max=16, tap_min=-16,
                                          tap_step_percent=0.625, tap_pos=0, tap_phase_shifter=False,
                                          name='Regulator 3')

    # Regulator 4
    pp.create_transformer_from_parameters(net, bus_160, bus_67, sn_mva=1., vn_hv_kv=4.16,
                                          vn_lv_kv=4.16, vkr_percent=2.0288, vk_percent=4.6252,
                                          pfe_kw=0.0, i0_percent=0.0, shift_degree=0.0,
                                          tap_side='lv', tap_neutral=0, tap_max=2, tap_min=-2,
                                          tap_step_percent=2.5, tap_pos=0, tap_phase_shifter=False,
                                          name='Regulator 4')

    # Loads
    pp.create_load(net, bus_1,   p_mw=0.040, q_mvar=0.020, name='Load 1')
    pp.create_load(net, bus_2,   p_mw=0.020, q_mvar=0.010, name='Load 2')
    pp.create_load(net, bus_4,   p_mw=0.040, q_mvar=0.020, name='Load 4')
    pp.create_load(net, bus_5,   p_mw=0.020, q_mvar=0.010, name='Load 5')
    pp.create_load(net, bus_6,   p_mw=0.040, q_mvar=0.020, name='Load 5')
    pp.create_load(net, bus_7,   p_mw=0.020, q_mvar=0.010, name='Load 7')
    pp.create_load(net, bus_9,   p_mw=0.040, q_mvar=0.020, name='Load 9')
    pp.create_load(net, bus_10,  p_mw=0.020, q_mvar=0.010, name='Load 10')
    pp.create_load(net, bus_11,  p_mw=0.040, q_mvar=0.020, name='Load 11')
    pp.create_load(net, bus_12,  p_mw=0.020, q_mvar=0.010, name='Load 12')
    pp.create_load(net, bus_16,  p_mw=0.040, q_mvar=0.020, name='Load 16')
    pp.create_load(net, bus_17,  p_mw=0.020, q_mvar=0.010, name='Load 17')
    pp.create_load(net, bus_19,  p_mw=0.040, q_mvar=0.020, name='Load 19')
    pp.create_load(net, bus_20,  p_mw=0.040, q_mvar=0.020, name='Load 20')
    pp.create_load(net, bus_22,  p_mw=0.040, q_mvar=0.020, name='Load 22')
    pp.create_load(net, bus_24,  p_mw=0.040, q_mvar=0.020, name='Load 24')
    pp.create_load(net, bus_28,  p_mw=0.040, q_mvar=0.020, name='Load 28')
    pp.create_load(net, bus_29,  p_mw=0.040, q_mvar=0.020, name='Load 29')
    pp.create_load(net, bus_30,  p_mw=0.040, q_mvar=0.020, name='Load 30')
    pp.create_load(net, bus_31,  p_mw=0.020, q_mvar=0.010, name='Load 31')
    pp.create_load(net, bus_32,  p_mw=0.020, q_mvar=0.010, name='Load 32')
    pp.create_load(net, bus_33,  p_mw=0.040, q_mvar=0.020, name='Load 33')
    pp.create_load(net, bus_34,  p_mw=0.040, q_mvar=0.020, name='Load 34')
    pp.create_load(net, bus_35,  p_mw=0.040, q_mvar=0.020, name='Load 35')
    pp.create_load(net, bus_37,  p_mw=0.040, q_mvar=0.020, name='Load 37')
    pp.create_load(net, bus_38,  p_mw=0.020, q_mvar=0.010, name='Load 38')
    pp.create_load(net, bus_39,  p_mw=0.020, q_mvar=0.010, name='Load 39')
    pp.create_load(net, bus_41,  p_mw=0.020, q_mvar=0.010, name='Load 41')
    pp.create_load(net, bus_42,  p_mw=0.020, q_mvar=0.010, name='Load 42')
    pp.create_load(net, bus_43,  p_mw=0.040, q_mvar=0.020, name='Load 43')
    pp.create_load(net, bus_45,  p_mw=0.020, q_mvar=0.010, name='Load 45')
    pp.create_load(net, bus_46,  p_mw=0.020, q_mvar=0.010, name='Load 46')
    pp.create_load(net, bus_47,  p_mw=0.105, q_mvar=0.075, name='Load 47')
    pp.create_load(net, bus_48,  p_mw=0.210, q_mvar=0.150, name='Load 48')
    pp.create_load(net, bus_49,  p_mw=0.140, q_mvar=0.095, name='Load 49')
    pp.create_load(net, bus_50,  p_mw=0.040, q_mvar=0.020, name='Load 50')
    pp.create_load(net, bus_51,  p_mw=0.020, q_mvar=0.010, name='Load 51')
    pp.create_load(net, bus_52,  p_mw=0.040, q_mvar=0.020, name='Load 52')
    pp.create_load(net, bus_53,  p_mw=0.040, q_mvar=0.020, name='Load 53')
    pp.create_load(net, bus_55,  p_mw=0.020, q_mvar=0.010, name='Load 55')
    pp.create_load(net, bus_56,  p_mw=0.020, q_mvar=0.010, name='Load 56')
    pp.create_load(net, bus_58,  p_mw=0.020, q_mvar=0.010, name='Load 58')
    pp.create_load(net, bus_59,  p_mw=0.020, q_mvar=0.010, name='Load 59')
    pp.create_load(net, bus_60,  p_mw=0.020, q_mvar=0.010, name='Load 60')
    pp.create_load(net, bus_62,  p_mw=0.040, q_mvar=0.020, name='Load 62')
    pp.create_load(net, bus_63,  p_mw=0.040, q_mvar=0.020, name='Load 63')
    pp.create_load(net, bus_64,  p_mw=0.075, q_mvar=0.035, name='Load 64')
    pp.create_load(net, bus_65,  p_mw=0.140, q_mvar=0.095, name='Load 65')
    pp.create_load(net, bus_66,  p_mw=0.075, q_mvar=0.035, name='Load 66')
    pp.create_load(net, bus_68,  p_mw=0.020, q_mvar=0.010, name='Load 68')
    pp.create_load(net, bus_69,  p_mw=0.040, q_mvar=0.020, name='Load 69')
    pp.create_load(net, bus_70,  p_mw=0.020, q_mvar=0.010, name='Load 70')
    pp.create_load(net, bus_71,  p_mw=0.040, q_mvar=0.020, name='Load 71')
    pp.create_load(net, bus_73,  p_mw=0.040, q_mvar=0.020, name='Load 73')
    pp.create_load(net, bus_74,  p_mw=0.040, q_mvar=0.020, name='Load 74')
    pp.create_load(net, bus_75,  p_mw=0.040, q_mvar=0.020, name='Load 75')
    pp.create_load(net, bus_76,  p_mw=0.245, q_mvar=0.180, name='Load 76')
    pp.create_load(net, bus_77,  p_mw=0.040, q_mvar=0.020, name='Load 77')
    pp.create_load(net, bus_79,  p_mw=0.040, q_mvar=0.020, name='Load 79')
    pp.create_load(net, bus_80,  p_mw=0.040, q_mvar=0.020, name='Load 80')
    pp.create_load(net, bus_82,  p_mw=0.040, q_mvar=0.020, name='Load 82')
    pp.create_load(net, bus_83,  p_mw=0.020, q_mvar=0.010, name='Load 83')
    pp.create_load(net, bus_84,  p_mw=0.020, q_mvar=0.010, name='Load 84')
    pp.create_load(net, bus_85,  p_mw=0.040, q_mvar=0.020, name='Load 85')
    pp.create_load(net, bus_86,  p_mw=0.020, q_mvar=0.010, name='Load 86')
    pp.create_load(net, bus_87,  p_mw=0.040, q_mvar=0.020, name='Load 87')
    pp.create_load(net, bus_88,  p_mw=0.040, q_mvar=0.020, name='Load 88')
    pp.create_load(net, bus_90,  p_mw=0.040, q_mvar=0.020, name='Load 90')
    pp.create_load(net, bus_92,  p_mw=0.040, q_mvar=0.020, name='Load 92')
    pp.create_load(net, bus_94,  p_mw=0.040, q_mvar=0.020, name='Load 94')
    pp.create_load(net, bus_95,  p_mw=0.020, q_mvar=0.010, name='Load 95')
    pp.create_load(net, bus_96,  p_mw=0.020, q_mvar=0.010, name='Load 96')
    pp.create_load(net, bus_98,  p_mw=0.040, q_mvar=0.020, name='Load 98')
    pp.create_load(net, bus_99,  p_mw=0.040, q_mvar=0.020, name='Load 99')
    pp.create_load(net, bus_100, p_mw=0.040, q_mvar=0.020, name='Load 100')
    pp.create_load(net, bus_102, p_mw=0.020, q_mvar=0.010, name='Load 102')
    pp.create_load(net, bus_103, p_mw=0.040, q_mvar=0.020, name='Load 103')
    pp.create_load(net, bus_104, p_mw=0.040, q_mvar=0.020, name='Load 104')
    pp.create_load(net, bus_106, p_mw=0.040, q_mvar=0.020, name='Load 106')
    pp.create_load(net, bus_107, p_mw=0.040, q_mvar=0.020, name='Load 107')
    pp.create_load(net, bus_109, p_mw=0.040, q_mvar=0.020, name='Load 109')
    pp.create_load(net, bus_111, p_mw=0.020, q_mvar=0.010, name='Load 111')
    pp.create_load(net, bus_112, p_mw=0.020, q_mvar=0.010, name='Load 112')
    pp.create_load(net, bus_113, p_mw=0.040, q_mvar=0.020, name='Load 113')
    pp.create_load(net, bus_114, p_mw=0.020, q_mvar=0.010, name='Load 114')

    # External grid
    pp.create_ext_grid(net, bus_150, vm_pu=1.0, va_degree=0.0, s_sc_max_mva=5.0,
                       s_sc_min_mva=-5.0, rx_max=1, rx_min=1, r0x0_max=1, x0x_max=1)

    # Switches
    pp.create_switch(net, bus=bus_18,  element=bus_135, et='b', type="LS", z_ohm=0.1)
    pp.create_switch(net, bus=bus_13,  element=bus_152, et='b', type="LS", z_ohm=0.1)
    pp.create_switch(net, bus=bus_54,  element=bus_94,  et='b', type="LS", z_ohm=0.1)
    pp.create_switch(net, bus=bus_60,  element=bus_160, et='b', type="LS", z_ohm=0.1)
    pp.create_switch(net, bus=bus_97,  element=bus_197, et='b', type="LS", z_ohm=0.1)

    # Distributed generators
    pp.create_sgen(net, bus_24,  p_mw=0., q_mvar=0., name='DG 1', max_p_mw=0.66, min_p_mw=0, max_q_mvar=0.5, min_q_mvar=0)
    # pp.create_sgen(net, bus_41,  p_mw=0., q_mvar=0., name='DG 2', max_p_mw=0.66, min_p_mw=0, max_q_mvar=0.5, min_q_mvar=0)
    pp.create_sgen(net, bus_94,  p_mw=0., q_mvar=0., name='DG 3', max_p_mw=0.50, min_p_mw=0, max_q_mvar=0.375, min_q_mvar=0)
    pp.create_sgen(net, bus_71,  p_mw=0., q_mvar=0., name='DG 4', max_p_mw=0.50, min_p_mw=0, max_q_mvar=0.375, min_q_mvar=0)
    # pp.create_sgen(net, bus_114, p_mw=0., q_mvar=0., name='DG 5', max_p_mw=0.40, min_p_mw=0, max_q_mvar=0.3, min_q_mvar=0)
    pp.create_sgen(net, bus_22,  p_mw=0., type='PV', name='PV 1', max_p_mw=0.1, min_p_mw=0, max_q_mvar=0, min_q_mvar=0)
    pp.create_sgen(net, bus_250, p_mw=0., type='PV', name='PV 2', max_p_mw=0.1, min_p_mw=0, max_q_mvar=0, min_q_mvar=0)
    pp.create_sgen(net, bus_43,  p_mw=0., type='PV', name='PV 3', max_p_mw=0.1, min_p_mw=0, max_q_mvar=0, min_q_mvar=0)
    pp.create_sgen(net, bus_450, p_mw=0., type='PV', name='PV 4', max_p_mw=0.1, min_p_mw=0, max_q_mvar=0, min_q_mvar=0)
    pp.create_sgen(net, bus_39,  p_mw=0., type='PV', name='PV 5', max_p_mw=0.1, min_p_mw=0, max_q_mvar=0, min_q_mvar=0)
    pp.create_sgen(net, bus_4,   p_mw=0., type='WP', name='WP 1', max_p_mw=0.1, min_p_mw=0, max_q_mvar=0, min_q_mvar=0)
    pp.create_sgen(net, bus_59,  p_mw=0., type='WP', name='WP 2', max_p_mw=0.1, min_p_mw=0, max_q_mvar=0, min_q_mvar=0)
    pp.create_sgen(net, bus_46,  p_mw=0., type='WP', name='WP 3', max_p_mw=0.1, min_p_mw=0, max_q_mvar=0, min_q_mvar=0)
    pp.create_sgen(net, bus_75,  p_mw=0., type='WP', name='WP 4', max_p_mw=0.1, min_p_mw=0, max_q_mvar=0, min_q_mvar=0)
    pp.create_sgen(net, bus_83,  p_mw=0., type='WP', name='WP 5', max_p_mw=0.1, min_p_mw=0, max_q_mvar=0, min_q_mvar=0)

    # Shunt capacity bank
    pp.create_shunt_as_capacitor(net, bus_108, q_mvar=-0.3, name='SCB 1', step=0, max_step=4, loss_factor=0.0)
    pp.create_shunt_as_capacitor(net, bus_76,  q_mvar=-0.3, name='SCB 2', step=0, max_step=4, loss_factor=0.0)

    # storage
    pp.create_storage(net, bus_20, p_mw=0.0, max_e_mwh=2, sn_mva=1.0, soc_percent=50, min_e_mwh=0.2, name='Storage 1')
    pp.create_storage(net, bus_56, p_mw=0.0, max_e_mwh=2, sn_mva=1.0, soc_percent=50, min_e_mwh=0.2, name='Storage 2')
    #pp.create_storage(net, bus_113, p_mw=0.0, max_e_mwh=1, sn_mva=1.0, soc_percent=50, min_e_mwh=0.1, name='Storage 3')

    return net

# net = IEEE123Bus()
# pp.runpp(net)
# print(net.res_bus)
# print(net.res_ext_grid)
# print(net.trafo)
# print(net.res_load['p_mw'].values.sum())


from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import shutil
import sys
import glob

import numpy as np

import tflite_runtime.interpreter as tflite
from PIL import Image

from PyQt5.QtGui import *
from PyQt5.QtWidgets import *
from PyQt5.QtCore import *

input_height   = 299
input_width    = 299
input_mean     = 0
input_std      = 255
input_layer    = "input"
output_layer   = "InceptionV3/Predictions/Reshape_1"

def _most_recent_model():
    files = glob.glob("models/*.tflite")
    if len(files) > 0:
        files.sort(reverse=True, key=lambda name: name[-11:])
        labels = files[0].replace(".tflite", ".txt");
        return files[0], labels
    return None,None

def classify_image(file_name, return_filename=False):
    input_details = interpreter.get_input_details()
    output_details = interpreter.get_output_details()
    floating_model = input_details[0]['dtype'] == np.float32

    # Load the image
    height = input_details[0]['shape'][1]
    width = input_details[0]['shape'][2]
    img = Image.open(file_name).resize((width, height))

    input_data = np.expand_dims(img, axis=0)

    if floating_model:
        input_data = (np.float32(input_data) - input_mean) / input_std

    # run inference
    interpreter.set_tensor(input_details[0]['index'], input_data)
    interpreter.invoke()

    output_data = interpreter.get_tensor(output_details[0]['index'])
    results = np.squeeze(output_data)
    predictions = results.argsort()[-5:][::-1]

    if(return_filename):
        p = predictions[0]
        return "%s (%d%%)" % (species_labels[p], results[p] * 100)

    return_value = ""
    for i in predictions:
        return_value += "%s (%d%%)\n" % (species_labels[i], results[i] * 100)

    return return_value

def load_labels(label_file):
    with open(label_file, 'r') as f:
        return [line.strip() for line in f.readlines()]

class ClassifyDirectoryThread(QThread):
    count = pyqtSignal(int)
    display = pyqtSignal('QString')
    result = pyqtSignal('QString')
    progress = pyqtSignal(int)
    complete = pyqtSignal()

    def __init__(self, path):
        QThread.__init__(self)
        self.path = path
        self.running = True

    def __del__(self):
        self.wait()

    def stop(self):
        self.running = False

    def run(self):
        result_path = os.path.join(self.path, "results")
        if not os.path.isdir(result_path):
            os.mkdir(result_path)

        files = glob.glob(os.path.join(self.path, "*.bmp"))
        files.extend(glob.glob(os.path.join(self.path, "*.gif")))
        files.extend(glob.glob(os.path.join(self.path, "*.jpg")))
        files.extend(glob.glob(os.path.join(self.path, "*.jpeg")))
        files.extend(glob.glob(os.path.join(self.path, "*.png")))

        self.count.emit(len(files))

        i = 0
        for file in files:
            if not self.running:
                break;

            result = classify_image(file, True)
            self.display.emit(file)
            self.result.emit(result)
            self.progress.emit(i)

            file_name, file_ext = os.path.splitext(os.path.basename(file))
            destination = os.path.join(result_path, file_name + " - " + result + file_ext)
            shutil.copyfile(file, destination)
            i += 1

        self.complete.emit()

class ClassifyImageThread(QThread):
    result = pyqtSignal('QString')
    complete = pyqtSignal()

    def __init__(self, path):
        QThread.__init__(self)
        self.path = path

    def __del__(self):
        self.wait()

    def run(self):
        result = classify_image(self.path, False)
        self.result.emit(result)
        self.complete.emit()

class ImageLabel(QLabel):
    def __init__(self, image):
        super(ImageLabel, self).__init__("")
        self.setScaledContents(True)
        self.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
        self.setAlignment(Qt.AlignCenter)
        self.setImage(image)

    def setImage(self, image):
        self.pixmap = QPixmap(image).scaled(400, 400, Qt.KeepAspectRatio, Qt.SmoothTransformation)
        self.setPixmap(self.pixmap)

    def width(self):
        return self.pixmap.width()

    def height(self):
        return self.pixmap.height()

class Fitch(QMainWindow):
    progdialog = None
    thread = None

    def __init__(self):
        super(Fitch, self).__init__()
        self.setAcceptDrops(True)
        self.initUI()

    def closeEvent(self, e):
        self.quit()

    def dragEnterEvent(self, e):
        if e.mimeData().hasUrls:
            e.accept()
        else:
            e.ignore()

    def dropEvent(self, e):
        path = str(e.mimeData().urls()[0].toLocalFile())

        if os.path.isdir(path):
            self.classifyDirectory(path)
        else:
            self.displayAndClassifyImage(path)

    def quit(self):
        app.quit()
        sys.exit()

    def initUI(self):
        self.setWindowTitle("Fitch - " + os.path.basename(species_model))
        self.setWindowIcon(QIcon('fitch.ico'))
        widget = QWidget(self)
        self.setCentralWidget(widget)
        layout = QVBoxLayout()
        widget.setLayout(layout)

        fileAct = QAction('Classify &Image', self)
        fileAct.setShortcut('Ctrl+I')
        fileAct.setStatusTip('Classify a single image')
        fileAct.triggered.connect(self.classifyImage)

        directoryAct = QAction('Classify &Directory', self)
        directoryAct.setShortcut('Ctrl+D')
        directoryAct.setStatusTip('Classify all images in a directory')
        directoryAct.triggered.connect(self.menuClassifyDirectory)

        chooseAct = QAction('Choose &Model', self)
        chooseAct.setShortcut('Ctrl+M')
        chooseAct.setStatusTip('Choose a new model file for classifications.')
        chooseAct.triggered.connect(self.chooseModelFile)

        exitAct = QAction('&Exit', self)
        exitAct.setShortcut('Ctrl+Q')
        exitAct.setStatusTip('Exit application')
        exitAct.triggered.connect(self.quit)

        menubar = self.menuBar()
        fileMenu = menubar.addMenu('&File')
        fileMenu.addAction(fileAct)
        fileMenu.addAction(directoryAct)
        fileMenu.addSeparator()
        fileMenu.addAction(chooseAct)
        fileMenu.addSeparator()
        fileMenu.addAction(exitAct)

        self.label = ImageLabel("drop.png");
        layout.addWidget(self.label)

        self.text = QTextEdit()
        self.text.setReadOnly(True)
        self.text.setAcceptDrops(False)
        layout.addWidget(self.text)

        self.resizeUI()

        rect = self.frameGeometry()
        center = QDesktopWidget().availableGeometry().center()
        rect.moveCenter(center)
        self.move(rect.topLeft())

        self.show()

    def resizeUI(self):
        self.resize(self.label.width(), self.label.height() + self.menuBar().height() + + 150);

    def displayImage(self, path):
        self.label.setImage(path)
        self.resizeUI()

    def displayAndClassifyImage(self, path):
        self.displayImage(path)
        self.text.setText('')

        QApplication.setOverrideCursor(Qt.WaitCursor)

        self.thread = ClassifyImageThread(path)
        self.thread.result.connect(self.signalResult)
        self.thread.complete.connect(self.signalComplete)
        self.thread.start()
        return

    def signalCancel(self):
        if self.thread:
            self.thread.stop()
            self.thread.wait()
            self.thread = None
        if self.progdialog:
            self.progdialog.close()
            self.progdialog = None

    def signalImageCount(self, count):
        if self.progdialog:
            self.progdialog.setMaximum(count)

    def signalProgress(self, value):
        if self.progdialog:
            self.progdialog.setValue(value)

    def signalComplete(self):
        if self.progdialog:
            self.progdialog.close()
        QApplication.restoreOverrideCursor()

    def signalResult(self, result):
        self.text.setText(result)

    def classifyImage(self):
        path = QFileDialog.getOpenFileName(self, "Choose an image", "", "Images (*.bmp *.gif *.jpg *.jpeg *.png)")
        if(path[0]):
            self.displayAndClassifyImage(path[0])

    def classifyDirectory(self, path):
        self.progdialog = QProgressDialog("", "Cancel", 0, 100, self)
        self.progdialog.setWindowTitle("Classifying")
        self.progdialog.setWindowModality(Qt.WindowModal)
        self.progdialog.canceled.connect(self.signalCancel)
        self.progdialog.show()

        self.thread = ClassifyDirectoryThread(path)
        self.thread.count.connect(self.signalImageCount)
        self.thread.display.connect(self.displayImage)
        self.thread.result.connect(self.signalResult)
        self.thread.progress.connect(self.signalProgress)
        self.thread.complete.connect(self.signalComplete)
        self.thread.start()

    def menuClassifyDirectory(self):
        path = QFileDialog.getExistingDirectory(self, "Choose a directory", "", QFileDialog.ShowDirsOnly)
        if(path):
            self.classifyDirectory(path)

    def chooseModelFile(self):
        global species_model, species_graph, species_labels, interpreter
        path = QFileDialog.getOpenFileName(self, "Choose a model file", "./models", "Models (*.tflite)")
        if(path[0]):
            file_model  = path[0]
            file_labels = path[0].replace(".tflite", ".txt");
            if os.path.exists(file_labels):
                species_model = file_model
                interpreter = tflite.Interpreter(model_path=species_model)
                interpreter.allocate_tensors()
                species_labels = load_labels(file_labels)
                self.setWindowTitle("Moth ID - " + os.path.basename(species_model))
            else:
                QMessageBox.warning(self, "Error", "Label file not found.")

if __name__ == '__main__':
    app = QApplication(sys.argv)

    species_model, species_labels = _most_recent_model()
    if species_model == None:
        QMessageBox.critical(None, "Error", "Could not find default model file.")
        sys.exit(-1)

    interpreter = tflite.Interpreter(model_path=species_model)
    interpreter.allocate_tensors()

    species_labels = load_labels(species_labels)

    win = Fitch()
    sys.exit(app.exec_())


from tensorflow.keras.preprocessing.image import img_to_array
from tensorflow.keras.models import load_model
import numpy as np
import cv2
import os
import cvlib as cv

# load model
model = load_model('model.h5')

# open webcam
webcam = cv2.VideoCapture(0)

# loop through frames
while webcam.isOpened():

    # read frame from webcam
    status, frame = webcam.read()

    # apply face detection
    face, confidence = cv.detect_face(frame)

    #arrays to hold the values
    age_ = []
    gender_ = []

    # loop through detected faces
    for idx, f in enumerate(face):

        # get corner points of face rectangle
        (startX, startY) = f[0], f[1]
        (endX, endY) = f[2], f[3]

        # draw rectangle over face
        cv2.rectangle(frame, (startX,startY), (endX,endY), (0,255,0), 4)

        # crop the detected face region
        face_crop = np.copy(frame[startY:endY,startX:endX])

        if (face_crop.shape[0]) < 10 or (face_crop.shape[1]) < 10:
            continue

        # preprocessing for gender detection model
        face_crop = cv2.resize(face_crop, (200,200))
        face_crop = face_crop.astype("float") / 255.0
        face_crop = img_to_array(face_crop)
        face_crop = np.expand_dims(face_crop, axis=0)

        # apply gender and age detection on face
        predict = model.predict(np.array(face_crop).reshape(-1,200,200,3))
        age_.append(predict[0])
        gender_.append(np.argmax(predict[1]))
        gend = np.argmax(predict[1])
        if gend == 0:
            gend = 'Man'
            col = (255,0,0)
        else:
            gend = 'Woman'
            col = (203,12,255)

        #labels
        label =  "Age : " + str(int(predict[0])) + " / " + str(gend)

        # write label and confidence above face rectangle
        cv2.putText(frame, label, (startX, StartY),  cv2.FONT_HERSHEY_SIMPLEX,
                    0.7, col, 4)

    # display output
    cv2.imshow("age and gender detection", frame)

    # press "Q" to stop
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

# release resources
webcam.release()
cv2.destroyAllWindows()


#!/usr/bin/env python3
# coding: utf8

"""
Model of a prediction of a new song.
"""

__author__ = 'David Flury, Andreas Kaufmann, Raphael Müller'
__email__ = "info@unmix.io"

import glob
import h5py
import os
import math
import numpy as np
import progressbar
import librosa

from unmix.source.configuration import Configuration
from unmix.source.data.track import Track
from unmix.source.exceptions.dataerror import DataError
from unmix.source.helpers import converter
from unmix.source.logging.logger import Logger

class Prediction(object):

    def __init__(self, engine, sample_rate=22050, fft_window=1536, stereo=False):
        self.model = engine.model
        self.transformer = engine.transformer
        self.graph = engine.graph
        self.sample_rate = sample_rate
        self.sample_rate_origin = 0
        self.fft_window = fft_window
        self.stereo = stereo
        self.mix = np.empty(0)
        self.vocals = np.empty(0)
        self.instrumental = np.empty(0)
        self.progress = 0
        self.initialized = False
        self.length = 0

    def save(self, file, folder='', extension='wav'):
        'Saves the predicted vocals and instrumental track to separate audio files.'
        return self.save_vocals(file, folder, extension), self.save_instrumental(file, folder, extension)

    def save_vocals(self, file, folder='', extension='wav'):
        'Saves the predicted vocals track to an audio file.'
        return self.__save(self.vocals, 'vocals', file, folder, extension)

    def save_instrumental(self, file, folder='', extension='wav'):
        'Saves the predicted instrumental track to an audio file.'
        return self.__save(self.instrumental, 'instrumental', file, folder, extension)

    def __save(self, prediction, type, file, folder='', extension='wav'):
        name = os.path.splitext(os.path.basename(file))[0]
        file_name = converter.get_timestamp() + "_" + name + \
            '_predicted_%s.%s' % (type, extension)
        if folder:
            output_file = os.path.join(
                folder, file_name)
        else:
            output_file = os.path.join(os.path.dirname(file), file_name)

        if self.stereo:
            track = np.array([librosa.istft(channel)
                              for channel in prediction])
        else:
            track = librosa.istft(prediction[0])
        librosa.output.write_wav(
            output_file, track, self.sample_rate, norm=False)
        Logger.info("Output prediction file: %s" % output_file)
        return output_file

    def predict_part(self, i, part):
        with self.graph.as_default():
            predicted = self.model.predict(np.array([part]))[0]
        predicted_vocals, predicted_instrumental = \
            self.transformer.untransform_target(
                self.mix, predicted, i)
        if not self.initialized:
            self.__init_shapes(predicted_vocals.shape)

        self.__expand_track(predicted_vocals, self.vocals, i)
        self.__expand_track(predicted_instrumental, self.instrumental, i)
        self.progress += 1

    def __expand_track(self, prediction, track, i):
        # Transformers can have their own expander - use that if available
        if(getattr(self.transformer, "expand_track", None)):
            self.transformer.expand_track(prediction, track, i)
        else:
            left = prediction.shape[2] * i
            right = prediction.shape[2] * (i+1)
            size = track.shape[2]
            if size < right:
                track = np.append(track, np.zeros(
                    (track.shape[0], track.shape[1], right - size)), axis=2)
            track[:, :, left:right] = prediction

    def __init_shapes(self, shape):
        self.vocals = np.zeros(
            (shape[0], shape[1], self.transformer.step * self.length), np.complex)
        self.instrumental = np.zeros_like(self.vocals)
        self.initialized = True

    def unpad(self):
        mix_length = self.mix[0].shape[1]
        offset = self.transformer.size // 2
        self.vocals = self.vocals[:, :, offset:offset + mix_length]
        self.instrumental = self.instrumental[:, :, offset:offset + mix_length]


# -*- coding: utf-8 -*-
"""
Created on Wed Nov 28 15:13:08 2018

@author: User
"""

import numpy as np
from keras.models import load_model
import json
from DataGenerator import *

def accuracy_test(model_name, test=test):
    '''
    Оценка точности предсказанных нейросетью данных.
    Входные данные:
        model_name - str, путь к сохраненной модели
        test - str, путь к тестовой выборке
    Выходные данные:
        mean_acc - float, точность
    '''
    dg_test = DataGenerator(test)
    Num = 10  # число эпох генератора
    model = load_model(model_name)
    acc = []
    for i in range(Num):
        length = dg_test.__len__()
        for j in range(length):
            X, Y = dg_test.__getitem__(j)
            Y_pred = model.predict_classes(X)
            Y_not = []
            for y in Y[0]:
                Y_not.append(list(y).index(1.))
            Y_diff = Y_pred[0] - np.array(Y_not)
            count = 0
            for y in Y_diff:
                if y == 0:
                    count += 1
            acc.append(count / len(Y_diff))
        dg_test.on_epoch_end()
    return np.mean(acc)

def predict_single_sample(model_name, test=test):
    '''
    Предсказание на одном образце.
    Входные данные:
        model_name - str, путь к сохраненной модели
        test - str, путь к тестовой выборке
    '''
    sig, json_arr = mixGenerate(json.load(open(test)), duration=5,
                                              shuffle=True)
    matrix = stft(sig)
    target = notation(sig, json_arr)
    model = load_model(model_name)
    res = model.predict_classes(np.array([matrix]))
    print('Predicted notation:', res)
    print('Real notation:', target)
    diff = res[0] - np.array(target)
    count = 0
    for y in diff:
        if y == 0:
            count += 1
    acc = count / len(diff)
    print('Accuracy:', acc)

def main():
    model_name = '1_500_81_1_500_5_with_regs_0_6overlapPool_e-5_decay5e-4_abs_norm_peak_norm.600-1.36.hdf5'
    #print(accuracy_test(model_name='1_500_81_1_500_5_with_regs_0_6overlapPool_e-5_decay5e-4_abs_norm_peak_norm.600-1.36.hdf5'))
    predict_single_sample(model_name)
    return 0

if __name__ == '__main__':
    main()


import numpy as np
from xml.etree.ElementTree import Element
from .shape import Shape
from ..shape_util import find_corners, rectangle_to_box, poly_to_mask

class Rectangle(Shape):
    """Class for a single rectangle

    Parameters
    ----------
    data : np.ndarray
        Nx2 array of vertices specifying the shape.
    edge_width : float
        thickness of lines and edges.
    edge_color : str | tuple
        If string can be any color name recognized by vispy or hex value if
        starting with `#`. If array-like must be 1-dimensional array with 3 or
        4 elements.
    face_color : str | tuple
        If string can be any color name recognized by vispy or hex value if
        starting with `#`. If array-like must be 1-dimensional array with 3 or
        4 elements.
    opacity : float
        Opacity of the shape, must be between 0 and 1.
    z_index : int
        Specifier of z order priority. Shapes with higher z order are displayed
        ontop of others.
    """
    def __init__(self, data, *, edge_width=1, edge_color='black',
                 face_color='white', opacity=1, z_index=0):

        super().__init__(edge_width=edge_width, edge_color=edge_color,
                         face_color=face_color, opacity=opacity,
                         z_index=z_index)

        self._closed = True
        self.data = np.array(data)
        self.name = 'rectangle'

    @property
    def data(self):
        """np.ndarray: Nx2 array of vertices.
        """
        return self._data

    @data.setter
    def data(self, data):
        if len(data) == 2:
            data = find_corners(data)
        if len(data) != 4:
            raise ValueError("""Data shape does not match a rectangle.
                             Rectangle expects four corner vertices""")
        else:
            # Add four boundary lines and then two triangles for each
            self._set_meshes(data, face=False)
            self._face_vertices = data
            self._face_triangles = np.array([[0, 1, 2], [0, 2, 3]])
            self._box = rectangle_to_box(data)

        self._data = data

    def to_mask(self, mask_shape=None):
        """Converts the shape vertices to a boolean mask with `True` for points
        lying inside the shape.

        Parameters
        ----------
        mask_shape : np.ndarray | tuple | None
            1x2 array of shape of mask to be generated. If non specified, takes
            the max of the vertiecs

        Returns
        ----------
        mask : np.ndarray
            Boolean array with `True` for points inside the shape
        """
        if mask_shape is None:
            mask_shape = self.data.max(axis=0).astype('int')

        mask = poly_to_mask(mask_shape, self.data)

        return mask

    def to_xml(self):
        """Generates an xml element that defintes the shape according to the
        svg specification.

        Returns
        ----------
        element : xml.etree.ElementTree.Element
            xml element specifying the shape according to svg.
        """
        props = self.svg_props
        data = self.data[:, ::-1]

        offset = data[1] - data[0]
        angle = -np.arctan2(offset[0], -offset[1])
        if not angle == 0:
            # if shape has been rotated, shift to origin
            cen = data.mean(axis=0)
            coords = data - cen

            # rotate back to axis aligned
            c, s = np.cos(angle), np.sin(-angle)
            rotation = np.array([[c, s],
                                 [-s, c]])
            coords = coords @ rotation.T

            # shift back to center
            coords = coords + cen

            # define rotation around center
            transform = f'rotate({np.degrees(-angle)} {cen[0]} {cen[1]})'
            props['transform'] = transform
        else:
            coords = data

        x = str(coords.min(axis=0)[0])
        y = str(coords.min(axis=0)[1])
        size = abs(coords[2] - coords[0])
        width = str(size[0])
        height = str(size[1])

        element = Element('rect', x=x, y=y, width=width, height=height,
                          **props)
        return element


from agent_dir.agent import Agent
import numpy as np
import scipy
import random

import os
import keras
import tensorflow as tf
from keras.models import Sequential,load_model, Model
from keras.layers import Dense, Dropout, Flatten, Multiply
from keras.layers import Conv2D, MaxPooling2D, Input, Lambda
from keras.optimizers import Adam, Adamax, RMSprop
from keras import backend as K

MAX_TIMESTEP = 1000
MAX_EP = 3000
from keras.backend.tensorflow_backend import set_session
config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.1
set_session(tf.Session(config=config))

random.seed(2)
np.random.seed(2)
tf.set_random_seed(2)

# ref : https://github.com/MorvanZhou/Reinforcement-learning-with-tensorflow/

def categorical_crossentropy(target, output):
    _epsilon =  tf.convert_to_tensor(10e-8, dtype=output.dtype.base_dtype)
    output = tf.clip_by_value(output, _epsilon, 1. - _epsilon)
    return (- target * tf.log(output))

class Agent_ActorCritic(Agent):
    def __init__(self, env, args):
        super(Agent_ActorCritic,self).__init__(env)

        self.log_path = './actor_critic.log'

        self.env = env
        self.actions_avialbe = env.action_space.n
        self.feature_dim = env.observation_space.shape[0]
        self.t = 0
        self.prev_x = None
        self.actor_learning_rate  = 1e-3
        self.critic_learning_rate = 1e-3
        self.gamma = 0.9

        self.dummy_act_picked = np.zeros((1,self.actions_avialbe))

        # Actor
        input_frame  = Input(shape=(self.feature_dim,))
        act_picked = Input(shape=(self.actions_avialbe,))
        hidden_f = Dense(20,activation='relu')(input_frame)

        act_prob = Dense(self.actions_avialbe,activation='softmax')(hidden_f)
        selected_act_prob = Multiply()([act_prob,act_picked])
        selected_act_prob = Lambda(lambda x:K.sum(x, axis=-1, keepdims=True),output_shape=(1,))(selected_act_prob)

        model = Model(inputs=[input_frame,act_picked], outputs=[act_prob, selected_act_prob])

        opt = Adam(lr=self.actor_learning_rate)
        model.compile(loss=['mse',categorical_crossentropy], loss_weights=[0.0,1.0],optimizer=opt)
        self.actor = model

        # Critic
        model = Sequential()
        model.add(Dense(20,activation='relu',input_shape=(self.feature_dim,)))
        model.add(Dense(1))

        opt = Adam(lr=self.critic_learning_rate)
        model.compile(loss='mse', optimizer=opt)
        self.critic = model

    def init_game_setting(self):
        self.prev_x = None

    def train(self):
        # Init
        log = open(self.log_path,'w')
        log.write('reward,avg_reward\n')
        batch_size = 1
        frames, prob_actions, dlogps, drs =[], [], [], []
        tr_x, tr_y = [],[]
        reward_record = []
        avg_reward = []
        reward_sum = 0
        ep_number = 0
        ep_step = 0
        #explore_rate = 0
        observation = self.env.reset()
        # Training progress
        while True:
            act = np.random.choice(np.arange(self.actions_avialbe),
                    p=self.actor.predict([np.expand_dims(observation,axis=0),self.dummy_act_picked])[0].flatten())

            act_one_hot = np.zeros((1,self.actions_avialbe))
            act_one_hot[0,act]=1.0
            next_observation, reward, done, info = self.env.step(act)
            if done: reward = -20

            reward_sum += reward
            predict_reward = self.critic.predict(np.expand_dims(observation,axis=0))
            predict_next_reward = self.critic.predict(np.expand_dims(next_observation,axis=0))

            td_target = np.expand_dims(reward,axis=0) + self.gamma*predict_next_reward
            td_error = td_target - predict_reward

            self.critic.train_on_batch(np.expand_dims(observation,axis=0),td_target)
            self.actor.train_on_batch([np.expand_dims(observation,axis=0),act_one_hot],[self.dummy_act_picked,td_error])

            observation = next_observation

            self.t += 1
            ep_step += 1

            if done or ep_step>MAX_TIMESTEP:
                ep_number += 1

                avg_reward.append(float(reward_sum))
                if len(avg_reward)>30: avg_reward.pop(0)

                print('EPISODE: {0:6d} / TIMESTEP: {1:8d} / REWARD: {2:5d} / AVG_REWARD: {3:2.3f} '.format(
                            ep_number, self.t, int(reward_sum), np.mean(avg_reward)))
                print('{:.4f},{:.4f}'.format(reward_sum,np.mean(avg_reward)),end='\n',file=log,flush=True)

                observation = self.env.reset()
                reward_sum = 0.0
                ep_step = 0

            if ep_number >= MAX_EP:
                self.actor.save('actor.h5')
                self.critic.save('critictor.h5')
                break

    def make_action(self, observation, test=True):
        """
        Input:
            observation: np.array
                current RGB screen of game, shape: (210, 160, 3)

        Return:
            action: int
                the predicted action from trained model
        """
        pass