# Copyright 2024 Toyota Research Institute.  All rights reserved.

from collections import OrderedDict
from copy import deepcopy

import os
import torch
import torch.nn as nn
import numpy as np
import wandb

from argparse import Namespace
from termcolor import colored


# Simple utilities (borrowed from vidar to avoid hard dependency)

def cfg_has(*args):
    """Check if key is in configuration, optionally return default."""
    if len(args) == 2:
        cfg, name = args
        if isinstance(name, (list, tuple)):
            return all(n in cfg.__dict__ for n in name)
        return name in cfg.__dict__
    elif len(args) == 3:
        cfg, name, default = args
        return cfg.__dict__.get(name, default)
    raise ValueError(f'cfg_has expects 2 or 3 args, got {len(args)}')


def world_size():
    """Return distributed world size."""
    ws = os.environ.get('WORLD_SIZE', None)
    return int(ws) if ws is not None else 1


def pcolor(string, color, on_color=None, attrs=None):
    """Colored string for printing."""
    return colored(string, color, on_color, attrs)


def is_dict(data):
    return isinstance(data, (dict, torch.nn.ModuleDict))

def is_tensor(data):
    return type(data) == torch.Tensor

def is_seq(data):
    return isinstance(data, (list, tuple, torch.nn.ModuleList))

def is_namespace(data):
    return isinstance(data, Namespace)


def get_from_dict(data, key, key2=None):
    """Safe nested dict lookup."""
    if data is None:
        return None
    if key2 is not None:
        out = get_from_dict(data, key)
        return None if out is None else get_from_dict(out, key2)
    if isinstance(key, (list, tuple)):
        for k in key:
            if k in data:
                return data[k]
        return None
    if not is_dict(data):
        return data
    return data.get(key, None)


def ctx_str(ctx):
    replaces = [[' ', ''], ['(', ''], [')', ''], [',', '_']]
    ctx = str(ctx)
    for replace in replaces:
        ctx = ctx.replace(replace[0], replace[1])
    return ctx


def recursive_convert_config(cfg):
    cfg = cfg.__dict__
    for key, val in cfg.items():
        if is_namespace(val):
            cfg[key] = recursive_convert_config(val)
    return cfg

class WandbLogger:
    def __init__(self, cfg, verbose=False):
        super().__init__()
        self._name = cfg.name if cfg_has(cfg, 'name') else None
        self._entity = cfg.entity
        self._project = cfg.project
        self._dir = cfg.dir
        self._tags = cfg_has(cfg, 'tags', '')
        self._notes = cfg_has(cfg, 'notes', '')
        self._id = cfg_has(cfg, 'id', None)
        self._anonymous = None

        if 'TRAINING_JOB_NAME' in os.environ.keys():
            del os.environ['TRAINING_JOB_NAME']

        self._experiment = self._create_experiment()
        self._metrics = OrderedDict()

        cfg.name = self.run_name
        cfg.url = self.run_url

        if verbose:
            self.print()

    @staticmethod
    def finish():
        wandb.finish()

    def print(self):
        font_base = {'color': 'red', 'attrs': ('bold', 'dark')}
        font_name = {'color': 'red', 'attrs': ('bold',)}
        font_underline = {'color': 'red', 'attrs': ('underline',)}
        print(pcolor('#' * 60, **font_base))
        print(pcolor('### WandB: ', **font_base) + \
              pcolor('{}'.format(self.run_name), **font_name))
        print(pcolor('### ', **font_base) + \
              pcolor('{}'.format(self.run_url), **font_underline))
        print(pcolor('#' * 60, **font_base))

    def __getstate__(self):
        state = self.__dict__.copy()
        state['_id'] = self._experiment.id if self._experiment is not None else None
        state['_experiment'] = None
        return state

    def _create_experiment(self):
        experiment = wandb.init(
            name=self._name, dir=self._dir, project=self._project,
            anonymous=self._anonymous, reinit=True, id=self._id, notes=self._notes,
            resume='allow', tags=self._tags, entity=self._entity,
            settings=wandb.Settings(_service_wait=120),
        )
        # wandb <0.18: wandb.run.save() (no args) persists the run name/id.
        # wandb >=0.18: init() already persists; same call now requires glob_str.
        if tuple(int(x) for x in wandb.__version__.split('.')[:2]) < (0, 18):
            wandb.run.save()
        return experiment

    @property
    def experiment(self):
        if self._experiment is None:
            self._experiment = self._create_experiment()
        return self._experiment

    @property
    def run_name(self):
        return wandb.run.name if self._experiment else None

    @property
    def run_url(self):
        return f'https://app.wandb.ai/' \
               f'{wandb.run.entity}/' \
               f'{wandb.run.project}/runs/' \
               f'{wandb.run.id}' if self._experiment else None

    def log_config(self, cfg):
        cfg = recursive_convert_config(deepcopy(cfg))
        self.experiment.config.update(cfg, allow_val_change=True)

    def log_metrics(self, metrics):
        self._metrics.update(metrics)
        if 'epochs' in metrics or 'samples' in metrics:
            self.experiment.log(self._metrics)
            self._metrics.clear()


def prep_image(key, image):
    """Prepare image for logging"""
    if is_tensor(image):
        if image.dim() == 2:
            image = image.unsqueeze(0)
        if image.dim() == 4:
            image = image[0]
        image = image.detach().permute(1, 2, 0).cpu().numpy()
    image[np.isnan(image)] = 0.0
    image[np.isinf(image)] = 0.0
    return {key: wandb.Image(image)}

def prep_video(key, video):
    return {key: wandb.Video(video)}

def prep_points(key, points, rgb=None):
    if is_tensor(points):
        if points.dim() == 4:
            points = points[0]
            if rgb is not None:
                rgb = rgb[0]
        points = points.view(3, -1).permute(1, 0)
        if rgb is not None:
            rgb = rgb.view(3, -1).permute(1, 0) * 255
            points = torch.cat([points, rgb], -1)
        points = points.detach().cpu().numpy()
    return {key: wandb.Object3D(points)}