from __future__ import division
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Wedge
import poincareModel as pm
from mpmath import mp

# This file contains various operations related to Schottky groups
# generated by reflections in hyperbolic 2-space

# computes Cayley transform of z
def cayley_transform(z) :
    if z == np.inf :
        return 1
    return (z - 1j) / (z + 1j)

# computes inverse Cayley transform of w
def inv_cayley_transform(w) :
    if w == 1 :
        return np.inf
    return 1j*(1 + w) / (1 - w)

def rect_to_pol(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

# computes reflection through geodesic semicircle in upper half plane
def reflect(z, cent, rad) :
    if abs(z - cent) < 1e-10 :
        return np.inf
    return rad**2 / (z.real - z.imag*1j - cent) + cent

# We can define a symmetric Schottky group by
#   - N: number of circles
#   - thet: angle cut out by one circle in the disk domain
# This function takes (N, thet) as input and outputs the corresponding circles
# in the upper half plane model (as a list of [center, radius])
def get_UHP_circles(N, thet) :
    if thet > 2*np.pi/N :
        raise ValueError("theta must be <= 2pi/N")

    # circles[i, 0] will be the i^th center, and circles[i, 1] its radius
    circles = np.zeros((N, 2))
    for i in range(0, N) :
        center_angle = i*2*np.pi/N
        ep1 = inv_cayley_transform(np.exp(1j*(center_angle - thet/2)))
        ep2 = inv_cayley_transform(np.exp(1j*(center_angle + thet/2)))

        circles[i, 0] = (ep1 + ep2) / 2
        circles[i, 1] = abs(ep1 - ep2) / 2

    return circles

# Given a center and a radius of a geodesic, computes the two endpoints
def get_endpoints(center, radius) :
    ep1 = center - radius
    ep2 = center + radius
    return ep1, ep2

# Given a list of circles and radii, creates a FundamentalDomain object
# (for use in plotting)
def get_FD(circles) :
    # For each circle, form a geodesic
    geodesics = []
    for i in range(len(circles)) :
        ep1, ep2 = get_endpoints(circles[i, 0], circles[i, 1])
        geodesics.append(pm.Geodesic(ep1, ep2))

    return pm.FundamentalDomain(geodesics, 1j)

# Given a list of circles, plots the fundamental domain
def plot_FD(circles, x_lim, y_lim) :
    fd = get_FD(circles)
    ax = pm.setupFig(-x_lim, x_lim, 0, y_lim)
    fd.draw(ax, 'k', fill=1)
    return ax

#---------------------------------------------
# Assume 3 circles in the following functions!
#---------------------------------------------

# Pullback algorithm for fundamental domain of the symmetric Schottky group
# with 3 circles mapped to the upper half plane model
def pullback_Schottky(z, circles) :

    # loop until in fundamental domain
    i = 0
    while 1 :
        in_FD = 1

        # check if outside R1
        if abs(z - circles[0, 0]) > circles[0, 1] :
            # map by R1
            z = reflect(z, circles[0, 0], circles[0, 1])
            in_FD = 0

        # check if in R2
        if abs(z - circles[1, 0]) < circles[1, 1] :
            # map by R2
            z = reflect(z, circles[1, 0], circles[1, 1])
            in_FD = 0

        # check if in R3
        if abs(z - circles[2, 0]) < circles[2, 1] :
            # map by R3
            z = reflect(z, circles[2, 0], circles[2, 1])
            in_FD = 0

        if in_FD == 1:
            return z

        i = i + 1
        if i == 30 :
            print('Warning: z = ' + str(z) + ' took more than 30 steps in pullback algorithm.')
            return None

# Given two (center, radius) pairs (c, r) and (a, t), computes the fixed
# points of the Mobius transformation obtained from reflection across the
# first circle composed with reflection across the second circle
def fixed_points_refls(c, r, a, t) :
    A = (c**2 - a**2 + t**2 - r**2)/2/(c - a)
    B = np.sqrt(( (a - c)**2 - (r**2 + t**2) )**2 - 4*r**2*t**2 )/2/(c - a)
    return A - B, A + B

# Produces the geodesics cutting off the flares in the
# symmetric Schottky group with 3 circles of angle thet
def get_axes(thet) :
    # circles in upper half plane
    circles = get_UHP_circles(3, thet)

    # get geodesic fixed by R1R2
    ep1, ep2 = fixed_points_refls(circles[0, 0], circles[0, 1],
                                  circles[1, 0], circles[1, 1])
    g1 = pm.Geodesic(ep1, ep2)

    # get geodesic fixed by R1R3
    ep1, ep2 = fixed_points_refls(circles[0, 0], circles[0, 1],
                                  circles[2, 0], circles[2, 1])
    g2 = pm.Geodesic(ep1, ep2)

    # get geodesic fixed by R2R3
    ep1, ep2 = fixed_points_refls(circles[1, 0], circles[1, 1],
                                  circles[2, 0], circles[2, 1])
    g3 = pm.Geodesic(ep1, ep2)

    return g1, g2, g3

# Draws Schottky group with 3 circles of angle thet, plus the
# geodesics cutting off the flares
def draw_labeled_flares(thet, xwidth=5, ywidth=8) :
    # get circles and geodesics
    circles = get_UHP_circles(3, thet)
    g1, g2, g3 = get_axes(thet)

    # draw
    ax = plot_FD(circles, xwidth, ywidth)
    g1.draw(ax, 'k', ':')
    g2.draw(ax, 'k', ':')
    g3.draw(ax, 'k', ':')

    # label flares
    plt.text(-.28, .5, '$R_2R_3$', fontsize=12)
    plt.text(-3.2, 3.5, '$R_1R_2$', fontsize=12)
    plt.text(2.5, 3.5, '$R_1R_3$', fontsize=12)

    plt.show()

# Draws flare domain for one of the three flares in Schottky group
# with 3 circles of angle thet (all three domains are equivalent)
def draw_flare_domains(thet, X=5, Y=5) :
    # get geodesics
    circles = get_UHP_circles(3, thet)
    _, _, g = get_axes(thet)

    # get endpoints
    z1 = min(g.ep1, g.ep2)
    z2 = max(g.ep1, g.ep2)
    t = circles[1, 0] + circles[1, 1]

    # Define Mobius transformation and apply to whole fundamental domain
    mult = (t - z2)/(t - z1)
    U = np.array([ [mult, -z1*mult], [1, -z2] ])
    fd_temp = get_FD(circles)
    fd = pm.mobius(U, fd_temp)
    fd.referencePoint = 1j

    # And plot the fundamental domain!mult = (t - z2)/(t - z1)
    U = np.array([ [mult, -z1*mult], [1, -z2] ])
    fd_temp = get_FD(circles)
    fd = pm.mobius(U, fd_temp)
    ax = pm.setupFig(-X, X, 0, Y)
    fd.draw(ax, fill=1)
    return ax

# draws Schottky group plus flares expanded out to angle omega
def draw_expanded_flares(thet, omega, xwidth=5, ywidth=8) :
    # get geodesics
    circles = get_UHP_circles(3, thet)
    g1, g2, g3 = get_axes(thet)

    # draw domain
    ax = plot_FD(circles, xwidth, ywidth)

    # for each flare, draw flare at angle omega
    gs = [g1, g2, g3]
    for i in range(3) :
        # identify endpoints
        g = gs[i]
        z1 = min(g.ep1, g.ep2)
        z2 = max(g.ep1, g.ep2)

        # compute circle center
        r0 = (z2 - z1)/2
        beta = omega - np.pi/2
        r = r0/np.cos(beta)
        cent = z1 + r*np.exp(1j*beta)

        # compute angle in degrees
        ang = beta * 180 / np.pi

        # draw circle
        wedge = Wedge( (cent.real, cent.imag), r, -ang, 180 + ang, color='k', ls=':', fill=0, width=0)
        ax.add_patch(wedge)

    plt.show()

# draws flare domain plus the sector conatining the flare
def draw_flare_with_sector() :
    ax = draw_flare_domains(np.pi/2, 13, 6)

    g1 = pm.Geodesic(inv_cayley_transform(np.exp(2*np.pi*1j/3)).real,
                     inv_cayley_transform(-np.exp(2*np.pi*1j/3)).real)
    g2 = pm.Geodesic(inv_cayley_transform(np.exp(4*np.pi*1j/3)).real,
                     inv_cayley_transform(-np.exp(4*np.pi*1j/3)).real)

    # get geodesics
    thet = np.pi/2
    circles = get_UHP_circles(3, thet)
    _, _, g = get_axes(thet)

    # get endpoints
    z1 = min(g.ep1, g.ep2)
    z2 = max(g.ep1, g.ep2)
    t = circles[1, 0] + circles[1, 1]

    mult = (t - z2)/(t - z1)
    U = np.array([ [mult, -z1*mult], [1, -z2] ])
    fd_temp = pm.FundamentalDomain([g1, g2])
    fd = pm.mobius(U, fd_temp)

    fd.draw(ax, lstyle=':')

    plt.show()

def test_pullback(X=5, Y=5) :
    circles = get_UHP_circles(3, np.pi/2)

    #zs = np.arange(1000)/180 + 1 + 1j/10
    zs = np.array([-9.4 + 1j])
    ws = []
    for z in zs :
        ws.append(pullback(z, circles))
    ws = np.array(ws)

    ax = plot_FD(circles, X, Y)
    ax.plot(zs.real, zs.imag, 'k.')
    ax.plot(ws.real, ws.imag, 'b.')

    plt.show()

#----------------------------------
# Functions for checking pari code!
#----------------------------------

# map to flare (fixed by R2R3)
def map_to_flare(z, thet) :
    # get circles/geodesics
    circles = get_UHP_circles(3, thet)
    _, _, g = get_axes(thet)

    # RECALL: circles order is R1, R3R1(im_axis), R3(R1),
    #                          R2R1(im_axis), R2(R1)

    # depending on type, get appropriate endpoints
    z1 = min(g.ep1, g.ep2)
    z2 = max(g.ep1, g.ep2)
    t = circles[1, 0] + circles[1, 1]

    # Apply Mobius transformation to z
    mult = (t - z2)/(t - z1)
    return mult*(z - z1)/(z - z2)

# use mpmath to get flare expansion
def flare_Whitt(thet, m, s, kappa) :
    mu = -1/2 + 2*mp.pi*1j*m/mp.log(kappa)
    return mp.sqrt(mp.sin(thet))*mp.legenp(mu, -s + 1/2, mp.cos(thet))

# use mpmath to get disk expansion
def disk_Whitt(rho, m, s) :
    return (1 - rho**2)**s*rho**abs(m)*mp.hyp2f1(s, s + abs(m), 1 + abs(m), rho**2)

# print some calculations on the point z
def print_z_and_pb(z, circles) :
    print('z =', z)
    print('')
    w = pullback_shifted(z, circles)
    print('Pullback =', w)
    return w

def plot_test_points() :
    # read file with points
    with open('testPoints.txt') as f :
        line = f.readline().strip()

    # convert point strings to complex numbers
    pts_string = line.replace('I', 'j')
    pts_string = pts_string.replace('*', '').split(',')
    pts = []
    for p in pts_string :
        p_new = ''.join(p.split())
        try :
            pts.append(complex(p_new))
        except ValueError :
            print(p_new)
            return

    # get geodesics
    circles = get_UHP_circles(3, np.pi/2)
    g1, g2, g3 = get_axes(np.pi/2)

    # draw fundamental domain
    ax = plot_FD(circles, 5, 8)
    g1.draw(ax, 'k', ':')
    g2.draw(ax, 'k', ':')
    g3.draw(ax, 'k', ':')

    # draw points
    pts = np.array(pts)
    ax.plot(pts.real, pts.imag, '.k')

    plt.show()

if __name__ == '__main__' :
    plot_test_points()
