"""Provides examples and visualization of several meshless domains."""

import matplotlib.pyplot as plt
import torch
from mpl_toolkits import axes_grid1

from scimba_torch.domain.meshless_domain import (
    ArcCircle2D,
    Cube3D,
    Disk2D,
    Disk3D,
    Polygon2D,
    Segment1D,
    Square2D,
)

# sampler for visu
from scimba_torch.integration.monte_carlo import (
    DomainSampler,
    SurfacicSampler,
    VolumetricSampler,
)
from scimba_torch.utils.mapping import Mapping


# little utility to make cute colorbar find here https://stackoverflow.com/questions/18195758/set-matplotlib-colorbar-size-to-match-graph
def add_colorbar(im, aspect=20, pad_fraction=0.5, **kwargs):
    """Add a vertical color bar to an image plot."""
    divider = axes_grid1.make_axes_locatable(im.axes)
    width = axes_grid1.axes_size.AxesY(im.axes, aspect=1.0 / aspect)
    pad = axes_grid1.axes_size.Fraction(pad_fraction, width)
    current_ax = plt.gca()
    cax = divider.append_axes("right", size=width, pad=pad)
    plt.sca(current_ax)
    return im.axes.figure.colorbar(im, cax=cax, **kwargs)


cmap = plt.get_cmap("tab20")


def visu_segment_1d(low_high=(1.0, 2.0)):
    domain = Segment1D(low_high=low_high, is_main_domain=False)
    sampler = VolumetricSampler(domain)

    bc_domains = domain.full_bc_domain()
    bc_samplers = [SurfacicSampler(bc) for bc in bc_domains]

    pts_bc = []
    normals_bc = []
    for bc_sampler in bc_samplers:
        x_bc, n = bc_sampler.sample(5, compute_normals=True)
        assert x_bc.shape == (5, 1)
        assert n.shape == (5, 1)

        pts_bc.append(x_bc)
        normals_bc.append(n)

    N = 1000
    pts = sampler.sample(N)
    assert pts.shape == (N, 1)

    x = torch.linspace(low_high[0] - 1.0, low_high[1] + 1.0, N)
    sdf = domain.sdf(x)
    assert sdf.shape == (N, 1)

    fig = plt.figure(figsize=(5, 5))

    ax = fig.add_subplot(111)
    ax.set_title("Segment1D")
    ax.set_xlabel("x")
    ax.set_ylabel("SDF")
    ax.grid()

    # plot sampled points
    ax.scatter(pts[:, 0], torch.zeros(N), label="sampled points")

    # plot pts/normals of points on boundaries
    for i, (x_bc, n) in enumerate(zip(pts_bc, normals_bc)):
        ax.scatter(x_bc[:, 0], torch.zeros_like(x_bc[:, 0]), color=cmap(i + 1), s=10)
        ax.quiver(
            x_bc[:, 0],
            torch.zeros_like(x_bc[:, 0]),
            n[:, 0],
            torch.zeros_like(x_bc[:, 0]),
            color=cmap(i + 1),
        )

    # plot sdf
    ax.plot(x, sdf, label="sdf")

    # plot the segment
    h = (low_high[1] - low_high[0]) / 50
    ax.vlines([*low_high], -h, h, color="k")
    ax.plot([*low_high], [0, 0], "k")

    ax.legend()

    fig.tight_layout()
    plt.show()


def visu_segment_1d_hole(low_high=(-2.0, 2.0), low_high_hole=(-1.0, 1.0)):
    domain = Segment1D(low_high=low_high, is_main_domain=True)
    domain.add_hole(Segment1D(low_high=low_high_hole))
    sampler = VolumetricSampler(domain)

    N = 1000
    pts = sampler.sample(N)
    assert pts.shape == (N, 1)

    x = torch.linspace(low_high[0] - 1.0, low_high[1] + 1.0, N)
    sdf = domain.sdf(x)
    assert sdf.shape == (N, 1)

    fig = plt.figure(figsize=(5, 5))

    ax = fig.add_subplot(111)
    ax.set_title("Segment1D")
    ax.set_xlabel("x")
    ax.set_ylabel("SDF")
    ax.grid()

    # plot sampled points
    ax.scatter(pts[:, 0], torch.zeros(N), label="sampled points")

    # plot sdf
    ax.plot(x, sdf, label="sdf")

    # plot the segment
    h = (low_high[1] - low_high[0]) / 50
    ax.vlines(
        [low_high[0], low_high_hole[0], low_high_hole[1], low_high[1]], -h, h, color="k"
    )
    ax.plot([low_high[0], low_high_hole[0]], [0, 0], "k")
    ax.plot([low_high_hole[1], low_high[1]], [0, 0], "k")

    ax.legend()

    fig.tight_layout()
    plt.show()


def visu_square_2d(bounds=[(-1, 1), (2, 4)]):
    domain = Square2D(bounds)

    bc_domains = domain.full_bc_domain()

    sampler = VolumetricSampler(domain)
    bc_samplers = [SurfacicSampler(bc) for bc in bc_domains]

    _plot_2d(
        N=10_000,
        N_bc=10,
        sampler=sampler,
        bc_samplers=bc_samplers,
        bc_samplers_holes=[],
        sdf=domain.sdf,
        lim_x=(bounds[0][0] - 1, bounds[0][1] + 1),
        lim_y=(bounds[1][0] - 1, bounds[1][1] + 1),
    )
    plt.show()


def visu_disk_2d(center=torch.tensor([0.0, 0.0]), radius=2.0):
    domain = Disk2D(center=center, radius=radius)

    bc_domains = domain.full_bc_domain()

    sampler = VolumetricSampler(domain)
    bc_samplers = [SurfacicSampler(bc) for bc in bc_domains]

    _plot_2d(
        N=10_000,
        N_bc=40,
        sampler=sampler,
        bc_samplers=bc_samplers,
        bc_samplers_holes=[],
        sdf=domain.sdf,
        lim_x=(center[0] - radius - 1, center[0] + radius + 1),
        lim_y=(center[1] - radius - 1, center[1] + radius + 1),
    )
    plt.show()


def visu_polygon_2d(vertices=[(0, 0), (1, 0), (1.5, 0.5), (1, 1), (0, 1)]):
    domain = Polygon2D(vertices)

    bc_domains = domain.full_bc_domain()

    sampler = VolumetricSampler(domain)
    bc_samplers = [SurfacicSampler(bc) for bc in bc_domains]

    lim_x = (
        torch.tensor(vertices)[:, 0].min() - 0.1,
        torch.tensor(vertices)[:, 0].max() + 0.1,
    )
    lim_y = (
        torch.tensor(vertices)[:, 1].min() - 0.1,
        torch.tensor(vertices)[:, 1].max() + 0.1,
    )

    _plot_2d(
        N=10_000,
        N_bc=10,
        sampler=sampler,
        bc_samplers=bc_samplers,
        bc_samplers_holes=[],
        sdf=domain.sdf,
        lim_x=lim_x,
        lim_y=lim_y,
    )
    plt.show()


def visu_polygon_2d_rot(vertices=[(0, 0), (1, 0), (1.5, 0.5), (1, 1), (0, 1)]):
    domain = Polygon2D(vertices, is_main_domain=True)

    lim_x = (
        torch.tensor(vertices)[:, 0].min() - 0.1,
        torch.tensor(vertices)[:, 0].max() + 0.1,
    )
    lim_y = (
        torch.tensor(vertices)[:, 1].min() - 0.1,
        torch.tensor(vertices)[:, 1].max() + 0.1,
    )

    bounds = (min(lim_x[0], lim_y[0]), max(lim_x[1], lim_y[1]))
    domain.set_mapping(Mapping.rot_2d(torch.tensor(torch.pi / 3)), [bounds, bounds])

    bc_domains = domain.full_bc_domain()

    sampler = VolumetricSampler(domain)
    bc_samplers = [SurfacicSampler(bc) for bc in bc_domains]

    _plot_2d(
        N=10_000,
        N_bc=10,
        sampler=sampler,
        bc_samplers=bc_samplers,
        bc_samplers_holes=[],
        sdf=domain.sdf,
        lim_x=lim_x,
        lim_y=lim_y,
    )
    plt.show()


def visu_polygon_2d_hole():
    vertices = [(0, 0), (1, 0), (1.5, 0.5), (1, 1), (0, 1)]
    domain = Polygon2D(vertices, is_main_domain=True)

    domain.add_hole(Square2D([(0.3, 0.7), (0.3, 0.7)]))

    lim_x = (
        torch.tensor(vertices)[:, 0].min() - 0.1,
        torch.tensor(vertices)[:, 0].max() + 0.1,
    )
    lim_y = (
        torch.tensor(vertices)[:, 1].min() - 0.1,
        torch.tensor(vertices)[:, 1].max() + 0.1,
    )

    bc_domains = domain.full_bc_domain()

    sampler = VolumetricSampler(domain)
    bc_samplers = [SurfacicSampler(bc) for bc in bc_domains]

    bc_domains_holes = [
        bc for hole in domain.list_holes for bc in hole.full_bc_domain()
    ]
    bc_samplers_holes = [SurfacicSampler(bc) for bc in bc_domains_holes]

    _plot_2d(
        N=10_000,
        N_bc=10,
        sampler=sampler,
        bc_samplers=bc_samplers,
        bc_samplers_holes=bc_samplers_holes,
        sdf=domain.sdf,
        lim_x=lim_x,
        lim_y=lim_y,
    )
    plt.show()


def visu_polygon_2d_hole_rot():
    vertices = [(0, 0), (1, 0), (1.5, 0.5), (1, 1), (0, 1)]
    domain = Polygon2D(vertices, is_main_domain=True)

    domain.add_hole(Square2D([(0.3, 0.7), (0.3, 0.7)]))

    lim_x = (
        torch.tensor(vertices)[:, 0].min() - 0.1,
        torch.tensor(vertices)[:, 0].max() + 0.1,
    )
    lim_y = (
        torch.tensor(vertices)[:, 1].min() - 0.1,
        torch.tensor(vertices)[:, 1].max() + 0.1,
    )

    bounds = (min(lim_x[0], lim_y[0]), max(lim_x[1], lim_y[1]))
    domain.set_mapping(Mapping.rot_2d(torch.tensor(torch.pi / 3)), [bounds, bounds])

    bc_domains = domain.full_bc_domain()

    sampler = VolumetricSampler(domain)
    bc_samplers = [SurfacicSampler(bc) for bc in bc_domains]

    bc_domains_holes = [
        bc for hole in domain.list_holes for bc in hole.full_bc_domain()
    ]
    bc_samplers_holes = [SurfacicSampler(bc) for bc in bc_domains_holes]

    _plot_2d(
        N=10_000,
        N_bc=10,
        sampler=sampler,
        bc_samplers=bc_samplers,
        bc_samplers_holes=bc_samplers_holes,
        sdf=domain.sdf,
        lim_x=lim_x,
        lim_y=lim_y,
    )
    plt.show()


def visu_polygon_2d_fixedhole_rot():
    vertices = [(-1, -1), (1, -1), (2.0, 0.0), (1, 1), (-1, 1)]
    domain = Polygon2D(vertices, is_main_domain=True)

    lim_x = (
        torch.tensor(vertices)[:, 0].min() - 0.1,
        torch.tensor(vertices)[:, 0].max() + 0.1,
    )
    lim_y = (
        torch.tensor(vertices)[:, 1].min() - 0.1,
        torch.tensor(vertices)[:, 1].max() + 0.1,
    )

    bounds = (min(lim_x[0], lim_y[0]), max(lim_x[1], lim_y[1]))
    domain.set_mapping(
        Mapping.rot_2d(torch.tensor(torch.pi / 3)), [bounds, bounds], to_holes=False
    )

    # if the hole is added after, we add copy_mapping=False
    domain.add_hole(Square2D([(0.3, 0.7), (0.3, 0.7)]), copy_mapping=False)

    # # if we hade made in the other way,
    # # copy_mapping is not necessary anymore but to_holes=False when we set the mapping:
    # domain.add_hole(Square2D([(.3, .7), (.3, .7)]))
    # domain.set_mapping(Mapping.Rot2D(torch.tensor(torch.pi/3)), [bounds, bounds], to_holes=False)

    bc_domains = domain.full_bc_domain()

    sampler = VolumetricSampler(domain)
    bc_samplers = [SurfacicSampler(bc) for bc in bc_domains]

    bc_domains_holes = [
        bc for hole in domain.list_holes for bc in hole.full_bc_domain()
    ]
    bc_samplers_holes = [SurfacicSampler(bc) for bc in bc_domains_holes]

    _plot_2d(
        N=10_000,
        N_bc=10,
        sampler=sampler,
        bc_samplers=bc_samplers,
        bc_samplers_holes=bc_samplers_holes,
        sdf=domain.sdf,
        lim_x=lim_x,
        lim_y=lim_y,
    )
    plt.show()


def visu_square_2d_complex():
    bounds = torch.tensor([(-3.0, 3.0), (-3.0, 3.0)])
    bounds_postmap = [(-5, 5), (-5, 5)]
    domain = Square2D(bounds, is_main_domain=True)

    domain.set_mapping(Mapping.rot_2d(torch.tensor(torch.pi / 4)), bounds_postmap)

    center, radius = torch.tensor([0.0, 2]), 0.5
    subdomain1 = Disk2D(center, radius)
    subdomain1.add_bc_domain(ArcCircle2D(center, radius, 0, torch.pi))
    subdomain1.add_bc_domain(ArcCircle2D(center, radius, torch.pi, 2 * torch.pi))
    domain.add_hole(subdomain1, copy_mapping=False)

    subdomain2 = Square2D([(-2, 0), (-0.5, 0.5)])
    domain.add_hole(subdomain2)

    bc_domains = domain.full_bc_domain()

    sampler = VolumetricSampler(domain)
    bc_samplers = [SurfacicSampler(bc) for bc in bc_domains]

    bc_domains_holes = subdomain1.list_bc_domains.copy()
    bc_domains_holes += subdomain2.full_bc_domain()

    bc_samplers_holes = [SurfacicSampler(bc) for bc in bc_domains_holes]

    _plot_2d(
        N=10_000,
        N_bc=2,
        sampler=sampler,
        bc_samplers=bc_samplers,
        bc_samplers_holes=bc_samplers_holes,
        sdf=domain.sdf,
        lim_x=(bounds_postmap[0][0] - 1, bounds_postmap[0][1] + 1),
        lim_y=(bounds_postmap[1][0] - 1, bounds_postmap[1][1] + 1),
    )

    # change the mapping
    def map(x):
        return torch.stack((x[:, 0], x[:, 1] / (x[:, 0] ** 2 + 1)), dim=-1)

    domain.set_mapping(Mapping(2, 2, map), bounds_postmap, to_holes=False)

    bc_domains = domain.full_bc_domain()

    bc_samplers = [SurfacicSampler(bc) for bc in bc_domains]

    bc_domains_holes = subdomain1.list_bc_domains.copy()
    bc_domains_holes += subdomain2.full_bc_domain()

    bc_samplers_holes = [SurfacicSampler(bc) for bc in bc_domains_holes]

    _plot_2d(
        N=10_000,
        N_bc=2,
        sampler=sampler,
        bc_samplers=bc_samplers,
        bc_samplers_holes=bc_samplers_holes,
        sdf=domain.sdf,
        lim_x=(bounds_postmap[0][0] - 1, bounds_postmap[0][1] + 1),
        lim_y=(bounds_postmap[1][0] - 1, bounds_postmap[1][1] + 1),
    )
    plt.show()


def _plot_2d(N, N_bc, sampler, bc_samplers, bc_samplers_holes, sdf, lim_x, lim_y):
    # sample some points
    pts = sampler.sample(N)
    assert pts.shape == (N, 2)

    # sample some points on the boundaries
    pts_bc = []
    normals_bc = []
    for bc_sampler in bc_samplers:
        x, n = bc_sampler.sample(N_bc, compute_normals=True)
        assert x.shape == (N_bc, 2)
        assert n.shape == (N_bc, 2)

        pts_bc.append(x)
        normals_bc.append(n)

    # sample some points on the boundaries of the holes
    pts_bc_hole = []
    normals_bc_hole = []
    for bc_sampler in bc_samplers_holes:
        x, n = bc_sampler.sample(N_bc, compute_normals=True)
        assert x.shape == (N_bc, 2)
        assert n.shape == (N_bc, 2)

        pts_bc_hole.append(x)
        normals_bc_hole.append(n)

    ## compute sdf
    n_mesh = 100
    xx, yy = torch.meshgrid(
        torch.linspace(*lim_x, n_mesh), torch.linspace(*lim_y, n_mesh), indexing="ij"
    )
    x = torch.stack([xx.flatten(), yy.flatten()], dim=1)
    sdf = sdf(x)
    assert sdf.shape == (n_mesh * n_mesh, 1)
    sdf = sdf.reshape(xx.shape)

    fig = plt.figure(figsize=(10, 5))

    ax = fig.add_subplot(121)
    ax.set_title("Domain Sampling")
    ax.set_aspect("equal")
    ax.grid()

    # plot pts in the domain
    ax.scatter(pts[:, 0], pts[:, 1], s=1, alpha=0.5)

    # plot pts/normals of points on boundaries
    for i, (x, n) in enumerate(zip(pts_bc, normals_bc)):
        ax.scatter(x[:, 0], x[:, 1], color=cmap(i + 1), s=10)
        ax.quiver(x[:, 0], x[:, 1], n[:, 0], n[:, 1], color=cmap(i + 1))

    # plot pts/normals of points on holes boundaries
    for i, (x, n) in enumerate(zip(pts_bc_hole, normals_bc_hole)):
        ax.scatter(x[:, 0], x[:, 1], color=cmap(i + 1 + len(pts_bc)), s=10)
        # the normals of a hole need to be inverted to be pointing outwards the domain
        ax.quiver(x[:, 0], x[:, 1], -n[:, 0], -n[:, 1], color=cmap(i + 1 + len(pts_bc)))

    ax = fig.add_subplot(122)
    ax.set_title("SDF")
    ax.set_aspect("equal")
    ax.grid()

    # plot sdf
    ax.contour(xx, yy, sdf, levels=[0], colors="k")
    cf = ax.contourf(xx, yy, sdf, levels=100, cmap="jet")
    add_colorbar(cf, aspect=20, pad_fraction=0.5)

    fig.tight_layout()


def visu_cube_3d():
    lim_x = (-1, 1)
    lim_y = (2, 4)
    lim_z = (0, 1)
    domain = Cube3D([lim_x, lim_y, lim_z])

    sampler = VolumetricSampler(domain)
    bc_samplers = [SurfacicSampler(bc) for bc in domain.full_bc_domain()]

    fig = plt.figure(figsize=(5, 5))

    ax = fig.add_subplot(111, projection="3d")
    ax.set_title("Domain Sampling")
    ax.set_xlim(lim_x[0] - 1, lim_x[1] + 1)
    ax.set_ylim(lim_y[0] - 1, lim_y[1] + 1)
    ax.set_zlim(lim_z[0] - 1, lim_z[1] + 1)
    ax.set_xlabel("x")
    ax.set_ylabel("y")
    ax.set_zlabel("z")

    pts = sampler.sample(5_000)

    ax.scatter(pts[:, 0], pts[:, 1], pts[:, 2], s=1, alpha=0.5)
    for i, bc_sampler in enumerate(bc_samplers):
        pts, normals = bc_sampler.sample(10, compute_normals=True)
        ax.quiver(
            pts[:, 0],
            pts[:, 1],
            pts[:, 2],
            normals[:, 0],
            normals[:, 1],
            normals[:, 2],
            color=cmap(i + 1),
        )

    plt.show()


def visu_disk_3d():
    domain = Disk3D(torch.ones(3), radius=2)

    sampler = VolumetricSampler(domain)
    bc_samplers = [SurfacicSampler(bc) for bc in domain.full_bc_domain()]

    fig = plt.figure(figsize=(5, 5))

    ax = fig.add_subplot(111, projection="3d")
    ax.set_title("Domain Sampling")
    ax.set_aspect("equal")
    ax.set_xlabel("x")
    ax.set_ylabel("y")
    ax.set_zlabel("z")

    pts = sampler.sample(5_000)

    ax.scatter(pts[:, 0], pts[:, 1], pts[:, 2], s=1, alpha=0.5)
    for i, bc_sampler in enumerate(bc_samplers):
        pts, normals = bc_sampler.sample(40, compute_normals=True)
        ax.quiver(
            pts[:, 0],
            pts[:, 1],
            pts[:, 2],
            normals[:, 0],
            normals[:, 1],
            normals[:, 2],
            color=cmap(i + 1),
        )

    plt.show()


def visu_multi_domain_disk_2d(center=torch.tensor([0.0, 0.0]), radius=2.0):
    domain = Disk2D(center=center, radius=radius, is_main_domain=True)
    sub_domain = Disk2D(
        center=torch.tensor([0.2, 0.2]), radius=0.2, is_main_domain=False
    )
    hole = Disk2D(center=torch.tensor([-1.0, -1.0]), radius=0.25, is_main_domain=False)
    domain.add_subdomain(sub_domain)
    domain.add_hole(hole)

    arc1 = ArcCircle2D(
        center=torch.tensor([0.0, 0.0]), radius=2.0, theta1=0.0, theta2=torch.pi / 2.0
    )
    arc2 = ArcCircle2D(
        center=torch.tensor([0.0, 0.0]),
        radius=2.0,
        theta1=torch.pi / 2.0,
        theta2=2.0 * torch.pi,
    )
    domain.add_bc_domain(arc1)
    domain.add_bc_domain(arc2)

    sampler = DomainSampler(domain)
    data = sampler.sample(20000)
    bc_data, _ = sampler.bc_sample(20000)

    x, y = data.get_components(label=0)
    plt.scatter(x.detach().numpy(), y.detach().numpy(), s=1, alpha=0.5)
    x, y = data.get_components(label=1)
    plt.scatter(x.detach().numpy(), y.detach().numpy(), s=1, alpha=0.5, c="red")
    plt.show()

    bc_x, bc_y = bc_data.get_components(label=0)
    plt.scatter(
        bc_x.detach().numpy(), bc_y.detach().numpy(), s=1, alpha=0.5, c="yellow"
    )
    bc_x, bc_y = bc_data.get_components(label=1)
    plt.scatter(bc_x.detach().numpy(), bc_y.detach().numpy(), s=1, alpha=0.5, c="green")
    bc_x, bc_y = bc_data.get_components(label=2)
    plt.scatter(bc_x.detach().numpy(), bc_y.detach().numpy(), s=1, alpha=0.5, c="blue")
    bc_x, bc_y = bc_data.get_components(label=100)
    plt.scatter(bc_x.detach().numpy(), bc_y.detach().numpy(), s=1, alpha=0.5, c="gray")
    plt.show()


if __name__ == "__main__":
    visu_segment_1d()
    visu_segment_1d_hole()
    visu_square_2d()
    visu_disk_2d()
    visu_polygon_2d()
    visu_polygon_2d_rot()
    visu_polygon_2d_hole()
    visu_polygon_2d_hole_rot()
    visu_polygon_2d_fixedhole_rot()
    visu_square_2d_complex()
    visu_cube_3d()
    visu_disk_3d()
    visu_multi_domain_disk_2d()