import matplotlib.pyplot as plt
import matplotlib.cbook as cbook
import numpy as np

from medcouple import medcouple_1d

def boxhistplot(data, fig=None, rect=None, xticks=None, colour=None,
                bins=None, title=None, adjusted=False):
    data_stats = cbook.boxplot_stats(data)

    left, bottom, width, height = rect
    histheight = 0.75*height
    boxheight = 0.25*height

    # setup the figure and axes
    histAx = fig.add_axes([left, bottom, width, histheight])
    bpAx = fig.add_axes([left, bottom+histheight, width, boxheight])

    # plot stuff
    bpAx.bxp(data_stats, vert=False, flierprops={"marker": 'x'})
    histAx.hist(data, bins=bins, color=colour)


    xlims = np.array([bpAx.get_xlim(), histAx.get_xlim()])

    # Do an adjusted boxplot
    if adjusted:
        mc = medcouple_1d(data)
        print "MC = ", mc
        iqr = data_stats[0]['iqr']
        q1 =  data_stats[0]['q1']
        q3 =  data_stats[0]['q3']
        if mc > 0:
            whishi = 1.5*iqr*np.exp(4*mc)
            whislo = 1.5*iqr*np.exp(-3*mc)
        else:
            whishi = 1.5*iqr*np.exp(3*mc)
            whislo = 1.5*iqr*np.exp(-4*mc)
        data_stats[0]['whishi'] = q3 + whishi
        data_stats[0]['whislo'] = q1 - whislo

        # Recompute the outliers
        data_stats[0]['fliers'] = [
            flier for flier in data
            if (flier < data_stats[0]['whislo']  or
                flier > data_stats[0]['whishi'])
        ]
        bpAx.cla()
        bpAx.bxp(data_stats, vert=False, flierprops={"marker": 'x'})

    # confirm that the axes line up
    for ax in [bpAx, histAx]:
        ax.set_xlim([xlims.min(), xlims.max()])

    bpAx.set_xticklabels([])  # clear out overlapping xlabels
    bpAx.set_yticks([])  # don't need that 1 tick mark
    bpAx.set_title(title,fontsize=20)

    histAx.set_xticks(xticks)
    histAx.get_xaxis().tick_bottom()
    bpAx.get_xaxis().tick_top()

    print "%d outliers for %s" % (len(data_stats[0]['fliers']), title)

with open("../../data/men") as f:
    men = [float(x) for x in f.readlines()]

with open("../../data/women") as f:
    women = [float(x) for x in f.readlines()]

# curiousgnu's data

xticks = np.arange(5,105,5)
bins = 0.5 + np.arange(0,100)

boyrect = [0.05, 0.55, 0.9, 0.35]
boycolour = [0.3, 0.3, 1]

girlrect = [0.05, 0.1, 0.9, 0.35]
girlcolour = 'pink'

fig = plt.figure(figsize=(12,8))
boxhistplot(men, fig=fig, rect=boyrect, xticks=xticks, colour=boycolour,
            bins=bins, title="Ages of actors")
boxhistplot(women, fig=fig, rect=girlrect, xticks=xticks, colour=girlcolour,
            bins=bins, title="Ages of actresses")
plt.savefig("boys-and-girls.pdf")

fig = plt.figure(figsize=(12,8))
boxhistplot(men, fig=fig, rect=boyrect, xticks=xticks, colour=boycolour,
            bins=bins, title="Ages of actors", adjusted=True)
boxhistplot(women, fig=fig, rect=girlrect, xticks=xticks, colour=girlcolour,
            bins=bins, title="Ages of actresses", adjusted=True)
plt.savefig("boys-and-girls-adjusted.pdf")

# Geometric distro

np.random.seed(0)
geo = np.random.geometric(0.20, size=10000)

plt.close('all')
plt.plot(geo, [1]*geo.size, 'x')
plt.yticks([])
plt.savefig("geometric-points.pdf")

plt.close('all')
plt.boxplot(geo, vert=False, flierprops=dict(marker='x'))
plt.savefig("geometric-boxplot.pdf")

plt.axis('off')
plt.savefig("geometric-boxplot-bare.pdf")

fig = plt.figure(figsize=(12,8))
boxhistplot(geo, fig=fig, rect=[0.05, 0.1, 0.9, 0.8], colour='yellow',
            bins=20, title="MS lesion counts (simulated)",
            xticks=[0, 10, 20, 30, 40, 50])
plt.savefig("geometric-boxhistplot.pdf")

fig = plt.figure(figsize=(12,8))
boxhistplot(geo, fig=fig, rect=[0.05, 0.1, 0.9, 0.8], colour='yellow',
            bins=20, title="MS lesion counts (simulated)",
            xticks=[0, 10, 20, 30, 40, 50], adjusted=True)
plt.savefig("geometric-boxhistplot-adjusted.pdf")

# Normal distro
N = np.random.normal(size=1000)

plt.close('all')
plt.plot(N, [1]*N.size, 'x')
plt.yticks([])
plt.savefig("normal-points.pdf")

plt.close('all')
plt.boxplot(N, vert=False, flierprops=dict(marker='x'))
plt.savefig("normal-boxplot.pdf")

plt.axis('off')
plt.savefig("normal-boxplot-bare.pdf")

fig = plt.figure(figsize=(12,8))
boxhistplot(N, fig=fig, rect=[0.05, 0.1, 0.9, 0.8], colour='green',
            bins=20, title="Brain volumes (simulated)",
            xticks=range(-4,5))
plt.savefig("normal-boxhistplot.pdf")

fig = plt.figure(figsize=(12,8))
boxhistplot(N, fig=fig, rect=[0.05, 0.1, 0.9, 0.8], colour='green',
            bins=20, title="Brain volumes (simulated)",
            xticks=range(-4,5), adjusted=True)
plt.savefig("normal-boxhistplot-adjusted.pdf")