import numpy as np
import math
import sys
import os
import IPython

"""
Explanation: CosmicFish PyLib: introduction and examples
This is a small iPython notebook to illustrate, with a set of examples, the use of the
CosmicFish Python Library.
Examples will be pretty basic and cover a small subset of the tools implemented in the library. To have a complete view of the structure of the library and all tools implemented you can have a look at the documentation.
Basic import and dependencies:
End of explanation
"""


cosmicfish_pylib_path = './python'
sys.path.insert(0, os.path.normpath(cosmicfish_pylib_path))

"""
Explanation: Then we get the path of the CosmicFish library and add it to the Python path:
End of explanation
"""


import cosmicfish_pylib.utilities            as fu
import cosmicfish_pylib.colors               as fc
import cosmicfish_pylib.fisher_matrix        as fm
import cosmicfish_pylib.fisher_derived       as fd
import cosmicfish_pylib.fisher_operations    as fo
import cosmicfish_pylib.fisher_plot_settings as fps
import cosmicfish_pylib.fisher_plot_analysis as fpa
import cosmicfish_pylib.fisher_plot          as fp

"""
Explanation: And now we can import the CosmicFish library:
End of explanation
"""


fu.CosmicFish_write_header(' Introduction and examples')

"""
Explanation: The first thing that we might want to do, when designing an application, is printing to screen the CosmicFish header:
End of explanation
"""


matrix = [ [ 1.0, 0.0 ],
           [ 0.0, 2.0 ] ]

"""
Explanation: Fisher matrices:
Let's start with the fisher_matrix class. This contains all the necessary tools to define and operate a Fisher matrix.
If you want the full details have a look at the fisher_matrix documentation.
We start by defining a matrix that we want to make a Fisher matrix:
End of explanation
"""


param_names = ['par1', 'par2']

"""
Explanation: Now we define the parameter names for our Fisher matrix.
End of explanation
"""


fisher = fm.fisher_matrix( fisher_matrix=matrix, param_names=param_names )

"""
Explanation: Now we have all the ingredients to define a simple Fisher matrix. We can do it by just calling the fisher_matrix class constructor. This will automatically make sure that the input matrix is well behaved and is really a Fisher matrix.
End of explanation
"""


print fisher.inverse_fisher_matrix()

"""
Explanation: Once we have initialized our Fisher matrix we have a set of useful operations defined on it. 
The inverse:
End of explanation
"""


print fisher.determinant()

"""
Explanation: The determinant:
End of explanation
"""


(eigenvalues,eigenvectors)=fisher.PCA()

print eigenvalues

print eigenvectors

"""
Explanation: Its eigenvalues and eigenvectors:
End of explanation
"""


print fisher.get_confidence_bounds( confidence_level=0.68 )
print fisher.get_confidence_bounds( confidence_level=0.95 )

"""
Explanation: And errors estimate based on the specified confidence level:
End of explanation
"""


fisher_2 = fisher
fisher_3 = fisher + fisher_2
print fisher_3.get_fisher_matrix()

"""
Explanation: We also have at our disposal some of the standard operations defined on the Fisher matrix object.
We can easily sum Fisher matrices:
End of explanation
"""


print fisher_2==fisher

print fisher_3!=fisher

"""
Explanation: And check for equality/disequality:
End of explanation
"""


derived_matrix = [[3.0],[1.0]]

"""
Explanation: Derived Fisher matrices:
Several tools are also included to define derived Fisher matrices. 
With this name we mean the Jacobian of the transformation that projects the Fisher matrix onto a new set 
of parameters. 
Also here we will cover the most basic usage and for the details refer to the fisher_derived documentation.
Let's start by defining a derived Fisher matrix that takes two parameters and projects them into one:
End of explanation
"""


base_parameters = param_names
derived_parameters = ['par3']

"""
Explanation: We need to supply also the names of the parameters (otherwise dumb defaults will be used) to make sure that the code can
interpret well what is going on:
End of explanation
"""


fisher_derived = fd.fisher_derived( derived_matrix=derived_matrix, 
                                    param_names=base_parameters, 
                                    derived_param_names=derived_parameters )

"""
Explanation: Now define the derived matrix:
End of explanation
"""


new_fisher = fisher_derived.add_derived( fisher )
print new_fisher.get_fisher_matrix()

"""
Explanation: And use it along our base Fisher matrix:
End of explanation
"""


new_fisher = fisher_derived.add_derived( fisher, preserve_input=True )
print new_fisher.get_fisher_matrix()

"""
Explanation: Preserving the input if we want to:
End of explanation
"""


params = ['par1']
fisher_marg = fo.marginalise_over(fisher, params)
print 'Fisher matrix:', fisher_marg.get_fisher_matrix()
print 'Parameter names:', fisher_marg.get_param_names()

"""
Explanation: Notice that the result of this operation will be a Fisher matrix with the new parameters.
Operations on Fisher matrices:
We can also perform several operations on Fisher matrices. 
Several operations are implemented in fisher_operations. We cover here the most important and refer to the documentation for all the others.
We start with marginalization. If we choose a set of parameters we can marginalize over those, so that the resulting Fisher matrix does not contain them:
End of explanation
"""


params = ['par1']
fisher_marg = fo.marginalise(fisher, params)
print 'Fisher matrix:', fisher_marg.get_fisher_matrix()
print 'Parameter names:', fisher_marg.get_param_names()

"""
Explanation: Or we can marginalise over the other parameters:
End of explanation
"""


params = ['par2','par1']
fisher_new = fo.reshuffle(fisher, params)
print 'Fisher matrix:'
print fisher_new.get_fisher_matrix()
print 'Parameter names:', fisher_new.get_param_names()

"""
Explanation: In some occasions we might want to restructure a Fisher matrix, reordering the parametes or eliminating some, without marginalizing. This is the goal of the reshuffle function. 
If we want to reorder parameters:
End of explanation
"""


params = ['par2']
fisher_new = fo.reshuffle(fisher, params)
print 'Fisher matrix:', fisher_new.get_fisher_matrix()
print 'Parameter names:', fisher_new.get_param_names()

"""
Explanation: If we want to delete a parameter:
End of explanation
"""


fisher_1 = fm.fisher_matrix( [[1.0,0.0],[0.0,2.0]], param_names = ['par1','par2'] )
fisher_2 = fm.fisher_matrix( [[1.0,0.5],[0.5,2.0]], param_names = ['par1','par2'] )
fisher_3 = fm.fisher_matrix( [[0.8,0.2],[0.2,3.0]], param_names = ['par1','par3'] )

"""
Explanation: Notice that the result is the same as the marginalization because in the example Fisher matrix the parameters are assumed independent. This will not be the case if parameters are correlated.
Analysis of a set of Fisher matrices:
So far we have only considered one single Fisher matrix. In many real life applications, however, we might want to work with a set of several Fisher matrices. This is the purpose of the fisher_plot_analysis class. This handles set of Fisher matrices and has been engeneered to be extremely flexible in its usage. Here we cover its basic usage and refer the documentation for all the details. 
We start by defining some Fisher matrices:
End of explanation
"""


fisher_1.name = 'fisher1'
fisher_2.name = 'fisher2'
fisher_3.name = 'fisher3'

"""
Explanation: In a Fisher matrix set the name of the Fisher matrix is used as the unique identifier of the object. We want to rename all our Fisher matrices to make them different:
End of explanation
"""


fisher_list = fpa.CosmicFish_FisherAnalysis()

"""
Explanation: We then create an empty list of Fisher matrices:
End of explanation
"""


fisher_list.add_fisher_matrix(fisher_1)
fisher_list.add_fisher_matrix([fisher_2, fisher_3])

"""
Explanation: We then add the Fisher matrix previously defined:
End of explanation
"""


print 'Fisher matrices in the list:', fisher_list.get_fisher_name_list()

"""
Explanation: We can now query the list to see what are the names of the Fisher matrices stored inside:
End of explanation
"""


print fisher_list.get_fisher_list()

"""
Explanation: Or we can directly ask for the Fisher matrices themselves:
End of explanation
"""


fisher_list.delete_fisher_matrix('fisher1')
print 'Fisher matrices in the list:', fisher_list.get_fisher_name_list()
fisher_list.add_fisher_matrix(fisher_1)

"""
Explanation: We can delete a Fisher matrix from the list, using its name as its identifier:
End of explanation
"""


print 'Parameters in the full set:',fisher_list.get_parameter_list()
print 'Parameters in a sub set:',fisher_list.get_parameter_list(['fisher1','fisher2'])

"""
Explanation: We can then query, for example, for all the parameters that are in the set of Fisher matrix or a sub set of it:
End of explanation
"""


fisher_plotter = fp.CosmicFishPlotter(fishers=fisher_list)

"""
Explanation: Many other things are defined to handle efficiently and with extreme flexibility our set of Fisher matrices. In addition this class contains a powerfull function that will search for files containing Fisher matrices.
Plotting:
The last main part of the library contains plot tools. The idea is that the user defines a set of Fisher matrices, as explained above, and then plots them. As before, we show here the basic usage, that might be enough for a user that just wants to get the plots, and we refer to the documentation of the plotting facilities and the documentation to the plot setting module for further details.
Here we start by defining a plotter with the Fisher matrix set that we previously defined:
End of explanation
"""


fisher_plotter = fp.CosmicFishPlotter(fishers=fisher_list,legend_ncol=3)

"""
Explanation: When we initialize the plotter we can pass settings, as described in the documentation to change the appearence of the plot. Here we just slightly change the legend format:
End of explanation
"""


fisher_plotter.new_plot()
fisher_plotter.plot1D()
fisher_plotter.export('test.png', dpi=100)
IPython.display.Image(filename='test.png',unconfined=True) 

"""
Explanation: Notice that in the following we have to go for a rather dumb trick to avoid IPython backend messing up with figures...
Plot 1D:
After the plotter is initialized we can produce marginalized 1D plots as follows:
End of explanation
"""


fisher_plotter.new_plot()
fisher_plotter.plot1D(['par2','par3'],['fisher1'])
fisher_plotter.export('test.png', dpi=100)
IPython.display.Image(filename='test.png',unconfined=True) 

fisher_plotter.new_plot()
fisher_plotter.plot1D(['par2','par3'])
fisher_plotter.export('test.png', dpi=100)
IPython.display.Image(filename='test.png',unconfined=True) 

"""
Explanation: Notice that calling the plotter without any argument will make him plot all the Fisher matrices and parameters in the set.
If we want to plot just some parameters or some Fisher matrices we can easily select them:
End of explanation
"""


fisher_plotter.new_plot()
fisher_plotter.plot2D()
fisher_plotter.export('test.png', dpi=100)
IPython.display.Image(filename='test.png',unconfined=True) 

"""
Explanation: Notice that the colors, labels, sizes and legend are consistent in all plots done with the same plotter.
Plot 2D:
In the same way we can do marginalized 2D plots as follows. Beware that calling plot2D without arguments will result in a plot of all the parameter's couples and this might take a while in real life.
End of explanation
"""


fisher_plotter.new_plot()
fisher_plotter.plot2D([['par1','par2']],legend_ncol=1)
fisher_plotter.export('test.png', dpi=100)
IPython.display.Image(filename='test.png',unconfined=True) 

"""
Explanation: Notice that the empty plot is signaling that there is no Fisher matrix with both par2 and par3.
As before we can select parameters by passing parameters couples:
End of explanation
"""


fisher_plotter.new_plot()
fisher_plotter.plot_tri(legend_ncol=1)
fisher_plotter.export('test.png', dpi=100)
IPython.display.Image(filename='test.png',unconfined=True) 

fisher_plotter.new_plot()
fisher_plotter.plot_tri(['par1','par2'],legend_ncol=1)
fisher_plotter.export('test.png', dpi=100)
IPython.display.Image(filename='test.png',unconfined=True) 

"""
Explanation: Notice that we changed a setting on the flight. Other settings can be changed on the flight as well and will hold only for the selected plot.
Triangular plot:
Precisely in the same way we can do triangular plots. Beware that calling plot_tri without arguments will result in a full triangular plot and this might take a while in real applications.
End of explanation
"""


files = './documentation/python/fisher_test/'

fisher_list = fpa.CosmicFish_FisherAnalysis(fisher_path=files)

fisher_list = fpa.CosmicFish_FisherAnalysis(fisher_path=[files+'Planck_Pre_Launch_fisher_matrix_cls.dat' ])

"""
Explanation: A realistic example:
So far we have seen just simple, unrealistic, examples. Now we turn to a realistic situation. 
We shall now take a couple of Fisher matrices that were produced with the CosmicFish code and show how to perform simple operations.
The CosmicFish python library comes with a set of applications that will facilitate every of these operations for real situations. Instead of doing a lot of Python code check their documentation.
We start by importing the Fisher matrices. This can be done in two ways. By letting CosmicFish_FisherAnalysis search a set of paths or point him directly to the Fisher matrices. Here we go showing both:
End of explanation
"""


names    = fisher_list.get_fisher_matrix('Planck_Pre_Launch_fisher_matrix_cls')[0].get_param_names()
fiducial = fisher_list.get_fisher_matrix('Planck_Pre_Launch_fisher_matrix_cls')[0].get_param_fiducial()
bounds68 = fisher_list.get_fisher_matrix('Planck_Pre_Launch_fisher_matrix_cls')[0].get_confidence_bounds()

"""
Explanation: Now we might want the bounds on some parameters and we can simply do:
End of explanation
"""


bounds68 = fu.v_nice_number(bounds68)
fiducial = [ fu.significant_digits((i,j)) for i,j  in zip(fiducial,bounds68) ]

"""
Explanation: In this format the results will contain a lot of useless digits. We use some utility to get rid of them:
End of explanation
"""


fu.print_table( [names,fiducial, bounds68] )

"""
Explanation: And print them in a nicely formatted table:
End of explanation
"""


fisher_plotter = fp.CosmicFishPlotter(fishers=fisher_list)

fisher_plotter.new_plot()

fisher_plotter.new_plot()
fisher_plotter.plot1D(['omegabh2','omegach2'])
fisher_plotter.export('test.png', dpi=100)
IPython.display.Image(filename='test.png',unconfined=True) 

fisher_plotter.new_plot()
fisher_plotter.plot2D([['omegabh2','omegach2'],['ns','r']])
fisher_plotter.export('test.png', dpi=100)
IPython.display.Image(filename='test.png',unconfined=True) 

fisher_plotter.new_plot()
fisher_plotter.plot_tri(['omegabh2','logA','tau'])
fisher_plotter.export('test.png', dpi=100)
IPython.display.Image(filename='test.png',unconfined=True) 

"""
Explanation: And now we do some plot:
End of explanation
"""


