import tensorflow as tf
import selfsup
import os
from .base import Method
from autocolorize.tensorflow.sparse_extractor import sparse_extractor
from autocolorize.extraction import calc_rgb_from_hue_chroma
import functools
from collections import OrderedDict
import numpy as np

def kl_divergence(labels, logits):
    #return tf.contrib.distributions.kl(p, q)
    return (tf.nn.softmax_cross_entropy_with_logits(labels=labels, logits=logits) -
            tf.nn.softmax_cross_entropy_with_logits(labels=labels, logits=tf.log(labels+1e-5)))


class ColorizeHypercolumn3(Method):
    def __init__(self, name, basenet, loader, use_scalers=False, version=1):
        self.name = name
        self._loader = loader
        self.basenet = basenet
        self.locations = 128
        self.edge_buffer = 32
        self._bins = 8
        self._use_scalers = use_scalers
        self._version = version
        self._hist_win_radius = 3
        self._hist_win_size = self._hist_win_radius * 2 + 1

    @property
    def basenet_settings(self):
        return {'convolutional': True}

    def batch(self):
        #input_size = self.basenet.canonical_input_size
        x, _ = self._loader.batch()

        input_size = x.get_shape().as_list()[1]

        self.mult = 1
        self.centroids = tf.to_float(tf.random_uniform(
                shape=[self._loader.batch_size, self.locations, 2],
                minval=self.edge_buffer // self.mult,
                maxval=(input_size - self.edge_buffer) // self.mult,
                name='centroids',
                dtype=tf.int32))

        gray_single = tf.reduce_mean(x, 3, keep_dims=True)
        gray = tf.tile(gray_single, [1, 1, 1, 3])

        w = input_size // self.mult
        if self.mult > 1:
            x4 = tf.image.resize_bilinear(x, (w, w))
        else:
            x4 = x

        hsv = tf.image.rgb_to_hsv(x4)

        self.hue = tf.expand_dims(hsv[..., 0], -1)
        self.chr = tf.reduce_max(x4, 3, keep_dims=True) - tf.reduce_min(x4, 3, keep_dims=True)

        #self.chr = tf.reduce_mean(x, 3, keep_dims=True)

        self.huechr = tf.concat([self.hue, self.chr], 3)

        if self._version == 4:
            self.huechr_onehot = tf.one_hot(tf.to_int32(0.9999 * self._bins * self.huechr), self._bins)
            self.hue_onehot, self.chr_onehot = tf.unstack(self.huechr_onehot, axis=3)

        elif self._version == 3:
            self.huechr_onehot = tf.one_hot(tf.to_int32(0.9999 * self._bins * self.huechr), self._bins)
            self.hue_onehot, self.chr_onehot = tf.unstack(self.huechr_onehot, axis=3)

            #self.hue_onehot = tf.zeros([8, 227, 227, 32], dtype=tf.float32)
            #self.chr_onehot = tf.zeros([8, 227, 227, 32], dtype=tf.float32)

            self.hue_integral = tf.cumsum(tf.cumsum(self.hue_onehot, 1), 2)
            self.chr_integral = tf.cumsum(tf.cumsum(self.chr_onehot, 1), 2)

            #import ipdb; ipdb.set_trace()

            #self.hue_integral = self.hue_onehot
            #self.chr_integral = self.chr_onehot

            #self.hue_integral = tf.zeros([8, 227, 227, 32], dtype=tf.float32)
            #self.chr_integral = tf.zeros([8, 227, 227, 32], dtype=tf.float32)

            #import ipdb; ipdb.set_trace()

            """
            Wa = tf.ones([self._hist_win_size, self._hist_win_size, self._bins, 1]) / self._hist_win_size**2
            Wb = tf.expand_dims(tf.expand_dims(tf.eye(self._bins), 0), 0)
            self.hue_bin = tf.nn.separable_conv2d(self.hue_onehot,
                    Wa,
                    Wb,
                    strides=[1, 1, 1, 1],
                    padding='SAME')


            #self.chr_bin = tf.nn.conv2d(self.chr_onehot, W, strides=[1, 1, 1, 1], padding='SAME')
            self.chr_bin = tf.nn.separable_conv2d(self.chr_onehot,
                    Wa,
                    Wb,
                    strides=[1, 1, 1, 1],
                    padding='SAME')
            #self.hue_bin = self.hue_onehot
            #self.chr_bin = self.chr_onehot

            self.huechr_bin = tf.concat([
                tf.expand_dims(self.hue_bin, 3),
                tf.expand_dims(self.chr_bin, 3)
                ], 3)

            sh = self.huechr_bin.get_shape().as_list()[:3] + [-1]
            self.huechr_bin = tf.reshape(self.huechr_bin, sh)

            #self.hue = tf.nn.conv2d(self.hue_onehot, W0, strides=[1, 1, 1, 1], padding='SAME')
            #self.hue = tf.nn.conv2d(self.hue, W1, strides=[1, 1, 1, 1], padding='SAME')
            #self.chr = tf.nn.conv2d(self.chr_onehot, W0, strides=[1, 1, 1, 1], padding='SAME')
            #self.chr = tf.nn.conv2d(self.chr, W1, strides=[1, 1, 1, 1], padding='SAME')
            """

        return gray, dict(grayscale=gray_single, hsv=hsv), #, imgshape, imgname

    def build_network(self, basenet_info, extra, phase_test, global_step):
        info = selfsup.info.create(scale_summary=True)

        hyper = []

        dropout = functools.partial(selfsup.ops.dropout, phase_test=phase_test, info=info)

        with tf.name_scope('hypercolumn'):
            layers = self.basenet.hypercolumn_layers
            if self._version == 2:
                layers = layers[4:]
            for name, scale in layers:
                sparse_layer = sparse_extractor(self.centroids,
                                                basenet_info['activations'][name],
                                                scale, [0.0, 0.0])
                if self._use_scalers:
                    sparse_layer = selfsup.ops.scale(sparse_layer, name=name+'/scale')
                hyper.append(sparse_layer)

            flat_x = tf.concat(hyper, 1, name='concat')
            info['activations']['hypercolumn'] = flat_x

        flat_h = selfsup.ops.inner(flat_x, 1024, stddev=0.0001, info=info, name='pre_h_fc1')
        flat_h = dropout(flat_h, 0.5, name='h_fc1')

        z_hue = selfsup.ops.inner(flat_h, self._bins, activation=None, info=info, name='hue')
        z_chr = selfsup.ops.inner(flat_h, self._bins, activation=None, info=info, name='chroma')

        with tf.name_scope('sparse_y'):
            if self._version == 3:
                w = self._hist_win_radius

                cc = tf.concat([
                    self.centroids + [w+1, w+1],
                    self.centroids + [w+1, -w],
                    self.centroids + [-w, w+1],
                    self.centroids + [-w, -w],
                ], 1)

                yy = sparse_extractor(cc, self.hue_integral, 1.0, [0.0, 0.0])
                #yy = tf.zeros([4096, 32], dtype=tf.float32)
                yy = tf.reshape(yy, [self.hue_integral.shape[0].value, self.centroids.shape[1].value*4, self._bins])
                y11, y10, y01, y00 = tf.split(yy, 4, axis=1)

                y_hue = tf.reshape((y11 - y10 - y01 + y00) / self._hist_win_size**2, [-1, self._bins])

                yy = sparse_extractor(cc, self.chr_integral, 1.0, [0.0, 0.0])
                #yy = tf.zeros([4096, 32], dtype=tf.float32)
                yy = tf.reshape(yy, [self.chr_integral.shape[0].value, self.centroids.shape[1].value*4, self._bins])
                y11, y10, y01, y00 = tf.split(yy, 4, axis=1)

                y_chr = tf.reshape((y11 - y10 - y01 + y00) / self._hist_win_size**2, [-1, self._bins])
            else:
                #y_huechr = sparse_extractor(self.centroids, self.huechr, self.mult, [0.0, 0.0])
                y_huechr = sparse_extractor(self.centroids, self.huechr, self.mult, [0.0, 0.0])
                y_hue, y_chr = tf.unstack(tf.to_int32(y_huechr * self._bins * 0.99999), axis=1)

        with tf.variable_scope('primary_loss'):
            if self._version == 3:
                loss_chr_each = kl_divergence(labels=y_chr, logits=z_chr)
                loss_hue_each = kl_divergence(labels=y_hue, logits=z_hue)
            #y_chr = sparse_extractor(
            else:
                loss_chr_each = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=y_chr, logits=z_chr)
                loss_hue_each = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=y_hue, logits=z_hue)
            #loss_hue_each = tf.multiply(5 * y_huechr[..., 1], loss_hue_each)
            hue_loss = tf.reduce_mean(loss_hue_each)
            chr_loss = tf.reduce_mean(loss_chr_each)
            primary_loss = hue_loss + chr_loss

        l2_layers = [
            info['weights']['pre_h_fc1:weights'],
            info['weights']['hue:weights'],
            info['weights']['chroma:weights'],
        ]
        with tf.name_scope('weight_decay'):
            wd = 1e-6
            l2_loss = tf.add_n([
                tf.nn.l2_loss(v) for v in l2_layers
            ])
            weight_decay = wd * l2_loss

        with tf.name_scope('loss'):
            loss = weight_decay + primary_loss


        self._losses = OrderedDict([
            ('hue', hue_loss),
            ('chr', chr_loss),
            ('+weight_decay', weight_decay),
        ])
        self._primary_loss = primary_loss
        self._loss = loss

        variables = info['vars']

        info['activations']['y_hue'] = y_hue
        info['activations']['y_chr'] = y_chr
        info['activations']['z_hue'] = z_hue
        info['activations']['z_chr'] = z_chr

        self.feedback_variables = [
            #self.hue_integral,
            #self.hue,
            #self.centroids,
            #y_hue,
            #extra['grayscale'],
            #tf.nn.softmax(z_hue),
            #tf.nn.softmax(z_chr),
        ]

        info['activations']['hue'] = self.huechr[..., 0]
        info['activations']['chr'] = self.huechr[..., 1]
        info['activations']['centroids'] = self.centroids
        info['activations']['primary_loss'] = primary_loss
        info['activations']['loss'] = loss
        info['activations']['weight_decay'] = weight_decay
        return info

    def feedback(self, variables, iteration):
        """
        rgb = calc_rgb_from_hue_chroma(variables[0], variables[1], variables[2])

        vzlog.ColorImageGrid([
            rgb[:3]
        ], rows=1, vmin=0, vmax=1).save(self.make_path('hue', 'png'))
        """

    @property
    def losses(self):
        return self._losses

    @property
    def primary_loss(self):
        return self._primary_loss

    @property
    def loss(self):
        return self._loss


import tensorflow as tf
import selfsup
import os
from .base import Method
from autocolorize.tensorflow.sparse_extractor import sparse_extractor
from autocolorize.extraction import calc_rgb_from_hue_chroma
import functools
from collections import OrderedDict
import numpy as np

def kl_divergence(labels, logits):
    #return tf.contrib.distributions.kl(p, q)
    return (tf.nn.softmax_cross_entropy_with_logits(labels=labels, logits=logits) -
            tf.nn.softmax_cross_entropy_with_logits(labels=labels, logits=tf.log(labels+1e-5)))


class ColorizeHypercolumn3(Method):
    def __init__(self, name, basenet, loader, use_scalers=False, version=1):
        self.name = name
        self._loader = loader
        self.basenet = basenet
        self.locations = 128
        self.edge_buffer = 32
        self._bins = 8
        self._use_scalers = use_scalers
        self._version = version
        self._hist_win_radius = 3
        self._hist_win_size = self._hist_win_radius * 2 + 1

    @property
    def basenet_settings(self):
        return {'convolutional': True}

    def batch(self):
        #input_size = self.basenet.canonical_input_size
        x, _ = self._loader.batch()

        input_size = x.get_shape().as_list()[1]

        self.mult = 1
        self.centroids = tf.to_float(tf.random_uniform(
                shape=[self._loader.batch_size, self.locations, 2],
                minval=self.edge_buffer // self.mult,
                maxval=(input_size - self.edge_buffer) // self.mult,
                name='centroids',
                dtype=tf.int32))

        gray_single = tf.reduce_mean(x, 3, keep_dims=True)
        gray = tf.tile(gray_single, [1, 1, 1, 3])

        w = input_size // self.mult
        if self.mult > 1:
            x4 = tf.image.resize_bilinear(x, (w, w))
        else:
            x4 = x

        hsv = tf.image.rgb_to_hsv(x4)

        self.hue = tf.expand_dims(hsv[..., 0], -1)
        self.chr = tf.reduce_max(x4, 3, keep_dims=True) - tf.reduce_min(x4, 3, keep_dims=True)

        #self.chr = tf.reduce_mean(x, 3, keep_dims=True)

        self.huechr = tf.concat([self.hue, self.chr], 3)

        if self._version == 4:
            self.huechr_onehot = tf.one_hot(tf.to_int32(0.9999 * self._bins * self.huechr), self._bins)
            self.hue_onehot, self.chr_onehot = tf.unstack(self.huechr_onehot, axis=3)

        elif self._version == 3:
            self.huechr_onehot = tf.one_hot(tf.to_int32(0.9999 * self._bins * self.huechr), self._bins)
            self.hue_onehot, self.chr_onehot = tf.unstack(self.huechr_onehot, axis=3)

            #self.hue_onehot = tf.zeros([8, 227, 227, 32], dtype=tf.float32)
            #self.chr_onehot = tf.zeros([8, 227, 227, 32], dtype=tf.float32)

            self.hue_integral = tf.cumsum(tf.cumsum(self.hue_onehot, 1), 2)
            self.chr_integral = tf.cumsum(tf.cumsum(self.chr_onehot, 1), 2)

            #import ipdb; ipdb.set_trace()

            #self.hue_integral = self.hue_onehot
            #self.chr_integral = self.chr_onehot

            #self.hue_integral = tf.zeros([8, 227, 227, 32], dtype=tf.float32)
            #self.chr_integral = tf.zeros([8, 227, 227, 32], dtype=tf.float32)

            #import ipdb; ipdb.set_trace()

            """
            Wa = tf.ones([self._hist_win_size, self._hist_win_size, self._bins, 1]) / self._hist_win_size**2
            Wb = tf.expand_dims(tf.expand_dims(tf.eye(self._bins), 0), 0)
            self.hue_bin = tf.nn.separable_conv2d(self.hue_onehot,
                    Wa,
                    Wb,
                    strides=[1, 1, 1, 1],
                    padding='SAME')


            #self.chr_bin = tf.nn.conv2d(self.chr_onehot, W, strides=[1, 1, 1, 1], padding='SAME')
            self.chr_bin = tf.nn.separable_conv2d(self.chr_onehot,
                    Wa,
                    Wb,
                    strides=[1, 1, 1, 1],
                    padding='SAME')
            #self.hue_bin = self.hue_onehot
            #self.chr_bin = self.chr_onehot

            self.huechr_bin = tf.concat([
                tf.expand_dims(self.hue_bin, 3),
                tf.expand_dims(self.chr_bin, 3)
                ], 3)

            sh = self.huechr_bin.get_shape().as_list()[:3] + [-1]
            self.huechr_bin = tf.reshape(self.huechr_bin, sh)

            #self.hue = tf.nn.conv2d(self.hue_onehot, W0, strides=[1, 1, 1, 1], padding='SAME')
            #self.hue = tf.nn.conv2d(self.hue, W1, strides=[1, 1, 1, 1], padding='SAME')
            #self.chr = tf.nn.conv2d(self.chr_onehot, W0, strides=[1, 1, 1, 1], padding='SAME')
            #self.chr = tf.nn.conv2d(self.chr, W1, strides=[1, 1, 1, 1], padding='SAME')
            """

        return gray, dict(grayscale=gray_single, hsv=hsv), #, imgshape, imgname

    def build_network(self, basenet_info, extra, phase_test, global_step):
        info = selfsup.info.create(scale_summary=True)

        hyper = []

        dropout = functools.partial(selfsup.ops.dropout, phase_test=phase_test, info=info)

        with tf.name_scope('hypercolumn'):
            layers = self.basenet.hypercolumn_layers
            if self._version == 2:
                layers = layers[4:]
            for name, scale in layers:
                sparse_layer = sparse_extractor(self.centroids,
                                                basenet_info['activations'][name],
                                                scale, [0.0, 0.0])
                if self._use_scalers:
                    sparse_layer = selfsup.ops.scale(sparse_layer, name=name+'/scale')
                hyper.append(sparse_layer)

            flat_x = tf.concat(hyper, 1, name='concat')
            info['activations']['hypercolumn'] = flat_x

        flat_h = selfsup.ops.inner(flat_x, 1024, stddev=0.0001, info=info, name='pre_h_fc1')
        flat_h = dropout(flat_h, 0.5, name='h_fc1')

        z_hue = selfsup.ops.inner(flat_h, self._bins, activation=None, info=info, name='hue')
        z_chr = selfsup.ops.inner(flat_h, self._bins, activation=None, info=info, name='chroma')

        with tf.name_scope('sparse_y'):
            if self._version == 3:
                w = self._hist_win_radius

                cc = tf.concat([
                    self.centroids + [w+1, w+1],
                    self.centroids + [w+1, -w],
                    self.centroids + [-w, w+1],
                    self.centroids + [-w, -w],
                ], 1)

                yy = sparse_extractor(cc, self.hue_integral, 1.0, [0.0, 0.0])
                #yy = tf.zeros([4096, 32], dtype=tf.float32)
                yy = tf.reshape(yy, [self.hue_integral.shape[0].value, self.centroids.shape[1].value*4, self._bins])
                y11, y10, y01, y00 = tf.split(yy, 4, axis=1)

                y_hue = tf.reshape((y11 - y10 - y01 + y00) / self._hist_win_size**2, [-1, self._bins])

                yy = sparse_extractor(cc, self.chr_integral, 1.0, [0.0, 0.0])
                #yy = tf.zeros([4096, 32], dtype=tf.float32)
                yy = tf.reshape(yy, [self.chr_integral.shape[0].value, self.centroids.shape[1].value*4, self._bins])
                y11, y10, y01, y00 = tf.split(yy, 4, axis=1)

                y_chr = tf.reshape((y11 - y10 - y01 + y00) / self._hist_win_size**2, [-1, self._bins])
            else:
                #y_huechr = sparse_extractor(self.centroids, self.huechr, self.mult, [0.0, 0.0])
                y_huechr = sparse_extractor(self.centroids, self.huechr, self.mult, [0.0, 0.0])
                y_hue, y_chr = tf.unstack(tf.to_int32(y_huechr * self._bins * 0.99999), axis=1)

        with tf.variable_scope('primary_loss'):
            if self._version == 3:
                loss_chr_each = kl_divergence(labels=y_chr, logits=z_chr)
                loss_hue_each = kl_divergence(labels=y_hue, logits=z_hue)
            #y_chr = sparse_extractor(
            else:
                loss_chr_each = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=y_chr, logits=z_chr)
                loss_hue_each = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=y_hue, logits=z_hue)
            #loss_hue_each = tf.multiply(5 * y_huechr[..., 1], loss_hue_each)
            hue_loss = tf.reduce_mean(loss_hue_each)
            chr_loss = tf.reduce_mean(loss_chr_each)
            primary_loss = hue_loss + chr_loss

        l2_layers = [
            info['weights']['pre_h_fc1:weights'],
            info['weights']['hue:weights'],
            info['weights']['chroma:weights'],
        ]
        with tf.name_scope('weight_decay'):
            wd = 1e-6
            l2_loss = tf.add_n([
                tf.nn.l2_loss(v) for v in l2_layers
            ])
            weight_decay = wd * l2_loss

        with tf.name_scope('loss'):
            loss = weight_decay + primary_loss


        self._losses = OrderedDict([
            ('hue', hue_loss),
            ('chr', chr_loss),
            ('+weight_decay', weight_decay),
        ])
        self._primary_loss = primary_loss
        self._loss = loss

        variables = info['vars']

        info['activations']['y_hue'] = y_hue
        info['activations']['y_chr'] = y_chr
        info['activations']['z_hue'] = z_hue
        info['activations']['z_chr'] = z_chr

        self.feedback_variables = [
            #self.hue_integral,
            #self.hue,
            #self.centroids,
            #y_hue,
            #extra['grayscale'],
            #tf.nn.softmax(z_hue),
            #tf.nn.softmax(z_chr),
        ]

        info['activations']['hue'] = self.huechr[..., 0]
        info['activations']['chr'] = self.huechr[..., 1]
        info['activations']['centroids'] = self.centroids
        info['activations']['primary_loss'] = primary_loss
        info['activations']['loss'] = loss
        info['activations']['weight_decay'] = weight_decay
        return info

    def feedback(self, variables, iteration):
        """
        rgb = calc_rgb_from_hue_chroma(variables[0], variables[1], variables[2])

        vzlog.ColorImageGrid([
            rgb[:3]
        ], rows=1, vmin=0, vmax=1).save(self.make_path('hue', 'png'))
        """

    @property
    def losses(self):
        return self._losses

    @property
    def primary_loss(self):
        return self._primary_loss

    @property
    def loss(self):
        return self._loss


#return tf.contrib.distributions.kl(p, q) #input_size = self.basenet.canonical_input_size #self.chr = tf.reduce_mean(x, 3, keep_dims=True) #self.hue_onehot = tf.zeros([8, 227, 227, 32], dtype=tf.float32) #self.chr_onehot = tf.zeros([8, 227, 227, 32], dtype=tf.float32) #import ipdb; ipdb.set_trace() #self.hue_integral = self.hue_onehot #self.chr_integral = self.chr_onehot #self.hue_integral = tf.zeros([8, 227, 227, 32], dtype=tf.float32) #self.chr_integral = tf.zeros([8, 227, 227, 32], dtype=tf.float32) #import ipdb; ipdb.set_trace() Wa = tf.ones([self._hist_win_size, self._hist_win_size, self._bins, 1]) / self._hist_win_size**2             Wb = tf.expand_dims(tf.expand_dims(tf.eye(self._bins), 0), 0)             self.hue_bin = tf.nn.separable_conv2d(self.hue_onehot,                     Wa,                     Wb,                     strides=[1, 1, 1, 1],                     padding='SAME')               #self.chr_bin = tf.nn.conv2d(self.chr_onehot, W, strides=[1, 1, 1, 1], padding='SAME')             self.chr_bin = tf.nn.separable_conv2d(self.chr_onehot,                     Wa,                     Wb,                     strides=[1, 1, 1, 1],                     padding='SAME')             #self.hue_bin = self.hue_onehot             #self.chr_bin = self.chr_onehot              self.huechr_bin = tf.concat([                 tf.expand_dims(self.hue_bin, 3),                 tf.expand_dims(self.chr_bin, 3)                 ], 3)              sh = self.huechr_bin.get_shape().as_list()[:3] + [-1]             self.huechr_bin = tf.reshape(self.huechr_bin, sh)              #self.hue = tf.nn.conv2d(self.hue_onehot, W0, strides=[1, 1, 1, 1], padding='SAME')             #self.hue = tf.nn.conv2d(self.hue, W1, strides=[1, 1, 1, 1], padding='SAME')             #self.chr = tf.nn.conv2d(self.chr_onehot, W0, strides=[1, 1, 1, 1], padding='SAME')             #self.chr = tf.nn.conv2d(self.chr, W1, strides=[1, 1, 1, 1], padding='SAME') #, imgshape, imgname #yy = tf.zeros([4096, 32], dtype=tf.float32) #yy = tf.zeros([4096, 32], dtype=tf.float32) #y_huechr = sparse_extractor(self.centroids, self.huechr, self.mult, [0.0, 0.0]) #y_chr = sparse_extractor( #loss_hue_each = tf.multiply(5 * y_huechr[..., 1], loss_hue_each) #self.hue_integral, #self.hue, #self.centroids, #y_hue, #extra['grayscale'], #tf.nn.softmax(z_hue), #tf.nn.softmax(z_chr), rgb = calc_rgb_from_hue_chroma(variables[0], variables[1], variables[2])          vzlog.ColorImageGrid([             rgb[:3]         ], rows=1, vmin=0, vmax=1).save(self.make_path('hue', 'png'))