import os

import numpy as np
import torch
from matplotlib import pyplot as plt

from libs.vision_libs.utils import draw_bounding_boxes
from models.wirenet.postprocess import postprocess
from torchvision import transforms


def save_latest_model(model, save_path, epoch, optimizer=None):
    os.makedirs(os.path.dirname(save_path), exist_ok=True)

    checkpoint = {
        'epoch': epoch,
        'model_state_dict': model.state_dict(),
    }

    if optimizer is not None:
        checkpoint['optimizer_state_dict'] = optimizer.state_dict()

    torch.save(checkpoint, save_path)

def save_best_model(model, save_path, epoch, current_loss, best_loss, optimizer=None):
    os.makedirs(os.path.dirname(save_path), exist_ok=True)

    if current_loss < best_loss:
        checkpoint = {
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'loss': current_loss
        }

        if optimizer is not None:
            checkpoint['optimizer_state_dict'] = optimizer.state_dict()

        torch.save(checkpoint, save_path)
        print(f"Saved best model at epoch {epoch} with loss {current_loss:.4f}")

        return current_loss

    return best_loss
def show_line(img, pred, epoch, writer):
    im = img.permute(1, 2, 0)
    writer.add_image("ori", im, epoch, dataformats="HWC")

    boxed_image = draw_bounding_boxes((img * 255).to(torch.uint8), pred[0]["boxes"],
                                      colors="yellow", width=1)
    writer.add_image("boxes", boxed_image.permute(1, 2, 0), epoch, dataformats="HWC")

    PLTOPTS = {"color": "#33FFFF", "s": 15, "edgecolors": "none", "zorder": 5}
    # print(f'pred[1]:{pred[1]}')
    H = pred[-1]['wires']
    lines = H["lines"][0].cpu().numpy() / 128 * im.shape[:2]
    scores = H["score"][0].cpu().numpy()
    for i in range(1, len(lines)):
        if (lines[i] == lines[0]).all():
            lines = lines[:i]
            scores = scores[:i]
            break

    # postprocess lines to remove overlapped lines
    diag = (im.shape[0] ** 2 + im.shape[1] ** 2) ** 0.5
    nlines, nscores = postprocess(lines, scores, diag * 0.01, 0, False)

    for i, t in enumerate([0.85]):
        plt.gca().set_axis_off()
        plt.subplots_adjust(top=1, bottom=0, right=1, left=0, hspace=0, wspace=0)
        plt.margins(0, 0)
        for (a, b), s in zip(nlines, nscores):
            if s < t:
                continue
            plt.plot([a[1], b[1]], [a[0], b[0]], c=c(s), linewidth=2, zorder=s)
            plt.scatter(a[1], a[0], **PLTOPTS)
            plt.scatter(b[1], b[0], **PLTOPTS)
        plt.gca().xaxis.set_major_locator(plt.NullLocator())
        plt.gca().yaxis.set_major_locator(plt.NullLocator())
        plt.imshow(im)
        plt.tight_layout()
        fig = plt.gcf()
        fig.canvas.draw()
        image_from_plot = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8).reshape(
            fig.canvas.get_width_height()[::-1] + (3,))
        plt.close()
        img2 = transforms.ToTensor()(image_from_plot)

        writer.add_image("output", img2, epoch)