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


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) after each of
        the 'total' parallel profiles is projected, 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,
                **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)

            if progress_callback is not None:
                progress_callback(ndx, num_profilers)

        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)

        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:
            map_axes = [Axis(trend, plunge) for trend, plunge in kwargs['individual_axes_values']]
        else:
            raise Exception(f"Mapping method is {projection_method}")

        projected_points = defaultdict(list)

        # 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
        segment_profilers = list(line_profiler)
        segment_planes = [sp.segment_as_vertical_plane() for sp in segment_profilers]

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

            previous_segments_offset = 0

            for ndx, (segment_profiler, section_plane) in enumerate(
                    zip(segment_profilers, segment_planes), 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=map_axes if not isinstance(map_axes, list) else map_axes[ndx],
                    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)

                previous_segments_offset += segment_profiler.length()

        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) after each of
        the 'total' parallel profiles is projected, 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,
                **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)

            if progress_callback is not None:
                progress_callback(ndx, num_profilers)

        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)

        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:
            map_axes = [Axis(trend, plunge) for trend, plunge in kwargs['individual_axes_values']]
        else:
            raise Exception(f"Mapping method is {projection_method}")

        """
        if mapping_method['method'] not in ('nearest', 'common axis', 'individual axes', 'attitude'):
            return None, Error(True, caller_name(), Exception(
                f"Mapping method is '{mapping_method['method']}'. One of 'nearest', 'common axis', 'individual axes', 'attitude' expected"))

        if mapping_method['method'] == 'individual axes' and len(mapping_method['individual_axes_values']) != len(attitudes):
            return None, Error(True, caller_name(), Exception(
                f"Individual axes values are {len(mapping_method['individual_axes_values'])} but attitudes are {len(attitudes)}"))

        if mapping_method['method'] == 'nearest':
            map_axes = None
        elif mapping_method['method'] == 'common axis':
            map_axes = Axis(mapping_method['trend'], mapping_method['plunge'])
        else:
            map_axes = [Axis(trend, plunge) for trend, plunge in mapping_method['individual_axes_values']]
        """

        projected_attitudes = defaultdict(list)

        # 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
        segment_profilers = list(line_profiler)
        segment_planes = [sp.segment_as_vertical_plane() for sp in segment_profilers]

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

            previous_segments_offset = 0

            for ndx, (segment_profiler, section_plane) in enumerate(
                    zip(segment_profilers, segment_planes), 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=map_axes if not isinstance(map_axes, list) else map_axes[ndx],
                    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)

                previous_segments_offset += segment_profiler.length()

        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())
