import math
from typing import List, Optional, Tuple

import numpy as np


def leggauss_2d(order: int) -> Optional[Tuple[np.ndarray, np.ndarray]]:
    """Generate Gauss-Legendre points and weights of the given order on the reference square [-1,1].
    """
    if math.floor(math.sqrt(order))**2 != order:
        print('order has to be a square')
        return None
    points_1d, weights_1d = np.polynomial.legendre.leggauss(int(math.sqrt(order)))
    xv, yv = np.meshgrid(points_1d, points_1d)
    points = np.stack((xv, yv), axis=2).reshape(order, 2)
    weights = np.outer(weights_1d, weights_1d).flatten()
    return points, weights

def ref_to_element_1d(x1: float, x2: float, num_points: int) -> Tuple[np.ndarray, np.ndarray]:
    ref_points, ref_weights = np.polynomial.legendre.leggauss(num_points)
    points: np.ndarray = ref_points*0.5*abs(x2 - x1) + 0.5*(x2 + x1)
    weights: np.ndarray = ref_weights*0.5*(x2 - x1)
    return points, weights

def ref_to_element_2d(x1: float, x2: float, y1: float, y2: float, num_points: int) -> Tuple[np.ndarray, np.ndarray]:
    ref_element: Optional[Tuple[np.ndarray, np.ndarray]] = leggauss_2d(num_points)
    assert ref_element is not None
    ref_points, ref_weights = ref_element
    points: np.ndarray = np.array([[x*(x2 - x1)*0.5 + (x1 + x2)*0.5, y*(y2 - y1)*0.5 + (y1 + y2)*0.5] for x, y in ref_points])
    weights: np.ndarray = ref_weights*(x2 - x1)*(y2 - y1)*0.25
    return points, weights

def generate_points_weights(elements: List[List[float]], num_points: int, dim: int) -> Tuple[np.ndarray, np.ndarray]:
    points_list: List[np.ndarray] = []
    weights_list: List[np.ndarray] = []
    for element in elements:
        if dim == 1:
            points, weights = ref_to_element_1d(*element, num_points)
        elif dim == 2:
            points, weights = ref_to_element_2d(*element, num_points)
        points_list.append(points)
        weights_list.append(weights)
    return np.concatenate(points_list), np.concatenate(weights_list)

def quadtree(x_start, x_end, y_start, y_end, elements, nodes) -> None:
    pass
