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}

[((804, 836), 'scipy.interpolate.interp1d', 'interp1d', (['x', 'y'], {'kind': 'interp_type'}), '(x, y, kind=interp_type)\n', (812, 836), False, 'from scipy.interpolate import interp1d\n'), ((1129, 1158), 'numpy.linspace', 'np.linspace', (['(0)', 'np.pi', 'ntheta'], {}), '(0, np.pi, ntheta)\n', (1140, 1158), True, 'import numpy as np\n'), ((1167, 1198), 'numpy.linspace', 'np.linspace', (['(0)', '(2 * np.pi)', 'nphi'], {}), '(0, 2 * np.pi, nphi)\n', (1178, 1198), True, 'import numpy as np\n'), ((1425, 1458), 'numpy.linspace', 'np.linspace', (['(0)', '(2 * np.pi)', 'ntheta'], {}), '(0, 2 * np.pi, ntheta)\n', (1436, 1458), True, 'import numpy as np\n'), ((1466, 1487), 'numpy.linspace', 'np.linspace', (['(0)', 'h', 'nz'], {}), '(0, h, nz)\n', (1477, 1487), True, 'import numpy as np\n'), ((1501, 1518), 'numpy.meshgrid', 'np.meshgrid', (['u', 'z'], {}), '(u, z)\n', (1512, 1518), True, 'import numpy as np\n'), ((2139, 2173), 'numpy.linspace', 'np.linspace', (['thetaMin', 'thetaMax', 'N'], {}), '(thetaMin, thetaMax, N)\n', (2150, 2173), True, 'import numpy as np\n'), ((2389, 2423), 'numpy.linspace', 'np.linspace', (['thetaMin', 'thetaMax', 'N'], {}), '(thetaMin, thetaMax, N)\n', (2400, 2423), True, 'import numpy as np\n'), ((2615, 2651), 'numpy.cos', 'np.cos', (['(view_degrees * np.pi / 180.0)'], {}), '(view_degrees * np.pi / 180.0)\n', (2621, 2651), True, 'import numpy as np\n'), ((2660, 2696), 'numpy.sin', 'np.sin', (['(view_degrees * np.pi / 180.0)'], {}), '(view_degrees * np.pi / 180.0)\n', (2666, 2696), True, 'import numpy as np\n'), ((3911, 3927), 'numpy.zeros_like', 'np.zeros_like', (['x'], {}), '(x)\n', (3924, 3927), True, 'import numpy as np\n'), ((1528, 1538), 'numpy.cos', 'np.cos', (['UU'], {}), '(UU)\n', (1534, 1538), True, 'import numpy as np\n'), ((1550, 1560), 'numpy.sin', 'np.sin', (['UU'], {}), '(UU)\n', (1556, 1560), True, 'import numpy as np\n'), ((2186, 2200), 'numpy.sin', 'np.sin', (['thetas'], {}), '(thetas)\n', (2192, 2200), True, 'import numpy as np\n'), ((2214, 2228), 'numpy.cos', 'np.cos', (['thetas'], {}), '(thetas)\n', (2220, 2228), True, 'import numpy as np\n'), ((2436, 2450), 'numpy.sin', 'np.sin', (['thetas'], {}), '(thetas)\n', (2442, 2450), True, 'import numpy as np\n'), ((2465, 2479), 'numpy.cos', 'np.cos', (['thetas'], {}), '(thetas)\n', (2471, 2479), True, 'import numpy as np\n'), ((2704, 2719), 'numpy.ones', 'np.ones', (['(2, 1)'], {}), '((2, 1))\n', (2711, 2719), True, 'import numpy as np\n'), ((2744, 2760), 'numpy.array', 'np.array', (['[0, L]'], {}), '([0, L])\n', (2752, 2760), True, 'import numpy as np\n'), ((754, 771), 'numpy.array', 'np.array', (['[-x, x]'], {}), '([-x, x])\n', (762, 771), True, 'import numpy as np\n'), ((870, 880), 'numpy.abs', 'np.abs', (['xs'], {}), '(xs)\n', (876, 880), True, 'import numpy as np\n'), ((1217, 1226), 'numpy.sin', 'np.sin', (['u'], {}), '(u)\n', (1223, 1226), True, 'import numpy as np\n'), ((1228, 1237), 'numpy.sin', 'np.sin', (['v'], {}), '(v)\n', (1234, 1237), True, 'import numpy as np\n'), ((1258, 1267), 'numpy.sin', 'np.sin', (['u'], {}), '(u)\n', (1264, 1267), True, 'import numpy as np\n'), ((1269, 1278), 'numpy.cos', 'np.cos', (['v'], {}), '(v)\n', (1275, 1278), True, 'import numpy as np\n'), ((1299, 1308), 'numpy.cos', 'np.cos', (['u'], {}), '(u)\n', (1305, 1308), True, 'import numpy as np\n'), ((1310, 1325), 'numpy.ones_like', 'np.ones_like', (['v'], {}), '(v)\n', (1322, 1325), True, 'import numpy as np\n'), ((1720, 1742), 'numpy.linspace', 'np.linspace', (['x0', 'x1', 'N'], {}), '(x0, x1, N)\n', (1731, 1742), True, 'import numpy as np\n'), ((1762, 1784), 'numpy.linspace', 'np.linspace', (['y0', 'y1', 'N'], {}), '(y0, y1, N)\n', (1773, 1784), True, 'import numpy as np\n'), ((1804, 1826), 'numpy.linspace', 'np.linspace', (['z0', 'z1', 'N'], {}), '(z0, z1, N)\n', (1815, 1826), True, 'import numpy as np\n'), ((1846, 1861), 'numpy.ones', 'np.ones', (['(N, 1)'], {}), '((N, 1))\n', (1853, 1861), True, 'import numpy as np\n'), ((3560, 3569), 'numpy.cos', 'np.cos', (['a'], {}), '(a)\n', (3566, 3569), True, 'import numpy as np\n'), ((3572, 3581), 'numpy.sin', 'np.sin', (['a'], {}), '(a)\n', (3578, 3581), True, 'import numpy as np\n'), ((3593, 3602), 'numpy.cos', 'np.cos', (['a'], {}), '(a)\n', (3599, 3602), True, 'import numpy as np\n'), ((3605, 3614), 'numpy.sin', 'np.sin', (['a'], {}), '(a)\n', (3611, 3614), True, 'import numpy as np\n'), ((367, 381), 'numpy.log', 'np.log', (['p[idx]'], {}), '(p[idx])\n', (373, 381), True, 'import numpy as np\n'), ((396, 405), 'numpy.log', 'np.log', (['x'], {}), '(x)\n', (402, 405), True, 'import numpy as np\n')]