from __future__ import print_function, division

from sympy import lambdify
from sympy.core.numbers import pi
from sympy.functions import sin, cos
from sympy.plotting.pygletplot.plot_curve import PlotCurve
from sympy.plotting.pygletplot.plot_surface import PlotSurface

from math import sin as p_sin
from math import cos as p_cos


def float_vec3(f):
    def inner(*args):
        v = f(*args)
        return float(v[0]), float(v[1]), float(v[2])
    return inner


class Cartesian2D(PlotCurve):
    i_vars, d_vars = 'x', 'y'
    intervals = [[-5, 5, 100]]
    aliases = ['cartesian']
    is_default = True

    def _get_sympy_evaluator(self):
        fy = self.d_vars[0]
        x = self.t_interval.v

        @float_vec3
        def e(_x):
            return (_x, fy.subs(x, _x), 0.0)
        return e

    def _get_lambda_evaluator(self):
        fy = self.d_vars[0]
        x = self.t_interval.v
        return lambdify([x], [x, fy, 0.0])


class Cartesian3D(PlotSurface):
    i_vars, d_vars = 'xy', 'z'
    intervals = [[-1, 1, 40], [-1, 1, 40]]
    aliases = ['cartesian', 'monge']
    is_default = True

    def _get_sympy_evaluator(self):
        fz = self.d_vars[0]
        x = self.u_interval.v
        y = self.v_interval.v

        @float_vec3
        def e(_x, _y):
            return (_x, _y, fz.subs(x, _x).subs(y, _y))
        return e

    def _get_lambda_evaluator(self):
        fz = self.d_vars[0]
        x = self.u_interval.v
        y = self.v_interval.v
        return lambdify([x, y], [x, y, fz])


class ParametricCurve2D(PlotCurve):
    i_vars, d_vars = 't', 'xy'
    intervals = [[0, 2*pi, 100]]
    aliases = ['parametric']
    is_default = True

    def _get_sympy_evaluator(self):
        fx, fy = self.d_vars
        t = self.t_interval.v

        @float_vec3
        def e(_t):
            return (fx.subs(t, _t), fy.subs(t, _t), 0.0)
        return e

    def _get_lambda_evaluator(self):
        fx, fy = self.d_vars
        t = self.t_interval.v
        return lambdify([t], [fx, fy, 0.0])


class ParametricCurve3D(PlotCurve):
    i_vars, d_vars = 't', 'xyz'
    intervals = [[0, 2*pi, 100]]
    aliases = ['parametric']
    is_default = True

    def _get_sympy_evaluator(self):
        fx, fy, fz = self.d_vars
        t = self.t_interval.v

        @float_vec3
        def e(_t):
            return (fx.subs(t, _t), fy.subs(t, _t), fz.subs(t, _t))
        return e

    def _get_lambda_evaluator(self):
        fx, fy, fz = self.d_vars
        t = self.t_interval.v
        return lambdify([t], [fx, fy, fz])


class ParametricSurface(PlotSurface):
    i_vars, d_vars = 'uv', 'xyz'
    intervals = [[-1, 1, 40], [-1, 1, 40]]
    aliases = ['parametric']
    is_default = True

    def _get_sympy_evaluator(self):
        fx, fy, fz = self.d_vars
        u = self.u_interval.v
        v = self.v_interval.v

        @float_vec3
        def e(_u, _v):
            return (fx.subs(u, _u).subs(v, _v),
                    fy.subs(u, _u).subs(v, _v),
                    fz.subs(u, _u).subs(v, _v))
        return e

    def _get_lambda_evaluator(self):
        fx, fy, fz = self.d_vars
        u = self.u_interval.v
        v = self.v_interval.v