import numpy as np
from fatiando.gravmag import talwani
from fatiando.mesher import Polygon
from fatiando.inversion.base import Misfit

class PolyBasin2D(Misfit):
    def __init__(self, x, z, gz, height, npoints, density):
        super(PolyBasin2D, self).__init__(
            data=gz, nparams=npoints, islinear=False,
            positional=dict(x=x, z=z), 
            model=dict(height=height, 
                       props={'density': density}))
    
    def p2vertices(self, p):
        x, h = self.positional['x'], self.model['height']
        n = self.nparams + 2
        verts = np.empty((n, 2))
        verts[:, 0] = np.linspace(x.min(), x.max(), n)[::-1]
        verts[:, 1] = np.concatenate([[h], p, [h]])
        return verts
    
    def _get_predicted(self, p):
        x, z = self.positional['x'], self.positional['z']
        verts = self.p2vertices(p)
        poly = Polygon(verts, self.model['props'])
        return talwani.gz(x, z, [poly])
        
    def _get_jacobian(self, p):    
        x, z = self.positional['x'], self.positional['z']
        verts = self.p2vertices(p)
        delta = np.array([0, 10])
        jac = np.zeros((self.ndata, self.nparams))
        for i in xrange(self.nparams):
            diff = Polygon([verts[i + 2], verts[i + 1] - delta, 
                            verts[i], verts[i + 1] + delta],
                           self.model['props'])
            jac[:, i] = talwani.gz(x, z, [diff])/(2 * delta[1])
        return jac
    
    def fit(self):
        super(PolyBasin2D, self).fit()
        self._estimate = Polygon(self.p2vertices(self.p_), 
                                 self.model['props'])
        return self

def confidence(obj, n=10):
    "Get a 99% confidence polygon from a Misfit object that has an estimate."
    dof = obj.ndata - obj.nparams
    # Estimate the variance factor
    variance = np.linalg.norm(obj.residuals(obj.p_))**2/dof
    noisefree = obj.predicted()
    backup = obj.data
    # Contaminate the data with Gaussian random noise
    # run the inverseion and store the solutions
    solutions = np.empty((n, obj.nparams))
    for i in range(n):
        obj.data = utils.contaminate(noisefree, np.sqrt(variance), seed=i)
        solutions[i, :] = np.copy(obj.fit().p_)
    obj.data = backup
    # Estimate the covariance matrix from the solutions
    A = solutions - np.mean(solutions, axis=0)
    cov = np.dot(A.T, A)/n
    # Make a 3sigma polygon
    std = np.sqrt(np.diagonal(cov))
    verts = np.concatenate([obj.p2vertices(obj.p_ - 3*std),
                            obj.p2vertices(obj.p_ + 3*std)[1:-1][::-1]])
    return Polygon(verts)

%matplotlib inline

from fatiando.vis import mpl
from fatiando import utils

x_sketch = np.linspace(-10000, 110000, 100)
verts_sketch = np.empty((100, 2))
verts_sketch[:, 0] = np.linspace(0, 100000, 100)[::-1]
verts_sketch[:, 1] = -1e-15*(verts_sketch[:, 0] - 50000)**4 + 8000
verts_sketch[:, 1] -= verts_sketch[:, 1].min()
sketch = Polygon(verts_sketch, {'density':-500})

fontsize = 50
labelsize = 28
linewidth = 6
tickpad = 15
markersize = 15
savefigrc = {'transparent': True, 'dpi': 600}

with mpl.xkcd():
    fig = mpl.figure(figsize=(13, 8))
    ax = fig.add_subplot(2, 1, 1)
    mpl.plot(x_sketch, talwani.gz(x_sketch, -100*np.ones_like(x_sketch), [sketch]), '-k', linewidth=linewidth)
    mpl.title('Gravity anomaly', fontsize=fontsize)
    mpl.xlabel("Earth surface", fontsize=fontsize) 
    mpl.ylim(-130, 0)
    mpl.xlim(x_sketch.min(), x_sketch.max())
    ax.set_yticklabels(['-%.0f' % (abs(i)) for i in ax.get_yticks()])
    ax.set_xticklabels([''])
    ax.set_xticks([])
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.spines['bottom'].set_visible(False)
    ax.spines['left'].set_linewidth(linewidth)
    ax.yaxis.set_ticks_position('left')
    ax.tick_params(labelsize=labelsize, pad=tickpad)
    
    ax = fig.add_subplot(2, 1, 2)
    mpl.polygon(sketch, fill='grey', alpha=0.4, linewidth=linewidth)
    mpl.text(50000, 4500, 'Sedimentary rocks\n<low density>', fontsize=fontsize, horizontalalignment='center')
    mpl.text(50000, 9000, 'Basement rocks <high density>', fontsize=fontsize, horizontalalignment='center')
    mpl.ylim(0, 10000)
    mpl.xlim(x_sketch.min(), x_sketch.max())
    mpl.m2km()
    ax.invert_yaxis()
    ax.spines['bottom'].set_visible(False)  
    ax.spines['right'].set_visible(False)    
    ax.spines['left'].set_linewidth(linewidth)
    ax.spines['top'].set_linewidth(linewidth)
    ax.yaxis.set_ticks_position('left')
    ax.set_xticks([])
    ax.tick_params(labelsize=labelsize, pad=tickpad)
    mpl.ylabel('depth', fontsize=labelsize)
    
    mpl.tight_layout()
    fig.savefig('sketch.png', **savefigrc)

params_sketch = np.concatenate([verts_sketch[::20], [verts_sketch[-1]]]) + 100

with mpl.xkcd():
    mpl.figure(figsize=(13, 4))
    mpl.polygon(sketch, fill='grey', alpha=0.4, linewidth=linewidth)
    lines = mpl.polygon(Polygon(params_sketch), 'o-r', linewidth=1.5*linewidth)
    lines.set_markersize(1.5*markersize)
    for i, v in enumerate(params_sketch[1:-1]):
        xc, zc = v
        mpl.arrow(xc, 0, 0, zc - 300, color='k', width=100*linewidth, head_length=500, head_width=2000, 
                  length_includes_head=True, zorder=10)
        mpl.text(xc + 1000, 2500, r'$p_%d$' % (i + 1), fontsize=1.5*fontsize)
    mpl.ylim(0, 7000)
    mpl.xlim(x_sketch.min(), x_sketch.max())
    mpl.m2km()
    ax = mpl.gca()
    ax.invert_yaxis()
    ax.spines['bottom'].set_visible(False)  
    ax.spines['right'].set_visible(False)
    ax.spines['left'].set_visible(False)
    ax.spines['top'].set_linewidth(linewidth)    
    ax.spines['top'].zorder = 0
    ax.yaxis.set_ticks_position('left')
    ax.set_xticks([])
    ax.set_yticks([])
    ax.tick_params(labelsize=fontsize, pad=tickpad)
    mpl.tight_layout()
    mpl.savefig('parametrization.png', **savefigrc)

u = np.zeros_like(params_sketch)
i = 3
u[i, 1] = 2000
size = (8, 2.5)
def boilerplate():
    mpl.ylim(0, 10000)
    mpl.xlim(x_sketch.min(), x_sketch.max())
    mpl.m2km()
    ax = mpl.gca()
    ax.invert_yaxis()
    ax.spines['bottom'].set_visible(False)  
    ax.spines['right'].set_visible(False)
    ax.spines['left'].set_visible(False)
    ax.spines['left'].set_linewidth(linewidth)
    ax.spines['top'].set_linewidth(linewidth)
    ax.spines['top'].zorder = 0
    ax.yaxis.set_ticks_position('left')
    ax.set_xticks([])
    ax.set_yticks([])
    ax.tick_params(labelsize=fontsize, pad=tickpad)
    mpl.tight_layout()

with mpl.xkcd():
    mpl.figure(figsize=size)
    lines = mpl.polygon(Polygon(params_sketch), 'o-k', linewidth=linewidth)
    lines.set_markersize(markersize)
    lines = mpl.polygon(Polygon(params_sketch + u), '.-g', fill='green', alpha=0.5, linewidth=linewidth)
    lines.set_markersize(markersize)
    mpl.text(50000, 2500,  r'$\mathbf{p} + \mathbf{u}^j \Delta$', 
             horizontalalignment='center', verticalalignment='center', fontsize=0.75*fontsize)
    boilerplate()
    mpl.savefig('jacobian-1.png', **savefigrc)
    
    mpl.figure(figsize=size)
    lines = mpl.polygon(Polygon(params_sketch), 'o-k', linewidth=linewidth)
    lines.set_markersize(markersize)
    lines = mpl.polygon(Polygon(params_sketch - u), '.-b', fill='blue', alpha=0.5, linewidth=linewidth)
    lines.set_markersize(markersize)
    mpl.text(50000, 2000,  r'$\mathbf{p}\, - \mathbf{u}^j \Delta$', 
             horizontalalignment='center', verticalalignment='center', fontsize=0.75*fontsize)
    boilerplate()
    mpl.savefig('jacobian-2.png', **savefigrc)
                
    mpl.figure(figsize=size)
    lines = mpl.polygon(Polygon(params_sketch), 'o-k', linewidth=linewidth)
    lines.set_markersize(markersize)
    lines = mpl.polygon(Polygon(params_sketch + u), '.-g', fill='green', alpha=0.5, linewidth=linewidth)
    lines.set_markersize(markersize)
    lines = mpl.polygon(Polygon(params_sketch - u), '.-b', fill='blue', alpha=0.5, linewidth=linewidth)
    lines.set_markersize(markersize)
    boilerplate()
    mpl.savefig('jacobian-3.png', **savefigrc)
        
    mpl.figure(figsize=size)
    lines = mpl.polygon(Polygon(params_sketch), 'o-k', linewidth=linewidth)
    lines.set_markersize(markersize)
    lines = mpl.polygon(Polygon([params_sketch[i + 1], params_sketch[i] - u[i], 
                                 params_sketch[i - 1], params_sketch[i] + u[i]]), 
                        '.-r', fill='red', alpha=0.5, linewidth=linewidth)
    lines.set_markersize(markersize)    
    boilerplate()
    mpl.savefig('jacobian-4.png', **savefigrc)

x = np.linspace(0, 100000, 100)
z = -100*np.ones_like(x)
xs = np.linspace(0, 100000, 100)[::-1]
depths = (-1e-15*(xs - 50000)**4 + 8000 
          - 3000*np.exp(-(xs - 70000)**2/(10000**2)))
depths -= depths.min()
density = -500
model = Polygon(np.transpose([xs, depths]), {'density':density})
gz = utils.contaminate(talwani.gz(x, z, [model]), 0.5, seed=10000)

# Makes a figure of the model and data in a way that fits nicely on the poster
def plot_model():
    mpl.figure(figsize=(14, 6))
    mpl.subplot(2, 1, 1)
    mpl.ylabel('mGal', fontsize=fontsize)
    mpl.plot(x*0.001, gz, 'ok', markersize=markersize, linewidth=linewidth)
    mpl.ylabel('GA (mGal)', fontsize=labelsize)
    mpl.ylim(-130, 0)
    mpl.xlabel('x (km)', fontsize=labelsize)
    ax = mpl.gca()
    ax.xaxis.labelpad = -30
    ax.spines['left'].set_linewidth(linewidth)
    ax.spines['right'].set_visible(False)
    ax.spines['top'].set_visible(False)
    ax.spines['bottom'].set_linewidth(linewidth)
    ax.yaxis.set_ticks_position('left')
    ax.xaxis.set_ticks_position('bottom')
    ax.set_yticks(ax.get_yticks()[::2])
    ax.tick_params(labelsize=labelsize, pad=tickpad, width=linewidth, length=2*linewidth)
    
    mpl.subplot(2, 1, 2)
    lines = mpl.polygon(model, fill='grey', alpha=0.5, linewidth=linewidth)
    lines.set_markersize(markersize)
    mpl.ylim(0, 12000)
    mpl.m2km()
    ax = mpl.gca()
    ax.invert_yaxis()
    ax.set_xticklabels([''])
    mpl.ylabel('Depth (km)', fontsize=labelsize)
    ax.spines['left'].set_linewidth(linewidth)
    ax.spines['right'].set_visible(False)
    ax.spines['top'].set_linewidth(linewidth)
    ax.spines['bottom'].set_visible(False)
    ax.yaxis.set_ticks_position('left')
    ax.xaxis.set_ticks_position('top')
    ax.tick_params(labelsize=labelsize, pad=tickpad, width=linewidth, length=2*linewidth)
    
    mpl.tight_layout()
    mpl.subplots_adjust(hspace=0.28)

plot_model()
mpl.savefig('synthetic.png', **savefigrc)

misfit = PolyBasin2D(x, z, gz, 0, 30, density)
initial = 3000*np.ones(misfit.nparams)
misfit.config('levmarq', initial=initial).fit()

misfit_conf = confidence(misfit, n=50)

def plot_estimate(obj, conf):
    plot_model()
    mpl.subplot(2, 1, 1)
    mpl.plot(x*0.001, obj.predicted(), '-r', linewidth=linewidth)
    mpl.subplot(2, 1, 2)
    mpl.polygon(conf, '-b', linealpha=0, fill='blue', alpha=0.8, linewidth=linewidth)
    lines = mpl.polygon(obj.estimate_, 'o-r', linewidth=linewidth)
    lines.set_markersize(markersize)

plot_estimate(misfit, misfit_conf)
mpl.text(70000, 11000, r'99% confidence', fontsize=labelsize, color='b')
mpl.arrow(69000, 10500, -11000, 0, width=500, head_width=1000, head_length=2000, 
          length_includes_head=True, color='k')
mpl.savefig('estimate.png', **savefigrc)

from fatiando.inversion.regularization import Smoothness1D
solver = misfit + 1e-8*Smoothness1D(misfit.nparams)
solver.config('levmarq', initial=initial).fit()

solver_conf = confidence(solver, n=50)

plot_estimate(solver, solver_conf)
mpl.text(60000, 11000, "99% confidence\n(Believe me, it's there)", fontsize=labelsize, color='b')
mpl.arrow(59000, 9500, -10000, -3000, width=500, head_width=1000, head_length=2000, 
          length_includes_head=True, color='k')
mpl.savefig('estimate-regularized.png', **savefigrc)

from fatiando.inversion.regularization import LCurve
lc = LCurve(misfit, Smoothness1D(misfit.nparams),
            [10**i for i in xrange(-15, -5)]) # List of regularization parameters to try
lc.config('levmarq', initial=initial).fit()

print(lc.regul_param_)
lc.plot_lcurve()

# This time I'll have to cheat and make a nicer plot of the L-curve by hand.
mpl.figure(figsize=(8, 6))
mpl.loglog(lc.dnorm, lc.mnorm, 'o-k', linewidth=linewidth, markersize=markersize)
mpl.plot(lc.dnorm[lc.corner_], lc.mnorm[lc.corner_], '*r', markersize=2*markersize)
mpl.ylabel('Smoothness', fontsize=labelsize)
mpl.xlabel('Misfit', fontsize=labelsize)
mpl.text(0.6, 1e8, r'$\mu = 10^{-8}$', fontsize=fontsize)
mpl.set_area([1e-1, 2, 2e6, 1e9])

ax = mpl.gca()            
ax.spines['right'].set_visible(False)
ax.spines['top'].set_visible(False)
ax.spines['bottom'].set_linewidth(linewidth)
ax.spines['left'].set_linewidth(linewidth)
ax.spines['right'].set_linewidth(linewidth)
ax.spines['top'].set_linewidth(linewidth)
ax.yaxis.set_ticks_position('left')
ax.xaxis.set_ticks_position('bottom')
ax.tick_params(labelsize=labelsize, pad=tickpad, width=linewidth, length=2*linewidth)
ax.xaxis.labelpad = -10
mpl.tight_layout()
mpl.savefig('lcurve.png', **savefigrc)

lc_conf = confidence(lc.objectives[lc.corner_], n=50)

plot_estimate(lc.objectives[lc.corner_], lc_conf)