import numpy as np
from scipy.interpolate import interp1d
from pyTools import *

################################################################################
#~~~~~~~~~Log ops
################################################################################
def logPolyVal(p,x):
    ord = p.order()
    logs = []
    for idx in xrange(ord+1):
        logs.append( np.log( p[idx] ) + (ord-idx)*np.log(x) )
    return logs

################################################################################
#~~~~~~~~~Symmeterize data
################################################################################
def symmeterize( x, y, interp_type='cubic' ):
    if x.min() <= 0:
        raise ValueError('x.min() must be greater than zero.')

    xs = np.array([-x,x]).flatten()
    xs.sort()

    f = interp1d( x , y , kind=interp_type )

    return { 'x':xs , 'y':f(np.abs(xs)) }

################################################################################
#~~~~~~~~~3D Shapes
################################################################################
def makeSphere(x0=0,y0=0,z0=0,r=1,ntheta=30,nphi=30):
    u = np.linspace(0, np.pi, ntheta)
    v = np.linspace(0, 2 * np.pi, nphi)

    x = np.outer(np.sin(u), np.sin(v))*r
    y = np.outer(np.sin(u), np.cos(v))*r
    z = np.outer(np.cos(u), np.ones_like(v))*r
    return x+x0, y+y0, z+z0

def makeCylinder(x0=0,y0=0,z0=0,r=1,h=10,ntheta=30,nz=30):
    u = np.linspace(0, 2*np.pi, ntheta)
    z =  np.linspace(0, h, nz)

    UU,ZZ = np.meshgrid(u,z)

    XX = np.cos(UU)*r
    YY = np.sin(UU)*r

    # ax.plot_wireframe(x, y, z)
    return XX+x0, YY+y0, ZZ+z0

def generateLine3D( x0=0, x1=1, y0=0, y1=1, z0=0, z1=0, N=2 ):
    return {'line':{'xData':np.linspace(x0,x1,N),
            'yData':np.linspace(y0,y1,N),
            'zData':np.linspace(z0,z1,N),
            'cData':np.ones((N,1))}}

################################################################################
#~~~~~~~~~2D Shapes
################################################################################
def generateCircle(R=1, X0=0, Y0=0, N = 60, thetaMin = 0, thetaMax = 2*np.pi ):
    thetas = np.linspace( thetaMin , thetaMax , N)
    uY = np.sin( thetas )*R
    uX = np.cos( thetas )*R
    return {'circle':{'xData':uX+X0, 'yData':uY+Y0}}

def generateEllipse( RX=2, RY=1, X0=0, Y0=0, N = 60, thetaMin = 0, thetaMax = 2*np.pi ):
    thetas = np.linspace( thetaMin , thetaMax , N)
    uY = np.sin( thetas )*RY
    uX = np.cos( thetas )*RX
    return {'ellipse':{'xData':uX+X0, 'yData':uY+Y0}}

def makeCylinder2D( L = 10., R = 1., N=60, view_degrees=30. ):

    yFac = np.cos(view_degrees * np.pi/180.)
    zFac = np.sin(view_degrees * np.pi/180.)

    xL = np.ones((2,1))*-R
    xR = -xL
    y  = np.array([0,L])*yFac

    cylinder = { 'leftSide':{'xData':xL, 'yData':y},
                 'rightSide':{'xData':xR, 'yData':y},
                 'upperEllipse':generateEllipse(RX = R, RY=R*zFac, Y0=L*yFac,N=N)['ellipse'],
                 'lowerHalfEllipse':generateEllipse(RX = R, RY=R*zFac, thetaMin=np.pi, thetaMax=2*np.pi, N=int(N/2.))['ellipse']}
    return cylinder

################################################################################
#~~~~~~~~~Rotations
################################################################################
def rotateObject(x,y,z,ax=None,ay=None,az=None):
    if ax is not None:
        y,z = rotateAt(y,z,ax)
    if ay is not None:
        x,z = rotateAt(x,z,-ay)
    if az is not None:
        x,y = rotateAt(x,y,az)
    return x,y,z

def rotateAt(x,y,a):
    xp = np.cos(a)*x-np.sin(a)*y
    yp = np.cos(a)*y+np.sin(a)*x
    return xp, yp

def rotateObj2D( obj_in, degrees ):
    obj = obj_in.copy()

    keys = obj.keys()
    for key in keys:
        obj[key] = rotate2D( degrees=degrees, **obj[key] )

    return obj

def rotate2D( xData, yData, degrees ):
    x = xData.flatten()
    y = yData.flatten()
    z = np.zeros_like(x)
    x,y,z = rotateObject( x, y, z, az=float(degrees)/180.*np.pi )
    return {'xData':x, 'yData':y}


def rotateObj3D( obj_in, gamma, theta, phi ):
    obj = obj_in.copy()

    keys = obj.keys()
    for key in keys:
        obj[key] = rotate3D( gamma=gamma, theta=theta, phi=phi, **obj[key] )

    return obj

def rotate3D( xData, yData, zData, gamma, theta, phi, kwargs_toggle=True, **kwargs ):
    ignore_kwargs(kwargs, toggle=kwargs_toggle)
    x = xData.flatten()
    y = yData.flatten()
    z = zData.flatten()

    x,y,z = rotateObject( x, y, z, az=float(gamma)/180.*np.pi )
    x,y,z = rotateObject( x, y, z, ay=float(theta)/180.*np.pi )
    x,y,z = rotateObject( x, y, z, az=float(phi)/180.*np.pi )

    return {'xData':x, 'yData':y, 'zData':z}

[((25, 8, 25, 44), 'scipy.interpolate.interp1d', 'interp1d', (), '', False, 'from scipy.interpolate import interp1d\n'), ((33, 8, 33, 37), 'numpy.linspace', 'np.linspace', ({(33, 20, 33, 21): '0', (33, 23, 33, 28): 'np.pi', (33, 30, 33, 36): 'ntheta'}, {}), '(0, np.pi, ntheta)', True, 'import numpy as np\n'), ((34, 8, 34, 39), 'numpy.linspace', 'np.linspace', ({(34, 20, 34, 21): '0', (34, 23, 34, 32): '2 * np.pi', (34, 34, 34, 38): 'nphi'}, {}), '(0, 2 * np.pi, nphi)', True, 'import numpy as np\n'), ((42, 8, 42, 39), 'numpy.linspace', 'np.linspace', ({(42, 20, 42, 21): '0', (42, 23, 42, 30): '2 * np.pi', (42, 32, 42, 38): 'ntheta'}, {}), '(0, 2 * np.pi, ntheta)', True, 'import numpy as np\n'), ((43, 9, 43, 30), 'numpy.linspace', 'np.linspace', ({(43, 21, 43, 22): '0', (43, 24, 43, 25): 'h', (43, 27, 43, 29): 'nz'}, {}), '(0, h, nz)', True, 'import numpy as np\n'), ((45, 12, 45, 28), 'numpy.meshgrid', 'np.meshgrid', ({(45, 24, 45, 25): 'u', (45, 26, 45, 27): 'z'}, {}), '(u, z)', True, 'import numpy as np\n'), ((63, 13, 63, 50), 'numpy.linspace', 'np.linspace', ({(63, 26, 63, 34): 'thetaMin', (63, 37, 63, 45): 'thetaMax', (63, 48, 63, 49): 'N'}, {}), '(thetaMin, thetaMax, N)', True, 'import numpy as np\n'), ((69, 13, 69, 50), 'numpy.linspace', 'np.linspace', ({(69, 26, 69, 34): 'thetaMin', (69, 37, 69, 45): 'thetaMax', (69, 48, 69, 49): 'N'}, {}), '(thetaMin, thetaMax, N)', True, 'import numpy as np\n'), ((76, 11, 76, 44), 'numpy.cos', 'np.cos', ({(76, 18, 76, 43): 'view_degrees * np.pi / 180.0'}, {}), '(view_degrees * np.pi / 180.0)', True, 'import numpy as np\n'), ((77, 11, 77, 44), 'numpy.sin', 'np.sin', ({(77, 18, 77, 43): 'view_degrees * np.pi / 180.0'}, {}), '(view_degrees * np.pi / 180.0)', True, 'import numpy as np\n'), ((118, 8, 118, 24), 'numpy.zeros_like', 'np.zeros_like', ({(118, 22, 118, 23): 'x'}, {}), '(x)', True, 'import numpy as np\n'), ((47, 9, 47, 19), 'numpy.cos', 'np.cos', ({(47, 16, 47, 18): 'UU'}, {}), '(UU)', True, 'import numpy as np\n'), ((48, 9, 48, 19), 'numpy.sin', 'np.sin', ({(48, 16, 48, 18): 'UU'}, {}), '(UU)', True, 'import numpy as np\n'), ((64, 9, 64, 25), 'numpy.sin', 'np.sin', ({(64, 17, 64, 23): 'thetas'}, {}), '(thetas)', True, 'import numpy as np\n'), ((65, 9, 65, 25), 'numpy.cos', 'np.cos', ({(65, 17, 65, 23): 'thetas'}, {}), '(thetas)', True, 'import numpy as np\n'), ((70, 9, 70, 25), 'numpy.sin', 'np.sin', ({(70, 17, 70, 23): 'thetas'}, {}), '(thetas)', True, 'import numpy as np\n'), ((71, 9, 71, 25), 'numpy.cos', 'np.cos', ({(71, 17, 71, 23): 'thetas'}, {}), '(thetas)', True, 'import numpy as np\n'), ((79, 9, 79, 23), 'numpy.ones', 'np.ones', ({(79, 17, 79, 22): '(2, 1)'}, {}), '((2, 1))', True, 'import numpy as np\n'), ((81, 9, 81, 24), 'numpy.array', 'np.array', ({(81, 18, 81, 23): '[0, L]'}, {}), '([0, L])', True, 'import numpy as np\n'), ((22, 9, 22, 25), 'numpy.array', 'np.array', ({(22, 18, 22, 24): '[-x, x]'}, {}), '([-x, x])', True, 'import numpy as np\n'), ((27, 28, 27, 38), 'numpy.abs', 'np.abs', ({(27, 35, 27, 37): 'xs'}, {}), '(xs)', True, 'import numpy as np\n'), ((36, 17, 36, 26), 'numpy.sin', 'np.sin', ({(36, 24, 36, 25): 'u'}, {}), '(u)', True, 'import numpy as np\n'), ((36, 28, 36, 37), 'numpy.sin', 'np.sin', ({(36, 35, 36, 36): 'v'}, {}), '(v)', True, 'import numpy as np\n'), ((37, 17, 37, 26), 'numpy.sin', 'np.sin', ({(37, 24, 37, 25): 'u'}, {}), '(u)', True, 'import numpy as np\n'), ((37, 28, 37, 37), 'numpy.cos', 'np.cos', ({(37, 35, 37, 36): 'v'}, {}), '(v)', True, 'import numpy as np\n'), ((38, 17, 38, 26), 'numpy.cos', 'np.cos', ({(38, 24, 38, 25): 'u'}, {}), '(u)', True, 'import numpy as np\n'), ((38, 28, 38, 43), 'numpy.ones_like', 'np.ones_like', ({(38, 41, 38, 42): 'v'}, {}), '(v)', True, 'import numpy as np\n'), ((54, 28, 54, 48), 'numpy.linspace', 'np.linspace', ({(54, 40, 54, 42): 'x0', (54, 43, 54, 45): 'x1', (54, 46, 54, 47): 'N'}, {}), '(x0, x1, N)', True, 'import numpy as np\n'), ((55, 20, 55, 40), 'numpy.linspace', 'np.linspace', ({(55, 32, 55, 34): 'y0', (55, 35, 55, 37): 'y1', (55, 38, 55, 39): 'N'}, {}), '(y0, y1, N)', True, 'import numpy as np\n'), ((56, 20, 56, 40), 'numpy.linspace', 'np.linspace', ({(56, 32, 56, 34): 'z0', (56, 35, 56, 37): 'z1', (56, 38, 56, 39): 'N'}, {}), '(z0, z1, N)', True, 'import numpy as np\n'), ((57, 20, 57, 34), 'numpy.ones', 'np.ones', ({(57, 28, 57, 33): '(N, 1)'}, {}), '((N, 1))', True, 'import numpy as np\n'), ((102, 9, 102, 18), 'numpy.cos', 'np.cos', ({(102, 16, 102, 17): 'a'}, {}), '(a)', True, 'import numpy as np\n'), ((102, 21, 102, 30), 'numpy.sin', 'np.sin', ({(102, 28, 102, 29): 'a'}, {}), '(a)', True, 'import numpy as np\n'), ((103, 9, 103, 18), 'numpy.cos', 'np.cos', ({(103, 16, 103, 17): 'a'}, {}), '(a)', True, 'import numpy as np\n'), ((103, 21, 103, 30), 'numpy.sin', 'np.sin', ({(103, 28, 103, 29): 'a'}, {}), '(a)', True, 'import numpy as np\n'), ((12, 21, 12, 37), 'numpy.log', 'np.log', ({(12, 29, 12, 35): 'p[idx]'}, {}), '(p[idx])', True, 'import numpy as np\n'), ((12, 50, 12, 59), 'numpy.log', 'np.log', ({(12, 57, 12, 58): 'x'}, {}), '(x)', True, 'import numpy as np\n')]