import numpy as np
import numpy.typing as npt
from dataclasses import dataclass

from chromatic_tda.utils.floating_point_utils import FloatingPointUtils
from chromatic_tda.utils.linear_algebra_utils import LinAlgUtils
from chromatic_tda.utils.timing import TimingUtils


@dataclass
class AffineSpace:
    shift: npt.NDArray
    vector_space: npt.NDArray
    dimension: int

    def __init__(self, shift: npt.NDArray, vector_space: npt.NDArray):
        if not (len(vector_space.shape) == 2 and len(shift.shape) == 1 and vector_space.shape[1] == shift.shape[0]):
            raise ValueError("Shape mismatch: `vector_space` needs to be a 2D array, "
                             "and `shift` a 1D array of the same dimension as the rows of `vector_space`")
        self.shift = shift
        self.vector_space = vector_space
        self.dimension = len(vector_space)


class GeometryUtils:
    @staticmethod
    def construct_equispace(*point_sets: npt.NDArray) -> AffineSpace:
        """Return the 'point + vector space' representation of the affine space A of points
        equidistant to all points of each argument. That is for a in A and for each argument P,
        all distances ||p - a|| for p in P are the same.

        :return: center `z`, list of generating vectors `ker`"""
        TimingUtils().start("Geom :: Construct Equispace")
        if len(point_sets) < 1:
            raise TypeError("At least one point set expected.")
        dim: int = point_sets[0].shape[1]
        k: int = len(point_sets)

        for point_set in point_sets[1:]:
            if point_set.shape[1] != dim:
                raise ValueError('All points need to be from the same dimension')

        a_mat_blocks: list[npt.NDArray] = []
        b_vec_blocks: list[npt.NDArray] = []
        for i, points in enumerate(point_sets):
            a_mat_blocks.append(GeometryUtils.one_hot_embedding(k, i, points))
            b_vec_blocks.append(np.array([(pt ** 2).sum() / 2 for pt in points]))
        a_mat = np.concatenate(a_mat_blocks)
        b_vec = np.concatenate(b_vec_blocks)
        z, ker = LinAlgUtils.solve(a_mat, b_vec)

        equispace = AffineSpace(z[:dim], ker[:, :dim])

        TimingUtils().stop("Geom :: Construct Equispace")
        return equispace

    @staticmethod
    def one_hot_embedding(number_of_categories: int, category: int, points: npt.NDArray) -> npt.NDArray:
        TimingUtils().start("Geom :: One Hot Embedding")
        if category >= number_of_categories:
            raise ValueError(f'Category out of range: need category < number_of_categories')
        suffix = np.zeros((len(points), number_of_categories))
        suffix[:, category] = 1
        one_hot_embedding = np.concatenate([points, suffix], axis=1)
        TimingUtils().stop("Geom :: One Hot Embedding")
        return one_hot_embedding

    @staticmethod
    def orthogonal_projection_onto_affine_space(affine_space: AffineSpace, *points: npt.NDArray) -> npt.NDArray:
        TimingUtils().start("Geom :: Project Points To Affine Space")
        u_mat_t = LinAlgUtils.orthogonalize_rows(affine_space.vector_space)  # matrix with orthogonal rows
        trans_mat = u_mat_t.transpose() @ u_mat_t
        projected_points = np.array([trans_mat @ (pt - affine_space.shift) + affine_space.shift for pt in points])
        TimingUtils().stop("Geom :: Project Points To Affine Space")
        return projected_points

    @staticmethod
    def reflect_points_through_affine_space(affine_space: AffineSpace, *points: npt.NDArray):
        TimingUtils().start("Geom :: Reflect Points Through Affine Space")
        projected_points = GeometryUtils.orthogonal_projection_onto_affine_space(affine_space, *points)
        reflected_points = 2 * projected_points - points  # == points + 2 * (projected_points - points)
        TimingUtils().stop("Geom :: Reflect Points Through Affine Space")
        return reflected_points

    @staticmethod
    def circumsphere_of_weighted_points(points: npt.NDArray, weights: npt.NDArray) -> tuple[npt.NDArray, float]:
        """Return center and power distance radius of the circumsphere of given weighted points. It lies in affine space
        generated by the points."""
        TimingUtils().start("Geom :: Circumsphere Of Weighted Points")
        if not (len(points.shape) == 2 and len(weights.shape) == 1 and points.shape[0] == weights.shape[0]):
            raise ValueError("Shape mismatch: points need to be a 2D array, and weights a 1D array of the same length")
        if len(points) == 0:
            raise ValueError("At least one point needs to be given.")
        if len(points) == 1:
            return points[0], weights[0]

        a_mat_t = points[1:] - points[0]  # difference vectors in row
        a_mat = a_mat_t.transpose()  # difference vectors in columns
        b_vec = (np.square(a_mat).sum(axis=0) + weights[1:] - weights[0]) / 2  # i-th entry: (||pi-p0||^2 + wi - w0)/2
        eq_mat = a_mat_t @ a_mat
        try:
            x = np.linalg.solve(eq_mat, b_vec)
        except np.linalg.LinAlgError:
            x, *_ = np.linalg.lstsq(eq_mat, b_vec)
            if not FloatingPointUtils.is_all_close(eq_mat @ x, b_vec):
                raise np.linalg.LinAlgError('Linear equation has no solution; check for degeneracies in the point set.'
                                            ' Try passing a small number, e.g., 1e-7, as `point_perturbation` argument')
        y = a_mat @ x
        rad2 = np.square(y).sum() + weights[0]
        z = y + points[0]

        TimingUtils().stop("Geom :: Circumsphere Of Weighted Points")
        return z, rad2
