%matplotlib notebook
from mpl_tools import is_notebook_or_qt as nb

import dapper as dpr
import dapper.da_methods as da

"""
Explanation: Illustrate usage of DAPPER to (interactively) run a synthetic ("twin") experiment.
Imports
<b>NB:</b> If you're on <mark><b>Gooble Colab</b></mark>,
then replace %matplotlib notebook below by
!python -m pip install git+https://github.com/nansencenter/DAPPER.git .
Also note that liveplotting does not work on Colab.
End of explanation
"""


from dapper.mods.Lorenz63.sakov2012 import HMM

"""
Explanation: Load experiment setup: the hidden Markov model (HMM)
End of explanation
"""


seed = dpr.set_seed(3000)

"""
Explanation: Generate the same random numbers each time this script is run
End of explanation
"""


HMM.tseq.T = 30  # shorten experiment
xx, yy = HMM.simulate()

"""
Explanation: Simulate synthetic truth (xx) and noisy obs (yy)
End of explanation
"""


# xp = da.OptInterp()
# xp = da.Var3D()
# xp = da.ExtKF(infl=90)
xp = da.EnKF('Sqrt', N=10, infl=1.02, rot=True)
# xp = da.PartFilt(N=100, reg=2.4, NER=0.3)

"""
Explanation: Specify a DA method configuration ("xp" is short for "experiment")
End of explanation
"""


xp.assimilate(HMM, xx, yy, liveplots=not nb)

"""
Explanation: Assimilate yy, knowing the HMM; xx is used to assess the performance
End of explanation
"""


# print(xp.stats)  # ⇒ long printout
xp.stats.average_in_time()

"""
Explanation: Average the time series of various statistics
End of explanation
"""


# print(xp.avrgs)  # ⇒ long printout
print(xp.avrgs.tabulate(['rmse.a', 'rmv.a']))

"""
Explanation: Print some of these time-averages
End of explanation
"""


xp.stats.replay(
    # speed=.6
)

"""
Explanation: Replay liveplotters
End of explanation
"""


if nb:
    import dapper.tools.viz as viz
    viz.plot_rank_histogram(xp.stats)
    viz.plot_err_components(xp.stats)
    viz.plot_hovmoller(xx)

"""
Explanation: Further diagnostic plots
End of explanation
"""


if nb:
    print(xp)

if nb:
    print(HMM)

"""
Explanation: Explore objects
End of explanation
"""


