from __future__ import division
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Arc

C = np.array( [ [1, -1j], [1, 1j] ] )
Cinv = np.array( [ [1/2, 1/2], [1/2*1j, -1/2*1j] ] )
Tstar = np.array( [ [1+1/2*1j, -1/2*1j], [1/2*1j, 1-1/2*1j] ] )
Sstar = np.array( [ [-1j, 0], [0, 1j] ] )

# draw unit circle
def drawUnitCirc(axis) :
    unitCirc = Arc( (0, 0), 2, 2, theta1=0, theta2=360, color='k')
    axis.add_patch(unitCirc)

# given invertible 2x2 matrix A, returns A(z) where action is by Mobius transformation
def mobiusTransform(A, z) :
    if z == np.inf :
        if A[1, 0] == 0.0 :
            return np.inf
        return A[0, 0] / A[1, 0]
    elif A[1, 0] * z + A[1, 1] == 0.0 :
        return np.inf
    return (A[0, 0] * z + A[0, 1]) / (A[1, 0] * z + A[1, 1])

# given point in upper half plane, computes the corresponding point in standard
# fundamental domain
#   Assuming group is generated by z maps to z+1 and z maps to -R^2/z
def pullbackToF(z, R) :
    # check if |re(z)| <= 1/2
    x = z.real
    if x <= 0.5 and x >= -0.5 :
        # if norm is bigger than R, done
        if abs(z) >= R :
            return z
        # otherwise, map out of circle of radius R
        else :
            return pullbackToF(-R**2/z, R)

    # otherwise, need to map into |re(z)| <= 1/2
    else :
        return pullbackToF(x - round(x) + z.imag*1j, R)

# Returns theta and P values of F-pullback
#   (Does this by mapping to upper half plane model)
def pullbackList(thetajs, P, R) :
    numPoints = len(thetajs)
    thetajStars = np.zeros(numPoints)
    pjStars = np.zeros(numPoints)
    for j in range(numPoints) :
        # convert polar to rectangular
        tau = P*np.cos(thetajs[j]) + 1j*P*np.sin(thetajs[j])

        # map unit disk to upper half plane
        z = 1j*(1+tau)/(1-tau)

        # pull back to standard fundamental domain
        zStar = pullbackToF(z, R)

        # map back to unit disk
        tauStar = (zStar - 1j)/(zStar + 1j)

        # convert rectangular to polar
        x = tauStar.real
        y = tauStar.imag
        p = np.sqrt(x**2 + y**2)
        theta = np.arccos(x/p)
        if y < 0 :
            theta = 2*np.pi - theta

        # store results
        thetajStars[j] = theta
        pjStars[j] = p

    return thetajStars, pjStars

# converts rectangular (x, y) coordinates to polar (r, theta) coordinates
def rectToPol(x, y) :
    r = np.sqrt(x**2 + y**2)
    theta = np.arccos(x/r)
    if y < 0 :
        theta = 2*np.pi - theta
    return r, theta

# converts polar (r, theta) coordinates to rectangular (x, y) coordinates
def polToRect(r, theta) :
    x = r*np.cos(theta)
    y = r*np.sin(theta)
    return x, y

# plot list of points given in polar coordinates
def plotPolarList(rs, thetas, optionStr, axis) :
    # get x and y coordinates
    N = len(rs)
    xs = np.zeros(N)
    ys = np.zeros(N)
    for i in range(N) :
        x, y = polToRect(rs[i], thetas[i])
        xs[i] = x
        ys[i] = y

    axis.plot( xs, ys, optionStr)

# plot geodesic defined by two endpoints
#   - endpoints are on unit circle, so are described by angles
#   - assumes thetas are input in range [0, 2pi)
#   - geodesics are circles tangent to the boundary
def plotGeodesic(theta1, theta2, axis) :
    if theta1 == theta2 :
        raise Exception('theta1 cannot equal theta2')

    # find acute angle between thetas
    ang = np.abs(theta1 - theta2)
    wasObtuse = 0
    if ang == np.pi :
        # plot line between them
        x1, y1 = polToRect(1, theta1)
        x2, y2 = polToRect(1, theta2)
        axis.plot( [x1, x2], [y1, y2], color='k')
        return
    elif ang > np.pi :
        ang = 2*np.pi - ang
        wasObtuse = 1

    # Assume we rotate to be symmetric across x-axis and in 1st/4th quadrant
    theta = ang / 2

    # get coordinates of center
    rC = 1 / np.cos(theta)
    thetaC = (theta1 + theta2) / 2
    if wasObtuse :
        thetaC += np.pi
    xC, yC = polToRect(rC, thetaC)

    # get diameter
    xP, yP = polToRect(1, theta)
    diam = 2*np.sqrt( (xP - rC)**2 + yP**2 )

    # plot arc
    ang1 = np.pi/2 + theta + thetaC
    ang2 = 3*np.pi/2 - theta + thetaC
    arc = Arc( (xC, yC), diam, diam, theta1=180*ang1/np.pi, theta2=180*ang2/np.pi, color='k')
    axis.add_patch(arc)

# gets 6 endpoints defining the standard fundamental domain for a Hecke group with parameter R
def getHeckeEndpoints(R) :
    endPts = [0, 0, 0, 0, 0, 0]

    # end points are images of -1/2, inf, 1/2, inf, -R, R under Cayley transform
    endPts[0] = mobiusTransform(C, -1/2)
    endPts[1] = mobiusTransform(C, np.inf)
    endPts[2] = mobiusTransform(C, 1/2)
    endPts[3] = mobiusTransform(C, np.inf)
    endPts[4] = mobiusTransform(C, -R)
    endPts[5] = mobiusTransform(C, R)

    return endPts

# plots the standard fundamental domain for a Hecke group with parameter R
def heckeFundDomain(R, axis) :
    # draw unit circle
    drawUnitCirc(axis)

    # get end points
    endPts = getHeckeEndpoints(R)

    # for each pair, plot geodesic
    for i in range(3) :
        e1 = endPts[2*i]
        e2 = endPts[2*i+1]

        r, theta1 = rectToPol(e1.real, e1.imag)
        r, theta2 = rectToPol(e2.real, e2.imag)

        plotGeodesic(theta1, theta2, axis)

# given six starting points and a Mobius transformation, plots fundamental domain
# defined by the image of those points under the map
#   optional argument: a point in the domain to be plotted
def plotShiftedDomain(pts, A, axis, ptInDomain=None) :
    # for each pair, find image under A then plot geodesic
    for i in range(3) :
        e1 = pts[2*i]
        e2 = pts[2*i+1]

        Ae1 = mobiusTransform(A, e1)
        Ae2 = mobiusTransform(A, e2)

        r, theta1 = rectToPol(Ae1.real, Ae1.imag)
        r, theta2 = rectToPol(Ae2.real, Ae2.imag)

        plotGeodesic(theta1, theta2, axis)

    # if optional argument is included, also plot that shifted point
    if ptInDomain != None :
        P = mobiusTransform(A, ptInDomain)
        axis.plot(P.real, P.imag, 'k.')

# given radius rho, plots N points in that circle after map back to upper half plane
def plotCircInH(rho, N) :
    # make circle of thetas
    thetas = 2*np.pi*np.arange(N) / N

    # for each theta, get x, y vals of pts in H
    xs = np.zeros(N)
    ys = np.zeros(N)
    for i in range(N) :
        x, y = polToRect(rho, thetas[i])
        z = mobiusTransform(Cinv, x + 1j*y)
        xs[i] = z.real
        ys[i] = z.imag

    # plot the points
    plt.plot(xs, ys, 'k.')

# set up figure with disk drawn to plot on
def setupDisk() :
    fig = plt.figure()
    ax = plt.axes( xlim=(-1.1, 1.1), ylim=(-1.1, 1.1), aspect='equal')
    drawUnitCirc(ax)
    return fig, ax
