"""Tracking projects experiments with **WandB**"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs-dev/01_wandb.ipynb.

# %% auto 0
__all__ = ['login_nb', 'WandbRun', 'entity_projects', 'get_project', 'print_entity_project_list', 'project_artifacts',
           'run_name_exists', 'unique_run_name']

# %% ../nbs-dev/01_wandb.ipynb 2
# Imports all dependencies
import numpy as np
import pandas as pd
import os
# import tensorflow as tf
import wandb

from datetime import datetime
from pathlib import Path
from typing import Callable, Tuple

# %% ../nbs-dev/01_wandb.ipynb 8
def login_nb(
    nb_file: str|Path   # name of the notebook (str) or path to the notebook (Path)
    ):
    """First step to setup WandB from notebook. Logs in and logs passed notebook as source of code"""

    # Validate nb_file
    if nb_file is None:
        raise TypeError('the file name of the current nb is required to allow code tracking')
    if isinstance(nb_file, str):
        if nb_file[-6:] != '.ipynb': nb_file = f"{nb_file}.ipynb"
        nb_file = Path.cwd() / nb_file
    elif not isinstance(nb_file, Path):
        raise TypeError('nb_file must me a `str` or a `Path`')
    
    if not nb_file.is_file():
        raise ValueError(f"{nb_file.name} is not a file, please correct the file name")

    # Registers notebook as WandB code
    os.environ['WANDB_NOTEBOOK_NAME'] = str(nb_file.absolute())
    print(f"Logging in from notebook: {os.environ['WANDB_NOTEBOOK_NAME']}")

    wandb.login(relogin=False)    

# %% ../nbs-dev/01_wandb.ipynb 18
class WandbRun():
    """Manages a WandB run and all logged actions performed while run is active. Close run with .finish()"""
    
    def __init__(
        self,
        entity: str='',             # user or organization under which the run will be logged. Default: `metagenomics_sh` 
        project: str='',            # name of the WandB project under which the run will be logged 
        run_name: str='',           # unique name for the run,
        job_type: str='',           # e.g.: `load_datasets`, `train_exp`, ... 
        notes: str='',              # any text description or additional information to store with the run 
        logs_dir: str|Path|None= None,   # default is project_root/wandb-logs if None, or uses the passed Path
        testing: bool=False         # (optional) If True, will not create a run on WandB. Use for local testing
        ) :
        """Validates metadata inputs and initialize the wandb run, unless testing is set to True"""
        
        # Validate inputs
        for k,v in [key_val for key_val in locals().items() if key_val[0] not in ['self', 'notes', 'testing', 'logs_dir']]:
            if v == '': raise ValueError(f"{k} may not be an empty string. Please provide a value")

        for k,v in [key_val for key_val in locals().items() if key_val[0] not in ['self', 'testing', 'logs_dir']]:
            if not isinstance(v, str): raise TypeError(f"{k} must be a string, not a {type(v)}")

        self.entity = entity
        self.project = project
        self.run_name = run_name
        self.job_type = job_type
        self.notes = notes
        
        if logs_dir is None:
            self.wandb_logs = self._get_wandb_logs_dir()
        elif isinstance(logs_dir, str):
            self.wandb_logs = Path(logs_dir)
        elif isinstance(logs_dir, Path):
            self.wandb_logs = logs_dir
        else:
            raise ValueError(f"logs_dir must be None, a str or a Path, not a {type(logs_dir)}")
        
        if not testing:
            self.run = wandb.init(
                entity=entity, 
                project=project, 
                name=run_name, 
                job_type=job_type, 
                notes=notes, 
                save_code=True,
                dir= self.wandb_logs
            )

    def finish(self):
        """End the run"""
        self.run.finish()
    
    @staticmethod
    def _get_wandb_logs_dir():
        cur_dir_parents = Path().absolute().parents
        wandb_logs_dir = [p for p in cur_dir_parents if 'nbs' not in p.name][0] / 'wandb-logs'
        if not wandb_logs_dir.is_dir():
            raise ValueError(f"Cannot find the wandb-logs directory. Please specify the correct path ")
        return wandb_logs_dir
        
    def upload_dataset(
        self, 
        ds_path: str,                 # path to the file or directory to load as dataset artifact 
        ds_name: str,                 # name for the dataset
        ds_type: str,                 # type of dataset: e.g. raw_data, processed_data, ...
        ds_descr: str,                # short description of the dataset
        ds_metadata: dict,            # keys/values for metadata on the dataset, eg. nb_samples, ...
        load_type:str = 'file',       # `file` to load a single file, `dir` to load all files in a directory
        wait_completion: bool = False # when True, wait completion of the logging before returning artifact
        ):
        """Load a dataset from a file as WandB artifact, with associated information and metadata"""
        
        # validate ds_path
        if load_type not in ['file', 'dir']:
            raise ValueError(f"load_type must be 'file' or 'dir'")
        if load_type == 'file' and not Path(ds_path).is_file():
            raise ValueError(f"No file found as {ds_path}. Please check path or load type")
        if load_type == 'dir' and not Path(ds_path).is_dir():
            raise ValueError(f"No directory found as {ds_path}. Please check path or load type")

        artifact = wandb.Artifact(name=ds_name, type=ds_type, description=ds_descr, metadata=ds_metadata)

        if load_type == 'file':
            artifact.add_file(ds_path, ds_name)
        if load_type == 'dir':
            artifact.add_dir(ds_path, ds_name)
        
        self.run.log_artifact(artifact)
        
        print(f"Dataset {ds_name} is being logged as artifact ...")
        
        if wait_completion:
            artifact.wait()
            print(f"Dataset {ds_name} logging completed")
            print(f"Artifact state: {artifact.state}")
        
        return artifact

# %% ../nbs-dev/01_wandb.ipynb 45
def entity_projects(
    entity: str # name of the entity from which the projects will be retrieved
    ) -> wandb.apis.public.Projects : # Projects iterator
    """Returns all projects under 'entity', as an iterable collection"""
    api = wandb.Api()
    projects = api.projects(entity=entity)
    return projects

# %% ../nbs-dev/01_wandb.ipynb 48
def get_project(
    entity: str,        # name of the entity from which the project will be retrieved 
    project_name:str,   # name of the project to retrieve
    ) -> wandb.apis.public.Project : # Project object
    """Returns project object defined by entity and project name"""
    api = wandb.Api()
    return api.from_path(f"{entity}/{project_name}")

# %% ../nbs-dev/01_wandb.ipynb 50
def print_entity_project_list(entity):
    """Print the name and url of all projects in entity"""
    projects = entity_projects(entity)
    print(f"List of projects under entity <{entity}>")
    for i, p in enumerate(projects):
        print(f" {i:2d}. {p.name:30s} (url: {p.url})")

# %% ../nbs-dev/01_wandb.ipynb 52
def project_artifacts(
    entity: str,                     # name of the entity from which to retrieve the artifacts 
    project_name: str,               # name of the project from which to retrieve the artifacts 
    by_alias: str='latest',          # name of the alias to filter by
    by_type: str=None,               # name of the artifact type to filter by (optional)
    by_version: str=None             # version to filter by (optional)
    ) -> Tuple[pd.DataFrame, list]:  # df w/ all artifacts and related info; list of artifact types in the project
    """Returns all artifacts in project, w/ key info, filtered by alias, types and version + list of artifact types"""
    api = wandb.Api()
    project = api.from_path(f"{entity}/{project_name}")
    atx_types = project.artifacts_types()
    runs = api.runs(path=f"{entity}/{project_name}")

    # validate by_type parameter
    if by_type is not None and by_type not in [t.name for t in atx_types]:
        raise ValueError(f"{by_type} is not an artifact type in {entity}/{project_name}")

    # create a df where each row corresponds to one artifact logged during one run in this project
    # some artifact may be duplicated when linked to more than one run. Those duplicate need to be filtered out
    cols = 'atx_name atx_type atx_id atx_state atx_version atx_aliases file_count created updated'.split(' ')
    artifacts_df = pd.DataFrame(columns=cols)
    
    for r in runs:
        for atx in r.logged_artifacts():
            metadata = [atx.name, atx.type, atx.id, atx.state, atx.version, atx.aliases, atx.file_count, atx.created_at, atx.updated_at]
            row = pd.DataFrame({k:v for k, v in zip(cols, metadata)})
            artifacts_df = pd.concat((artifacts_df, row), axis=0, ignore_index=True)
    artifacts_df = artifacts_df.loc[~artifacts_df.duplicated(subset=['atx_id'], keep='first'), :]

    cols2show = 'atx_name atx_version atx_type atx_aliases file_count created updated atx_id'.split(' ')
    # filtering by passed alias and type:
    #   if by_xxx is not None:    filter is a boolean vector
    #   if by_xxx is None:        filter is an array of 'True'
    nbr_rows = artifacts_df.shape[0]
    alias_filter = artifacts_df.atx_aliases==by_alias if by_alias is not None else np.ones(shape=(nbr_rows,), dtype=bool)
    type_filter = artifacts_df.atx_type==by_type if by_type is not None else np.ones(shape=(nbr_rows,), dtype=bool)
    version_filter = artifacts_df.atx_version==by_version if by_version is not None else np.ones(shape=(nbr_rows,), dtype=bool)

    row_filter = alias_filter * type_filter * version_filter

    latest = artifacts_df.loc[row_filter, cols2show].sort_values(by='created').reset_index(drop=True)
    return latest, [t.name for t in atx_types]

# %% ../nbs-dev/01_wandb.ipynb 58
def run_name_exists(
    run_name: str,      # name of the run to check 
    entity: str,        # name of the entity from which to retrieve the artifacts 
    project_name: str,  # name of the project from which to retrieve the artifacts 
    ) -> bool:          # True if a run exists with the name run_name, False otherwise
    """Check whether a run with name `run_name` already exists in `entity/project_name`"""
    api = wandb.Api()
    runs = api.runs(path=f"{entity}/{project_name}")
    run_matches = [run_name == r.name for r in runs]
    return any(run_matches)

# %% ../nbs-dev/01_wandb.ipynb 61
def unique_run_name(
    name_seed:str     # Run name to which a timestamp will be added
    ):
    """Create a unique run name by adding a timestamp to the passed seed"""
    timestamp = datetime.now().strftime('%y%m%d-%H%M')
    return f"{name_seed}-{timestamp}"
