"""
2D Posterior analysis of Resolve inference
------------------------------------------

All plotting in GeoBIPy can be carried out using the 3D inference class

"""

import matplotlib.pyplot as plt
import numpy as np
from geobipy import Model
from geobipy import Inference2D

def plot_2d_summary(folder, data_type, model_type):
   #%%
   # Inference for a line of inferences
   # ++++++++++++++++++++++++++++++++++
   #
   # We can instantiate the inference handler by providing a path to the directory containing
   # HDF5 files generated by GeoBIPy.
   #
   # The InfereceXD classes are low memory.  They only read information from the HDF5 files
   # as and when it is needed.
   #
   # The first time you use these classes to create plots, expect longer initial processing times.
   # I precompute expensive properties and store them in the HDF5 files for later use.

   from numpy.random import Generator
   from numpy.random import PCG64DXSM
   generator = PCG64DXSM(seed=0)
   prng = Generator(generator)

   #%%
   results_2d = Inference2D.fromHdf('{}/{}/{}/0.0.h5'.format(folder, data_type, model_type), prng=prng)

   kwargs = {
         "log" : 10,
         "cmap" : 'jet'
         }

   true_model = Model.create_synthetic_model(model_type)
   true_model.mesh.y_edges = true_model.mesh.y_edges / 10.0

   kwargs['vmin'] = np.log10(np.min(true_model.values))
   kwargs['vmax'] = np.log10(np.max(true_model.values))

   t = results_2d.inference_1d(index=7)

   plt.figure()
   t.model.plot_posteriors(edges_kwargs={
                           'transpose': True,
                           'trim': False},
                           values_kwargs={
                           'colorbar': False,
                           'flipY': True,
                           'xscale': 'log',
                           'credible_interval_kwargs': {
                           }})

   tmp = t.model.values.posterior


   plt.show(block=True)
   # plt.savefig('{}_{}.png'.format(data_type, model_type), dpi=300)

if __name__ == '__main__':
   models = ['glacial']#, 'saline_clay', 'resistive_dolomites', 'resistive_basement', 'coastal_salt_water', 'ice_over_salt_water']

   for model in models:
      # try:
         plot_2d_summary("../../../Parallel_Inference/", "resolve", model)
      # except Exception as e:
      #    print(model)
      #    print(e)
      #    pass
