from __future__ import absolute_import

import numpy as np

from constants import *

from animation.animation import Animation
from animation.movement import Homotopy
from animation.creation import ShowPartial
from animation.transform import Transform
from mobject.mobject import Group
from mobject.mobject import Mobject
from mobject.types.vectorized_mobject import VMobject
from mobject.geometry import Circle
from mobject.geometry import Dot
from utils.config_ops import digest_config
from utils.rate_functions import squish_rate_func
from utils.rate_functions import there_and_back


class FocusOn(Transform):
    CONFIG = {
        "opacity" : 0.2,
        "color" : GREY,
        "run_time" : 2,
        "remover" : True,
    }
    def __init__(self, mobject_or_point, **kwargs):
        digest_config(self, kwargs)
        big_dot = Dot(
            radius = FRAME_X_RADIUS+FRAME_Y_RADIUS,
            stroke_width = 0,
            fill_color = self.color,
            fill_opacity = 0,
        )
        little_dot = Dot(radius = 0)
        little_dot.set_fill(self.color, opacity = self.opacity)
        little_dot.move_to(mobject_or_point)

        Transform.__init__(self, big_dot, little_dot, **kwargs)

class Indicate(Transform):
    CONFIG = {
        "rate_func" : there_and_back,
        "scale_factor" : 1.2,
        "color" : YELLOW,
    }
    def __init__(self, mobject, **kwargs):
        digest_config(self, kwargs)
        target = mobject.copy()
        target.scale_in_place(self.scale_factor)
        target.set_color(self.color)
        Transform.__init__(self, mobject, target, **kwargs)

class CircleIndicate(Indicate):
    CONFIG = {
        "rate_func" : squish_rate_func(there_and_back, 0, 0.8),
        "remover" : True
    }
    def __init__(self, mobject, **kwargs):
        digest_config(self, kwargs)
        circle = Circle(color = self.color, **kwargs)
        circle.surround(mobject)
        Indicate.__init__(self, circle, **kwargs)

class ShowPassingFlash(ShowPartial):
    CONFIG = {
        "time_width" : 0.1,
        "remover" : True,
    }
    def get_bounds(self, alpha):
        alpha *= (1+self.time_width)
        alpha -= self.time_width/2.0
        lower = max(0, alpha - self.time_width/2.0)
        upper = min(1, alpha + self.time_width/2.0)
        return (lower, upper)

    def clean_up(self, *args, **kwargs):
        ShowPartial.clean_up(self, *args, **kwargs)
        for submob, start_submob in self.get_all_families_zipped():
            submob.pointwise_become_partial(start_submob, 0, 1)

class ShowCreationThenDestruction(ShowPassingFlash):
    CONFIG = {
        "time_width" : 2.0,
        "run_time" : 1,
    }

class ApplyWave(Homotopy):
    CONFIG = {
        "direction" : DOWN,
        "amplitude" : 0.2,
        "run_time" : 1,
        "apply_function_kwargs" : {
            "maintain_smoothness" : False,
        },
    }
    def __init__(self, mobject, **kwargs):
        digest_config(self, kwargs, locals())
        left_x = mobject.get_left()[0]
        right_x = mobject.get_right()[0]
        vect = self.amplitude*self.direction
        def homotopy(x, y, z, t):
            start_point = np.array([x, y, z])
            alpha = (x-left_x)/(right_x-left_x)
            power = np.exp(2*(alpha-0.5))
            nudge = there_and_back(t**power)
            return np.array([x, y, z]) + nudge*vect
        Homotopy.__init__(self, homotopy, mobject, **kwargs)

class WiggleOutThenIn(Animation):
    CONFIG = {
        "scale_value" : 1.1,
        "rotation_angle" : 0.01*TAU,
        "n_wiggles" : 6,
        "run_time" : 2,
        "scale_about_point" : None,
        "rotate_about_point" : None,
    }
    def __init__(self, mobject, **kwargs):
        digest_config(self, kwargs)
        if self.scale_about_point is None:
            self.scale_about_point = mobject.get_center()
        if self.rotate_about_point is None:
            self.rotate_about_point = mobject.get_center()
        Animation.__init__(self, mobject, **kwargs)

    def update_submobject(self, submobject, starting_sumobject, alpha):
        submobject.points[:,:] = starting_sumobject.points
        submobject.scale(
            interpolate(1, self.scale_value, there_and_back(alpha)),
            about_point = self.scale_about_point
        )
        submobject.rotate(
            wiggle(alpha, self.n_wiggles)*self.rotation_angle,
            about_point = self.rotate_about_point
        )

class Vibrate(Animation):
    CONFIG = {
        "spatial_period"  : 6,
        "temporal_period" : 1,
        "overtones"       : 4,
        "amplitude"       : 0.5,
        "radius"          : FRAME_X_RADIUS/2,
        "run_time"        : 3.0,
        "rate_func"      : None
    }
    def __init__(self, mobject = None, **kwargs):
        if mobject is None:
            mobject = Line(3*LEFT, 3*RIGHT)
        Animation.__init__(self, mobject, **kwargs)

    def wave_function(self, x, t):
        return sum([
            reduce(op.mul, [
                self.amplitude/(k**2), #Amplitude
                np.sin(2*np.pi*(k**1.5)*t/self.temporal_period), #Frequency
                np.sin(2*np.pi*k*x/self.spatial_period) #Number of waves
            ])
            for k in range(1, self.overtones+1)
        ])


    def update_mobject(self, alpha):
        time = alpha*self.run_time
        families = map(
            Mobject.submobject_family,
            [self.mobject, self.starting_mobject]
        )
        for mob, start in zip(*families):
            mob.points = np.apply_along_axis(
                lambda (x, y, z) : (x, y + self.wave_function(x, time), z),
                1, start.points
            )

class TurnInsideOut(Transform):
    CONFIG = {
        "path_arc" : TAU/4,
    }
    def __init__(self, mobject, **kwargs):
        mobject.sort_points(np.linalg.norm)
        mob_copy = mobject.copy()
        mob_copy.sort_points(lambda p : -np.linalg.norm(p))
        Transform.__init__(self, mobject, mob_copy, **kwargs)