import numpy as np
import matplotlib as mpl

import common

fontsize = 16

def braket(ax, x, y, text):
    ax.text(
        x,
        y,
        common.braket(text),
        verticalalignment='center',
        horizontalalignment='center',
        fontdict=dict(size=fontsize, weight='bold')
    )

def plot(plt, params):
    color = 'black'
    linewidth = 2
    fig, ax = plt.subplots()
    draw_width = 1.0
    ax.set_aspect(1)
    n = 9
    dx = draw_width/n
    draw_height = 2*dx
    ax.set_xlim(-1.2*dx, draw_width+1.2*dx)
    ax.set_ylim(draw_height * -0.75, draw_height * 1.5)
    line_length = dx*0.8;
    arrow_head_length = 6
    name_map = {
        n//4: ['+', 'i'],
        n//2: ['1', '1'],
        3*(n//4): ['-', '-i'],
    };
    for i in range(n):
        translate_x = draw_width * i/(n-1)
        rotate = np.pi*i/(n - 1)

        def draw_double_arrow(a, b, middleline):
            '''
            This is a mess.
            https://stackoverflow.com/questions/25761717/matplotlib-simple-and-two-head-arrows#comment114130205_64368063
            https://stackoverflow.com/questions/25761717/matplotlib-simple-and-two-head-arrows#comment114130705_55191333
            '''
            for forward in [True, False]:
                if forward:
                    a2 = a
                    b2 = b
                else:
                    a2 = b
                    b2 = a
                arrow = mpl.patches.FancyArrowPatch(
                    a2,
                    b2,
                    color=color,
                    linewidth=0,
                    shrinkA=arrow_head_length,
                    shrinkB=0,
                    arrowstyle=mpl.patches.ArrowStyle.Simple(
                        head_length=arrow_head_length,
                        head_width=15,
                        tail_width=linewidth
                    ),
                )
                t = mpl.transforms.Affine2D().rotate(rotate) + \
                    mpl.transforms.Affine2D().translate(translate_x, draw_height*middleline) + \
                    ax.transData
                arrow.set_transform(t)
                ax.add_artist(arrow)

        # Linear.
        draw_double_arrow((0, -line_length/2), (0,  line_length/2), 0.75)

        # Circular.
        middleline = .25
        if i in [0, n//2, n - 1]:
            draw_double_arrow((0, -line_length/2), (0,  line_length/2), middleline)
        else:
            ellipse_height = np.abs(line_length*np.cos(rotate))
            ellipse = mpl.patches.Ellipse(
                (0, 0),
                np.abs(line_length*np.sin(rotate)),
                ellipse_height,
                fill=False,
                linewidth=linewidth
            )
            t = mpl.transforms.Affine2D().translate(translate_x, draw_height*middleline) + \
                ax.transData
            ellipse.set_transform(t)
            ax.add_artist(ellipse)
            ellipse_top = draw_height*middleline + ellipse_height/2
            marker = ">" if i < n//2 else "<"
            ax.plot(translate_x, ellipse_top, ls="", marker=marker, markersize=8, color="k",
                    clip_on=False)


        # Annotations top and bottom.
        if i in name_map:
            braket(ax, translate_x, draw_height* 1.25, name_map[i][0])
            braket(ax, translate_x, draw_height*-0.25, name_map[i][1])

    # Annotations outside of loop.
    braket(ax, -dx, draw_height*0.5, '0')
    braket(ax, draw_width+dx, draw_height*0.5, '0')
    ax.text(
        draw_width/2,
        draw_height*-0.75,
        'Increasing θ',
        verticalalignment='center',
        horizontalalignment='center',
        fontdict=dict(size=fontsize)
    )
    def draw_axis_arrow(height):
        ax.add_artist(mpl.patches.FancyArrowPatch(
            (draw_width*0.0, height),
            (draw_width*1.0, height),
            color='black',
            linewidth=0,
            arrowstyle=mpl.patches.ArrowStyle.Simple(
                head_length=15,
                head_width=15,
                tail_width=linewidth
            ),
        ))
    draw_axis_arrow(draw_height*-0.55)

    debug = False
    if debug:
        ax.grid()
    else:
        ax.set_axis_off()
    return {'height': 500}