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

np.random.seed(12345)
nbins = 10
shp = (int(1e2), 3)
WID = 0.7
colors = ['red', 'blue', 'forestgreen']
ALIGN = ['left', 'mid', 'right']

# THIS WORKS FINE:
# XSCALE = 'linear'
# YSCALE = 'linear'

# THIS DOES NOT WORK:
XSCALE = 'log'
YSCALE = 'log'

fig, axes = plt.subplots(figsize=[12, 6], ncols=4, nrows=2)
plt.subplots_adjust(left=0.05, right=0.95, bottom=0.05, top=0.9)

# Generate Data
# ------------------------
xx = np.random.lognormal(1.0, 2.0, shp)
yy = np.random.lognormal(1.0, 2.0, shp)

for ii in range(shp[1]):
    aa = 4.0
    bb = (ii + 1)
    xx[:, ii] += np.maximum(np.random.normal(aa, bb, shp[0]), 0.0)
    yy[:, ii] += np.maximum(np.random.normal(aa, bb, shp[0]), 0.0)

bads = (xx < np.percentile(xx, 1)) | (xx > np.percentile(xx, 99))
xx[bads] = np.random.choice(xx[~bads], bads.sum())

bads = (yy < np.percentile(yy, 1)) | (yy > np.percentile(yy, 99))
yy[bads] = np.random.choice(yy[~bads], bads.sum())

xlim = [np.min(xx), np.max(xx)]
ylim = [np.min(yy), np.max(yy)]


# Plot Data
# -----------------------

def plot_scatter_hist(axes, align_flag):
    align = [aa if align_flag else 'mid' for aa in ALIGN]

    # Scatter
    axes[0, 1].axis('off')
    ax = axes[1, 0]
    ax.set(xscale=XSCALE, yscale=YSCALE, xlim=xlim, ylim=ylim)
    for ii in range(3):
        ax.scatter(xx[:, ii], yy[:, ii], alpha=0.2, color=colors[ii])


    # Marginalize y
    if XSCALE.startswith('log'):
        edges = np.logspace(*np.log10(xlim), nbins)
    else:
        edges = np.linspace(*xlim, nbins)
        
    ax = axes[0, 0]
    ax.set(xscale=XSCALE, yscale=YSCALE, xlim=xlim)
    for ii in range(3):
        ax.hist(xx[:, ii], bins=edges, color=colors[ii], rwidth=WID, alpha=0.5, align=align[ii])

    # Marginalize x
    if YSCALE.startswith('log'):
        edges = np.logspace(*np.log10(ylim), nbins)
    else:
        edges = np.linspace(*ylim, nbins)

    ax = axes[1, 1]
    ax.set(xscale=XSCALE, yscale=YSCALE, ylim=ylim)
    for ii in range(3):
        ax.hist(yy[:, ii], bins=edges, color=colors[ii], rwidth=WID, alpha=0.5,
                align=align[ii], orientation='horizontal')

    text = "Align: {}".format(align) + "\n" + "Matplotlib version: {}".format(mpl.__version__)
    axes[0, 0].set_title(text)
    

plot_scatter_hist(axes[:, :2], False)
plot_scatter_hist(axes[:, 2:], True)
    
fig.savefig('scatter-hist-align_x-{}_y-{}_mpl-{}.png'.format(XSCALE, YSCALE, mpl.__version__))