from __future__ import annotations

from time import sleep

import pydiffvg
import torch
import random
from random import uniform

from gmtypes import GradientMesh, Quad, Patch, Point, join_quads


def quads():
    return [
        Quad.random(),
        Quad.random(),
        Quad.random(),
        Quad.random(),
    ]


def rand_quad_test(filename='random_quad.png', width=256, height=256,
                   degree=4, num_control_points=2):
    pydiffvg.set_use_gpu(torch.cuda.is_available())
    render = pydiffvg.RenderFunction.apply

    patch = Patch.random()

    shape_groups = [patch.as_shape_group()]
    shapes = [patch.as_path(width, height)]

    scene_args = pydiffvg.RenderFunction.serialize_scene(width, height,
                                                         shapes, shape_groups)

    img = render(width, height, 2, 2, 0, None, *scene_args)
    pydiffvg.imwrite(img.cpu(), f"test_data/{filename}", gamma=2.2)

    return img


def mult_quad_test(filename='multiple_quads.png', width=1024,
                   height=1024, num_control_points=None, mask=None, seed=None):
    random.seed(seed)
    mask = mask or [1, 1, 1, 1]
    pydiffvg.set_use_gpu(torch.cuda.is_available())
    render = pydiffvg.RenderFunction.apply

    a, b, c, d = quads()
    join_quads(a, b, c, d)

    to_render = [a, b, c, d]
    to_render = [x for x in to_render if mask[to_render.index(x)]]

    shape_groups = [patch.as_shape_group(color=(
        uniform(0, 1),
        uniform(0, 1),
        uniform(0, 1),
        0.8
    )) for patch in to_render]

    for i in range(len(to_render)):
        shape_groups[i].shape_ids = torch.tensor([i])
    shapes = [patch.as_path(width, height) for patch in to_render]

    scene_args = pydiffvg.RenderFunction.serialize_scene(width, height,
                                                         shapes, shape_groups)

    img = render(width, height, 2, 2, 0, None, *scene_args)
    pydiffvg.imwrite(img.cpu(), f"test_data/{filename}", gamma=2.2)
    return img.clone()


def om():
    filename = 'optimize_test.png'
    pydiffvg.set_use_gpu(torch.cuda.is_available())
    render = pydiffvg.RenderFunction.apply

    target = mult_quad_test(width=256, height=256)

    squad = quads()

    join_quads(*squad)

    gm = GradientMesh(*squad)

    points_n = []
    for s in squad:
        out = []
        for pt in s.points:
            out.append([pt.x, pt.y])
            for cpt in pt.controls:
                out.append([cpt.x, cpt.y])
        points_n.append(out)

    points_n = torch.tensor(points_n, requires_grad=True)
    color = torch.tensor([s.color for s in squad], requires_grad=True)

    paths = [s.as_path() for s in squad]
    path_groups = [pydiffvg.ShapeGroup(shape_ids=torch.tensor([i]),
                                       fill_color=torch.tensor(squad[i].color))
                   for i in range(len(squad))]
    scene_args = pydiffvg.RenderFunction.serialize_scene(
        256, 256, paths, path_groups
    )
    img = render(256,  # width
                 256,  # height
                 2,   # num_samples_x
                 2,   # num_samples_y
                 1,   # seed
                 None,
                 *scene_args)

    points, controls, color = [torch.tensor(x, requires_grad=True)
                               for x in gm.to_numbers()]

    optimizer = torch.optim.Adam([points, color, points_n], lr=1e-2)

    for t in range(180):
        print(f"iteration {t}")
        optimizer.zero_grad()

        points_n.data = torch.tensor(
            GradientMesh.from_path_points(points_n, color).to_path_points()
        )

        for i in range(len(paths)):
            paths[i].points = points_n[i] * 256

        for i in range(len(path_groups)):
            path_groups[i].fill_color = color[i]

        scene_args = pydiffvg.RenderFunction.serialize_scene(
            256, 256, paths, path_groups)

        img = render(256,   # width
                     256,   # height
                     2,     # num_samples_x
                     2,     # num_samples_y
                     t+1,   # seed
                     None,
                     *scene_args)

        pydiffvg.imwrite(img.cpu(),
                         f'test_data/test_curve/iter_{filename}_'
                         f'{str(t).zfill(5)}.png',
                         gamma=2.2)

        loss = (img - target).pow(2).sum()

        loss.backward()

        print(f'loss: {loss}')
        print(f'points.grad {points.grad}')
        print(f'color.grad {color.grad}')

        optimizer.step()


def slideshow(n=30, s=1, do_mask=False):
    mask = None
    for i in range(n):
        if do_mask:
            mask = [1] * 4
            print(i % n)
            mask[i % 4] = 0
            print(mask)

        mult_quad_test(mask=mask)
        sleep(s)


def get_mesh():
    a, b, c, d = quads()
    join_quads(a,b,c,d)

    gm = GradientMesh(a, b, c, d)
    return gm