%matplotlib inline

# here the usual imports. If any of the imports fails, 
# make sure that pynoddy is installed
# properly, ideally with 'python setup.py develop' 
# or 'python setup.py install'
import sys, os
import matplotlib.pyplot as plt
import numpy as np
# adjust some settings for matplotlib
from matplotlib import rcParams
# print rcParams
rcParams['font.size'] = 15
# determine path of repository to set paths corretly below
repo_path = os.path.realpath('../..')
import pynoddy.history
import pynoddy.experiment
reload(pynoddy.experiment)
rcParams.update({'font.size': 15})

# From notebook 4/ Traning Set example 1:
reload(pynoddy.history)
reload(pynoddy.events)
nm = pynoddy.history.NoddyHistory()
# add stratigraphy
strati_options = {'num_layers' : 3,
                  'layer_names' : ['layer 1', 'layer 2', 'layer 3'],
                  'layer_thickness' : [1500, 500, 1500]}
nm.add_event('stratigraphy', strati_options )

# The following options define the fault geometry:
fault_options = {'name' : 'Fault_E',
                 'pos' : (4000, 0, 5000),
                 'dip_dir' : 90.,
                 'dip' : 60,
                 'slip' : 1000}

nm.add_event('fault', fault_options)
history = 'normal_fault.his'
output_name = 'normal_fault_out'
nm.write_history(history)

"""
Explanation: Generate Training sets
Based on "Reproducible Experiments" notebook
End of explanation
"""


reload(pynoddy.history)
reload(pynoddy.experiment)

from pynoddy.experiment import monte_carlo
ue = pynoddy.experiment.Experiment(history)

ue.change_cube_size(100)
ue.plot_section('y')

"""
Explanation: Initiate experiment with this input file:
End of explanation
"""


ue.freeze()

"""
Explanation: Before we start to draw random realisations of the model, we should first store the base state of the model for later reference. This is simply possibel with the freeze() method which stores the current state of the model as the "base-state":
End of explanation
"""


ue.set_random_seed(12345)

"""
Explanation: We now intialise the random generator. We can directly assign a random seed to simplify reproducibility (note that this is not essential, as it would be for the definition in a script function: the random state is preserved within the model and could be retrieved at a later stage, as well!):
End of explanation
"""


ue.info(events_only = True)

ev2 = ue.events[2]

ev2.properties

"""
Explanation: The next step is to define probability distributions to the relevant event parameters. Let's first look at the different events:
End of explanation
"""


param_stats = [{'event' : 2, 
              'parameter': 'Slip',
              'stdev': 300.0,
              'type': 'normal'},
              {'event' : 2, 
              'parameter': 'Dip',
              'stdev': 10.0,
              'type': 'normal'},]

ue.set_parameter_statistics(param_stats)

resolution = 100
ue.change_cube_size(resolution)
tmp = ue.get_section('y')
prob_2 = np.zeros_like(tmp.block[:,:,:])
n_draws = 100


for i in range(n_draws):
    ue.random_draw()
    tmp = ue.get_section('y', resolution = resolution)
    prob_2 += (tmp.block[:,:,:] == 2)

# Normalise
prob_2 = prob_2 / float(n_draws)

fig = plt.figure(figsize = (12,8))
ax = fig.add_subplot(111)
ax.imshow(prob_2.transpose()[:,0,:], 
           origin = 'lower left',
           interpolation = 'none')
plt.title("Estimated probability of unit 4")
plt.xlabel("x (E-W)")
plt.ylabel("z")

"""
Explanation: Next, we define the probability distributions for the uncertain input parameters:
End of explanation
"""


ue.random_draw()
s1 = ue.get_section('y')
s1.block.shape
s1.block[np.where(s1.block == 3)] = 1
s1.plot_section('y', cmap='Greys')

"""
Explanation: This example shows how the base module for reproducible experiments with kinematics can be used. For further specification, child classes of Experiment can be defined, and we show examples of this type of extension in the next sections.
Adjustments to generate training set
First step: generate more layers and randomly select layers to visualise:
End of explanation
"""


nm = pynoddy.history.NoddyHistory()
# add stratigraphy

n_layers = 8

strati_options['num_layers'] = n_layers
strati_options['layer_names'] = []
strati_options['layer_thickness'] = []

for n in range(n_layers):

    strati_options['layer_names'].append("layer %d" % n)
    strati_options['layer_thickness'].append(5000./n_layers)

nm.add_event('stratigraphy', strati_options )

# The following options define the fault geometry:
fault_options = {'name' : 'Fault_E',
                 'pos' : (1000, 0, 5000),
                 'dip_dir' : 90.,
                 'dip' : 60,
                 'slip' : 500}

nm.add_event('fault', fault_options)
history = 'normal_fault.his'
output_name = 'normal_fault_out'
nm.write_history(history)

reload(pynoddy.history)
reload(pynoddy.experiment)

from pynoddy.experiment import monte_carlo
ue = pynoddy.experiment.Experiment(history)
ue.freeze()
ue.set_random_seed(12345)
ue.set_extent(2800, 100, 2800)

ue.change_cube_size(50)
ue.plot_section('y')

param_stats = [{'event' : 2, 
              'parameter': 'Slip',
              'stdev': 100.0,
              'type': 'lognormal'},
              {'event' : 2, 
              'parameter': 'Dip',
              'stdev': 10.0,
              'type': 'normal'},
#              {'event' : 2, 
#              'parameter': 'Y',
#              'stdev': 150.0,
#              'type': 'normal'},
              {'event' : 2, 
              'parameter': 'X',
              'stdev': 150.0,
              'type': 'normal'},]

ue.set_parameter_statistics(param_stats)

# randomly select layers:
ue.random_draw()

s1 = ue.get_section('y')

# create "feature" model:
f1 = s1.block.copy()

# randomly select layers:
f1 = np.squeeze(f1)
# n_featuers: number of "features" -> gray values in image
n_features = 5
vals = np.random.randint(0,255,size=n_features)
for n in range(n_layers):
    f1[f1 == n] = np.random.choice(vals)

f1.shape

plt.imshow(f1.T, origin='lower_left', cmap='Greys', interpolation='nearest')

# blur image
from scipy import ndimage
f2 = ndimage.filters.gaussian_filter(f1, 1, mode='nearest')


plt.imshow(f2.T, origin='lower_left', cmap='Greys', interpolation='nearest', vmin=0, vmax=255)

# randomly swap image
if np.random.randint(2) == 1:
    f2 = f2[::-1,:]

plt.imshow(f2.T, origin='lower_left', cmap='Greys', interpolation='nearest', vmin=0, vmax=255)

"""
Explanation: Idea: generate many layers, then randomly extract a couple of these and also assign different density/ color values:
End of explanation
"""


# back to before: re-initialise model:
nm = pynoddy.history.NoddyHistory()
# add stratigraphy

n_layers = 18

strati_options['num_layers'] = n_layers
strati_options['layer_names'] = []
strati_options['layer_thickness'] = []

for n in range(n_layers):

    strati_options['layer_names'].append("layer %d" % n)
    strati_options['layer_thickness'].append(5000./n_layers)

nm.add_event('stratigraphy', strati_options )

# The following options define the fault geometry:
fault_options = {'name' : 'Fault_E',
                 'pos' : (1000, 0, 5000),
                 'dip_dir' : 90.,
                 'dip' : 60,
                 'slip' : 500}

nm.add_event('fault', fault_options)
history = 'normal_fault.his'
output_name = 'normal_fault_out'
nm.write_history(history)

reload(pynoddy.history)
reload(pynoddy.experiment)

from pynoddy.experiment import monte_carlo
ue = pynoddy.experiment.Experiment(history)
ue.freeze()
ue.set_random_seed(12345)
ue.set_extent(2800, 100, 2800)
ue.change_cube_size(50)

param_stats = [{'event' : 2, 
              'parameter': 'Slip',
              'stdev': 100.0,
              'type': 'lognormal'},
              {'event' : 2, 
              'parameter': 'Dip',
              'stdev': 10.0,
              'type': 'normal'},
#              {'event' : 2, 
#              'parameter': 'Y',
#              'stdev': 150.0,
#              'type': 'normal'},
              {'event' : 2, 
              'parameter': 'X',
              'stdev': 150.0,
              'type': 'normal'},]

ue.set_parameter_statistics(param_stats)

"""
Explanation: All in one function
Generate images for normal faults
End of explanation
"""


n_train = 10000
F_train = np.empty((n_train, 28*28))

ue.change_cube_size(100)

for i in range(n_train):
    # randomly select layers:
    ue.random_draw()
    s1 = ue.get_section('y')
    # create "feature" model:
    f1 = s1.block.copy()
    # randomly select layers:
    f1 = np.squeeze(f1)
    # n_featuers: number of "features" -> gray values in image
    n_features = 4
    vals = np.random.randint(0,255,size=n_features)
    for n in range(n_layers):
        f1[f1 == n+1] = np.random.choice(vals)
    f1 = f1.T
    f2 = ndimage.filters.gaussian_filter(f1, 0, mode='nearest')
    # scale image
    f2 = f2 - np.min(f2)
    if np.max(f2) != 0:
        f2 = f2/np.max(f2)*255
    # randomly swap image
    if np.random.randint(2) == 1:
        f2 = f2[::-1,:]
    F_train[i] = f2.flatten().T


plt.imshow(f2, origin='lower_left', cmap='Greys', interpolation='nearest', vmin=0, vmax=255)

import matplotlib.pyplot as plt
%matplotlib inline

fig, ax = plt.subplots(nrows=2, ncols=5, sharex=True, sharey=True, figsize=(12,6))
ax = ax.flatten()
for i in range(10):
    img = F_train[i].reshape(28, 28)
    ax[i].imshow(img, cmap='Greys', interpolation='nearest')

ax[0].set_xticks([])
ax[0].set_yticks([])
plt.tight_layout()
# plt.savefig('./figures/mnist_all.png', dpi=300)
plt.show()

import pickle


pickle.dump(F_train, open("f_train_normal.pkl", 'w'))

"""
Explanation: Generate training set for normal faults:
End of explanation
"""


# back to before: re-initialise model:
nm = pynoddy.history.NoddyHistory()
# add stratigraphy

n_layers = 18

strati_options['num_layers'] = n_layers
strati_options['layer_names'] = []
strati_options['layer_thickness'] = []

for n in range(n_layers):

    strati_options['layer_names'].append("layer %d" % n)
    strati_options['layer_thickness'].append(5000./n_layers)

nm.add_event('stratigraphy', strati_options )

# The following options define the fault geometry:
fault_options = {'name' : 'Fault_E',
                 'pos' : (1000, 0, 5000),
                 'dip_dir' : 90.,
                 'dip' : 60,
                 'slip' : -500}

nm.add_event('fault', fault_options)
history = 'normal_fault.his'
output_name = 'normal_fault_out'
nm.write_history(history)

reload(pynoddy.history)
reload(pynoddy.experiment)

from pynoddy.experiment import monte_carlo
ue = pynoddy.experiment.Experiment(history)
ue.freeze()
ue.set_random_seed(12345)
ue.set_extent(2800, 100, 2800)
ue.change_cube_size(50)

param_stats = [{'event' : 2, 
              'parameter': 'Slip',
              'stdev': -100.0,
              'type': 'lognormal'},
              {'event' : 2, 
              'parameter': 'Dip',
              'stdev': 10.0,
              'type': 'normal'},
#              {'event' : 2, 
#              'parameter': 'Y',
#              'stdev': 150.0,
#              'type': 'normal'},
              {'event' : 2, 
              'parameter': 'X',
              'stdev': 150.0,
              'type': 'normal'},]

ue.set_parameter_statistics(param_stats)

n_train = 10000
F_train_rev = np.empty((n_train, 28*28))

ue.change_cube_size(100)

for i in range(n_train):
    # randomly select layers:
    ue.random_draw()
    s1 = ue.get_section('y')
    # create "feature" model:
    f1 = s1.block.copy()
    # randomly select layers:
    f1 = np.squeeze(f1)
    # n_featuers: number of "features" -> gray values in image
    n_features = 4
    vals = np.random.randint(0,255,size=n_features)
    for n in range(n_layers):
        f1[f1 == n+1] = np.random.choice(vals)
    f1 = f1.T
    f2 = ndimage.filters.gaussian_filter(f1, 0, mode='nearest')
    # scale image
    f2 = f2 - np.min(f2)
    if np.max(f2) != 0:
        f2 = f2/np.max(f2)*255
    # randomly swap image
    if np.random.randint(2) == 1:
        f2 = f2[::-1,:]
    F_train_rev[i] = f2.flatten().T



fig, ax = plt.subplots(nrows=2, ncols=5, sharex=True, sharey=True, figsize=(12,6))
ax = ax.flatten()
for i in range(10):
    img = F_train_rev[i].reshape(28, 28)
    ax[i].imshow(img, cmap='Greys', interpolation='nearest')

ax[0].set_xticks([])
ax[0].set_yticks([])
plt.tight_layout()
# plt.savefig('./figures/mnist_all.png', dpi=300)
plt.show()

pickle.dump(F_train_rev, open("f_train_reverse.pkl", 'w'))

"""
Explanation: Generate reverse faults
And now: the same for reverse faults:
End of explanation
"""


l1 = np.empty_like(s1.block[:,0,:])

n_layers = 18
for i in range(l1.shape[0]):
    l1[:,i] = i
l1_ori = np.floor(l1*n_layers/l1.shape[0])

F_train_line = np.empty((n_train, 28*28))



for i in range(n_train):
    n_features = 4
    vals = np.random.randint(0,255,size=n_features)
    l1 = l1_ori.copy()
    for n in range(n_layers):
        l1[l1 == n+1] = np.random.choice(vals)
    f1 = l1.T
    f2 = ndimage.filters.gaussian_filter(f1, 0, mode='nearest')
    # scale image
    f2 = f2 - np.min(f2)
    if np.max(f2) != 0:
        f2 = f2/np.max(f2)*255
    F_train_line[i] = f2.flatten().T

fig, ax = plt.subplots(nrows=2, ncols=5, sharex=True, sharey=True, figsize=(12,6))
ax = ax.flatten()
for i in range(10):
    img = F_train_line[i].reshape(28, 28)
    ax[i].imshow(img, cmap='Greys', interpolation='nearest')

ax[0].set_xticks([])
ax[0].set_yticks([])
plt.tight_layout()
# plt.savefig('./figures/mnist_all.png', dpi=300)
plt.show()

pickle.dump(F_train_line, open("f_train_line.pkl", 'w'))

"""
Explanation: Generate simple layer structure
No need for noddy, in this simple case - just adapt a numpy array:
End of explanation
"""


