# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/callback.simpleprofiler.ipynb (unless otherwise specified).

__all__ = ['SimpleProfilerPostCallback', 'SimpleProfilerCallback']

# Cell
import time
import pandas as pd
import numpy as np
from pathlib import Path
from packaging.version import parse

from fastcore.foundation import docs
from fastcore.basics import mk_class, noop, in_notebook

import fastai
from fastai.learner import Learner, Recorder
from fastai.callback.core import *

from ..imports import *

if in_notebook():
    from IPython.display import display

# Cell
if parse(fastai.__version__) >= parse('2.7.0'):
    _inner_loop = "before_draw before_batch after_pred after_loss before_backward after_cancel_backward after_backward before_step after_step after_cancel_batch after_batch".split()
else:
    _inner_loop = "before_draw before_batch after_pred after_loss before_backward before_step after_step after_cancel_batch after_batch".split()

# Cell
if parse(fastai.__version__) >= parse('2.7.0'):
    _events = L.split('after_create before_fit before_epoch before_train before_draw before_batch after_pred after_loss \
        before_backward after_cancel_backward after_backward before_step after_cancel_step after_step \
        after_cancel_batch after_batch after_cancel_train after_train before_validate after_cancel_validate \
        after_validate after_cancel_epoch after_epoch after_cancel_fit after_fit')
else:
    _events = L.split('after_create before_fit before_epoch before_train before_draw before_batch after_pred after_loss \
        before_backward before_step after_cancel_step after_step after_cancel_batch after_batch after_cancel_train \
        after_train before_validate after_cancel_validate after_validate after_cancel_epoch \
        after_epoch after_cancel_fit after_fit')

mk_class('event', **_events.map_dict(),
         doc="All possible events as attributes to get tab-completion and typo-proofing")

# Cell
@patch
def __call__(self:Callback, event_name):
    "Call `self.{event_name}` if it's defined"
    _run = (event_name not in _inner_loop or (self.run_train and getattr(self, 'training', True)) or
            (self.run_valid and not getattr(self, 'training', False)))
    res = None
    if self.run and _run: res = getattr(self, event_name, noop)()
    if event_name=='after_fit': self.run=True #Reset self.run to True at each end of fit
    return res

# Cell
@patch
def _call_one(self:Learner, event_name):
    if not hasattr(event, event_name): raise Exception(f'missing {event_name}')
    for cb in self.cbs.sorted('order'): cb(event_name)

# Cell
@patch
def all_batches(self:Learner):
    self.n_iter = len(self.dl)
    if hasattr(self, 'simple_profiler'):
        self.it = iter(self.dl)
        for i in range(self.n_iter):
            self("before_draw")
            self.one_batch(i, next(self.it))
        del(self.it)
    else:
        for o in enumerate(self.dl): self.one_batch(*o)

# Cell
_loop = ['Start Fit', 'before_fit', 'Start Epoch Loop', 'before_epoch', 'Start Train', 'before_train',
         'Start Batch Loop', 'before_draw', 'before_batch', 'after_pred', 'after_loss', 'before_backward',
         'before_step', 'after_step', 'after_cancel_batch', 'after_batch','End Batch Loop', 'End Train',
         'after_cancel_train', 'after_train', 'Start Valid', 'before_validate', 'Start Batch Loop',
         '**CBs same as train batch**', 'End Batch Loop', 'End Valid', 'after_cancel_validate',
         'after_validate', 'End Epoch Loop', 'after_cancel_epoch', 'after_epoch', 'End Fit',
         'after_cancel_fit', 'after_fit']

# Internal Cell
@patch
def show_training_loop(self:Learner):
    indent = 0
    for s in _loop:
        if s.startswith('Start'): print(f'{" "*indent}{s}'); indent += 2
        elif s.startswith('End'): indent -= 2; print(f'{" "*indent}{s}')
        else: print(f'{" "*indent} - {s:15}:', self.ordered_cbs(s))

# Internal Cell
_all = ['fit', 'epoch', 'train', 'validate']
_train = ['backward', 'step', 'zero_grad']
_multiple = ['draw', 'batch', 'pred', 'loss']

# Cell
class SimpleProfilerPostCallback(Callback):
    "Pair with `SimpleProfilerCallback` to profile training performance. Removes itself after training is over."
    order,remove_on_fetch = Recorder.order-1,True

    def before_fit(self):
        self.profiler = self.learn.simple_profiler

    def after_train(self):
        self.profiler._raw_time['train'].append(time.monotonic() - self.profiler._train_start)

    def after_validate(self):
        self.profiler._raw_time['validate'].append(time.monotonic() - self.profiler._validate_start)

    def after_pred(self):
        if self.training: self.profiler._raw_time['train_pred'].append(time.monotonic() - self.profiler._train_batch_start)
        else: self.profiler._raw_time['valid_pred'].append(time.monotonic() - self.profiler._valid_batch_start)

        if self.training: self.profiler._train_loss_start = time.monotonic()
        else: self.profiler._valid_loss_start = time.monotonic()

    def after_loss(self):
        if self.training: self.profiler._raw_time['train_loss'].append(time.monotonic() - self.profiler._train_loss_start)
        else: self.profiler._raw_time['valid_loss'].append(time.monotonic() - self.profiler._valid_loss_start)

    def after_step(self):
        self.profiler._raw_time['step'].append(time.monotonic() - self.profiler._step_start)
        self.profiler._zero_start = time.monotonic()

    def after_batch(self):
        if self.training: self.profiler._raw_time['train_batch'].append(time.monotonic() - self.profiler._train_draw_start)
        else: self.profiler._raw_time['valid_batch'].append(time.monotonic() - self.profiler._valid_draw_start)

    def after_epoch(self):
        self.profiler._raw_time['epoch'].append(time.monotonic() - self.profiler._epoch_start)

    def after_fit(self):
        self.profiler._raw_time['fit'].append(time.monotonic() - self.profiler._fit_start)
        self.profiler._generate_report()
        if not hasattr(self.learn, 'lr_finder'):
            self.profiler._display_report()
            self.learn.remove_cbs([SimpleProfilerCallback, SimpleProfilerPostCallback])

# Cell
class SimpleProfilerCallback(Callback):
    """
    Adds a simple profiler to the fastai `Learner`. Optionally showing formatted report or saving unformatted results as csv.

    Pair with SimpleProfilerPostCallback to profile training performance.

    Post fit, access report & results via `Learner.simple_profile_report` & `Learner.simple_profile_results`.
    """
    order,remove_on_fetch = TrainEvalCallback.order+1,True
    def __init__(self, show_report=True, plain=False, markdown=False,
                 save_csv=False, csv_name='simple_profile.csv'):
        store_attr()
        self.csv_name = Path(csv_name)

    def before_fit(self):
        self._raw_time = dict()
        for c in _all + _train:
            self._raw_time[c] = []
        for p in ['train', 'valid']:
            for c in _multiple:
                self._raw_time[f'{p}_{c}'] = []

        self._fit_start = time.monotonic()

    def before_epoch(self):
        self._epoch_start = time.monotonic()

    def before_train(self):
        self._train_start = time.monotonic()

    def before_validate(self):
        self._validate_start = time.monotonic()

    def before_draw(self):
        if self.training: self._train_draw_start = time.monotonic()
        else: self._valid_draw_start = time.monotonic()

    def before_batch(self):
        if self.training: self._raw_time['train_draw'].append(time.monotonic() - self._train_draw_start)
        else: self._raw_time['valid_draw'].append(time.monotonic() - self._valid_draw_start)

        if self.training: self._train_batch_start = time.monotonic()
        else: self._valid_batch_start = time.monotonic()

    def before_backward(self):
        self._backward_start = time.monotonic()

    def before_step(self):
        self._raw_time['backward'].append(time.monotonic() - self._backward_start)
        self._step_start = time.monotonic()

    def after_batch(self):
        if self.training: self._raw_time['zero_grad'].append(time.monotonic() - self._zero_start)

    def _generate_report(self):
        total_time = self._raw_time['fit'][0]
        self.report = pd.DataFrame(columns=['Phase', 'Action', 'Mean Duration', 'Duration Std Dev',
                                               'Number of Calls', 'Total Time', 'Percent of Total'])

        for c in _all:
            if c == 'fit':
                self._append_to_df(['fit', c, 0, 0, 1, total_time, f'{self._calc_percent(total_time):.0%}'])
            else: self._append_to_df(self._create_row('fit', c, self._raw_time[c]))

        for c in _train: self._append_to_df(self._create_row('train', f'{c}', self._raw_time[c]))

        for p in ['train', 'valid']:
            for c in _multiple: self._append_to_df(self._create_row(p, f'{c}', self._raw_time[f'{p}_{c}']))

        self.report = self.report.sort_values(['Phase','Total Time'], ascending=[True,False])
        self.learn.simple_profile_results = self.report.copy()
        for c in ['Mean Duration', 'Duration Std Dev', 'Total Time']:
            self.report[c] = self.report[c].apply(self._scale)
        self.report['Phase'] = self.report['Phase'].where(~self.report['Phase'].duplicated(), '')

        self.learn.simple_profile_report = self.report

    def _display_report(self):
        if self.show_report:
            if self.markdown: print(self.report.to_markdown(index=False))
            else:
                if in_notebook() and not self.plain:
                    with pd.option_context('display.max_rows', len(self.report.index)):
                        s = self.report.style.set_caption("Simple Profiler Results").hide_index()
                        display(s)
                else:
                    print('Simple Profiler Results')
                    print(self.report.to_string(index=False))

        if self.save_csv:
            self.path.parent.mkdir(parents=True, exist_ok=True)
            self.learn.simple_profile_results.to_csv(self.path/self.csv_name, index=False)

    def _append_to_df(self, row):
        self.report.loc[len(self.report.index)] = row

    def _calc_percent(self, time):
        return time / self._raw_time['fit'][0]

    def _create_row(self, phase, name, input):
        return [phase, name, np.mean(input), np.std(input), len(input),
                np.sum(input), f'{self._calc_percent(np.sum(input)):.0%}']

    # modified from https://github.com/thomasbrandon/mish-cuda/blob/master/test/perftest.py
    def _scale(self, val, spec="#0.4G"):
        if val == 0: return '-'
        PREFIXES = np.array([c for c in u"yzafpnµm kMGTPEZY"])
        exp = np.int8(np.log10(np.abs(val)) // 3 * 3 * np.sign(val))
        val /= 10.**exp
        prefix = PREFIXES[exp//3 + len(PREFIXES)//2]
        return f"{val:{spec}}{prefix}s"

# Cell
@patch
def profile(self:Learner, show_report=True, plain=False, markdown=False,
            save_csv=False, csv_name='simple_profile.csv'):
    "Run Simple Profiler when training. Simple Profiler removes itself when finished."
    self.add_cbs([SimpleProfilerCallback(show_report, plain, markdown, save_csv, csv_name),
                  SimpleProfilerPostCallback()])
    return self

# Internal Cell
try:
    import wandb
    from fastai.callback.wandb import WandbCallback

    if not hasattr(WandbCallback,'_orig_before_fit'): WandbCallback._orig_before_fit = WandbCallback.before_fit
    if not hasattr(WandbCallback,'_orig_after_fit'): WandbCallback._orig_after_fit = WandbCallback.after_fit

    @patch
    def before_fit(self:WandbCallback):
        if not hasattr(self.learn, 'lr_finder') and hasattr(self.learn, 'simple_profiler'):
            self.log_simple_profiler = True
        else:
            self.log_simple_profiler = False

        self._orig_before_fit()

    @patch
    def after_fit(self:WandbCallback):
        self._orig_after_fit()

        if self.log_simple_profiler:
            report = wandb.Table(dataframe=self.learn.simple_profile_report)
            results = wandb.Table(dataframe=self.learn.simple_profile_results)

            wandb.log({"simple_profile_report": report})
            wandb.log({"simple_profile_results": results})
            wandb.log({}) # ensure sync

            self.log_simple_profiler = False
except:
    pass