import math

def nearPose((x1,y1,th1),(x2,y2,th2), distEps, angleEps):
   return nearPose2((x1,y1),(x2,y2), distEps) and nearAngle(th1, th2, angleEps)

def nearPose2(p1, p2, eps):
   return dist(p1, p2) < eps

def dist((x1,y1),(x2,y2)):
   return math.sqrt((x1 - x2)**2 + (y1 - y2)**2) 

def nearAngle(a1,a2,eps):
   return abs(fixAnglePlusMinusPi(a1-a2)) < eps

def diffPose((x1,y1,th1),(x2,y2,th2)):
   return ((x1-x2),(y1-y2),fixAnglePlusMinusPi(th1-th2))

def diffPose2((x1,y1),(x2,y2)):
   return (x1-x2,y1-y2)

def nearlyEqual(x,y):
    return abs(x-y)<.0001


########   New transform stuff 

def poseToTransform((ax, ay, atheta)):
    cos = math.cos(atheta)
    sin = math.sin(atheta)
    return [[cos, -sin, ax],[sin, cos, ay], [0, 0, 1]]

def inverseTransform(((c, ms, x),(s, c2, y), (z1, z2, o))):
    return [[c, s, (-c*x)-(s*y)],
            [-s, c, (s*x)-(c*y)],
            [0, 0, 1]]

def composeTransforms(t1, t2):
   return mm(t2, t1)

def mm(t1, t2):   
   result = make2DArray(3, 3, 0)
   for i in range(3):
      for j in range(3):
         for k in range(3):
            result[i][j] += t1[i][k]*t2[k][j]
   return result

def mvmult(m, v):
   result = makeVector(3, 0)
   for i in range(3):
      for j in range(3):
         result[i] += m[i][j]*v[j]
   return result

def vmmult(v, m):
   result = makeVector(3, 0)
   for i in range(3):
      for j in range(3):
         result[i] += v[i] * m[i][j]
   return result

from operator import add
def vplus(v1, v2):
   return map(add, v1, v2)
   
def transformToPose(t):
   theta = math.atan2(t[1][0],t[0][0])
   return (t[0][2], t[1][2], theta)

def transformPose(pose, transform):
   return transformToPose(composeTransforms(poseToTransform(pose), transform))

#####################
#
# Gaussian stuff

def logGaussian(x, mu, sigma):
   return -((x-mu)**2 / (2*sigma**2)) - math.log(sigma*math.sqrt(2*math.pi))

def gaussian(x, mu, sigma):
   return math.exp(-((x-mu)**2 / (2*sigma**2))) / (sigma*math.sqrt(2*math.pi))   

#####################

def localToGlobal((rx,ry,rtheta), (hx,hy)):
    cos = math.cos(rtheta)
    sin = math.sin(rtheta)
    return [rx + cos*hx - sin*hy, ry + sin*hx + cos*hy]

# Given robot's pose in a global frame and a point in the global frame
# return coordinates of point in local frame
def globalToLocal((rx,ry,rtheta), (gx,gy)):
    # Translate
    ax = gx - rx
    ay = gy -ry
    # Rotate
    cos = math.cos(-rtheta)
    sin = math.sin(-rtheta)
    return [ax*cos - ay*sin, ay*cos + ax*sin]

def globalPoseToLocalPose((rx, ry, rtheta), (gx, gy, gtheta)):
   (x,y) = globalToLocal((rx, ry, rtheta), (gx, gy))
   return (x, y, fixAnglePlusMinusPi(gtheta - rtheta))

def localPoseToGlobalPose((rx, ry, rtheta), (gx, gy, gtheta)):
   (x,y) = localToGlobal((rx, ry, rtheta), (gx, gy))
   return (x, y, fixAnglePlusMinusPi(gtheta + rtheta))

# Given robot's pose in a global frame an a point in the global frame
# return coordinates of point in local frame
def globalDeltaToLocal((rx,ry,rtheta), (dx,dy)):
    # Translate
    # Rotate
    cos = math.cos(-rtheta)
    sin = math.sin(-rtheta)
    return [dx*cos - dy*sin, dy*cos + dx*sin]

def sub(x,y): return x - y
def add(x,y): return x + y
def times(a,b): return a*b

# Put betwen minus pi and pi
def fixAnglePlusMinusPi(a):
   if a > math.pi:
      return fixAnglePlusMinusPi(a - 2.0* math.pi)
   elif a < -math.pi:
      return fixAnglePlusMinusPi(a + 2.0*math.pi)
   else:
      return a

# Put an angle between 0 and 2 pi
def fixAngle02Pi(a):
    if a < 0:
        return fixAngle02Pi(a + 2*math.pi)
    elif a > 2*math.pi:
        return fixAngle02Pi(a - 2*math.pi)
    else:
        return a

# Return angle in radians of vector from p to g
def angleToPoint((px,py),(gx,gy)):
   if abs(gx - px) < .0000001 and gy > py:
      return math.pi/2.0
   elif abs(gx - px) < .0000001 and py > gy:
      return -math.pi/2.0
   elif abs(gy - py) < .0000001 and gx > px:
      return 0.0
   elif abs(gy - py) < .0000001 and px > gx:
      return - math.pi
   else:
      signx = sign(gx-px)
      signy = sign(gy-py)
      if signx == 1:
         return math.atan((gy - py) / (gx - px))
      else:
         return fixAnglePlusMinusPi(math.atan((gy - py) / (gx - px)) + math.pi)

def argmax(seq, fn):
    best = seq[0]; best_score = fn(best)
    for x in seq:
        x_score = fn(x)
        if x_score > best_score:
            best, best_score = x, x_score
    return best

def randomMultinomial(dist):
    r = random.random()
    for i in range(len(dist)):
        r = r - dist[i]
        if r < 0.0:
            return i
    return "weird"

def sign(x):
    if x > 0.0:
        return 1
    elif x == 0.0:
        return 0
    else:
        return -1

def make3DArray(dim1, dim2, dim3, initValue):
   result = []
   for i in range(dim1):
      result = result + [make2DArray(dim2, dim3, initValue)]
   return result

def make2DArray(dim1, dim2, initValue):
   result = []
   for i in range(dim1):
      result = result + [makeVector(dim2, initValue)]
   return result

def makeVector(dim, initValue):
   return [initValue]*dim

def mapArray3D(arr, f):
   for i1 in range(len(arr)):
      for i2 in range(len(arr[0])):
         for i3 in range(len(arr[0][0])):
            arr[i1][i2][i3] = f(arr[i1][i2][i3])

def sumArray3D(arr, f):
   theSum = 0
   for slice1 in arr:
      for slice2 in slice1:
         for elt in slice2:
            theSum += f(elt)
   return theSum

      
