from dataclasses import dataclass
from typing import List

import numpy as np
from matplotlib import pyplot as plt

from geometry.affine import AffineTransform
from geometry.vector import Vector


def draw_line_2d(axis, v1: Vector, v2: Vector, **kwargs):
  axis.plot(
    [v1.a[0, 0], v2.a[0, 0]],
    [v1.a[1, 0], v2.a[1, 0]],
    **kwargs)


def draw_line_3d(axis, v1: Vector, v2: Vector, **kwargs):
  axis.plot(
    [v1.a[0, 0], v2.a[0, 0]],
    [v1.a[1, 0], v2.a[1, 0]],
    [v1.a[2, 0], v2.a[2, 0]],
    **kwargs)


class Painter:
  axis: plt.Axes

  def __init__(self, axis):
    self.axis = axis

  def plot3d(self, vs: List[Vector], **kwargs):
    V = np.hstack([v.a for v in vs])

    self.axis.plot(V[0, :], V[1, :], V[2, :], **kwargs)

  def line3d(self, a: Vector, b: Vector, **kwargs):
    self.plot3d([a, b], **kwargs)


class TransformedPainter:
  painter: Painter
  transform: AffineTransform

  def __init__(self, painter, transform):
    self.painter = painter
    self.transform = transform

  def plot3d(self, vs: List[Vector], **kwargs):
    vs_ = list(map(self.transform.invert_shift, vs))

    self.painter.plot3d(vs_, **kwargs)

  def line3d(self, a: Vector, b: Vector, **kwargs):
    self.plot3d([a, b], **kwargs)
