import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision.utils import make_grid, save_image
import pytorch_lightning as pl
from pytorch_lightning.metrics import (
    PSNR,
    SSIM,
    MeanAbsoluteError,
    MetricCollection,
    Precision,
    Recall,
    F1,
)

import numpy as np
from collections import OrderedDict
import os

from .networks import load_network
from .loss import TverskyLoss, PerceptualLoss, GradientLoss
from .utils import get_scheduler, remove_key_prefix
from .frozen_models import MaskRefineNet, EdgeCompleteNet


class MaskRefineModel(pl.LightningModule):
    def __init__(self, hparams):
        super(MaskRefineModel, self).__init__()
        self.save_hyperparameters()

        # Hyperparameters
        self.hparams = hparams
        self.lr = hparams.lr

        # Modules
        self.net = load_network(hparams)
        self.criterion = TverskyLoss(hparams.tversky_alpha, hparams.tversky_beta)

        # Metrics
        self.train_metrics = MetricCollection([Precision(), Recall(), F1(2)])
        self.val_metrics = MetricCollection([Precision(), Recall(), F1(2)])

    def forward(self, img, bboxes):
        return self.net(torch.cat((img, bboxes), 1))

    def training_step(self, batch, batch_idx):
        img, bboxes, target, _ = batch

        # Pass through model
        pred = self(img, bboxes)

        # Calculate loss and metrics
        loss = self.criterion(pred, target)
        metrics = self.train_metrics(pred, target.long())

        self.log("train_loss", loss)
        self.log("train_prec", metrics["Precision"])
        self.log("train_rec", metrics["Recall"])
        self.log("train_f1", metrics["F1"])
        self.log("lr", self.trainer.optimizers[0].param_groups[0]["lr"], prog_bar=True)

        return loss

    def validation_step(self, batch, batch_idx):
        img, bboxes, target, _ = batch

        # Pass through model
        pred = self(img, bboxes)

        # Calculate loss and metrics
        loss = self.criterion(pred, target)
        metrics = self.val_metrics(pred, target.long())

        self.log("val_loss", loss)
        self.log("val_prec", metrics["Precision"], on_epoch=True)
        self.log("val_rec", metrics["Recall"], on_epoch=True)
        self.log("val_f1", metrics["F1"], on_epoch=True)

        return OrderedDict(
            {
                "val_loss": loss,
                "val_prec": metrics["Precision"],
                "val_rec": metrics["Recall"],
                "val_f1": metrics["F1"],
                "val_sample": torch.stack([img[0], target[0], pred[0]])
                if batch_idx % (250 // self.hparams.batch_size) == 0
                else None,  # Save some samples for visualization
            }
        )

    def validation_epoch_end(self, outputs):
        # Visualize validation samples
        imgs = torch.cat(
            [x["val_sample"] for x in outputs if x["val_sample"] is not None], 0
        )
        grid = make_grid(imgs, nrow=3, normalize=True)

        # Log to tensorboard
        tensorboard = self.logger.experiment[0]
        tensorboard.add_image("val_samples", grid, self.current_epoch)

        # Print validation results
        avg_loss = torch.stack([x["val_loss"] for x in outputs]).mean()
        avg_prec = torch.stack([x["val_prec"] for x in outputs]).mean()
        avg_rec = torch.stack([x["val_rec"] for x in outputs]).mean()
        avg_f1 = (2 * avg_prec * avg_rec) / (avg_prec + avg_rec)
        print(
            f"\nVal Loss: {avg_loss:.5f} Val Precision: {avg_prec:.4f}",
            f"Val Recall: {avg_rec:.4f} Val F1 {avg_f1:.4f}",
        )

    def test_step(self, batch, batch_idx):
        img, bboxes, target, name = batch

        # Pass through model
        pred = self(img, bboxes)

        # Calculate loss and metrics
        loss = self.criterion(pred, target)
        metrics = self.val_metrics(pred, target.long())

        return OrderedDict(
            {
                "test_loss": loss,
                "test_prec": metrics["Precision"],
                "test_rec": metrics["Recall"],
                "test_f1": metrics["F1"],
                "test_sample": torch.stack([img[0], target[0], pred[0]])
                if batch_idx % 100 == 0
                else None,
                "test_name": name[0],
            }
        )

    def test_epoch_end(self, outputs):
        # Save outputs
        sample_dir = os.path.join(self.hparams.output_path, "test_results")
        if not os.path.exists(sample_dir):
            os.makedirs(sample_dir)
        for out in outputs:
            if out["test_sample"] is not None:
                grid = make_grid(out["test_sample"], nrow=3, normalize=True)
                save_image(grid, os.path.join(sample_dir, out["test_name"]))

        # Print results
        avg_loss = torch.stack([x["test_loss"] for x in outputs]).mean()
        avg_prec = torch.stack([x["test_prec"] for x in outputs]).mean()
        avg_rec = torch.stack([x["test_rec"] for x in outputs]).mean()
        avg_f1 = (2 * avg_prec * avg_rec) / (avg_prec + avg_rec)
        print(
            f"\nTest Loss: {avg_loss:.5f} Test Precision: {avg_prec:.4f}",
            f"Test Recall: {avg_rec:.4f} Test F1 {avg_f1:.4f}",
        )

    def configure_optimizers(self):
        optimizer = optim.Adam(
            self.parameters(),
            lr=self.lr,
            betas=(self.hparams.beta1, self.hparams.beta2),
        )
        if self.hparams.schedule != "none":
            scheduler = get_scheduler(optimizer, self.lr, self.hparams)
            return [optimizer], [scheduler]
        else:
            print(f"Using no LR schedule lr={self.lr}")
            return [optimizer]


class EdgeCompleteModel(pl.LightningModule):
    def __init__(self, hparams):
        super(EdgeCompleteModel, self).__init__()
        self.save_hyperparameters()

        # Hyperparameters
        self.hparams = hparams
        self.lr = hparams.lr

        # Modules
        self.mask_refine_net = MaskRefineNet(hparams.mask_refine_weights)
        self.mask_refine_net.freeze()
        self.net = load_network(hparams)
        self.tv_loss = TverskyLoss(hparams.tversky_alpha, hparams.tversky_beta)

        # Metrics
        self.train_metrics = MetricCollection([Precision(), Recall(), F1(2)])
        self.val_metrics = MetricCollection([Precision(), Recall(), F1(2)])

    def forward(self, img, mask, edge_input):
        return self.net(torch.cat((img, mask, edge_input), 1))

    def training_step(self, batch, batch_idx):
        img, bboxes, edge_input, edge_real, _ = batch

        # Mask out edges from text in image
        mask = self.mask_refine_net(img, bboxes).round()
        mask = nn.MaxPool2d(3, stride=1, padding=1)(mask)  # Dilate mask
        edge_input = edge_input * (1 - mask)

        # Pass through model
        edge_fake = self(img, mask, edge_input)

        # Tversky loss
        loss = self.tv_loss(edge_fake * mask, edge_real * mask)

        # Calculate metrics
        metrics = self.train_metrics(edge_real * mask, (edge_fake * mask).long())

        self.log("train_loss", loss)
        self.log("train_prec", metrics["Precision"])
        self.log("train_rec", metrics["Recall"])
        self.log("train_f1", metrics["F1"])
        self.log("lr", self.trainer.optimizers[0].param_groups[0]["lr"], prog_bar=True)

        return loss

    def validation_step(self, batch, batch_idx):
        img, bboxes, edge_input, edge_real, _ = batch

        # Mask out edges from text in image
        mask = self.mask_refine_net(img, bboxes).round()
        mask = nn.MaxPool2d(3, stride=1, padding=1)(mask)  # Dilate mask
        edge_input = edge_input * (1 - mask)

        # Pass through model
        edge_fake = self(img, mask, edge_input)

        # Calculate loss
        loss_tv = self.tv_loss(edge_fake * mask, edge_real * mask)

        # Calculate metrics
        metrics = self.val_metrics(edge_real * mask, (edge_fake * mask).long())

        self.log("val_loss", loss_tv)
        self.log("val_prec", metrics["Precision"], on_epoch=True)
        self.log("val_rec", metrics["Recall"], on_epoch=True)
        self.log("val_f1", metrics["F1"], on_epoch=True)

        return OrderedDict(
            {
                "val_loss": loss_tv,
                "val_prec": metrics["Precision"],
                "val_rec": metrics["Recall"],
                "val_f1": metrics["F1"],
                "val_sample": torch.stack(
                    [
                        img[0],
                        edge_real[0],
                        edge_input[0],
                        edge_input[0] * (1 - mask[0]) + edge_fake[0] * mask[0],
                    ]
                )
                if batch_idx % (250 // self.hparams.batch_size) == 0
                else None,  # Save some samples for visualization
            }
        )

    def validation_epoch_end(self, outputs):
        # Visualize validation samples
        imgs = torch.cat(
            [x["val_sample"] for x in outputs if x["val_sample"] is not None], 0
        )
        grid = make_grid(imgs, nrow=4, normalize=True)

        # Log to tensorboard
        tensorboard = self.logger.experiment[0]
        tensorboard.add_image("val_samples", grid, self.current_epoch)

        # Print validation results
        avg_loss = torch.stack([x["val_loss"] for x in outputs]).mean()
        avg_prec = torch.stack([x["val_prec"] for x in outputs]).mean()
        avg_rec = torch.stack([x["val_rec"] for x in outputs]).mean()
        avg_f1 = (2 * avg_prec * avg_rec) / (avg_prec + avg_rec)
        print(
            f"\nVal Tversky Loss: {avg_loss:.5f} Val Precision: {avg_prec:.4f}",
            f"Val Recall: {avg_rec:.4f} Val F1 {avg_f1:.4f}",
        )

    def test_step(self, batch, batch_idx):
        img, bboxes, edge_input, edge_target, _ = batch

        # Mask out edges from text in image
        mask = self.mask_refine_net(img, bboxes).round()
        mask = nn.MaxPool2d(3, stride=1, padding=1)(mask)  # Dilate mask
        edge_input = edge_input * (1 - mask)

        # Pass through model
        edge_fake = self(img, mask, edge_input)
        edge_real = edge_target

        # Calculate loss and metrics
        loss_tv = self.tv_loss(edge_fake, edge_real)

        # Calculate metrics
        metrics = self.val_metrics(edge_real * mask, (edge_fake * mask).long())

        return OrderedDict(
            {
                "test_loss": loss_tv,
                "test_prec": metrics["Precision"],
                "test_rec": metrics["Recall"],
                "test_f1": metrics["F1"],
                "test_sample": torch.stack([img[0], edge_real[0], edge_fake[0]])
                if batch_idx % 100 == 0
                else None,
                "test_name": name[0],
            }
        )

    def test_epoch_end(self, outputs):
        # Save outputs
        sample_dir = os.path.join(self.hparams.output_path, "test_results")
        if not os.path.exists(sample_dir):
            os.makedirs(sample_dir)
        for out in outputs:
            if out["test_sample"] is not None:
                grid = make_grid(out["test_sample"], nrow=3, normalize=True)
                save_image(grid, os.path.join(sample_dir, out["test_name"]))

        # Print results
        avg_loss = torch.stack([x["test_loss"] for x in outputs]).mean()
        avg_prec = torch.stack([x["test_prec"] for x in outputs]).mean()
        avg_rec = torch.stack([x["test_rec"] for x in outputs]).mean()
        avg_f1 = (2 * avg_prec * avg_rec) / (avg_prec + avg_rec)
        print(
            f"\nTest Loss: {avg_loss:.5f} Test Precision: {avg_prec:.4f}",
            f"Test Recall: {avg_rec:.4f} Test F1 {avg_f1:.4f}",
        )

    def configure_optimizers(self):
        optimizer = optim.Adam(
            self.parameters(),
            lr=self.lr,
            betas=(self.hparams.beta1, self.hparams.beta2),
        )

        if self.hparams.schedule != "none":
            scheduler = get_scheduler(optimizer, self.lr, self.hparams)
            return [optimizer], [scheduler]
        else:
            print(f"Using no LR schedule lr={self.lr}")
            return [optimizer]


class TextRemovalModel(pl.LightningModule):
    def __init__(self, hparams):
        super(TextRemovalModel, self).__init__()
        self.save_hyperparameters()

        # Hyperparameters
        self.hparams = hparams
        self.lr = hparams.lr

        # Networks
        self.mask_refine_net = MaskRefineNet(hparams.mask_refine_weights)
        self.mask_refine_net.freeze()
        self.edge_complete_net = EdgeCompleteNet(hparams.edge_complete_weights)
        self.edge_complete_net.freeze()
        self.net = load_network(hparams)

        # Losses
        self.l1_loss = nn.L1Loss()
        self.gradient_loss = GradientLoss(type=hparams.gradient_loss_type)
        self.perceptual_loss = PerceptualLoss()

        # Metrics
        self.train_psnr = PSNR()
        self.val_metrics = MetricCollection([PSNR(),  MeanAbsoluteError()])

    def forward(self, img, mask, edge):
        return self.net(torch.cat((img, mask, edge), 1))

    def training_step(self, batch, batch_idx):
        img, img_target, bboxes, edge_input, _ = batch

        # Generate text mask
        mask = self.mask_refine_net(img, bboxes).round()
        mask = nn.MaxPool2d(3, stride=1, padding=1)(mask)  # Dilate mask

        # Complete masked out edges
        edge_input = edge_input * (1 - mask)
        completed_edges = self.edge_complete_net(img, mask, edge_input)
        edge = completed_edges * mask + edge_input * (1 - mask)

        # Inpaint
        img_masked = img * (1 - mask)
        img_pred = self(img_masked, mask, edge)
        img_pred = img_pred * mask + img_masked

        # Calculate losses
        loss_l1 = self.l1_loss(img_pred, img_target) / torch.mean(mask)
        loss_grad = self.gradient_loss(img_pred, img_target) / torch.mean(mask)
        loss_perceptual = self.perceptual_loss(img_pred, img_target)

        # Combine losses
        loss = (
            loss_l1 * self.hparams.l1_loss_weight
            + (loss_perceptual * self.hparams.percept_loss_weight)
            + (loss_grad * self.hparams.gradient_loss_weight)
        )

        # Calculate metrics
        psnr = self.train_psnr(img_pred, img_target)

        self.log("train_loss", loss)
        self.log("train_loss_l1", loss_l1, prog_bar=True)
        self.log("train_loss_perceptual", loss_perceptual, prog_bar=True)
        self.log("train_loss_grad", loss_grad, prog_bar=True)
        self.log("train_psnr", psnr)
        self.log("lr", self.trainer.optimizers[0].param_groups[0]["lr"], prog_bar=True)

        return loss

    def validation_step(self, batch, batch_idx):
        img, img_target, bboxes, edge_input, _ = batch

        # Generate text mask
        mask = self.mask_refine_net(img, bboxes).round()
        mask = nn.MaxPool2d(3, stride=1, padding=1)(mask)  # Dilate mask

        # Complete masked out edges
        edge_input = edge_input * (1 - mask)
        completed_edges = self.edge_complete_net(img, mask, edge_input)
        edge = completed_edges * mask + edge_input * (1 - mask)

        # Inpaint
        img_masked = img * (1 - mask)
        img_pred = self(img_masked, mask, edge)
        img_pred = img_pred * mask + img * (1 - mask)

        # Calculate metrics
        metrics = self.val_metrics(img_pred, img_target)

        self.log("val_psnr", metrics["PSNR"], on_epoch=True)
        #self.log("val_ssim", metrics["SSIM"], on_epoch=True)
        self.log("val_mae", metrics["MeanAbsoluteError"], on_epoch=True)

        return OrderedDict(
            {
                "val_psnr": metrics["PSNR"],
                #"val_ssim": metrics["SSIM"],
                "val_mae": metrics["MeanAbsoluteError"],
                "val_sample": torch.stack(
                    [
                        img[0],
                        img_target[0],
                        img_masked[0],
                        img_pred[0],
                    ]
                )
                if batch_idx % (250 // self.hparams.batch_size) == 0
                else None,  # Save some samples for visualization
            }
        )

    def validation_epoch_end(self, outputs):
        # Visualize validation samples
        imgs = torch.cat(
            [x["val_sample"] for x in outputs if x["val_sample"] is not None], 0
        )
        grid = make_grid(imgs, nrow=4, normalize=True)

        # Log to tensorboard
        tensorboard = self.logger.experiment[0]
        tensorboard.add_image("val_samples", grid, self.current_epoch)

        # Print validation results
        avg_psnr = torch.stack([x["val_psnr"] for x in outputs]).mean()
        #avg_ssim = torch.stack([x["val_ssim"] for x in outputs]).mean()
        avg_mae = torch.stack([x["val_mae"] for x in outputs]).mean()
        print(
            f"\nVal PSNR: {avg_psnr:.4f} Val MAE: {avg_mae:.4f}",
        )

    def test_step(self, batch, batch_idx):
        img, img_target, bboxes, edge_input, name = batch

        # Generate text mask
        mask = self.mask_refine_net(img, bboxes).round()
        mask = nn.MaxPool2d(3, stride=1, padding=1)(mask)  # Dilate mask

        # Complete masked out edges
        edge_input = edge_input * (1 - mask)
        completed_edges = self.edge_complete_net(img, mask, edge_input)
        edge = completed_edges * mask + edge_input * (1 - mask)

        # Inpaint
        img_masked = img * (1 - mask)
        img_pred = self(img_masked, mask, edge)
        img_pred = img_pred * mask + img * (1 - mask)

        # Calculate metrics
        metrics = self.val_metrics(img_pred, img_target)

        return OrderedDict(
            {
                "test_psnr": metrics["PSNR"],
                #"test_ssim": metrics["SSIM"],
                "test_mae": metrics["MeanAbsoluteError"],
                "test_sample": torch.stack(
                    [
                        img[0],
                        img_target[0],
                        img_masked[0],
                        img_pred[0],
                    ]
                )
                if batch_idx % (100 // self.hparams.batch_size) == 0
                else None,  # Save some samples for visualization
                "test_name": name[0],
            }
        )

    def test_epoch_end(self, outputs):
        # Save outputs
        sample_dir = os.path.join(self.hparams.output_path, "test_results")
        if not os.path.exists(sample_dir):
            os.makedirs(sample_dir)
        for out in outputs:
            if out["test_sample"] is not None:
                grid = make_grid(out["test_sample"], nrow=4, normalize=True)
                save_image(grid, os.path.join(sample_dir, out["test_name"]))

        # Print results
        avg_psnr = torch.stack([x["test_psnr"] for x in outputs]).mean()
        #avg_ssim = torch.stack([x["test_ssim"] for x in outputs]).mean()
        avg_mae = torch.stack([x["test_mae"] for x in outputs]).mean()
        print(
            f"\nTest PSNR: {avg_psnr:.4f}  Test MAE: {avg_mae:.4f}",
        )

    def configure_optimizers(self):
        optimizer = optim.Adam(
            self.net.parameters(),
            lr=self.lr,
            betas=(self.hparams.beta1, self.hparams.beta2),
        )

        if self.hparams.schedule != "none":
            scheduler = get_scheduler(optimizer, self.lr, self.hparams)
            return [optimizer], [scheduler]
        else:
            print(f"Using no LR schedule lr={self.lr}")
            return [optimizer]

    """
    def on_save_checkpoint(self, checkpoint):
        new_state_dict = OrderedDict()
        for key, val in checkpoint['state_dict'].items():
          if not key.startswith('perceptual_loss'):
            new_state_dict[key] = val
        checkpoint['state_dict'] = new_state_dict
    """
