from pyhande.data_preparing.hande_ccmc_fciqmc import PrepHandeCcmcFciqmc
from pyhande.extracting.extractor import Extractor
from pyhande.error_analysing.blocker import Blocker
from pyhande.error_analysing.hybrid_ana import HybridAna
from pyhande.results_viewer.get_results import get_results, analyse_data

"""
Explanation: This demonstration shows how CCMC [1] data (analysis) results can be analysed in a more customised way. 
This applies to FCIQMC [2] as well.
End of explanation
"""


results = get_results(["data/0.01_ccsd.out.gz", "data/0.002_ccsd.out.gz"], merge_type='no', analyser='hybrid', start_its='mser')

"""
Explanation: Part A
For now, still using the default, quick get_results function but this time specify merge_type to not merging
(no effect here as calculations are independent, the default is merge using UUIDs btw),
the analyser to hybrid [3] (not blocking [4] as by default) and while we don't specify analysis start MC iterations,
we specify that the MSER find starting iteration function should be used to automatically find them
(the default is 'blocking' for the blocking find starting iteration function).
End of explanation
"""


results.summary

"""
Explanation: The summary table shows the analysed data by the analyser.
The hybrid analyser analyses the instantaneous projected energy (as prepared by the preparator object).
End of explanation
"""


results.analyser.opt_block

print(results.analyser.start_its) # Used starting iterations, found using MSER find starting iteration function.
print(results.analyser.end_its) # Used end iterations, the last iteration by default.

"""
Explanation: The hybrid analyser's output can be viewed.
End of explanation
"""


extra = Extractor(merge={'type': 'legacy', 'md_shift': ['qmc:shift_damping'], 'shift_key': 'Shift'})

"""
Explanation: Part B
Now, we don't use get_results to get the results object but define the extractor, preparator and analyser objects ourselves.
Even though it doesn't have an effect here as there is no calculation to merge, we state that we want to merge using
the 'legacy' way, i.e. don't use UUID for merging but simply determine whether iterations from one output file to the next
(order matters here) are consecutive. If shift is already varying across that continuation, don't merge if 'shift_damping'
differs from one output file to the next ('md_shift' specifies that this restriction only applies when shift is already varying,
otherwise use 'md_always' for this restriction to always hold).
Since no merge is possibly, these options are ignored and just shown here for demonstration purposes.
End of explanation
"""


prep = PrepHandeCcmcFciqmc()

"""
Explanation: Define preparator object.  It contains the hard coded mapping of column name meaning to column name, i.e. 'ref_key' : 'N_0,
for the case of HANDE CCMC/FCIQMC.  If you use a different package, you'll need to create your own preparator class.
End of explanation
"""


ana = HybridAna.inst_hande_ccmc_fciqmc(start_its = 'blocking', find_start_kw_args={'show_graph': True})

"""
Explanation: Define analyser.  Use class method inst_hande_ccmc_fciqmc to pre-set what should be analysed (inst. projected energy), name
of iteration key ('iterations'), etc.
Use 'blocking' start iteration finder and specify that a graph should be shown by the start iteration finder.
End of explanation
"""


results2 = analyse_data(["data/0.01_ccsd.out.gz", "data/0.002_ccsd.out.gz"], extra, prep, ana)

"""
Explanation: Now, we can execute those three objects.  'analyse_data' is a handy helper to call their .exe() methods.
For each calculation, a graph is shown by the find starting iteration method.
End of explanation
"""


results2.analyser.start_its

"""
Explanation: Have used different starting iteration finder, so these will be different.
End of explanation
"""


results2.summary_pretty

"""
Explanation: But results are comparable.
End of explanation
"""


ana2 = HybridAna('iterations', 'Shift', 'replica id')

results3 = analyse_data(["data/0.01_ccsd.out.gz", "data/0.002_ccsd.out.gz"], extra, prep, ana2)

results3.summary_pretty

"""
Explanation: But what if we want to analyse the shift instead of the instantaneous projected energy with hybrid analysis?
-> BEWARE this is untested. Only used for illustration here!
Don't use class method for analyser instantiation anymore.
Keep default settings (find start iterations using 'mser' etc).
Note that when doing blocking [4], not hybrid [3], the order is a bit different, the columns to be analysed are 'cols'
for blocking [4] and 'hybrid_col' for hybrid analysis [3]. You might need to define both for a given analyser if you
are using the starting iteration function of the other type ('blocking' with start_its='mser' or 'hybrid' with
start_its='blocking').  Consult the docstring.
End of explanation
"""


