"""
Composed ALSSM Windows (Weighting Window) [ex106.0]
===================================================

Each ALSSM cost term from :mod:`lmlib.statespace.cost` is weighted by its own window (or window weighting function).
This example demonstrates the application of the basic exponentially decaying window of finite and infinite support.
A more complex, symmetric windows (cf. :class:`lmlib.statespace.cost.Segment`) is generated by composing multiple
exponential windows.

See also:
:meth:`lmlib.statespace.model.Segment.windows`,

"""

import numpy as np
import matplotlib.pyplot as plt
import lmlib as lm

K = 1000  # signal length
ks = [500]  # (arbitrary) window location
k = range(K)

segment_left_infinite = lm.Segment(a=-np.inf, b=-40, direction=lm.FORWARD, g=100, delta=-40)  # left decaying window
segment_left_finite = lm.Segment(a=-39, b=-1, direction=lm.FORWARD, g=1e6, delta=0)  # (nearly) rectangular window
segment_right_finite = lm.Segment(a=0, b=39, direction=lm.BACKWARD, g=1e6, delta=0)  # (nearly) rectangular window
segment_right_infinite = lm.Segment(a=40, b=np.inf, direction=lm.BACKWARD, g=100, delta=40)  # right decaying window

cost = lm.CompositeCost((lm.AlssmPoly(3),),
                        [segment_left_infinite, segment_left_finite, segment_right_finite, segment_right_infinite],
                        F=[[1, 1, 1, 1]])

# Generating Windows - for illustrative and plotting purpose only
# ---------------------------------------------------------------

# Defines the minimum window weight (threshold) that is displayed as a non-zero value.
# only relevant for exponentially decaying windows with infinite support, as these never reach exactly 0
display_thd = 0.1
# display without window thresholding
wins = lm.map_windows(cost.windows([0, 1, 2, 3]), ks, K, merge_ks=True,  merge_seg=False)
# display with window thresholding
wins_all_no_thd = lm.map_windows(cost.windows([0, 1, 2, 3], display_thd), ks, K, merge_ks=True, merge_seg=True)
# display without window thresholding
wins_all = lm.map_windows(cost.windows([0, 1, 2, 3]), ks, K, merge_ks=True, merge_seg=True)

# Plot
# ----
plt.figure(figsize=(6, 2))
for p, win in enumerate(wins):
    line = plt.plot(k, win, '-', lw=0.3, label=f'Segment {p}')
    plt.fill_between(k, win, color=line[0].get_color(), alpha=0.3)
plt.plot(k, wins_all, '--k', lw=1.0, label=f'Overall window \n with threshold')
plt.plot(k, wins_all_no_thd, '-k', lw=1.0, label=f'Overall window \n without threshold')

plt.axhline(display_thd, lw=0.6, c='k', ls='--', label='')
plt.xlabel('time index $k$')
plt.ylabel('window weight(s)')
plt.title(f'(Composed) Window localized at index $k={ks[0]}$')
plt.legend(loc=1, fontsize=8)
plt.show()
