from pathlib import Path

import numpy as np
import torch
import torchvision.utils as vutils
from addict import Dict
from PIL import Image
from torch.nn.functional import interpolate, sigmoid

from climategan.data import decode_segmap_merged_labels
from climategan.tutils import (
    all_texts_to_tensors,
    decode_bucketed_depth,
    normalize_tensor,
    write_architecture,
)
from climategan.utils import flatten_opts


class Logger:
    def __init__(self, trainer):
        self.losses = Dict()
        self.time = Dict()
        self.trainer = trainer
        self.global_step = 0
        self.epoch = 0

    def log_comet_images(self, mode, domain, minimal=False, all_only=False):
        trainer = self.trainer
        save_images = {}
        all_images = []
        n_all_ims = None
        all_legends = ["Input"]
        task_legends = {}

        if domain not in trainer.display_images[mode]:
            return

        # --------------------
        # -----  Masker  -----
        # --------------------
        n_ims = len(trainer.display_images[mode][domain])
        print(" " * 60, end="\r")
        if domain != "rf":
            for j, display_dict in enumerate(trainer.display_images[mode][domain]):

                print(f"Inferring sample {mode} {domain} {j+1}/{n_ims}", end="\r")

                x = display_dict["data"]["x"].unsqueeze(0).to(trainer.device)
                z = trainer.G.encode(x)

                s_pred = decoded_s_pred = d_pred = z_depth = None
                for k, task in enumerate(["d", "s", "m"]):

                    if (
                        task not in display_dict["data"]
                        or task not in trainer.opts.tasks
                    ):
                        continue

                    task_legend = ["Input"]
                    target = display_dict["data"][task]
                    target = target.unsqueeze(0).to(trainer.device)
                    task_saves = []

                    if task not in save_images:
                        save_images[task] = []

                    prediction = None
                    if task == "m":
                        cond = None
                        if s_pred is not None and d_pred is not None:
                            cond = trainer.G.make_m_cond(d_pred, s_pred, x)

                        prediction = trainer.G.decoders[task](z, cond, z_depth)
                    elif task == "d":
                        prediction, z_depth = trainer.G.decoders[task](z)
                    elif task == "s":
                        prediction = trainer.G.decoders[task](z, z_depth)

                    if task == "s":
                        # Log fire
                        wildfire_tens = trainer.compute_fire(x, prediction)
                        task_saves.append(wildfire_tens)
                        task_legend.append("Wildfire")
                        # Log seg output
                        s_pred = prediction.clone()
                        target = (
                            decode_segmap_merged_labels(target, domain, True)
                            .float()
                            .to(trainer.device)
                        )
                        prediction = (
                            decode_segmap_merged_labels(prediction, domain, False)
                            .float()
                            .to(trainer.device)
                        )
                        decoded_s_pred = prediction
                        task_saves.append(target)
                        task_legend.append("Target Segmentation")

                    elif task == "m":
                        prediction = sigmoid(prediction).repeat(1, 3, 1, 1)
                        task_saves.append(x * (1.0 - prediction))
                        if not minimal:
                            task_saves.append(
                                x * (1.0 - (prediction > 0.1).to(torch.int))
                            )
                            task_saves.append(
                                x * (1.0 - (prediction > 0.5).to(torch.int))
                            )

                        task_saves.append(x * (1.0 - target.repeat(1, 3, 1, 1)))
                        task_legend.append("Masked input")

                        if not minimal:
                            task_legend.append("Masked input (>0.1)")
                            task_legend.append("Masked input (>0.5)")

                        task_legend.append("Masked input (target)")
                        # dummy pixels to fool scaling and preserve mask range
                        prediction[:, :, 0, 0] = 1.0
                        prediction[:, :, -1, -1] = 0.0

                    elif task == "d":
                        # prediction is a log depth tensor
                        d_pred = prediction
                        target = normalize_tensor(target) * 255
                        if prediction.shape[1] > 1:
                            prediction = decode_bucketed_depth(
                                prediction, self.trainer.opts
                            )
                        smogged = self.trainer.compute_smog(
                            x, d=prediction, s=decoded_s_pred, use_sky_seg=False
                        )
                        prediction = normalize_tensor(prediction)
                        prediction = prediction.repeat(1, 3, 1, 1)
                        task_saves.append(smogged)
                        task_legend.append("Smogged")
                        task_saves.append(target.repeat(1, 3, 1, 1))
                        task_legend.append("Depth target")

                    task_saves.append(prediction)
                    task_legend.append(f"Predicted {task}")

                    save_images[task].append(x.cpu().detach())
                    if k == 0:
                        all_images.append(save_images[task][-1])

                    task_legends[task] = task_legend
                    if j == 0:
                        all_legends += task_legend[1:]

                    for im in task_saves:
                        save_images[task].append(im.cpu().detach())
                        all_images.append(save_images[task][-1])

                if j == 0:
                    n_all_ims = len(all_images)

            if not all_only:
                for task in save_images.keys():
                    # Write images:
                    self.upload_images(
                        image_outputs=save_images[task],
                        mode=mode,
                        domain=domain,
                        task=task,
                        im_per_row=trainer.opts.comet.im_per_row.get(task, 4),
                        rows_per_log=trainer.opts.comet.get("rows_per_log", 5),
                        legends=task_legends[task],
                    )

            if len(save_images) > 1:
                self.upload_images(
                    image_outputs=all_images,
                    mode=mode,
                    domain=domain,
                    task="all",
                    im_per_row=n_all_ims,
                    rows_per_log=trainer.opts.comet.get("rows_per_log", 5),
                    legends=all_legends,
                )
        # ---------------------
        # -----  Painter  -----
        # ---------------------
        else:
            # in the rf domain display_size may be different from fid.n_images
            limit = trainer.opts.comet.display_size
            image_outputs = []
            legends = []
            for im_set in trainer.display_images[mode][domain][:limit]:
                x = im_set["data"]["x"].unsqueeze(0).to(trainer.device)
                m = im_set["data"]["m"].unsqueeze(0).to(trainer.device)

                prediction = trainer.G.paint(m, x)

                image_outputs.append(x * (1.0 - m))
                image_outputs.append(prediction)
                image_outputs.append(x)
                image_outputs.append(prediction * m)
                if not legends:
                    legends.append("Masked Input")
                    legends.append("Painted Input")
                    legends.append("Input")
                    legends.append("Isolated Water")
            # Write images
            self.upload_images(
                image_outputs=image_outputs,
                mode=mode,
                domain=domain,
                task="painter",
                im_per_row=trainer.opts.comet.im_per_row.get("p", 4),
                rows_per_log=trainer.opts.comet.get("rows_per_log", 5),
                legends=legends,
            )

        return 0

    def log_losses(self, model_to_update="G", mode="train"):
        """Logs metrics on comet.ml

        Args:
            model_to_update (str, optional): One of "G", "D". Defaults to "G".
        """
        trainer = self.trainer
        loss_names = {"G": "gen", "D": "disc"}

        if trainer.opts.train.log_level < 1:
            return

        if trainer.exp is None:
            return

        assert model_to_update in {
            "G",
            "D",
        }, "unknown model to log losses {}".format(model_to_update)

        loss_to_update = self.losses[loss_names[model_to_update]]

        losses = loss_to_update.copy()

        if trainer.opts.train.log_level == 1:
            # Only log aggregated losses: delete other keys in losses
            for k in loss_to_update:
                if k not in {"masker", "total_loss", "painter"}:
                    del losses[k]
        # convert losses into a single-level dictionnary

        losses = flatten_opts(losses)
        trainer.exp.log_metrics(
            losses, prefix=f"{model_to_update}_{mode}", step=self.global_step
        )

    def log_learning_rates(self):
        if self.trainer.exp is None:
            return
        lrs = {}
        trainer = self.trainer
        if trainer.g_scheduler is not None:
            for name, lr in zip(
                trainer.lr_names["G"], trainer.g_scheduler.get_last_lr()
            ):
                lrs[f"lr_G_{name}"] = lr
        if trainer.d_scheduler is not None:
            for name, lr in zip(
                trainer.lr_names["D"], trainer.d_scheduler.get_last_lr()
            ):
                lrs[f"lr_D_{name}"] = lr

        trainer.exp.log_metrics(lrs, step=self.global_step)

    def log_step_time(self, time):
        """Logs step-time on comet.ml

        Args:
            step_time (float): step-time in seconds
        """
        if self.trainer.exp:
            self.trainer.exp.log_metric(
                "step-time", time - self.time.step_start, step=self.global_step
            )

    def log_epoch_time(self, time):
        """Logs step-time on comet.ml

        Args:
            step_time (float): step-time in seconds
        """
        if self.trainer.exp:
            self.trainer.exp.log_metric(
                "epoch-time", time - self.time.epoch_start, step=self.global_step
            )

    def log_comet_combined_images(self, mode, domain):

        trainer = self.trainer
        image_outputs = []
        legends = []
        im_per_row = 0
        for i, im_set in enumerate(trainer.display_images[mode][domain]):
            x = im_set["data"]["x"].unsqueeze(0).to(trainer.device)
            # m = im_set["data"]["m"].unsqueeze(0).to(trainer.device)

            m = trainer.G.mask(x=x)
            m_bin = (m > 0.5).to(m.dtype)
            prediction = trainer.G.paint(m, x)
            prediction_bin = trainer.G.paint(m_bin, x)

            image_outputs.append(x)
            legends.append("Input")
            image_outputs.append(x * (1.0 - m))
            legends.append("Soft Masked Input")
            image_outputs.append(prediction)
            legends.append("Painted")
            image_outputs.append(prediction * m)
            legends.append("Soft Masked Painted")
            image_outputs.append(x * (1.0 - m_bin))
            legends.append("Binary (0.5) Masked Input")
            image_outputs.append(prediction_bin)
            legends.append("Binary (0.5) Painted")
            image_outputs.append(prediction_bin * m_bin)
            legends.append("Binary (0.5) Masked Painted")

            if i == 0:
                im_per_row = len(image_outputs)
        # Upload images
        self.upload_images(
            image_outputs=image_outputs,
            mode=mode,
            domain=domain,
            task="combined",
            im_per_row=im_per_row or 7,
            rows_per_log=trainer.opts.comet.get("rows_per_log", 5),
            legends=legends,
        )

        return 0

    def upload_images(
        self,
        image_outputs,
        mode,
        domain,
        task,
        im_per_row=3,
        rows_per_log=5,
        legends=[],
    ):
        """
        Save output image

        Args:
            image_outputs (list(torch.Tensor)): all the images to log
            mode (str): train or val
            domain (str): current domain
            task (str): current task
            im_per_row (int, optional): umber of images to be displayed per row.
                Typically, for a given task: 3 because [input prediction, target].
                Defaults to 3.
            rows_per_log (int, optional): Number of rows (=samples) per uploaded image.
                Defaults to 5.
            comet_exp (comet_ml.Experiment, optional): experiment to use.
                Defaults to None.
        """
        trainer = self.trainer
        if trainer.exp is None:
            return
        curr_iter = self.global_step
        nb_per_log = im_per_row * rows_per_log
        n_logs = len(image_outputs) // nb_per_log + 1

        header = None
        if len(legends) == im_per_row and all(isinstance(t, str) for t in legends):
            header_width = max(im.shape[-1] for im in image_outputs)
            headers = all_texts_to_tensors(legends, width=header_width)
            header = torch.cat(headers, dim=-1)

        for logidx in range(n_logs):
            print(" " * 100, end="\r", flush=True)
            print(
                "Uploading images for {} {} {} {}/{}".format(
                    mode, domain, task, logidx + 1, n_logs
                ),
                end="...",
                flush=True,
            )
            ims = image_outputs[logidx * nb_per_log : (logidx + 1) * nb_per_log]
            if not ims:
                continue

            ims = self.upsample(ims)
            ims = torch.stack([im.squeeze() for im in ims]).squeeze()
            image_grid = vutils.make_grid(
                ims, nrow=im_per_row, normalize=True, scale_each=True, padding=0
            )

            if header is not None:
                image_grid = torch.cat(
                    [header.to(image_grid.device), image_grid], dim=1
                )

            image_grid = image_grid.permute(1, 2, 0).cpu().numpy()
            trainer.exp.log_image(
                Image.fromarray((image_grid * 255).astype(np.uint8)),
                name=f"{mode}_{domain}_{task}_{str(curr_iter)}_#{logidx}",
                step=curr_iter,
            )

    def upsample(self, ims):
        h = max(im.shape[-2] for im in ims)
        w = max(im.shape[-1] for im in ims)
        new_ims = []
        for im in ims:
            im = interpolate(im, (h, w), mode="bilinear")
            new_ims.append(im)
        return new_ims

    def padd(self, ims):
        h = max(im.shape[-2] for im in ims)
        w = max(im.shape[-1] for im in ims)
        new_ims = []
        for im in ims:
            ih = im.shape[-2]
            iw = im.shape[-1]
            if ih != h or iw != w:
                padded = torch.zeros(im.shape[-3], h, w)
                padded[
                    :, (h - ih) // 2 : (h + ih) // 2, (w - iw) // 2 : (w + iw) // 2
                ] = im
                new_ims.append(padded)
            else:
                new_ims.append(im)

        return new_ims

    def log_architecture(self):
        write_architecture(self.trainer)

        if self.trainer.exp is None:
            return

        for f in Path(self.trainer.opts.output_path).glob("archi*.txt"):
            self.trainer.exp.log_asset(str(f), overwrite=True)