from IPython.core.pylabtools import print_figure
from IPython.html import widgets
from IPython.display import Image, display
from fatiando.gravmag import talwani
from fatiando.gravmag.basin2d import PolygonalBasinGravity
from fatiando.gravmag.interactive import Moulder
from fatiando.inversion.regularization import Damping, Smoothness1D, Smoothness2D
from fatiando.vis import mpl
from fatiando import utils
import fatiando
import numpy as np

print('Versão do Fatiando a Terra: {}'.format(fatiando.__version__))

area = [0, 100e3, 0, 10e3]
x = np.linspace(area[0], area[1], 80)
z = -np.ones_like(x)

meu_modelo = Moulder(area, x, z, density_range=[-1000, 0])

meu_modelo.run()

meu_modelo.save('modelo_bacia')

modelo = Moulder.load('modelo_bacia')
modelo.plot(figsize=(12, 6))

dado = modelo.predicted
bacia = modelo.model[0]
densidade = bacia.props.copy()

pontos = 40
topo = bacia.y.min()
p0 = 1000*np.ones(pontos)
xlim = [bacia.x.min(), bacia.x.max()]
misfit = PolygonalBasinGravity(x, z, dado, pontos, bacia.props, top=topo, xlim=xlim).config('levmarq', initial=p0)

def plota_resultado(solver):
    fig = mpl.figure(figsize=(12, 4), facecolor='white')
    ax = mpl.subplot(2, 1, 1)
    mpl.plot(x, solver.data, 'ok')
    mpl.plot(x, solver.predicted(), '-r', linewidth=2)
    ax.set_ylim(1.2*dado.min(), 1.1*dado.max())
    ax.grid(True)
    ax.set_xticklabels([])
    ax.set_ylabel('Anomalia (mGal)')
    ax = mpl.subplot(2, 1, 2)
    mpl.polygon(bacia, style='o-k', fill='gray', alpha=0.5, linewidth=2)
    mpl.polygon(solver.estimate_, style='o-r', linewidth=2)
    ax.set_xlabel('x (km)')
    ax.set_ylabel('z (km)')
    mpl.set_area(area)
    mpl.m2km()
    ax.grid(True)
    ax.invert_yaxis()
    mpl.tight_layout(pad=0)
    return fig

def interativo_sem_vinculo(erro, seed):
    misfit.data = utils.contaminate(dado, erro, seed=seed)
    misfit.fit()
    fig = plota_resultado(misfit)
    mpl.close(fig)    
    display(Image(print_figure(fig, dpi=80)))
    misfit.data = dado
    
widgets.interact(interativo_sem_vinculo, 
                 erro=widgets.FloatSliderWidget(min=0, max=1, step=0.1, value=0), 
                 seed=widgets.IntSliderWidget(min=0, max=10, step=1, value=0))

def interativo_damping(erro, seed, mi):
    misfit.data = utils.contaminate(dado, erro, seed=seed)
    solver = misfit + 10**mi*Damping(misfit.nparams)
    solver.fit()
    fig = plota_resultado(solver)
    mpl.close(fig)    
    display(Image(print_figure(fig, dpi=80)))
    misfit.data = dado
    
widgets.interact(interativo_damping, 
                 erro=widgets.FloatSliderWidget(min=0, max=1, step=0.1, value=0), 
                 seed=widgets.IntSliderWidget(min=0, max=10, step=1, value=0),
                 mi=widgets.FloatSliderWidget(min=-15, max=0, step=0.5, value=-15))

def interativo_suave(erro, seed, mu):
    misfit.data = utils.contaminate(dado, erro, seed=seed)
    solver = misfit + 10**mu*Smoothness1D(misfit.nparams)
    solver.fit()
    fig = plota_resultado(solver)
    mpl.close(fig)    
    display(Image(print_figure(fig, dpi=80)))
    misfit.data = dado

widgets.interact(interativo_suave, 
                 erro=widgets.FloatSliderWidget(min=0, max=5, step=0.1, value=0), 
                 seed=widgets.IntSliderWidget(min=0, max=10, step=1, value=0),
                 mu=widgets.FloatSliderWidget(min=-15, max=0, step=0.5, value=-15))

def interativo_densidade(seed, mu, densidade):
    misfit.data = utils.contaminate(dado, 0.5, seed=seed)
    misfit.model['props']['density'] = densidade
    solver = misfit + 10**mu*Smoothness1D(misfit.nparams)
    solver.fit()
    fig = plota_resultado(solver)
    mpl.close(fig)    
    display(Image(print_figure(fig, dpi=80)))
    misfit.data = dado
    misfit.model['props']['density'] = bacia.props['density']

widgets.interact(interativo_densidade, 
                 seed=widgets.IntSliderWidget(min=0, max=10, step=1, value=0),
                 mu=widgets.FloatSliderWidget(min=-10, max=0, step=0.5, value=-10),
                 densidade=widgets.FloatSliderWidget(min=-1000, max=-50, step=20, value=-50))

print('Densidade correta = {} kg/m3'.format(densidade['density']))