# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/06f_memory.tensorboard.ipynb.

# %% auto 0
__all__ = ['create_event_accumulator', 'diff_tags', 'scalars2value', 'scalar2line', 'images2animation', 'figures_to_html',
           'create_experience_replay_fig']

# %% ../nbs/06f_memory.tensorboard.ipynb 3
# Python native modules
import os
from fnmatch import fnmatch
from warnings import warn
# Third party libs
from PIL import Image
from IPython.display import HTML
from io import BytesIO
import plotly.express as px
import numpy as np
import pandas as pd
from fastcore.all import *
from tensorboard.backend.event_processing import event_accumulator
# Local modules
from ..core import *

# %% ../nbs/06f_memory.tensorboard.ipynb 6
def create_event_accumulator(
        p:Optional[Path]=None, # Path or in-complete path to a tensorboard event directory. ref `pattern`
        pattern:str='*', # Will filter paths at `cwd` if `p` is None, or filter at `p`
        verbose:bool=False,
        images=None, # Reference EventAccumulator `size_guidance`
        scalars=None, # Reference EventAccumulator `size_guidance`
        audio=None, # Reference EventAccumulator `size_guidance`
        histograms=None, # Reference EventAccumulator `size_guidance`
        compressed_histograms=None, # Reference EventAccumulator `size_guidance`
    )->event_accumulator.EventAccumulator:
    "Creates an event accumulator at the first directory at `p`."
    p=ifnone(p,Path.cwd())
    events=L([o for o in p.ls() if fnmatch(str(o),pattern)]).sorted()
    
    if not events:
        raise Exception(f'No events were found at {p}'+('' if pattern is None else f' and pattern {pattern}'))
    if len(events)>1 and verbose:
        warn(f'Got {len(events)} events. Using the most recent one.')
    if verbose:
        print('Using: ',str(events[-1]))
    ea=event_accumulator.EventAccumulator(str(events[-1]),
        size_guidance={ 
            event_accumulator.COMPRESSED_HISTOGRAMS: ifnone(compressed_histograms,500),
            event_accumulator.IMAGES: ifnone(images,100),
            event_accumulator.AUDIO: ifnone(audio,4),
            event_accumulator.SCALARS: ifnone(scalars,0),
            event_accumulator.HISTOGRAMS: ifnone(histograms,1),
        })
    ea.Reload()
    return ea

# %% ../nbs/06f_memory.tensorboard.ipynb 8
def diff_tags(a,b):
    return ''.join(set(a.split('/'))-set(b.split('/')))

# %% ../nbs/06f_memory.tensorboard.ipynb 11
def scalars2value(scalars:list): return list(map(lambda o:o.value,scalars))

def scalar2line(
        tag_pattern:str, # Tags that will match this pattern to be displayed.
        ea:event_accumulator.EventAccumulator, # The event accumulator to load the images
        start:int=0, # Where in the list of images to start loading
        end:Optional[int]=None, # Optionally where to stop loading images
        step:int=1 # Number of steps betwee image reads
    ):
    tags=[t for t in ea.Tags()['scalars'] if fnmatch(t,tag_pattern)]
    if not tags: warn(f'There are not tags with pattern {tag_pattern}, there are: \n{ea.Tags()["scalars"]}')
    # slice the scalars
    sliced={t:scalars2value(ea.Scalars(t)[start:end:step]) for t in tags}
    # Create the data frame
    df=pd.DataFrame(
        data={'values':np.array(list(sliced.values())).reshape(-1,),
              'tags':np.array([[diff_tags(k,tag_pattern)]*len(v) 
                               for k,v in sliced.items()]).reshape(-1,),
              'steps':np.array([np.arange(len(v)) for v in sliced.values()]).reshape(-1,)}
    )
    
    value_array=np.array(list(sliced.values()))
    return px.line(data_frame=df,x='steps',y='values',animation_frame="tags",
                   title=tag_pattern,
            range_y=[value_array.min(),value_array.max()])

# %% ../nbs/06f_memory.tensorboard.ipynb 14
def _td_gif2np(img):
    return np.array(Image.open(BytesIO(img.encoded_image_string)).convert('RGB')).copy()

def images2animation(
        tag_pattern:str, # Tags that will match this pattern to be displayed.
        ea:event_accumulator.EventAccumulator, # The event accumulator to load the images
        start:int=0, # Where in the list of images to start loading
        end:Optional[int]=None, # Optionally where to stop loading images
        step:int=1 # Number of steps betwee image reads
    ):
    tags=[t for t in ea.Tags()['images'] if fnmatch(t,tag_pattern)]
    if not tags: 
        warn(f'There are no images with pattern {tag_pattern} given tags: '+'\n'.join(ea.Tags()['images']))
        return None
    
    imgs=ea.Images(tags[-1])[start:end:step]
    if not imgs: 
        warn(f'There are no images in {tags[-1]}')
        return None
    
    np_imgs=[_td_gif2np(o) for o in imgs]
    return px.imshow(np.array(np_imgs),animation_frame=0)

# %% ../nbs/06f_memory.tensorboard.ipynb 16
def figures_to_html(*figs):    
    html="<html><head></head><body>" + "\n"
    for fig in figs: html+=fig.to_html().split('<body>')[1].split('</body>')[0]
    html+="</body></html>" + "\n"
    display(HTML(html))

# %% ../nbs/06f_memory.tensorboard.ipynb 17
@delegates(create_event_accumulator)
def create_experience_replay_fig(
        tag_pattern_expected_reward:Union[str,list]='experience_replay/*/expected_reward/action_dim_0', # Tags about expected reward
        tag_pattern_td_error:str='experience_replay/*/td_error', # Tags about td_error
        tag_pattern_images:str='*', # Tags about images. 
        start:int=0, # Where in the list of images to start loading
        end:Optional[int]=None, # Optionally where to stop loading images
        step:int=1, # Number of steps betwee image reads
        **kwargs):
    ea=create_event_accumulator(**kwargs)
    
    common_args=dict(start=start,end=end,step=step)
    if isinstance(tag_pattern_expected_reward,str):
        expected_reward_line=(scalar2line(tag_pattern_expected_reward,ea=ea,**common_args),)
    else:
        expected_reward_line=tuple(scalar2line(tper,ea=ea,**common_args) 
                                   for tper in tag_pattern_expected_reward)
    td_error_line=scalar2line(tag_pattern_td_error,ea=ea,**common_args)
    
    if tag_pattern_images is not None:
        images_animation=images2animation(tag_pattern_images,ea=ea,**common_args)
    
        return figures_to_html(
            td_error_line,
            images_animation,
            *expected_reward_line
        )
    else:
        return figures_to_html(
            td_error_line,
            *expected_reward_line
        )

