#!/usr/bin/env python
# -*- coding: utf-8 -*-

import sys
import os.path

sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
from helpers import *

from mobject.tex_mobject import TexMobject
from mobject import Mobject, Group
from mobject.image_mobject import ImageMobject
from mobject.vectorized_mobject import *

from animation.animation import Animation
from animation.transform import *
from animation.simple_animations import *
from animation.compositions import *
from animation.playground import *
from animation.continual_animation import *
from topics.geometry import *
from topics.characters import *
from topics.functions import *
from topics.fractals import *
from topics.number_line import *
from topics.combinatorics import *
from topics.numerals import *
from topics.three_dimensions import *
from topics.objects import *
from topics.probability import *
from topics.complex_numbers import *
from scene import Scene
from scene.reconfigurable_scene import ReconfigurableScene
from scene.zoomed_scene import *
from camera import Camera
from mobject.svg_mobject import *
from mobject.tex_mobject import *

from nn.network import *
from nn.part1 import *


class Test(Scene):
    def construct(self):
        network = get_pretrained_network()
        training_data, validation_data, test_data = load_data_wrapper()
        self.show_weight_rows(network, index = 0)
        # self.show_maximizing_inputs(network)
        # self.show_all_activation_images(network, test_data)

        # group = Group()
        # for k in range(10):
        #     v = np.zeros((10, 1))
        #     v[k] = 1
        #     h_group = Group()
        #     for W, b in reversed(zip(network.weights, network.biases)):
        #         h_group.add(MNistMobject(v))
        #         v = np.dot(W.T, sigmoid_inverse(v) - b)
        #         v = sigmoid(v)
        #     h_group.add(MNistMobject(v))
        #     h_group.arrange_submobjects(LEFT)
        #     group.add(h_group)
        # group.arrange_submobjects(DOWN)
        # group.scale_to_fit_height(2*SPACE_Y_RADIUS - 1)
        # self.add(group)


    def show_random_results(self):
        group = Group(*[
            Group(*[
                MNistMobject(a)
                for a in network.get_activation_of_all_layers(
                    np.random.randn(784, 1)
                )
            ]).arrange_submobjects(RIGHT)
            for x in range(10)
        ]).arrange_submobjects(DOWN)
        group.scale_to_fit_height(2*SPACE_Y_RADIUS - 1)
        self.add(group)

    def show_weight_rows(self, network, index):
        group = VGroup()
        for row in network.weights[index]:
            mob = PixelsFromVect(np.zeros(row.size))
            for n, pixel in zip(row, mob):
                color = GREEN if n > 0 else RED
                opacity = 2*(sigmoid(abs(n)) - 0.5)
                pixel.set_fill(color, opacity = opacity)
            group.add(mob)
        group.arrange_submobjects_in_grid()
        group.scale_to_fit_height(2*SPACE_Y_RADIUS - 1)
        self.add(group)

    def show_all_activation_images(self, network, test_data):
        image_samples = Group(*[
            self.get_activation_images(digit, network, test_data)
            for digit in range(10)
        ])
        image_samples.arrange_submobjects_in_grid(
            n_rows = 2, buff = LARGE_BUFF
        )
        image_samples.scale_to_fit_height(2*SPACE_Y_RADIUS - 1)
        self.add(image_samples)

    def get_activation_images(self, digit, network, test_data, n_examples = 8):
        input_vectors = [
            data[0] 
            for data in test_data
            if data[1] == digit
        ]
        activation_iamges = Group(*[
            Group(*[
                MNistMobject(a)
                for a in network.get_activation_of_all_layers(vect)
            ]).arrange_submobjects(RIGHT)
            for vect in input_vectors[:n_examples]
        ]).arrange_submobjects(DOWN)
        activation_iamges.scale_to_fit_height(2*SPACE_Y_RADIUS - 1)
        return activation_iamges

    def show_two_blend(self):
        training_data, validation_data, test_data = load_data_wrapper()
        vects = [
            data[0]
            for data in training_data[:30]
            if np.argmax(data[1]) == 2
        ]
        mean_vect = reduce(op.add, vects)/len(vects)
        self.add(MNistMobject(mean_vect))

    def show_maximizing_inputs(self, network):
        training_data, validation_data, test_data = load_data_wrapper()
        layer = 1
        n_neurons = DEFAULT_LAYER_SIZES[layer]
        groups = Group()
        for k in range(n_neurons):
            out = np.zeros(n_neurons)
            out[k] = 1
            in_vect = maximizing_input(network, layer, out)
            new_out = network.get_activation_of_all_layers(in_vect)[layer]
            group = Group(*map(MNistMobject, [in_vect, new_out]))
            group.arrange_submobjects(DOWN+RIGHT, SMALL_BUFF)
            groups.add(group)
        groups.arrange_submobjects_in_grid()
        groups.scale_to_fit_height(2*SPACE_Y_RADIUS - 1)
        self.add(groups)

    def show_test_input(self, network):
        training_data, validation_data, test_data = load_data_wrapper()
        group = Group(*[
            self.get_set(network, test)
            for test in test_data[3:20]
            if test[1] in [4, 9]
        ])
        group.arrange_submobjects(DOWN, buff = MED_LARGE_BUFF)
        group.scale_to_fit_height(2*SPACE_Y_RADIUS - 1)
        self.play(FadeIn(group))

    def get_set(self, network, test):
        test_in, test_out = test
        activations = network.get_activation_of_all_layers(test_in)
        group = Group(*map(MNistMobject, activations))
        group.arrange_submobjects(RIGHT, buff = LARGE_BUFF)
        return group

    # def show_frame(self):
    #     pass


if __name__ == "__main__":
    save_pretrained_network()
    test_network()