from geogst.core.geology.profiles.geoprofiles import *
from geogst.core.geology.profiles.profilers import *
from geogst.core.geology.profiles.methods import PointsInput, ProjectionMethod


# How many points to project between progress reports. Reporting once per point
# would make the report itself (a GUI repaint, typically) dominate the run, while
# reporting once per profile leaves the caller with no feedback for as long as a
# whole profile takes — on a large dataset, minutes with a frozen interface.
PROGRESS_REPORT_POINT_BLOCK = 500


def _point_progress_relay(
    progress_callback,
    profile_ndx: numbers.Integral,
    num_profilers: numbers.Integral,
    num_points: numbers.Integral,
):
    """
    Adapts a per-point progress report from one profile into a report over the
    whole run, so that the caller sees a single steadily advancing count instead
    of one restarting per profile.

    :param progress_callback: the caller's progress_callback(done, total, unit),
        or None to report nothing.
    :param profile_ndx: the 1-based index of the profile being projected.
    :param num_profilers: the number of profiles in the run.
    :param num_points: the number of points projected onto each profile.
    :return: a progress_callback(done, total) for a single profile, or None.
    """

    if progress_callback is None:
        return None

    total_work = num_profilers * num_points
    profile_offset = (profile_ndx - 1) * num_points

    def _relay(done_points, _total_points):
        progress_callback(profile_offset + done_points, total_work, "point")

    return _relay


def _points_within_profile_buffer(
    points: List[Point],
    segment_profilers: List[SegmentProfiler],
    max_profile_distance: numbers.Real,
) -> np.ndarray:
    """
    Returns a boolean mask of the points lying within max_profile_distance of the
    profile in plan view, i.e. inside a buffer of that width around the profile
    trace.

    This is a broad-phase filter: cheap, vectorised, and used to skip the full
    per-segment projection for points that cannot possibly be accepted. On a
    regional dataset projected onto a few local profiles, most points fall outside
    the buffer, and each one would otherwise pay a plane projection plus an
    along-profile distance computation against every segment of every profile.

    It never discards a point the full projection would have kept. Acceptance
    requires both that the distance from the point to its projection is within
    max_profile_distance, and that the projection lands inside the segment. When
    it lands inside the segment, that landing point is itself a point of the
    segment lying no further from the source point, in plan, than the full
    (possibly inclined) projection distance. So the plan distance to the segment
    is bounded by the accepted distance, for perpendicular projection and for
    projection along an axis or an attitude alike: outside the buffer, no
    acceptance is possible.

    :param points: the points to test, in the profiler's CRS.
    :param segment_profilers: the profile segments to measure the distance to.
    :param max_profile_distance: the buffer half-width, i.e. the maximum allowed
        distance of a point from the profile.
    :return: a boolean mask, True where the point is inside the buffer.
    """

    num_points = len(points)

    if num_points == 0 or not segment_profilers:
        return np.zeros(num_points, dtype=bool)

    pts_x = np.fromiter((pt.x for pt in points), dtype=np.float64, count=num_points)
    pts_y = np.fromiter((pt.y for pt in points), dtype=np.float64, count=num_points)

    within = np.zeros(num_points, dtype=bool)

    for segment_profiler in segment_profilers:

        start_pt = segment_profiler.start_pt()
        end_pt = segment_profiler.end_pt()

        delta_x = end_pt.x - start_pt.x
        delta_y = end_pt.y - start_pt.y
        squared_length = delta_x * delta_x + delta_y * delta_y

        offsets_x = pts_x - start_pt.x
        offsets_y = pts_y - start_pt.y

        if squared_length == 0.0:
            # degenerate segment: distance to its single position
            distances = np.hypot(offsets_x, offsets_y)
        else:
            # clamped to [0, 1], so the distance is to the segment and not to its
            # infinite line: a point beyond a segment end is measured from that end
            positions = np.clip(
                (offsets_x * delta_x + offsets_y * delta_y) / squared_length,
                0.0,
                1.0,
            )
            distances = np.hypot(
                offsets_x - positions * delta_x,
                offsets_y - positions * delta_y,
            )

        np.logical_or(within, distances <= max_profile_distance, out=within)

        if within.all():
            break

    return within


def _batch_project_positions(
    positions: List[Point],
    src_crs_wkt,
    dest_crs_wkt,
) -> List[Point]:
    """
    Reproject a batch of points with a single Transformer construction, instead
    of one per point (as calling 'project(Point, ...)' inside a loop would do,
    which is a major cost when projecting many points/attitudes onto profiles).
    Preserves each point's original embedding dimension (e.g. z is kept as-is).
    """

    if not positions:
        return []

    num_pts = len(positions)
    xs = np.fromiter((pt.x for pt in positions), dtype=np.float64, count=num_pts)
    ys = np.fromiter((pt.y for pt in positions), dtype=np.float64, count=num_pts)

    proj_x, proj_y = project_xy_arrays(xs, ys, src_crs=src_crs_wkt, dest_crs=dest_crs_wkt)
    if proj_x is None:
        raise Exception("Unable to project points batch to profiler CRS")

    projected = []
    for pt, px, py in zip(positions, proj_x, proj_y):
        coords = np.copy(pt.coords)
        coords[0] = px
        coords[1] = py
        projected.append(Point(coords))

    return projected


@singledispatch
def project_points(
    obj,
    *args,
    **kwargs,
):

    return 1


@project_points.register(GeoProfiles)
def _(
    geoprofiles: GeoProfiles,
    data: Sequence,
    data_crs: Crs,
    max_profile_distance: numbers.Real,
    projection_method: ProjectionMethod = ProjectionMethod.NEAREST,
    input_type: PointsInput = PointsInput.POINTS,
    **kargs
) -> Union[Error, 'projections']:

    """
    Proiect points onto the geoprofiles.

    :param data: a list of point dataset.
    :param data_crs: the coordinate reference system of the data.
    :param max_profile_distance: the maximum point distance from the profile.
    :param projection_method: the projection method to use for the points.
    :param input_type: the type of the input dataset. Default is points.
    :param cat_key: the category name of the point dataset.
    :return: the error status.
    """

    try:

        data_crs = data_crs if isinstance(data_crs, Crs) else Crs(data_crs)

        if input_type in (PointsInput.POINTS, PointsInput.FOCAL_MECHANISMS):

            data_projected_onto_profiles = project_points(
                geoprofiles.profilers,
                data,
                data_crs,
                max_profile_distance,
                projection_method,
                input_type,
                **kargs
            )

        elif input_type == PointsInput.ATTITUDES:

            data_projected_onto_profiles = project_attitudes(
                geoprofiles.profilers,
                data,
                data_crs,
                max_profile_distance,
                projection_method,
                **kargs
            )

        elif input_type == PointsInput.FAULTS:

            data_projected_onto_profiles = Error(
                True,
                caller_name(),
                Exception(f"Faults not implemented"),
                traceback.format_exc())

        else:

            data_projected_onto_profiles = Error(
                True,
                caller_name(),
                Exception(f"Got not implemented point input type: {input_type}"),
                traceback.format_exc())

        return data_projected_onto_profiles

    except Exception as e:

        return Error(
            True,
            caller_name(),
            e,
            traceback.format_exc())


@project_points.register(Profilers)
def _(
    profilers: Profilers,
    data: List[Tuple[Category, Point]],
    data_crs: Crs,
    max_profile_distance: numbers.Real,
    projection_method: ProjectionMethod = ProjectionMethod.NEAREST,
    input_type: PointsInput = PointsInput.POINTS,
    **kwargs
) -> Union[Error, List[Dict[Category, List[PointTrace]]]]:
    """
    Projects a set of points onto the section profile.

    :param data: the set of points to be plotted onto the section.
    :param data_crs: the coordinate reference system of the data.
    :param max_profile_distance: the maximum allowed projection distance between the individual point and the profile
    :param projection_method: the method to project the points to the section.
    :param kwargs: the keyword arguments. If 'cat_key' is provided, each per-profile result is
        wrapped as {cat_key: profile_points}, matching the structure expected by GeoProfiles storage/plotting.
        If 'progress_callback' is provided, it is invoked as progress_callback(done, total, "point")
        every PROGRESS_REPORT_POINT_BLOCK points, counting over the whole run (every point
        against every profile), e.g. to drive a progress bar.
    :return: the parsed point traces and an error status.
    """

    try:

        data_crs = data_crs if isinstance(data_crs, Crs) else Crs(data_crs)

        cat_key = kwargs.pop('cat_key', None)
        progress_callback = kwargs.pop('progress_callback', None)

        point_traces = []

        num_profilers = len(profilers)

        for ndx, line_profiler in enumerate(profilers, start=1):

            profile_points = project_points(
                line_profiler,
                data=data,
                data_crs=data_crs,
                max_profile_distance=max_profile_distance,
                projection_method=projection_method,
                input_type=input_type,
                progress_callback=_point_progress_relay(
                    progress_callback, ndx, num_profilers, len(data)
                ),
                **kwargs
            )
            if isinstance(profile_points, Error):
                raise Exception(f"{profile_points!r}")

            point_traces.append({cat_key: profile_points} if cat_key is not None else profile_points)

        return point_traces

    except Exception as e:

        return Error(
            True,
            caller_name(),
            e,
            traceback.format_exc())


@project_points.register(LineProfiler)
def _(
    line_profiler: LineProfiler,
    data: List[Tuple[Category, Point, Fault]],
    data_crs: Crs,
    max_profile_distance: numbers.Real,
    projection_method: ProjectionMethod = ProjectionMethod.NEAREST,
    input_type: PointsInput = PointsInput.POINTS,
    **kwargs
) -> Union[Error, defaultdict[PointTrace]]:
    """
    Projects a set of 3D points onto the section profile.

    :param data: the set of 3D points to project onto the section. The third tuple element
        (a Fault) is only used when input_type is PointsInput.FOCAL_MECHANISMS, to carry the
        strike/dip/rake of the associated focal mechanism.
    :param data_crs: the coordinate reference system of the data.
    :param max_profile_distance: the maximum allowed projection distance between the points and the profile.
    :param projection_method: the method to use for projecting the points onto the section.
    :param input_type: PointsInput.POINTS for plain point projections, or PointsInput.FOCAL_MECHANISMS
        to also carry the focal mechanism (strike/dip/rake) needed for beachball plotting.
    :param kwargs: the keyword arguments.
    :return: dictionary storing projected points and error status.
    """

    try:

        data_crs = data_crs if isinstance(data_crs, Crs) else Crs(data_crs)

        progress_callback = kwargs.pop('progress_callback', None)

        if projection_method == ProjectionMethod.NEAREST:
            map_axes = None
        elif projection_method == ProjectionMethod.COMMON_AXIS:
            map_axes = Axis(kwargs['trend'], kwargs['plunge'])
        elif projection_method == ProjectionMethod.INDIVIDUAL_AXES:
            individual_axes_values = kwargs['individual_axes_values']
            # one axis per point, so a length mismatch is a caller error: checked
            # here, where it can be named, rather than surfacing later as an
            # IndexError or, worse, as points silently projected along the axis
            # meant for a different point
            if len(individual_axes_values) != len(data):
                raise Exception(
                    f"Individual axes values are {len(individual_axes_values)} "
                    f"but data points are {len(data)}"
                )
            map_axes = [Axis(trend, plunge) for trend, plunge in individual_axes_values]
        else:
            raise Exception(f"Mapping method is {projection_method}")

        projected_points = defaultdict(list)
        num_data_points = len(data)

        # batch-reproject all input positions once (single Transformer),
        # instead of once per point via project(Point, ...) inside the loop
        src_points_in_profiler_crs = _batch_project_positions(
            [position for _, position, _ in data],
            data_crs.to_wkt(),
            line_profiler.wkt_crs,
        )

        # materialize the segment profilers and their vertical planes once per
        # line profiler, instead of re-creating them (iterating line_profiler
        # builds a fresh SegmentProfiler per segment on every call) and
        # recomputing each segment's vertical plane on every single point.
        # Consecutive collinear segments are merged first: profile lines are
        # densified for DEM draping, and projecting against every sub-segment of
        # a straight run repeats the same work once per sampling step, only to
        # discard all but one result. Cost then scales with the number of
        # direction changes instead of the sampling step, with identical results.
        segment_profilers_with_offsets = line_profiler.direction_change_segment_profilers()
        segment_profilers = [sp for sp, _ in segment_profilers_with_offsets]
        segment_offsets = [offset for _, offset in segment_profilers_with_offsets]
        segment_planes = [sp.segment_as_vertical_plane() for sp in segment_profilers]

        # broad-phase filter: only the points inside the profile's buffer can be
        # accepted, so the rest skip the per-segment projection entirely
        near_profile = _points_within_profile_buffer(
            src_points_in_profiler_crs, segment_profilers, max_profile_distance
        )

        for ndx_super, ((category, position, fault), src_point_in_profiler_crs) in enumerate(
                zip(data, src_points_in_profiler_crs), start=1):

            if near_profile[ndx_super - 1]:

                # INDIVIDUAL_AXES supplies one axis per point, so it is picked
                # here, by point, and not inside the segment loop below: indexing
                # it by segment gave every point the axis of whichever segment it
                # happened to be tested against
                point_map_axis = map_axes[ndx_super - 1] if isinstance(map_axes, list) else map_axes

                for ndx, (segment_profiler, section_plane, previous_segments_offset) in enumerate(
                        zip(segment_profilers, segment_planes, segment_offsets), start=1):

                    intersection_point_in_profiler_crs, err = segment_profiler.project_point(
                        point=src_point_in_profiler_crs,
                        projection_method=projection_method,
                        min_profile_distance=MIN_DISTANCE_TOLERANCE,
                        map_axis=point_map_axis,
                        axis_angular_tolerance=MIN_DISORIENTATION_TOLERANCE,
                        precomputed_plane=section_plane,
                    )

                    if err:
                        raise Exception(f"{err!r}")

                    if intersection_point_in_profiler_crs is None:
                        continue

                    point_trace, err = segment_profiler.parse_single_point_projection_result(
                        source_pt=src_point_in_profiler_crs,
                        intersection_point=intersection_point_in_profiler_crs,
                        max_profile_distance=max_profile_distance,
                        offset=previous_segments_offset
                    )

                    if err:
                        raise Exception(f"{err!r}")

                    if point_trace is not None:

                        if input_type == PointsInput.FOCAL_MECHANISMS:
                            profile_versor = segment_profiler.versor()
                            point_trace = FocalMechanismTrace(
                                s=point_trace.s,
                                z=point_trace.z,
                                dist=point_trace.dist,
                                strike=fault.plane.rhr_strike,
                                dip=fault.plane.dipang,
                                rake=fault.rake(0) if fault.has_slickenlines() else 0.0,
                                profile_azimuth=Vect2D(profile_versor.x, profile_versor.y).azimuth_degr(),
                            )

                        projected_points[category].append(point_trace)

            if progress_callback is not None and (
                ndx_super % PROGRESS_REPORT_POINT_BLOCK == 0 or ndx_super == num_data_points
            ):
                progress_callback(ndx_super, num_data_points)

        best_points_projections = defaultdict(PointTrace)

        for category, solutions in projected_points.items():
            if len(solutions) == 1:
                best_points_projections[category] = solutions[0]
            else:
                min_distance = min([solution.dist for solution in solutions])
                for solution in solutions:
                    if solution.dist <= min_distance:
                        best_points_projections[category] = solution
                        break

        return best_points_projections

    except Exception as e:

        return Error(
            True,
            caller_name(),
            e,
            traceback.format_exc())


@singledispatch
def project_attitudes(
    obj,
    data: List[Tuple[Category, Point, Plane]],
    data_crs: Crs,
    max_profile_distance: numbers.Real,
    projection_method: ProjectionMethod = ProjectionMethod.NEAREST,
    **kwargs
) -> Tuple[Union[None, List[Dict[Category, PlaneTrace]]], Error]:
    """
    Projects a set of georeferenced space3d attitudes onto the section profile.

    :param data: the set of georeferenced space3d attitudes to plot on the section.
    :param data_crs: the coordinate reference system of the data.
    :param projection_method: the method to map the attitudes to the section.
    :param max_profile_distance: the maximum projection distance between the plane_attitude and the profile
    :return: an attitudes set and an error status.
    """

    return Error()


@project_attitudes.register(Profilers)
def _(
    profilers: Profilers,
    data: List[Tuple[Category, Point, Plane]],
    data_crs: Crs,
    max_profile_distance: numbers.Real,
    projection_method: ProjectionMethod = ProjectionMethod.NEAREST,
    **kwargs
) -> Union[None, Error, List[Dict[Category, PlaneTrace]]]:
    """
    Projects a set of georeferenced space3d attitudes onto the section profile.

    :param data: the set of georeferenced space3d attitudes to plot on the section.
    :param data_crs: the coordinate reference system of the data.
    :param projection_method: the method to map the attitudes to the section.
    :param max_profile_distance: the maximum projection distance between the plane_attitude and the profile
    :param kwargs: the keyword arguments. If 'cat_key' is provided, each per-profile result is
        wrapped as {cat_key: profile_attitudes}, matching the structure expected by GeoProfiles storage/plotting.
        If 'progress_callback' is provided, it is invoked as progress_callback(done, total, "point")
        every PROGRESS_REPORT_POINT_BLOCK attitudes, counting over the whole run (every attitude
        against every profile), e.g. to drive a progress bar.
    :return: an attitudes set and an error status.
    """

    try:

        data_crs = data_crs if isinstance(data_crs, Crs) else Crs(data_crs)

        cat_key = kwargs.pop('cat_key', None)
        progress_callback = kwargs.pop('progress_callback', None)

        plane_traces = []

        num_profilers = len(profilers)

        for ndx, line_profiler in enumerate(profilers, start=1):

            profile_attitudes = project_attitudes(
                line_profiler,
                data,
                data_crs,
                max_profile_distance,
                projection_method,
                progress_callback=_point_progress_relay(
                    progress_callback, ndx, num_profilers, len(data)
                ),
                **kwargs
            )
            if isinstance(profile_attitudes, Error):
                return profile_attitudes

            plane_traces.append({cat_key: profile_attitudes} if cat_key is not None else profile_attitudes)

        return plane_traces

    except Exception as e:

        return Error(
            True,
            caller_name(),
            e,
            traceback.format_exc())


@project_attitudes.register(LineProfiler)
def _(
    line_profiler: LineProfiler,
    attitudes: List[Tuple[Category, Point, Plane]],
    data_crs: Crs,
    max_profile_distance: numbers.Real,
    projection_method: ProjectionMethod,
    **kwargs
) -> Union[None, Dict[Category, PlaneTrace]]:
    """
    Projects a set of space3d attitudes onto the section profile.

    :param attitudes: the set of georeferenced space3d attitudes to plot on the section.
    :param data_crs: the coordinate reference system of the data.
    :param max_profile_distance: the maximum projection distance between the plane_attitude and the profile
    :param projection_method: the method to map the attitudes to the section.
    :return: list of PlaneTrace values.
    """

    try:

        data_crs = data_crs if isinstance(data_crs, Crs) else Crs(data_crs)

        progress_callback = kwargs.pop('progress_callback', None)

        if projection_method == ProjectionMethod.NEAREST:
            map_axes = None
        elif projection_method == ProjectionMethod.COMMON_AXIS:
            map_axes = Axis(kwargs['trend'], kwargs['plunge'])
        elif projection_method == ProjectionMethod.INDIVIDUAL_AXES:
            individual_axes_values = kwargs['individual_axes_values']
            # one axis per attitude: see the equivalent check in project_points
            if len(individual_axes_values) != len(attitudes):
                raise Exception(
                    f"Individual axes values are {len(individual_axes_values)} "
                    f"but attitudes are {len(attitudes)}"
                )
            map_axes = [Axis(trend, plunge) for trend, plunge in individual_axes_values]
        else:
            raise Exception(f"Mapping method is {projection_method}")

        projected_attitudes = defaultdict(list)
        num_attitudes = len(attitudes)

        # batch-reproject all input positions once (single Transformer),
        # instead of once per point via project(Point, ...) inside the loop
        src_points_in_profiler_crs = _batch_project_positions(
            [position for _, position, _ in attitudes],
            data_crs.to_wkt(),
            line_profiler.wkt_crs,
        )

        # materialize the segment profilers and their vertical planes once per
        # line profiler, instead of re-creating them (iterating line_profiler
        # builds a fresh SegmentProfiler per segment on every call) and
        # recomputing each segment's vertical plane on every single attitude.
        # Consecutive collinear segments are merged first, for the same reason as
        # in point projection: a densified straight run costs one projection per
        # sampling step and keeps a single result.
        segment_profilers_with_offsets = line_profiler.direction_change_segment_profilers()
        segment_profilers = [sp for sp, _ in segment_profilers_with_offsets]
        segment_offsets = [offset for _, offset in segment_profilers_with_offsets]
        segment_planes = [sp.segment_as_vertical_plane() for sp in segment_profilers]

        # broad-phase filter, as in point projection: attitudes outside the
        # profile's buffer cannot be accepted and skip the per-segment work
        near_profile = _points_within_profile_buffer(
            src_points_in_profiler_crs, segment_profilers, max_profile_distance
        )

        for ndx_super, ((attitude_id, position, attitude), src_point_in_profiler_crs) in enumerate(
                zip(attitudes, src_points_in_profiler_crs), start=1):

            if near_profile[ndx_super - 1]:

                # picked per attitude, not per segment: see project_points
                attitude_map_axis = map_axes[ndx_super - 1] if isinstance(map_axes, list) else map_axes

                for ndx, (segment_profiler, section_plane, previous_segments_offset) in enumerate(
                        zip(segment_profilers, segment_planes, segment_offsets), start=1):

                    result, err = segment_profiler.project_attitude(
                        attitude=attitude,
                        point=src_point_in_profiler_crs,
                        projection_method=projection_method,
                        min_profile_distance=MIN_DISTANCE_TOLERANCE,
                        map_axis=attitude_map_axis,
                        axis_angular_tolerance=MIN_DISORIENTATION_TOLERANCE,
                        precomputed_plane=section_plane,
                    )

                    if err:
                        raise Exception(f"{err!r}")

                    if result is None:
                        continue

                    intersection_vector, intersection_point_in_profiler_crs = result

                    profile_attitude = segment_profiler.parse_single_attitude_projection_result(
                        attitude=attitude,
                        attitude_intersection=intersection_vector,
                        source_pt=src_point_in_profiler_crs,
                        intersection_point=intersection_point_in_profiler_crs,
                        max_profile_distance=max_profile_distance,
                        offset=previous_segments_offset
                    )

                    if isinstance(profile_attitude, Error):
                        return profile_attitude

                    if profile_attitude is not None:
                        projected_attitudes[attitude_id].append(profile_attitude)

            if progress_callback is not None and (
                ndx_super % PROGRESS_REPORT_POINT_BLOCK == 0 or ndx_super == num_attitudes
            ):
                progress_callback(ndx_super, num_attitudes)

        best_attitudes_projections = defaultdict(PlaneTrace)

        for attitude_id, solutions in projected_attitudes.items():

            if len(solutions) == 1:
                best_attitudes_projections[attitude_id] = solutions[0]
            else:
                min_distance = min([solution.dist for solution in solutions])
                for solution in solutions:
                    if solution.dist <= min_distance:
                        best_attitudes_projections[attitude_id] = solution
                        break

        return best_attitudes_projections

    except Exception as e:

        return Error(
            True,
            caller_name(),
            e,
            traceback.format_exc())
