# -*- coding: utf-8 -*-
# Author: ximing
# Description: LIVE painter and optimizer
# Copyright (c) 2023, XiMing Xing.
# License: MIT License

import copy
import random

import omegaconf
from omegaconf import DictConfig

import cv2
import numpy as np
import pydiffvg
import torch
from torch.optim.lr_scheduler import LambdaLR

from pytorch_svgrender.diffvg_warp import DiffVGState


class Painter(DiffVGState):

    def __init__(
            self,
            target_img: torch.Tensor,
            diffvg_cfg: omegaconf.DictConfig,
            num_segments: int = 4,
            segment_init: str = 'random',
            radius: int = 5,
            canvas_size=240,
            trainable_bg: bool = False,
            stroke: bool = False,
            stroke_width: int = 3,
            device: torch.device = None,
    ):
        super(Painter, self).__init__(device, print_timing=diffvg_cfg.print_timing,
                                      canvas_width=canvas_size, canvas_height=canvas_size)

        self.target_img = target_img

        self.num_segments = num_segments
        self.segment_init = segment_init
        self.radius = radius
        self.train_stroke = stroke
        self.stroke_width = stroke_width

        self.points_vars = []
        self.stroke_width_vars = []
        self.stroke_color_vars = []
        self.color_vars = []

        self.strokes_counter = 0  # counts the number of calls to "get_path"

        # Background
        self.para_bg = torch.tensor([1., 1., 1.], requires_grad=trainable_bg, device=self.device)

        self.pos_init_method = None

    def component_wise_path_init(self, pred, init_type: str = 'sparse'):
        assert self.target_img is not None  # gt

        if init_type == 'random':
            self.pos_init_method = RandomCoordInit(self.canvas_height, self.canvas_width)
        elif init_type == 'sparse':
            # when initialized for the first time, the render result is None
            if pred is None:
                pred = self.para_bg.view(1, -1, 1, 1).repeat(1, 1, self.canvas_height, self.canvas_width)
            # then pred is the render result
            self.pos_init_method = SparseCoordInit(pred, self.target_img)
        elif init_type == 'naive':
            if pred is None:
                pred = self.para_bg.view(1, -1, 1, 1).repeat(1, 1, self.canvas_height, self.canvas_width)
            self.pos_init_method = NaiveCoordInit(pred, self.target_img)
        else:
            raise NotImplementedError(f"'{init_type}' is not support.")

    def init_image(self, num_paths=0):
        self.cur_shapes, self.cur_shape_groups = [], []

        for i in range(num_paths):
            path, color_ref = self.get_path()
            self.shapes.append(path)
            self.cur_shapes.append(path)

            wref, href = color_ref
            wref = max(0, min(int(wref), self.canvas_width - 1))
            href = max(0, min(int(href), self.canvas_height - 1))
            fill_color_init = list(self.target_img[0, :, href, wref]) + [1.]
            fill_color_init = torch.FloatTensor(fill_color_init)
            stroke_color_init = torch.FloatTensor(np.random.uniform(size=[4]))
            path_group = pydiffvg.ShapeGroup(
                shape_ids=torch.tensor([len(self.shapes) - 1]),
                fill_color=None if self.train_stroke else fill_color_init,
                stroke_color=stroke_color_init if self.train_stroke else None
            )
            self.shape_groups.append(path_group)
            self.cur_shape_groups.append(path_group)

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

    def get_image(self, step: int = 0):
        img = self.render_warp(seed=step)
        img = img[:, :, 3:4] * img[:, :, :3] + self.para_bg * (1 - img[:, :, 3:4])
        img = img.unsqueeze(0)  # convert img from HWC to NCHW
        img = img.permute(0, 3, 1, 2).to(self.device)  # NHWC -> NCHW
        return img

    def get_path(self):
        num_segments = self.num_segments
        num_control_points = [2] * num_segments

        points = []
        # init segment
        if self.segment_init == 'circle':
            radius = self.radius if self.radius is not None else np.random.uniform(0.5, 1)
            if self.pos_init_method is not None:
                center = self.pos_init_method()
            else:
                center = (random.random(), random.random())
            bias = center
            color_ref = copy.deepcopy(bias)

            avg_degree = 360 / (num_segments * 3)
            for i in range(0, num_segments * 3):
                point = (
                    np.cos(np.deg2rad(i * avg_degree)), np.sin(np.deg2rad(i * avg_degree))
                )
                points.append(point)

            points = torch.FloatTensor(points) * radius + torch.FloatTensor(bias).unsqueeze(dim=0)
        else:  # 'random' init
            p0 = self.pos_init_method()
            color_ref = copy.deepcopy(p0)
            points.append(p0)
            for j in range(num_segments):
                radius = self.radius
                p1 = (p0[0] + radius * np.random.uniform(-0.5, 0.5),
                      p0[1] + radius * np.random.uniform(-0.5, 0.5))
                p2 = (p1[0] + radius * np.random.uniform(-0.5, 0.5),
                      p1[1] + radius * np.random.uniform(-0.5, 0.5))
                p3 = (p2[0] + radius * np.random.uniform(-0.5, 0.5),
                      p2[1] + radius * np.random.uniform(-0.5, 0.5))
                points.append(p1)
                points.append(p2)
                if j < num_segments - 1:
                    points.append(p3)
                    p0 = p3
            points = torch.FloatTensor(points)

        path = pydiffvg.Path(
            num_control_points=torch.LongTensor(num_control_points),
            points=points,
            stroke_width=torch.tensor(float(self.stroke_width)) if self.train_stroke else torch.tensor(0.0),
            is_closed=True
        )

        self.strokes_counter += 1
        return path, color_ref

    def clip_curve_shape(self):
        for group in self.shape_groups:
            if self.train_stroke:
                group.stroke_color.data.clamp_(0.0, 1.0)
            else:
                group.fill_color.data.clamp_(0.0, 1.0)

    def calc_distance_weight(self, loss_weight_keep):
        shapes_forsdf = copy.deepcopy(self.cur_shapes)
        shape_groups_forsdf = copy.deepcopy(self.cur_shape_groups)
        for si in shapes_forsdf:
            si.stroke_width = torch.FloatTensor([0]).to(self.device)
        for sg_idx, sgi in enumerate(shape_groups_forsdf):
            sgi.fill_color = torch.FloatTensor([1, 1, 1, 1]).to(self.device)
            sgi.shape_ids = torch.LongTensor([sg_idx]).to(self.device)

        sargs_forsdf = pydiffvg.RenderFunction.serialize_scene(
            self.canvas_width, self.canvas_height, shapes_forsdf, shape_groups_forsdf
        )
        _render = pydiffvg.RenderFunction.apply
        with torch.no_grad():
            im_forsdf = _render(self.canvas_width,  # width
                                self.canvas_height,  # height
                                2,  # num_samples_x
                                2,  # num_samples_y
                                0,  # seed
                                None,
                                *sargs_forsdf)

        # use alpha channel is a trick to get 0-1 image
        im_forsdf = (im_forsdf[:, :, 3]).detach().cpu().numpy()
        loss_weight = get_sdf(im_forsdf, normalize='to1')
        loss_weight += loss_weight_keep
        loss_weight = np.clip(loss_weight, 0, 1)
        loss_weight = torch.FloatTensor(loss_weight).to(self.device)
        return loss_weight

    def set_parameters(self):
        # stroke`s location optimization
        self.points_vars = []
        for i, path in enumerate(self.cur_shapes):
            path.points.requires_grad = True
            self.points_vars.append(path.points)

            if self.train_stroke:
                path.stroke_width.requires_grad = True
                self.stroke_width_vars.append(path.stroke_width)

        # for stroke' color optimization
        self.color_vars = []
        for i, group in enumerate(self.cur_shape_groups):
            if self.train_stroke:
                group.stroke_color.requires_grad = True
                self.stroke_color_vars.append(group.stroke_color)
            else:
                group.fill_color.requires_grad = True
                self.color_vars.append(group.fill_color)

    def get_point_parameters(self):
        return self.points_vars

    def get_color_parameters(self):
        return self.color_vars

    def get_stroke_parameters(self):
        return self.stroke_width_vars, self.stroke_color_vars

    def get_bg_parameters(self):
        return self.para_bg

    def save_svg(self, fpath):
        pydiffvg.save_svg(f'{fpath}',
                          self.canvas_width,
                          self.canvas_height,
                          self.shapes,
                          self.shape_groups)


def get_sdf(phi, **kwargs):
    import skfmm  # local import

    phi = (phi - 0.5) * 2
    if (phi.max() <= 0) or (phi.min() >= 0):
        return np.zeros(phi.shape).astype(np.float32)
    sd = skfmm.distance(phi, dx=1)

    flip_negative = kwargs.get('flip_negative', True)
    if flip_negative:
        sd = np.abs(sd)

    truncate = kwargs.get('truncate', 10)
    sd = np.clip(sd, -truncate, truncate)
    # print(f"max sd value is: {sd.max()}")

    zero2max = kwargs.get('zero2max', True)
    if zero2max and flip_negative:
        sd = sd.max() - sd
    elif zero2max:
        raise ValueError

    normalize = kwargs.get('normalize', 'sum')
    if normalize == 'sum':
        sd /= sd.sum()
    elif normalize == 'to1':
        sd /= sd.max()
    return sd


class SparseCoordInit:

    def __init__(self, pred, gt, format='[bs x c x 2D]', quantile_interval=200, nodiff_thres=0.1):
        if torch.is_tensor(pred):
            pred = pred.detach().cpu().numpy()
        if torch.is_tensor(gt):
            gt = gt.detach().cpu().numpy()

        if format == '[bs x c x 2D]':
            self.map = ((pred[0] - gt[0]) ** 2).sum(0)
            self.reference_gt = copy.deepcopy(np.transpose(gt[0], (1, 2, 0)))
        elif format == ['[2D x c]']:
            self.map = (np.abs(pred - gt)).sum(-1)
            self.reference_gt = copy.deepcopy(gt[0])
        else:
            raise ValueError

        # OptionA: Zero too small errors to avoid the error too small deadloop
        self.map[self.map < nodiff_thres] = 0
        quantile_interval = np.linspace(0., 1., quantile_interval)
        quantized_interval = np.quantile(self.map, quantile_interval)
        # remove redundant
        quantized_interval = np.unique(quantized_interval)
        quantized_interval = sorted(quantized_interval[1:-1])
        self.map = np.digitize(self.map, quantized_interval, right=False)
        self.map = np.clip(self.map, 0, 255).astype(np.uint8)
        self.idcnt = {}
        for idi in sorted(np.unique(self.map)):
            self.idcnt[idi] = (self.map == idi).sum()
        # remove smallest one to remove the correct region
        self.idcnt.pop(min(self.idcnt.keys()))

    def __call__(self):
        if len(self.idcnt) == 0:
            h, w = self.map.shape
            return [np.random.uniform(0, 1) * w, np.random.uniform(0, 1) * h]

        target_id = max(self.idcnt, key=self.idcnt.get)
        _, component, cstats, ccenter = cv2.connectedComponentsWithStats(
            (self.map == target_id).astype(np.uint8),
            connectivity=4
        )
        # remove cid = 0, it is the invalid area
        csize = [ci[-1] for ci in cstats[1:]]
        target_cid = csize.index(max(csize)) + 1
        center = ccenter[target_cid][::-1]
        coord = np.stack(np.where(component == target_cid)).T
        dist = np.linalg.norm(coord - center, axis=1)
        target_coord_id = np.argmin(dist)
        coord_h, coord_w = coord[target_coord_id]

        # replace_sampling
        self.idcnt[target_id] -= max(csize)
        if self.idcnt[target_id] == 0:
            self.idcnt.pop(target_id)
        self.map[component == target_cid] = 0
        return [coord_w, coord_h]


class RandomCoordInit:
    def __init__(self, canvas_width, canvas_height):
        self.canvas_width, self.canvas_height = canvas_width, canvas_height

    def __call__(self):
        w, h = self.canvas_width, self.canvas_height
        return [np.random.uniform(0, 1) * w, np.random.uniform(0, 1) * h]


class NaiveCoordInit:
    def __init__(self, pred, gt, format='[bs x c x 2D]', replace_sampling=True):
        if isinstance(pred, torch.Tensor):
            pred = pred.detach().cpu().numpy()
        if isinstance(gt, torch.Tensor):
            gt = gt.detach().cpu().numpy()

        if format == '[bs x c x 2D]':
            self.map = ((pred[0] - gt[0]) ** 2).sum(0)
        elif format == ['[2D x c]']:
            self.map = ((pred - gt) ** 2).sum(-1)
        else:
            raise ValueError
        self.replace_sampling = replace_sampling

    def __call__(self):
        coord = np.where(self.map == self.map.max())
        coord_h, coord_w = coord[0][0], coord[1][0]
        if self.replace_sampling:
            self.map[coord_h, coord_w] = -1
        return [coord_w, coord_h]


class LinearDecayLR:

    def __init__(self, decay_every, decay_ratio):
        self.decay_every = decay_every
        self.decay_ratio = decay_ratio

    def __call__(self, n):
        decay_time = n // self.decay_every
        decay_step = n % self.decay_every
        lr_s = self.decay_ratio ** decay_time
        lr_e = self.decay_ratio ** (decay_time + 1)
        r = decay_step / self.decay_every
        lr = lr_s * (1 - r) + lr_e * r
        return lr


class PainterOptimizer:

    def __init__(self,
                 renderer: Painter,
                 num_iter: int,
                 lr_config: DictConfig,
                 trainable_stroke: bool = False,
                 trainable_bg: bool = False):
        self.renderer = renderer
        self.num_iter = num_iter
        self.trainable_stroke = trainable_stroke
        self.trainable_bg = trainable_bg

        self.lr_base = {
            'point': lr_config.point,
            'color': lr_config.color,
            'stroke_width': lr_config.stroke_width,
            'stroke_color': lr_config.stroke_color,
            'bg': lr_config.bg
        }

        self.learnable_params = []  # list[Dict]

        self.optimizer = None
        self.scheduler = None

    def init_optimizers(self):
        # optimizers
        params = {}
        self.renderer.set_parameters()
        params['point'] = self.renderer.get_point_parameters()
        if self.trainable_stroke:
            params['stroke_width'], params['stroke_color'] = self.renderer.get_stroke_parameters()
        else:
            params['color'] = self.renderer.get_color_parameters()

        if self.trainable_bg:
            params['bg'] = self.renderer.get_bg_parameters()

        self.learnable_params = [
            {'params': params[ki], 'lr': self.lr_base[ki]} for ki in sorted(params.keys())
        ]
        self.optimizer = torch.optim.Adam(self.learnable_params)
        # lr schedule
        lr_lambda_fn = LinearDecayLR(self.num_iter, 0.4)
        self.scheduler = LambdaLR(self.optimizer, lr_lambda=lr_lambda_fn, last_epoch=-1)

    def update_params(self, name: str, value: torch.tensor):
        for param_group in self.learnable_params:
            if param_group.get('_id') == name:
                param_group['params'] = value

    def update_lr(self):
        self.scheduler.step()

    def zero_grad_(self):
        self.optimizer.zero_grad()

    def step_(self):
        self.optimizer.step()

    def get_lr(self):
        return self.optimizer.param_groups[0]['lr']