# Path: frbpoppy/misc.py
# def pprint(*s, output=True):
#     """Hack to make for more informative print statements."""
#     f = inspect.stack()[1][1].split('/')[-1]
#     m = '{:13.13} |'.format(f)
# 
#     if output:
#         print(m, *s)
#     else:
#         lines = []
#         for e in s:
#             lines.append('\n'.join([f'{m} {f}' for f in e.split('\n')]))
#         return '\n'.join(lines)
# 
# Path: tests/convenience.py
# def plot_aa_style(cols=1):
#     """Use a plotting style specifically for Astronomy & Astrophyiscs."""
#     # Use a style file
#     dir = os.path.abspath(os.path.dirname(__file__))
#     style_file = os.path.join(dir, './aa.mplstyle')
#     plt.style.use(style_file)
# 
#     # Check whether interactive backend is present
#     if os.name == 'posix' and "DISPLAY" not in os.environ:
#         # On a cluster
#         plt.switch_backend('Agg')
#         # Use just the basic latex package allowing \text to be used
#         plt.rcParams['text.latex.preamble'] = r'\usepackage{amsmath}'
# 
#     if cols == 2:
#         plt.rcParams["figure.figsize"] = (5.75373, 3.556)
# 
# def rel_path(path):
#     """Return the relative path name for this directory."""
#     return os.path.join(os.path.abspath(os.path.dirname(__file__)), path)


import numpy as np
import matplotlib.pyplot as plt
import frbpoppy.time_dists as td
from matplotlib.colors import LogNorm
from frbpoppy import pprint
from tests.convenience import plot_aa_style, rel_path

def get_probabilities(normalise=False):
    """Get the number of bursts per maximum time.

    Args:
        normalise (type): Whether to normalise at each maximum time, such that
            there's a number of bursts at each time stamp having a maximum
            probability of one.

    Returns:
        array: Number of bursts per maximum time.

    """
    days = np.arange(1, N_DAYS, 1)
    m_bursts = np.arange(0, M_BURSTS, 1)
    xx, yy = np.meshgrid(m_bursts, days)
    prob = np.full((len(days), len(m_bursts)), np.nan)

    # Mask any frbs over the maximum time
    time = td.clustered(r=R, k=K, n_srcs=N_SRCS, n_days=N_DAYS, z=0)

    pprint('Masking days')
    for i, day in reversed(list(enumerate(days))):
        time[(time > day)] = np.nan
        m_bursts = (~np.isnan(time)).sum(1)
        unique, counts = np.unique(m_bursts, return_counts=True)

        try:
            prob[(i, unique)] = counts
        except IndexError:

# Path: frbpoppy/misc.py
# def pprint(*s, output=True):
#     """Hack to make for more informative print statements."""
#     f = inspect.stack()[1][1].split('/')[-1]
#     m = '{:13.13} |'.format(f)
# 
#     if output:
#         print(m, *s)
#     else:
#         lines = []
#         for e in s:
#             lines.append('\n'.join([f'{m} {f}' for f in e.split('\n')]))
#         return '\n'.join(lines)
# 
# Path: tests/convenience.py
# def plot_aa_style(cols=1):
#     """Use a plotting style specifically for Astronomy & Astrophyiscs."""
#     # Use a style file
#     dir = os.path.abspath(os.path.dirname(__file__))
#     style_file = os.path.join(dir, './aa.mplstyle')
#     plt.style.use(style_file)
# 
#     # Check whether interactive backend is present
#     if os.name == 'posix' and "DISPLAY" not in os.environ:
#         # On a cluster
#         plt.switch_backend('Agg')
#         # Use just the basic latex package allowing \text to be used
#         plt.rcParams['text.latex.preamble'] = r'\usepackage{amsmath}'
# 
#     if cols == 2:
#         plt.rcParams["figure.figsize"] = (5.75373, 3.556)
# 
# def rel_path(path):
#     """Return the relative path name for this directory."""
#     return os.path.join(os.path.abspath(os.path.dirname(__file__)), path)


# Path: tests/misc/prob_m_bursts_within_time_given_clustering.py
import numpy as np
import matplotlib.pyplot as plt
import frbpoppy.time_dists as td
from matplotlib.colors import LogNorm
from frbpoppy import pprint
from tests.convenience import plot_aa_style, rel_path


def get_probabilities(normalise=False):
    """Get the number of bursts per maximum time.

    Args:
        normalise (type): Whether to normalise at each maximum time, such that
            there's a number of bursts at each time stamp having a maximum
            probability of one.

    Returns:
        array: Number of bursts per maximum time.

    """
    days = np.arange(1, N_DAYS, 1)
    m_bursts = np.arange(0, M_BURSTS, 1)
    xx, yy = np.meshgrid(m_bursts, days)
    prob = np.full((len(days), len(m_bursts)), np.nan)

    # Mask any frbs over the maximum time
    time = td.clustered(r=R, k=K, n_srcs=N_SRCS, n_days=N_DAYS, z=0)

    pprint('Masking days')
    for i, day in reversed(list(enumerate(days))):
        time[(time > day)] = np.nan
        m_bursts = (~np.isnan(time)).sum(1)
        unique, counts = np.unique(m_bursts, return_counts=True)

        try:
            prob[(i, unique)] = counts
        except IndexError:

            pprint('Please ensure M is large enough.')