"""
logging.py

Log format convenience methods for training spinn.

"""

import numpy as np
from spinn.util.blocks import flatten
from spinn.util.misc import time_per_token


def train_accumulate(model, A):

    has_spinn = hasattr(model, 'spinn')
    has_transition_loss = hasattr(model, 'transition_loss')
    has_policy = has_spinn and hasattr(model, 'policy_loss')
    has_value = has_spinn and hasattr(model, 'value_loss')
    has_rae = has_spinn and hasattr(model.spinn, 'rae_loss')
    has_leaf = has_spinn and hasattr(model.spinn, 'leaf_loss')
    has_gen = has_spinn and hasattr(model.spinn, 'gen_loss')
    has_entropy = hasattr(model, 'avg_entropy')

    # Accumulate stats for transition accuracy.
    if has_transition_loss:
        preds = [m["t_preds"] for m in model.spinn.memories]
        truth = [m["t_given"] for m in model.spinn.memories]
        A.add('preds', preds)
        A.add('truth', truth)

    # Accumulate stats for leaf prediction accuracy.
    if has_leaf:
        A.add('leaf_acc', model.spinn.leaf_acc)

    # Accumulate stats for word prediction accuracy.
    if has_gen:
        A.add('gen_acc', model.spinn.gen_acc)

    if has_entropy:
        A.add('entropy', model.avg_entropy)


def train_stats(model, optimizer, A, step):

    has_spinn = hasattr(model, 'spinn')
    has_transition_loss = hasattr(model, 'transition_loss')
    has_policy = has_spinn and hasattr(model, 'policy_loss')
    has_value = has_spinn and hasattr(model, 'value_loss')
    has_rae = has_spinn and hasattr(model.spinn, 'rae_loss')
    has_leaf = has_spinn and hasattr(model.spinn, 'leaf_loss')
    has_gen = has_spinn and hasattr(model.spinn, 'gen_loss')
    has_epsilon = has_spinn and hasattr(model.spinn, "epsilon")
    has_entropy = hasattr(model, 'avg_entropy')

    if has_transition_loss:
        all_preds = np.array(flatten(A.get('preds')))
        all_truth = np.array(flatten(A.get('truth')))
        avg_trans_acc = (all_preds == all_truth).sum() / float(all_truth.shape[0])

    time_metric = time_per_token(A.get('total_tokens'), A.get('total_time'))

    ret = dict(
        step=step,
        class_acc=A.get_avg('class_acc'),
        transition_acc=avg_trans_acc if has_transition_loss else 0.0,
        xent_cost=A.get_avg('xent_cost'), # not actual mean
        transition_cost=model.transition_loss.data[0] if has_transition_loss else 0.0,
        l2_cost=A.get_avg('l2_cost'), # not actual mean
        policy_cost=model.policy_loss.data[0] if has_policy else 0.0,
        value_cost=model.spinn.value_loss.data[0] if has_value else 0.0,
        epsilon=model.spinn.epsilon if has_epsilon else 0.0,
        avg_entropy=A.get('avg_entropy') if has_entropy else 0.0,
        rae_cost=model.spinn.rae_loss.data[0] if has_rae else 0.0,
        leaf_acc=A.get_avg('leaf_acc') if has_leaf else 0.0,
        leaf_cost=model.spinn.leaf_loss.data[0] if has_leaf else 0.0,
        gen_acc=A.get_avg('gen_acc') if has_gen else 0.0,
        gen_cost=model.spinn.gen_loss.data[0] if has_gen else 0.0,
        learning_rate=optimizer.lr,
        time=time_metric,
    )

    total_cost = 0.0
    for key in ret.keys():
        if key == 'transition_cost' and has_transition_loss and model.optimize_transition_loss:
            total_cost += ret[key]
        elif 'cost' in key:
            total_cost += ret[key]
    ret['total_cost'] = total_cost

    return ret


def train_format(model):

    has_spinn = hasattr(model, 'spinn')

    stats_str = "Step: {step}"

    # Accuracy Component.
    stats_str += " Acc: {class_acc:.5f} {transition_acc:.5f}"
    if has_spinn and hasattr(model.spinn, 'leaf_loss'):
        stats_str += " leaf{leaf_acc:.5f}"
    if has_spinn and hasattr(model.spinn, 'gen_loss'):
        stats_str += " gen{gen_acc:.5f}"

    # Cost Component.
    stats_str += " Cost: {total_cost:.5f} {xent_cost:.5f} {transition_cost:.5f} {l2_cost:.5f}"
    if has_spinn and hasattr(model.spinn, 'policy_loss'):
        stats_str += " p{policy_cost:.5f}"
    if has_spinn and hasattr(model.spinn, 'value_loss'):
        stats_str += " v{value_cost:.5f}"
    if hasattr(model, 'avg_entropy'):
        stats_str += " e{avg_entropy:.5f}"
    if has_spinn and hasattr(model.spinn, 'rae_loss'):
        stats_str += " rae{rae_cost:.5f}"
    if has_spinn and hasattr(model.spinn, 'leaf_loss'):
        stats_str += " leaf{leaf_cost:.5f}"
    if has_spinn and hasattr(model.spinn, 'gen_loss'):
        stats_str += " gen{gen_cost:.5f}"

    # Time Component.
    stats_str += " Time: {time:.5f}"

    return stats_str


def train_extra_format(model):

    # Extra Component.
    extra_str = "Train Extra:"
    extra_str += " lr={learning_rate:.7f}"
    if hasattr(model, "spinn") and hasattr(model.spinn, "epsilon"):
        extra_str += " eps={epsilon:.7f}"

    return extra_str

@@ -3200,16 +3200,20 @@
 oss and 
+not 
 model.op
@@ -3248,38 +3248,20 @@
         
-total_cost += ret%5Bkey%5D
+pass
 %0A       

Remove transition_loss from total_loss when logging.