from pyproj import Transformer, CRS
from pyproj.transformer import TransformerGroup, AreaOfInterest

from geogst.core.geometries.shapes.polygons import *
from geogst.core.geometries.shapes.segments import Segment


def project_xy(
    x,
    y,
    src_crs: Union[numbers.Integral, CRS],
    dest_crs: Union[numbers.Integral, CRS],
) -> Union[None, Tuple[numbers.Real, numbers.Real]]:

    try:

        transformer = Transformer.from_crs(
            src_crs,
            dest_crs,
            always_xy=True,
        )
        return transformer.transform(x, y)

    except Exception as e:

        print(f"\n{e!r}")
        return None


def project_xy_arrays(
    x_arr: np.ndarray,
    y_arr: np.ndarray,
    src_crs: Union[numbers.Integral, CRS],
    dest_crs: Union[numbers.Integral, CRS],
) -> Union[None, Tuple[np.ndarray, np.ndarray]]:
    """
    Vectorized counterpart of 'project_xy', transforming a batch of x,y
    coordinates (as numpy arrays) with a single Transformer, instead of
    building one Transformer per point.

    :param x_arr: array of x coordinates in the source CRS.
    :param y_arr: array of y coordinates in the source CRS.
    :param src_crs: the source CRS.
    :param dest_crs: the destination CRS.
    :return: a pair of arrays with the transformed x,y coordinates.
    """

    try:

        transformer = Transformer.from_crs(
            src_crs,
            dest_crs,
            always_xy=True,
        )
        return transformer.transform(x_arr, y_arr)

    except Exception as e:

        print(f"\n{e!r}")
        return None


def project_xy_to_geographic(
    x,
    y,
    src_crs: Union[numbers.Integral, CRS],
) -> Union[None, Tuple[numbers.Real, numbers.Real]]:

    try:

        transformer = Transformer.from_crs(src_crs, "EPSG:4326", always_xy=True)
        return transformer.transform(x, y)

    except Exception as e:

        print(f"\n{e!r}")
        return None


def get_transformers(
    source_crs: Union[numbers.Integral, CRS],
    dest_crs: Union[numbers.Integral, CRS],
    area_of_interest: Tuple[ numbers.Real, numbers.Real, numbers.Real, numbers.Real,]
) -> Union[None, List[Transformer]]:


    # Crea un TransformerGroup con area_of_interest

    tg = TransformerGroup(
        source_crs,
        dest_crs,
        area_of_interest=AreaOfInterest(*area_of_interest)
    )

    return tg.transformers


@singledispatch
def project(
    shape: Union[Shape, Sequence[Shape]],
    source_crs: Union[numbers.Integral, CRS],
    dest_crs: Union[numbers.Integral, CRS],
) -> Union[None, Shape]:
    """

    :param obj:
    :param source_crs:
    :param dest_crs:
    :return:
    """

    return NotImplemented


@project.register(Point)
def _(
    point: Point,
    source_crs: Union[numbers.Integral, CRS],
    dest_crs: Union[numbers.Integral, CRS],
) -> Union[None, Point]:

    try:

        xy = project_xy(
            point.x,
            point.y,
            src_crs=source_crs,
            dest_crs=dest_crs,)

        if xy is None:
            return None

        x, y = xy
        proj_coords = np.copy(point.coords)
        proj_coords[:2] = x, y
        return Point(proj_coords)

    except Exception as e:

        print(f"\n{e!r}")
        return None


@project.register(Segment)
def _(
    segment: Segment,
    source_crs: Union[numbers.Integral, CRS],
    dest_crs: Union[numbers.Integral, CRS],
) -> Union[None, Segment]:

    try:

        start_point = project(
            segment.start_pt,
            source_crs,
            dest_crs,
        )

        end_point = project(
            segment.end_pt,
            source_crs,
            dest_crs,
        )

        if start_point is None or end_point is None:
            return None

        return Segment(start_point, end_point)

    except Exception as e:

        print(f"\n{e!r}")

        return None


@project.register(Ln)
def _(
    line: Ln,
    source_crs: Union[numbers.Integral, CRS],
    dest_crs: Union[numbers.Integral, CRS],
) -> Union[None, Ln]:

    try:

        coords = line.coords

        if coords is None or coords.shape[0] == 0:
            return Ln(coords)

        proj_xy = project_xy_arrays(
            coords[:, 0],
            coords[:, 1],
            src_crs=source_crs,
            dest_crs=dest_crs,
        )

        if proj_xy is None:
            return None

        proj_x, proj_y = proj_xy

        new_coords = np.copy(coords)
        new_coords[:, 0] = proj_x
        new_coords[:, 1] = proj_y

        return Ln(new_coords)

    except Exception as e:

        print(f"\n{e!r}")
        return None


def project_lines_batch(
    lines: List[Ln],
    source_crs: Union[numbers.Integral, CRS],
    dest_crs: Union[numbers.Integral, CRS],
) -> List[Union[None, Ln]]:
    """
    Vectorized counterpart of projecting a list of Ln one at a time:
    builds a single Transformer and reprojects every vertex of every
    line in one batched call, instead of one Transformer per line
    (and far worse, one per point, as the naive per-Point dispatch does).
    Line count and per-line vertex order are preserved.

    :param lines: the lines to project, all sharing source_crs.
    :param source_crs: the source CRS.
    :param dest_crs: the destination CRS.
    :return: the projected lines, in the same order as the input.
    """

    non_empty_idxs = [
        i for i, ln in enumerate(lines)
        if ln.coords is not None and ln.coords.shape[0] > 0
    ]

    if not non_empty_idxs:
        return [Ln(ln.coords) for ln in lines]

    counts = [lines[i].coords.shape[0] for i in non_empty_idxs]
    all_coords = np.vstack([lines[i].coords for i in non_empty_idxs])

    proj_xy = project_xy_arrays(
        all_coords[:, 0],
        all_coords[:, 1],
        src_crs=source_crs,
        dest_crs=dest_crs,
    )

    if proj_xy is None:
        return [None for _ in lines]

    proj_x, proj_y = proj_xy

    all_coords[:, 0] = proj_x
    all_coords[:, 1] = proj_y

    results: List[Union[None, Ln]] = [None] * len(lines)
    offset = 0
    for i, n in zip(non_empty_idxs, counts):
        results[i] = Ln(all_coords[offset:offset + n])
        offset += n

    for i, ln in enumerate(lines):
        if ln.coords is None or ln.coords.shape[0] == 0:
            results[i] = Ln(ln.coords)

    return results


def project_polygons_batch(
    polygons: List[Polygon],
    source_crs: Union[numbers.Integral, CRS],
    dest_crs: Union[numbers.Integral, CRS],
) -> List[Union[None, Polygon]]:
    """
    Vectorized counterpart of projecting a list of Polygon one at a time:
    reprojects every ring (outer plus inner holes) of every polygon in a
    single batched call, instead of building a Transformer per ring.
    Polygon count and ring order are preserved.

    :param polygons: the polygons to project, all sharing source_crs.
    :param source_crs: the source CRS.
    :param dest_crs: the destination CRS.
    :return: the projected polygons, in the same order as the input.
    """

    rings = []
    ring_counts = []

    for polygon in polygons:

        outer_line = polygon.outer_line()
        if outer_line is None:
            print(f"\nPolygon outer line is None")

        polygon_rings = [outer_line] + list(polygon.inner_lines() or [])
        polygon_rings = [ring for ring in polygon_rings if ring is not None]

        rings.extend(polygon_rings)
        ring_counts.append(len(polygon_rings))

    if not rings:
        return [Polygon() for _ in polygons]

    projected_rings = project_lines_batch(
        rings,
        source_crs,
        dest_crs,
    )

    projected_polygons = []
    offset = 0

    for num_rings in ring_counts:

        polygon_rings = projected_rings[offset:offset + num_rings]
        offset += num_rings

        if not polygon_rings or polygon_rings[0] is None:
            projected_polygons.append(Polygon())
            continue

        projected_polygons.append(
            Polygon(
                polygon_rings[0],
                [ring for ring in polygon_rings[1:] if ring is not None],
            )
        )

    return projected_polygons


@project.register(MultiLine)
def _(
    multiline: MultiLine,
    source_crs: Union[numbers.Integral, CRS],
    dest_crs: Union[numbers.Integral, CRS],
) -> Union[None, MultiLine]:

    try:

        projected_lines = []
        for line in multiline:
            projected_line = project(
                line,
                source_crs,
                dest_crs,
            )
            if projected_line is None:
                continue
            projected_lines.append(projected_line)

        return MultiLine(projected_lines)

    except Exception as e:

        print(f"\n{e!r}")
        return None


@project.register(Polygon)
def _(
    polygon: Polygon,
    source_crs: Union[numbers.Integral, CRS],
    dest_crs: Union[numbers.Integral, CRS],
) -> Union[None, Polygon]:

    try:

        outer_line = polygon.outer_line()
        if outer_line is None:
            print(f"\nPolygon outer line is None")

        projected_outer_line = project(
            outer_line,
            source_crs,
            dest_crs,
        )

        projected_inner_lines = []
        inner_lines = polygon.inner_lines()
        if inner_lines:
            for inner_line in inner_lines:
                projected_inner_line = project(
                    inner_line,
                    source_crs,
                    dest_crs,
                )
                if projected_inner_line is not None:
                    projected_inner_lines.append(projected_inner_line)

        return Polygon(projected_outer_line, projected_inner_lines)

    except Exception as e:

        print(f"\n{e!r}")
        return None


@project.register(list)
def _(
    features: list,
    source_crs: Union[numbers.Integral, CRS],
    dest_crs: Union[numbers.Integral, CRS],
) -> Union[None, list]:

    try:

        if features and all(isinstance(feature, Ln) for feature in features):
            # Batched fast path: one Transformer for the whole list instead
            # of one per line (or, worse, one per point).
            return project_lines_batch(
                features,
                source_crs,
                dest_crs,
            )

        if features and all(isinstance(feature, Polygon) for feature in features):
            # Same batched fast path, for the rings of a set of polygons.
            return project_polygons_batch(
                features,
                source_crs,
                dest_crs,
            )

        projected_features = []
        for feature in features:
            projected_feature = project(
                feature,
                source_crs,
                dest_crs,
            )
            projected_features.append(projected_feature)

        return projected_features

    except Exception as e:

        print(f"\n{e!r}")
        return None

