# Path: selfsup/multi/methods/base.py
# class Method:
#     def __init__(self, name, batch_size):
#         self.name = name
#         self.batch_size = batch_size
# 
#     def input_size(self):
#         return 0
# 
#     def batch_loader(self):
#         raise NotImplemented("Abstract base class")
# 
#     def build_network(network, y, phase_test):
#         raise NotImplemented("Abstract base class")
# 
#     @property
#     def basenet_settings():
#         return {}
# 
#     def make_path(self, name, ext, iteration):
#         return f'img/{self.name}_{name}_{iteration}.png'


import selfsup
import tensorflow as tf
import os
import deepdish as dd
import numpy as np
from .base import Method
from collections import OrderedDict
from __future__ import division, print_function, absolute_import

# http://stackoverflow.com/questions/36904298/how-to-implement-multi-class-hinge-loss-in-tensorflow/36928137#36928137
def multi_class_hinge_loss(logits, labels, n_classes):
    batch_size = logits.get_shape().as_list()[0]
    # get the correct logit
    flat_logits = tf.reshape(logits, (-1,))
    correct_id = tf.range(0, batch_size) * n_classes + labels
    correct_logit = tf.gather(flat_logits, correct_id)

    # get the wrong maximum logit
    max_label = tf.to_int32(tf.argmax(logits, 1))
    top2, _ = tf.nn.top_k(logits, k=2, sorted=True)
    first, second = tf.unstack(top2, axis=1)
    wrong_max_logit = tf.where(tf.equal(max_label, labels), second, first)

    # calculate multi-class hinge loss
    return tf.maximum(0., 1.0 + wrong_max_logit - correct_logit), wrong_max_logit, correct_logit



# Path: selfsup/multi/methods/base.py
# class Method:
#     def __init__(self, name, batch_size):
#         self.name = name
#         self.batch_size = batch_size
# 
#     def input_size(self):
#         return 0
# 
#     def batch_loader(self):
#         raise NotImplemented("Abstract base class")
# 
#     def build_network(network, y, phase_test):
#         raise NotImplemented("Abstract base class")
# 
#     @property
#     def basenet_settings():
#         return {}
# 
#     def make_path(self, name, ext, iteration):
#         return f'img/{self.name}_{name}_{iteration}.png'


# Path: selfsup/multi/methods/supervised.py
import selfsup
import tensorflow as tf
import os
import deepdish as dd
import numpy as np
from .base import Method
from collections import OrderedDict

from __future__ import division, print_function, absolute_import

# http://stackoverflow.com/questions/36904298/how-to-implement-multi-class-hinge-loss-in-tensorflow/36928137#36928137
def multi_class_hinge_loss(logits, labels, n_classes):
    batch_size = logits.get_shape().as_list()[0]
    # get the correct logit
    flat_logits = tf.reshape(logits, (-1,))
    correct_id = tf.range(0, batch_size) * n_classes + labels
    correct_logit = tf.gather(flat_logits, correct_id)

    # get the wrong maximum logit
    max_label = tf.to_int32(tf.argmax(logits, 1))
    top2, _ = tf.nn.top_k(logits, k=2, sorted=True)
    first, second = tf.unstack(top2, axis=1)
    wrong_max_logit = tf.where(tf.equal(max_label, labels), second, first)

    # calculate multi-class hinge loss
    return tf.maximum(0., 1.0 + wrong_max_logit - correct_logit), wrong_max_logit, correct_logit



class Supervised(Method):