import logging
import sys

import torch
import wandb

from .plot_color import Plot
from .metrics import Metric

class Logger:
    def __init__(self, args, mds):
        self.molecule = args.molecule
        self.save_dir = args.save_dir
        self.wandb = args.wandb
        self.plot = Plot(args, mds)
        self.metric = Metric(args, mds)
        self.rmsd = float("inf")

    def __call__(self, loss, rollout, policy):
        metrics = self.metric()
        if self.rmsd > metrics["rmsd"]:
            self.rmsd = metrics["rmsd"]
            torch.save(policy.state_dict(), f"{self.save_dir}/policy.pt")
        if self.wandb:
            if metrics["ets"] is not None:
                wandb.log({
                    "rmsd": metrics["rmsd"],
                    "rmsd_std": metrics["rmsd_std"],
                    "thp": metrics["thp"],
                    "ets": metrics["ets"],
                    "ets_std": metrics["ets_std"],
                    "loss": loss
                })
            else:
                wandb.log({
                    "rmsd": metrics["rmsd"],
                    "rmsd_std": metrics["rmsd_std"],
                    "thp": metrics["thp"],
                    "loss": loss
                })