# -*- coding: utf-8 -*-
# Copyright (c) XiMing Xing. All rights reserved.
# Author: XiMing Xing
# Description: SVG Painter and ist optimizer

from typing import Tuple

import omegaconf
import pydiffvg
import torch
import numpy as np

from pytorch_svgrender.diffvg_warp import DiffVGState
from pytorch_svgrender.utils import get_rgb_from_color


class Painter(DiffVGState):

    def __init__(self, device=None):
        super().__init__(device)
        self.device = device

        self.strokes_counter = 0  # num of paths

    def init_shapes(self, path_svg, reinit_cfg: omegaconf.DictConfig = None):
        print(f"-> init svg from `{path_svg}` ...")
        self.canvas_width, self.canvas_height, self.shapes, self.shape_groups = self.load_svg(path_svg)
        self.strokes_counter = len(self.shapes)

        """re-init font color"""
        if reinit_cfg is not None:
            self.color_init(reinit_cfg)

        img = self.render_warp()
        img = img[:, :, 3:4] * img[:, :, :3] + \
              torch.ones(img.shape[0], img.shape[1], 3, device=self.device) * (1 - img[:, :, 3:4])
        img = img[:, :, :3]
        # Convert img from HWC to NCHW
        img = img.unsqueeze(0)
        img = img.permute(0, 3, 1, 2).to(self.device)  # NHWC -> NCHW
        return img

    def color_init(self, reinit_cfg: omegaconf.DictConfig):
        if not reinit_cfg.reinit:
            return

        if reinit_cfg.reinit_color == 'randn':
            for i, group in enumerate(self.shape_groups):
                color_val = np.random.random(size=3).tolist() + [1.0]
                group.fill_color = torch.FloatTensor(color_val)
        elif reinit_cfg.reinit_color == 'randn_all':
            color_val = np.random.random(size=3).tolist() + [1.0]
            for i, group in enumerate(self.shape_groups):
                group.fill_color = torch.FloatTensor(color_val)
        else:
            rgb = get_rgb_from_color(str(reinit_cfg.reinit_color))
            color_val = list(rgb) + [1.0]
            for i, group in enumerate(self.shape_groups):
                group.fill_color = torch.FloatTensor(color_val)

    def clip_curve_shape(self):
        for group in self.shape_groups:
            group.fill_color.data.clamp_(0.0, 1.0)
            # force opacity
            group.fill_color.data[-1] = 1.0

    def get_image(self):
        img = self.render_warp()
        opacity = img[:, :, 3:4]
        img = opacity * img[:, :, :3] + torch.ones(img.shape[0], img.shape[1], 3, device=self.device) * (1 - opacity)
        img = img[:, :, :3]
        # Convert img from HWC to NCHW
        img = img.unsqueeze(0)
        img = img.permute(0, 3, 1, 2).to(self.device)  # NHWC -> NCHW
        return img

    def set_parameters(self):
        self.point_vars = []
        # the strokes point optimization
        for i, path in enumerate(self.shapes):
            path.points.requires_grad = True
            self.point_vars.append(path.points)

        # the strokes color optimization
        self.color_vars = []
        for i, group in enumerate(self.shape_groups):
            if group.fill_color is not None:
                group.fill_color.requires_grad = True
                self.color_vars.append(group.fill_color)
            if group.stroke_color is not None:
                group.stroke_color.requires_grad = True
                self.color_vars.append(group.stroke_color)

    def get_point_parameters(self):
        return self.point_vars

    def get_color_parameters(self):
        return self.color_vars

    def pretty_save_svg(self, filename, width=None, height=None, shapes=None, shape_groups=None):
        width = self.canvas_width if width is None else width
        height = self.canvas_height if height is None else height
        shapes = self.shapes if shapes is None else shapes
        shape_groups = self.shape_groups if shape_groups is None else shape_groups

        self.save_svg(filename, width, height, shapes, shape_groups, use_gamma=False, background=None)

    def load_svg(self, path_svg):
        canvas_width, canvas_height, shapes, shape_groups = pydiffvg.svg_to_scene(path_svg)
        return canvas_width, canvas_height, shapes, shape_groups


class PainterOptimizer:

    def __init__(self, renderer: Painter, lr_cfg: omegaconf.DictConfig):
        self.renderer = renderer
        self.point_lr = lr_cfg.point
        self.color_lr = lr_cfg.color
        self.point_optimizer, self.color_optimizer = None, None

    def init_optimizers(self):
        self.renderer.set_parameters()
        self.point_optimizer = torch.optim.Adam([
            {'params': self.renderer.get_point_parameters(), 'lr': self.point_lr}])
        self.color_optimizer = torch.optim.Adam([
            {'params': self.renderer.get_color_parameters(), 'lr': self.color_lr}])

    def update_lr(self, step):
        pass

    def zero_grad_(self):
        self.point_optimizer.zero_grad()
        self.color_optimizer.zero_grad()

    def step_(self):
        self.point_optimizer.step()
        self.color_optimizer.step()

    def get_lr(self) -> Tuple[float, float]:
        return self.point_optimizer.param_groups[0]['lr'], self.color_optimizer.param_groups[0]['lr']