import itertools
import numpy as np
import seaborn as sns
import matplotlib as mpl
import matplotlib.pyplot as plt

%matplotlib inline

sns.set(style="white")
mpl.rc("savefig", dpi=130)

fig = plt.figure(figsize=(6.85, 2.25))
ratio = 6.85 / 2.25

# Parameters to control screen axes location
left = .015
top = .92
w, h = .1625, .35
w_step, h_step = .115, .07

# Durations of the component trial parts
durations = ["4.0", "6.5", "1.5", "0.5", "1.0", "0.5", "2.0", "6.5"]

# Make an Axes for each "screen"
axes = []
poly_axes = []
for i in range(8):

    ax = fig.add_axes([left, top - h, w, h], axis_bgcolor="k")
    sns.despine(ax=ax, left=True, bottom=True)
    ax.plot([0, 0], [0, 1], "w", zorder=3, aa=False)
    ax.plot([0, 1], [1, 1], "w", zorder=3, aa=False)
    ax.set(xticks=[], yticks=[])
    axes.append(ax)

    # Add separate square axes to facilitate drawing the polygons
    if i in [3, 5]:
        ax = fig.add_axes([left + (w - (h / ratio)) / 2, top - h, h / ratio, h])
        ax.set_axis_off()
        poly_axes.append(ax)

    fig.text(left + w, top, durations[i] + " s", size=8, ha="right", va="bottom")
        
    left += w_step
    top -= h_step

# ============================================================================

# Add in the cue instructions
axes[0].text(.5, .5, "COLORS\nSAME\n[y n]", color="white",
             size=8, ha="center", va="center")
cross_kws = dict(size=20, ha="center", va="center")

# Add in the fixation crosses
axes[1].text(.5, .48, "+", color="w", **cross_kws)
axes[2].text(.5, .48, "+", color="r", **cross_kws)
axes[4].text(.5, .48, "+", color="w", **cross_kws)
axes[6].text(.5, .48, "+", color="g", **cross_kws)
axes[7].text(.5, .48, "+", color="w", **cross_kws)

# ============================================================================
# Draw the first stimulus polygon

ax = poly_axes[0]
poly = mpl.patches.CirclePolygon((.5, .5), .3, 7,
                                 facecolor="#8FB46B",
                                 edgecolor="white")

kws = dict(linewidth=1, color="#112801")
ax.plot([.40, .60], [.75, .75], **kws)
ax.plot([.32, .68], [.71, .71], **kws)

ax.plot([.23, .77], [.52, .52], **kws)
ax.plot([.22, .78], [.48, .48], **kws)

ax.plot([.30, .70], [.32, .32], **kws)
ax.plot([.33, .67], [.28, .28], **kws)

verts = poly.get_verts()
for (xa, ya), (xb, yb) in zip(verts[:-1], verts[1:]):
    ax.plot([xa, xb], [ya, yb], "white", lw=1)
ax.add_artist(poly)
ax.set(xlim=(0, 1), ylim=(0, 1))

# ============================================================================
# Draw the second stimulus polygon

ax = poly_axes[1]
poly = mpl.patches.CirclePolygon((.5, .5), .3, 5,
                                 facecolor="#6AF46B",
                                 edgecolor="white")

kws = dict(linewidth=1, color="#296006")
ax.plot([.39, .61], [.72, .72], **kws)

ax.plot([.23, .77], [.60, .60], **kws)

ax.plot([.25, .75], [.48, .48], **kws)

ax.plot([.29, .71], [.36, .36], **kws)

verts = poly.get_verts()
for (xa, ya), (xb, yb) in zip(verts[:-1], verts[1:]):
    ax.plot([xa, xb], [ya, yb], "white", lw=1)
ax.add_artist(poly)
ax.set(xlim=(0, 1), ylim=(0, 1))

# ============================================================================
# Draw an arrow to show the looping structure

ax = fig.add_axes([0, 0, 1, 1])
ax.set_axis_off()
arrow = mpl.patches.FancyArrowPatch((.82, .07), (.3, .4),
                                    color=".1",
                                    linewidth=2,
                                    arrowstyle="->,head_width=10,head_length=15",
                                    connectionstyle="arc3,rad=-.2")
ax.add_artist(arrow)
ax.text(.56, .1, "2x", size=10, color=".1")

# Save the figure
fig.savefig("figures/figure_1.pdf")
fig.savefig("figures/figure_1.tiff", dpi=300)