%matplotlib inline
import numpy
from fatiando.mesher import SquareMesh
from fatiando.seismic import ttime2d, srtomo
from fatiando.inversion.regularization import Smoothness2D
from fatiando.vis import mpl
from fatiando import utils

area = (0, 500000, 0, 500000)
shape = (30, 30)
model = SquareMesh(area, shape)
model.img2prop('logo.png', 4000, 10000, 'vp')
seed = 0 # Set the random seed so that points are the same everythime
src_loc = utils.random_points(area, 80, seed=seed)
rec_loc = utils.circular_points(area, 30, random=True, seed=seed)
srcs, recs = utils.connect_points(src_loc, rec_loc)
tts = ttime2d.straight(model, 'vp', srcs, recs)
tts, error = utils.contaminate(tts, 0.01, percent=True, return_stddev=True)

mesh = SquareMesh(area, shape)
datamisfit = srtomo.SRTomo(tts, srcs, recs, mesh)
regul = Smoothness2D(mesh.shape)

class LCurve(object):
    def __init__(self, datamisfit, regul, regul_params):
        self.regul_params = regul_params
        self.datamisfit = datamisfit
        self.regul = regul
        self.regul_param = None
        self.objectives = None
        self.corner = None
        self.estimate_ = None
        self.p_ = None
        self.dnorm = numpy.empty(len(regul_params))
        self.mnorm = numpy.empty(len(regul_params))
        
    def fit(self):
        self.objectives = []
        for i, mu in enumerate(self.regul_params):
            solver = (self.datamisfit + mu*self.regul).fit()            
            p = solver.p_
            self.dnorm[i] = self.datamisfit.value(p)
            self.mnorm[i] = self.regul.value(p)
            self.objectives.append(solver)
        corner = self.select_corner(self.dnorm, self.mnorm)
        self.corner = corner
        self.regul_param = self.regul_params[corner]
        self.p_ = self.objectives[corner].p_
        self.estimate_ = self.objectives[corner].estimate_
        return self
    
    def select_corner(self, dnorm, mnorm):
        # Uses http://www.sciencedirect.com/science/article/pii/S0168927401001799
        x, y = numpy.log(dnorm), numpy.log(mnorm)
        n = len(dnorm)
        corner = n - 1
        dist = lambda p1, p2: numpy.sqrt((p1[0] - p2[0])**2 + (p1[1] - p2[1])**2)
        cte = 7.*numpy.pi/8.
        angmin = None
        c = [x[-1], y[-1]]
        for k in xrange(0, n - 2):
            b = [x[k], y[k]]
            for j in xrange(k + 1, n - 1):
                a = [x[j], y[j]]
                ab = dist(a, b)
                ac = dist(a, c)
                bc = dist(b, c)
                cosa = (ab**2 + ac**2 - bc**2)/(2.*ab*ac)
                ang = numpy.arccos(cosa)
                area = 0.5*((b[0] - a[0])*(a[1] - c[1]) - (a[0] - c[0])*(b[1] - a[1]))
                # area is > 0 because in the paper C is index 0
                if area > 0 and (ang < cte and (angmin is None or ang < angmin)):
                    corner = j
                    angmin = ang
        return corner        
    
    def predicted(self, p=None):
        return self.objectives[self.corner].predicted(p)
        
    def residuals(self, p=None):
        return self.objectives[self.corner].residuals(p)
        
    def plot_lcurve(self):
        mpl.title('Corner regularization parameter: %g' % (self.regul_param))        
        mpl.loglog(self.dnorm, self.mnorm, '.-k')
        ax = mpl.gca()
        vmin, vmax = ax.get_ybound()     
        mpl.vlines(self.dnorm[self.corner], vmin, vmax)
        vmin, vmax = ax.get_xbound()     
        mpl.hlines(self.mnorm[self.corner], vmin, vmax)
        mpl.plot(self.dnorm[self.corner], self.mnorm[self.corner], '^b', markersize=10)
        mpl.xlabel('Data misfit')
        mpl.ylabel('Regularization')

tomo = LCurve(datamisfit, regul, [10**i for i in range(-10, 10, 1)]).fit()
tomo.plot_lcurve()
mpl.grid(True)

estimate = tomo.estimate_
residuals = tomo.residuals()
mesh.addprop('vp', estimate)

mpl.figure(figsize=(14, 5))
mpl.subplot(1, 2, 1)
mpl.axis('scaled')
mpl.title('Vp synthetic model of the Earth')
mpl.squaremesh(model, prop='vp', cmap=mpl.cm.seismic)
cb = mpl.colorbar()
cb.set_label('Velocity')
mpl.points(src_loc, '*y', label="Sources")
mpl.points(rec_loc, '^r', label="Receivers")
mpl.legend(loc='lower left', shadow=True, numpoints=1, prop={'size':10})
mpl.m2km()
mpl.subplot(1, 2, 2)
mpl.axis('scaled')
mpl.title('Tomography result (damped)')
mpl.squaremesh(mesh, prop='vp', vmin=4000, vmax=10000,
    cmap=mpl.cm.seismic)
cb = mpl.colorbar()
cb.set_label('Velocity')
mpl.m2km()
mpl.figure()
mpl.grid()
mpl.title('Residuals (data with %.4f s error)' % (error))
mpl.hist(residuals, color='gray', bins=10)
mpl.xlabel("seconds")