def arc_distance_theano_broadcast_prepare(dtype='float64'):
    """
    Calculates the pairwise arc distance between all points in vector a and b.
    """
    a = tensor.matrix(dtype=str(dtype))
    b = tensor.matrix(dtype=str(dtype))

    theta_1 = a[:, 0][None, :]
    theta_2 = b[:, 0][None, :]
    phi_1 = a[:, 1][:, None]
    phi_2 = b[:, 1][None, :]

    temp = (tensor.sin((theta_2 - theta_1) / 2)**2
            +
            tensor.cos(theta_1) * tensor.cos(theta_2)
            * tensor.sin((phi_2 - phi_1) / 2)**2)
    distance_matrix = 2 * (tensor.arctan2(tensor.sqrt(temp),
                                          tensor.sqrt(1 - temp)))
    name = "arc_distance_theano_broadcast"
    rval = theano.function([a, b],
                           distance_matrix,
                           name=name)
    rval.__name__ = name

    return rval

def poseDiff2D(startingPose, endingPose):
    x1, y1, theta1 = endingPose[0], endingPose[1], endingPose[2]
    x2, y2, theta2 = startingPose[0], startingPose[1], startingPose[2]
    dx = (x1 - x2)*T.cos(theta2) + (y1 - y2)*T.sin(theta2)
    dy = -(x1 - x2)*T.sin(theta2) + (y1 - y2)*T.cos(theta2)
    dtheta = normalizeAngle(theta1 - theta2)
    return dx, dy, dtheta

def arc_distance_theano_alloc_prepare(dtype='float64'):
    """
    Calculates the pairwise arc distance between all points in vector a and b.
    """
    a = tensor.matrix(dtype=str(dtype))
    b = tensor.matrix(dtype=str(dtype))
    # Theano don't implement all case of tile, so we do the equivalent with alloc.
    #theta_1 = tensor.tile(a[:, 0], (b.shape[0], 1)).T
    theta_1 = tensor.alloc(a[:, 0], b.shape[0], b.shape[0]).T
    phi_1 = tensor.alloc(a[:, 1], b.shape[0], b.shape[0]).T

    theta_2 = tensor.alloc(b[:, 0], a.shape[0], a.shape[0])
    phi_2 = tensor.alloc(b[:, 1], a.shape[0], a.shape[0])

    temp = (tensor.sin((theta_2 - theta_1) / 2)**2
            +
            tensor.cos(theta_1) * tensor.cos(theta_2)
            * tensor.sin((phi_2 - phi_1) / 2)**2)
    distance_matrix = 2 * (tensor.arctan2(tensor.sqrt(temp),
                                          tensor.sqrt(1 - temp)))
    name = "arc_distance_theano_alloc"
    rval = theano.function([a, b],
                           distance_matrix,
                           name=name)
    rval.__name__ = name

    return rval

    def get_celerite_matrices(self, x, diag):
        x = tt.as_tensor_variable(x)
        diag = tt.as_tensor_variable(diag)
        ar, cr, ac, bc, cc, dc = self.coefficients
        a = diag + tt.sum(ar) + tt.sum(ac)
        U = tt.concatenate((
            ar[None, :] + tt.zeros_like(x)[:, None],
            ac[None, :] * tt.cos(dc[None, :] * x[:, None])
            + bc[None, :] * tt.sin(dc[None, :] * x[:, None]),
            ac[None, :] * tt.sin(dc[None, :] * x[:, None])
            - bc[None, :] * tt.cos(dc[None, :] * x[:, None]),
        ), axis=1)

        V = tt.concatenate((
            tt.zeros_like(ar)[None, :] + tt.ones_like(x)[:, None],
            tt.cos(dc[None, :] * x[:, None]),
            tt.sin(dc[None, :] * x[:, None]),
        ), axis=1)

        dx = x[1:] - x[:-1]
        P = tt.concatenate((
            tt.exp(-cr[None, :] * dx[:, None]),
            tt.exp(-cc[None, :] * dx[:, None]),
            tt.exp(-cc[None, :] * dx[:, None]),
        ), axis=1)

        return a, U, V, P

def get_uhs_operator(uhs, depth, n_hidden, rhos):
    """

    :param uhs:
    :param depth:
    :param n_hidden:
    :param rhos: can be shared variable or constant of shape (depth, )!!
    :return:
    """
    # Will use a Fourier matrix (will be O(n^2)...)
    # Doesn't seem to slow things down much though!
    exp_phases = [T.cos(uhs), T.sin(uhs)]
    neg_exp_phases = [T.cos(uhs[:, ::-1]), -T.sin(uhs[:, ::-1])]
    ones_ = [T.ones((depth, 1), dtype=theano.config.floatX), T.zeros((depth, 1), dtype=theano.config.floatX)]

    rhos_reshaped = T.reshape(rhos, (depth, 1), ndim=2)
    rhos_reshaped = T.addbroadcast(rhos_reshaped, 1)

    eigvals_re = rhos_reshaped * T.concatenate((ones_[0], exp_phases[0], -ones_[0], neg_exp_phases[0]), axis=1)
    eigvals_im = rhos_reshaped * T.concatenate((ones_[1], exp_phases[1], -ones_[1], neg_exp_phases[1]), axis=1)
    phase_array = -2 * np.pi * np.outer(np.arange(n_hidden), np.arange(n_hidden)) / n_hidden
    f_array_re_val = np.cos(phase_array) / n_hidden
    f_array_im_val = np.sin(phase_array) / n_hidden
    f_array_re = theano.shared(f_array_re_val.astype(theano.config.floatX), name="f_arr_re")
    f_array_im = theano.shared(f_array_im_val.astype(theano.config.floatX), name="f_arr_im")

    a_k = T.dot(eigvals_re, f_array_re) + T.dot(eigvals_im, f_array_im)
    uhs_op = rep_vec(a_k, n_hidden, n_hidden)  # shape (depth, 2 * n_hidden - 1)

    return uhs_op

  def get_output_for(self, inputs, **kwargs):
    # see eq. (1) and sec 3.1 in [1]
    input, para = inputs
    num_batch, channels, height, width = input.shape
    _w = T.cast(width, dtype = self.dtype)
    _h = T.cast(height, dtype = self.dtype)
    mat = T.zeros((num_batch, 3, 3), dtype = self.dtype)
    mat = T.set_subtensor(mat[:, 0, 0], const(1.0))
    mat = T.set_subtensor(mat[:, 1, 1], const(1.0))
    mat = T.set_subtensor(mat[:, 2, 2], const(1.0))

    if self.method == 'perspective':
      mat = T.set_subtensor(mat[:, 2, 0], (para[:, 0] / 1e4 - 1e-3) * _w)
      mat = T.set_subtensor(mat[:, 2, 1], (para[:, 1] / 1e4 - 1e-3) * _h)
    elif self.method == 'angle':
      angle = T.cast(T.argmax(para, axis = 1), dtype = self.dtype) * np.pi / 90 - np.pi / 3.0
      # ss = np.sqrt(2.0)
      mat = T.set_subtensor(mat[:, :, :], T.stacklists([
        [T.cos(angle), T.sin(angle), -(T.cos(angle) * _w + T.sin(angle) * _h - _w) / (2.0 * _w)],
        [-T.sin(angle), T.cos(angle), -(-T.sin(angle) * _w + T.cos(angle) * _h - _h) / (2.0 * _h)],
        [constv(0, num_batch, self.dtype), constv(0, num_batch, self.dtype), constv(1, num_batch, self.dtype)]]).dimshuffle(2, 0, 1))
      # return [mat, _w, _h]
    elif self.method == 'all':
      mat = T.reshape(para, [-1, 3, 3])
      mat = T.set_subtensor(mat[:, 0, 2], mat[:, 0, 2] / T.cast(width, dtype))
      mat = T.set_subtensor(mat[:, 1, 2], mat[:, 1, 2] / T.cast(height, dtype))
      mat = T.set_subtensor(mat[:, 2, 0], mat[:, 2, 0] * T.cast(width, dtype))
      mat = T.set_subtensor(mat[:, 2, 1], mat[:, 2, 1] * T.cast(height, dtype))
    else:
      raise Exception('method not understood.')
    return transform_affine(mat, input, self.method, scale_factor = self.scale_factor)

    def grad(self, inputs, gradients):
        M, e = inputs
        E, f = self(M, e)

        bM = tt.zeros_like(M)
        be = tt.zeros_like(M)
        ecosE = e * tt.cos(E)

        if not isinstance(gradients[0].type, theano.gradient.DisconnectedType):
            # Backpropagate E_bar
            bM = gradients[0] / (1 - ecosE)
            be = tt.sin(E) * bM

        if not isinstance(gradients[1].type, theano.gradient.DisconnectedType):
            # Backpropagate f_bar
            sinf2 = tt.sin(0.5*f)
            cosf2 = tt.cos(0.5*f)
            tanf2 = sinf2 / cosf2
            e2 = e**2
            ome2 = 1 - e2
            ome = 1 - e
            ope = 1 + e
            cosf22 = cosf2**2
            twoecosf22 = 2 * e * cosf22
            factor = tt.sqrt(ope/ome)
            inner = (twoecosf22+ome) * tt.as_tensor_variable(gradients[1])

            bM += factor*(ome*tanf2**2+ope)*inner*cosf22/(ope*ome2)
            be += -2*cosf22*tanf2/ome2**2*inner*(ecosE-2+e2)

        return [bM, be]

 def _getrz(self,d):
     r=theano.shared(np.zeros((3,3),dtype='float32'))
     r=T.set_subtensor(r[2,2],1.0)
     r=T.set_subtensor(r[0,0],T.cos(d))
     r=T.set_subtensor(r[0,1],-T.sin(d))
     r=T.set_subtensor(r[1,0],T.sin(d))
     r=T.set_subtensor(r[1,1],T.cos(d))        
     return r

def xyz(R1, theta, R2, phi):
    theta_ = T.scalar('theta')
    phi_ = T.scalar('phi')

    x = R1 * T.cos(theta_) * R2 * T.cos(phi_)
    y = R2 * T.sin(theta_) * R2 * T.cos(phi_)
    z = R2 * T.sin(phi_)
    func = function([theta_, phi_], [x, y, z], allow_input_downcast=True)
    vals = func(theta, phi)
    return vals[0].flatten(), vals[1].flatten(), vals[2].flatten()

 def get_output(self, train=False):
     rnorm = self.omega / np.sqrt(2*np.pi)*self.kappa
     val = - self.omega**2 / (8 * self.kappa**2)
     dir1 = 4 * (self._outter(self.x, tensor.cos(self.theta)) +
                 self._outter(self.y, tensor.sin(self.theta)))**2
     dir2 = (-self._outter(self.x, tensor.sin(self.theta)) +
             self._outter(self.y, tensor.cos(self.theta)))**2
     ex = 1j * (self.omega * self._outter(tensor.cos(self.theta), self.x) +
                self.omega * self._outter(tensor.sin(self.theta), self.y))
     output = rnorm * tensor.exp(val * (dir1 + dir2)) * (tensor.exp(ex)
                                                         - tensor.exp(-self.kappa**2 / 2))
     return output

def xyz2(theta, a, b, m, n1, n2, n3, rho, a2, b2, m2, n4, n5, n6):
    theta_ = T.scalar('theta')
    rho_ = T.scalar('rho')
    R1 = R(theta, a, b, m, n1, n2, n3)
    R2 = R(rho, a2, b2, m2, n4, n5, n6)

    x = R1 * T.cos(theta_) * R2 * T.cos(rho_)
    y = R1 * T.sin(theta_) * R2 * T.cos(rho_)
    z = R2 * T.sin(rho_)
    func = function([theta_, rho_], [x, y, z], allow_input_downcast=True)
    vals = func(theta, rho)
    return vals[0].flatten(), vals[1].flatten(), vals[2].flatten()

def get_theano_func():
    a = tt.dvector("a")
    v = tt.dvector("v")
    freq = tt.dscalar("freq")
    t = tt.dscalar("t")
    dt = tt.dscalar("dt")
    tau = tt.dscalar("tau")
    return theano.function(
        [a, v, freq, t, dt, tau],
        a * tt.sin(2.0 * freq * pi * t)
        + b
        + v * tt.exp(-dt / tau)
        + (-a * tt.sin(2.0 * freq * pi * t) - b) * tt.exp(-dt / tau),
    )

    def theano_dynamics(self, x, u):
        G, L1, L2, M1, M2 = self.extract_constants()
        # TODO: this is just an approximation

        dydx = T.alloc(0.0, 4)

        dydx = T.set_subtensor(dydx[0], x[1])

        del_ = x[2]-x[0]

        den1 = (M1+M2)*L1 - M2*L1*T.cos(del_)*T.cos(del_)

        dydx = T.set_subtensor(dydx[1],
            (  M2*L1      *  x[1] * x[1] * T.sin(del_) * T.cos(del_)
               + M2*G       *  T.sin(x[2]) * T.cos(del_) +
                 M2*L2      *  x[3] * x[3] * T.sin(del_)
               - (M1+M2)*G  *  T.sin(x[0]))/den1 )


        dydx = T.set_subtensor(dydx[2], x[3])

        den2 = (L2/L1)*den1
        dydx = T.set_subtensor(dydx[3], (-M2*L2        *   x[3]*x[3]*T.sin(del_)*T.cos(del_)
                   + (M1+M2)*G   *   T.sin(x[0])*T.cos(del_)
                   - (M1+M2)*L1  *   x[1]*x[1]*T.sin(del_)
                   - (M1+M2)*G   *   T.sin(x[2]))/den2  + u )
        return x + dydx * self.dt

 def __init__(self, target, initial_phi, profile_s=None, A0=1.0):
     self.target = target
     self.n_pixels = int(target.shape[0] / 2)   # target should be 512x512, but SLM pattern calculated should be 256x256.
     self.intensity_calc = None
     
     self.cost = None   # placeholder for cost function.
     
     if profile_s is None:
         profile_s = np.ones((self.n_pixels, self.n_pixels))
         
     assert profile_s.shape == (self.n_pixels, self.n_pixels), 'profile_s is wrong shape, should be ({n},{n})'.format(n=self.n_pixels)
     self.profile_s_r = profile_s.real.astype('float64')
     self.profile_s_i = profile_s.imag.astype('float64')
     
     assert initial_phi.shape == (self.n_pixels**2,), "initial_phi must be a vector of phases of size N^2 (not (N,N)).  Shape is " + str(initial_phi.shape)
     
     self.A0 = A0
     
     # Set zeros matrix:
     self.zero_frame = np.zeros((2*self.n_pixels, 2*self.n_pixels), dtype='float64')
     
     # Phi and its momentum for use in gradient descent with momentum:
     self.phi    = theano.shared(value=initial_phi.astype('float64'),
                                 name='phi')
     self.phi_rate = theano.shared(value=np.zeros_like(initial_phi).astype('float64'),
                                   name='phi_rate')
     
     self.S_r = theano.shared(value=self.profile_s_r,
                              name='s_r')
     self.S_i = theano.shared(value=self.profile_s_i,
                              name='s_i')
     self.zero_matrix = theano.shared(value=self.zero_frame,
                                      name='zero_matrix')
     
     # E_in: (n_pixels**2)
     phi_reshaped = self.phi.reshape((self.n_pixels, self.n_pixels))
     self.E_in_r = self.A0 * (self.S_r*T.cos(phi_reshaped) - self.S_i*T.sin(phi_reshaped))
     self.E_in_i = self.A0 * (self.S_i*T.cos(phi_reshaped) + self.S_r*T.sin(phi_reshaped))
     
     # E_in padded: (4n_pixels**2)
     idx_0, idx_1 = get_centre_range(self.n_pixels)
     self.E_in_r_pad = T.set_subtensor(self.zero_matrix[idx_0:idx_1,idx_0:idx_1], self.E_in_r)
     self.E_in_i_pad = T.set_subtensor(self.zero_matrix[idx_0:idx_1,idx_0:idx_1], self.E_in_i)
     
     # E_out:
     self.E_out_r, self.E_out_i = fft(self.E_in_r_pad, self.E_in_i_pad)        
     
     # finally, the output intensity:
     self.E_out_2 = T.add(T.pow(self.E_out_r, 2), T.pow(self.E_out_i, 2))

def hdist(a, b):
    lat1 = a[:, 0] * deg2rad
    lon1 = a[:, 1] * deg2rad
    lat2 = b[:, 0] * deg2rad
    lon2 = b[:, 1] * deg2rad

    dlat = abs(lat1-lat2)
    dlon = abs(lon1-lon2)

    al = tensor.sin(dlat/2)**2  + tensor.cos(lat1) * tensor.cos(lat2) * (tensor.sin(dlon/2)**2)
    d = tensor.arctan2(tensor.sqrt(al), tensor.sqrt(const(1)-al))

    hd = const(2) * rearth * d

    return tensor.switch(tensor.eq(hd, float('nan')), (a-b).norm(2, axis=1), hd)

    def symbolic_call(self,x,u):

        u = TT.clip(u, -self.max_force, self.max_force) #pylint: disable=E1111

        dt = self.dt

        z = TT.take(x,0,axis=x.ndim-1)
        zdot = TT.take(x,1,axis=x.ndim-1)    
        th = TT.take(x,2,axis=x.ndim-1)
        thdot = TT.take(x,3,axis=x.ndim-1)
        u0 = TT.take(u,0,axis=u.ndim-1)

        th1 = np.pi - th

        g = 10.
        mc = 1. # mass of cart
        mp = .1 # mass of pole
        muc = .0005 # coeff friction of cart
        mup = .000002 # coeff friction of pole
        l = 1. # length of pole

        def sign(x):
            return TT.switch(x>0, 1, -1)

        thddot = -(-g*TT.sin(th1)
         + TT.cos(th1) * (-u0 - mp * l *thdot**2 * TT.sin(th1) + muc*sign(zdot))/(mc+mp)
          - mup*thdot / (mp*l))  \
        / (l*(4/3. - mp*TT.cos(th1)**2 / (mc + mp)))
        zddot = (u0 + mp*l*(thdot**2 * TT.sin(th1) - thddot * TT.cos(th1)) - muc*sign(zdot))  \
            / (mc+mp)

        newzdot = zdot + dt*zddot
        newz = z + dt*newzdot
        newthdot = thdot + dt*thddot
        newth = th + dt*newthdot

        done = (z > self.max_cart_pos) | (z < -self.max_cart_pos) | (th > self.max_pole_angle) | (th < -self.max_pole_angle) 

        ucost = 1e-5*(u**2).sum(axis=u.ndim-1)
        xcost = 1-TT.cos(th)
        # notdone = TT.neg(done) #pylint: disable=W0612,E1111
        notdone = 1-done
        costs = TT.stack((done-1)*10., notdone*xcost, notdone*ucost).T #pylint: disable=E1103


        newx = TT.stack(newz, newzdot, newth, newthdot).T #pylint: disable=E1103

        return [newx,newx,costs,done]

    def create_transport_gradient(self):
        HL2 = T.sqrt((self.search_direction ** 2).sum(axis=0,keepdims=True)) + np.spacing(np.single(1))

        transported = self.last_G - (self.last_G*self.search_direction).sum(axis=0,keepdims=True)/(HL2**2) * \
                               (self.X*HL2 * T.sin(self.t_prev*HL2) + self.search_direction * (1.0-T.cos(self.t_prev*HL2)))
        f = theano.function([],[],updates=[(self.last_G,transported)])
        return f

def rotate(angle, axis):
    """Returns a transform to represent a rotation"""

    angle = T.as_tensor_variable(angle)
    axis = T.as_tensor_variable(axis)
    a = axis

    radians = angle*np.pi/180.0
    s = T.sin(radians)
    c = T.cos(radians)

    m = T.alloc(0., 4, 4)

    m = T.set_subtensor(m[0,0], a[0] * a[0] + (1. - a[0] * a[0]) * c)
    m = T.set_subtensor(m[0,1], a[0] * a[1] * (1. - c) - a[2] * s)
    m = T.set_subtensor(m[0,2], a[0] * a[2] * (1. - c) + a[1] * s)

    m = T.set_subtensor(m[1,0], a[0] * a[1] * (1. - c) + a[2] * s)
    m = T.set_subtensor(m[1,1], a[1] * a[1] + (1. - a[1] * a[1]) * c)
    m = T.set_subtensor(m[1,2], a[1] * a[2] * (1. - c) - a[0] * s)

    m = T.set_subtensor(m[2,0], a[0] * a[2] * (1. - c) - a[1] * s)
    m = T.set_subtensor(m[2,1], a[1] * a[2] * (1. - c) + a[0] * s)
    m = T.set_subtensor(m[2,2], a[2] * a[2] + (1. - a[2] * a[2]) * c)

    m = T.set_subtensor(m[3,3], 1)

    return Transform(m, m.T)

 def For_MMD_Sub_class(self,target,data,omega,num_FF,Xlabel):
     
     Num=T.sum(Xlabel,0)
     D_num=Xlabel.shape[1]
     N=data.shape[0]
     
     F_times_Omega = T.dot(data, omega)#minibatch_size*n_rff
     Phi = (self.sf2**0.5 /num_FF**0.5 ) * T.concatenate([T.cos(F_times_Omega), T.sin(F_times_Omega)],1)
     
     #各RFFは２N_rffのたてベクトル
     Phi_total=T.sum(Phi.T,-1)/N
     
     #Domain_number*2N_rffの行列
     Phi_each_domain, updates = theano.scan(fn=lambda a,b: T.switch(T.neq(b,0), Phi.T*a/b, 0),
                           sequences=[Xlabel.T,Num])
     each_Phi=T.sum(Phi_each_domain,-1)
     #まず自分自身との内積 結果はD次元のベクトル
     each_domain_sum=T.sum(each_Phi*each_Phi,-1)
     
     #全体の内積
     tot_sum=T.dot(Phi_total,Phi_total)
     
     #全体とドメインのクロス内積
     tot_domain_sum, updates=theano.scan(fn=lambda a: a*Phi_total,
                           sequences=[each_Phi])
     
     #MMDの計算
     MMD_central=T.sum(each_domain_sum)+D_num*tot_sum-2*T.sum(tot_domain_sum)
     
     return MMD_central     

def beta_noise(rng, input):

    srng = theano.tensor.shared_randomstreams.RandomStreams(rng.randint(999999))
    
    # Beta(.5,.5) multiplicative noise                       
    mask = T.cast(T.sin( (np.pi / 2.0) * srng.uniform(size=input.shape, low=0.0, high=1.0) )**2, theano.config.floatX)
    return mask * input