from typing import Union

import numpy as np
from qgis.core import *
from qgis.gui import *


@qgsfunction(args="auto", group="Custom", referenced_columns=[])
def get_geometry_with_m_dimension(
    geometry: QgsGeometry,
    original_geometry: Union[int, str, QgsGeometry],
    layer_crs: int,
    desurvey_crs: int,
):
    """
    Return a 3D linestring, with M attribute storing cumulated 3D length.
    Args:
        - geometry : linestring geometry in database srid
        - original_geometry : linestring geometry in original srid or integer if it don't exist. Could be wkt in spatialite
        - collar_dict : dictionnary of collars
    """
    new_points = []
    segment_3Dlength = 0

    if isinstance(original_geometry, str):
        original_geometry = QgsGeometry.fromWkt(original_geometry)

    if isinstance(original_geometry, int):
        # reprojection in a planar system to get right 3Dlength
        proj_points = geometry.constGet().points()
        destCrs = QgsCoordinateReferenceSystem(desurvey_crs)
        sourceCrs = QgsCoordinateReferenceSystem(layer_crs)
        tr = QgsCoordinateTransform(sourceCrs, destCrs, QgsProject.instance())
        _ = [pt.transform(tr) for pt in proj_points]
    else:
        proj_points = original_geometry.constGet().points()

    vertices = geometry.constGet().points()
    for point in range(len(vertices) - 1):

        x = vertices[point].x()
        y = vertices[point].y()
        z = vertices[point].z()
        new_points.append(QgsPoint(x, y, z, segment_3Dlength))

        segment_3Dlength += QgsLineString(
            [proj_points[point], proj_points[point + 1]]
        ).length3D()

    # last vertex
    last_point = vertices[-1]
    new_points.append(
        QgsPoint(
            last_point.x(),
            last_point.y(),
            last_point.z(),
            QgsLineString(proj_points).length3D(),
        )
    )

    new_line = QgsLineString(new_points)
    new_geom = QgsGeometry(new_line)
    return new_geom


@qgsfunction(args="auto", group="Custom", referenced_columns=[])
def get_projected_length_of_3D_line(geom, length):
    """
    QGIS function line_interpolate_point need 2d distance, this function return 2d distance of 3d distance.
    """
    """
    QGIS function line_interpolate_point need 2d distance, this function return 2d distance of 3d distance.
    """
    vertices = geom.constGet().points()
    proj_length = 0
    total_3Dlength = 0
    for point in range(len(vertices) - 1):
        segment_3Dlength = QgsLineString(
            [vertices[point], vertices[point + 1]]
        ).length3D()
        if total_3Dlength + segment_3Dlength > length:
            proj_length += (
                (length - total_3Dlength) / segment_3Dlength
            ) * QgsLineString([vertices[point], vertices[point + 1]]).length()
            break
        total_3Dlength += segment_3Dlength
        proj_length += QgsLineString([vertices[point], vertices[point + 1]]).length()
    return proj_length


@qgsfunction(args="auto", group="Custom", referenced_columns=[])
def get_3d_distance_from_n_vertices(line, n_arcs, last_portion) -> float:
    """
    Compute length of a line part given number of arcs + remaining portion.
    """
    vertices = line.constGet().points()
    distance = 0
    scanned_arcs = 0

    for i in range(len(vertices) - 1):

        p1 = vertices[i]
        p2 = vertices[i + 1]

        z1 = p1.z()

        if z1 == np.nan:
            d = abs(p1.distance(p2))
        else:
            d = abs(p1.distance3D(p2))

        if scanned_arcs == n_arcs:
            distance += d * last_portion
            return distance

        distance += d
        scanned_arcs += 1


def get_n_vertices_from_z(line, target_z):
    """
    Return a tuple containing :
        - number of complete arcs between line start and intersection with z
        - lasting portion between last vertice and intersection point (expressed in % of intersected arc)

    """
    vertices = line.constGet().points()
    n_arcs = 0
    portion = 0

    for i in range(len(vertices) - 1):
        p1 = vertices[i]
        p2 = vertices[i + 1]

        # Obtenir les coordonnées XYZ
        x1, y1, z1 = p1.x(), p1.y(), p1.z()
        x2, y2, z2 = p2.x(), p2.y(), p2.z()

        # Vérifier si le Z cible se trouve entre ces deux points
        if z2 <= target_z <= z1:
            # Interpolation linéaire pour trouver l'exacte position
            ratio = (z1 - target_z) / (z1 - z2)
            x = x1 + ratio * (x2 - x1)
            y = y1 + ratio * (y2 - y1)

            d = (x1 - x) ** 2 + (y1 - y) ** 2 + (z1 - target_z) ** 2
            portion = (d**0.5) / abs(p1.distance3D(x2, y2, z2))
            return n_arcs, portion

        n_arcs += 1

    return None, None


def get_n_vertices_from_length(
    line,
    target_length,
    reproject: bool = True,
    layer_crs: int = None,
    desurvey_crs: int = None,
):
    """
    Return a tuple containing :
        - number of complete arcs between line start and intersection with length
        - lasting portion between last vertice and intersection point (expressed in % of intersected arc)

    """
    vertices = line.constGet().points()
    # transform to desurveying CRS
    if reproject and layer_crs is not None:
        destCrs = QgsCoordinateReferenceSystem(desurvey_crs)
        tr = QgsCoordinateTransform(layer_crs, destCrs, QgsProject.instance())
        _ = [point.transform(tr) for point in vertices]

    distance = 0
    n_arcs = 0
    portion = 0

    for i in range(len(vertices) - 1):
        p1 = vertices[i]
        p2 = vertices[i + 1]

        # arc length
        d = abs(p1.distance3D(p2))

        # Vérifier si le Z cible se trouve entre ces deux points
        if distance + d < target_length:
            distance += d
            n_arcs += 1
        else:

            delta = target_length - distance
            portion = delta / d
            return n_arcs, portion

    return None, None


def get_3d_length(
    line, reproject: bool = True, layer_crs: int = None, desurvey_crs: int = None
) -> float:
    """
    line: QgsGeometry
    """
    vertices = line.constGet().points()
    # transform to desurveying CRS
    if reproject and layer_crs is not None:
        destCrs = QgsCoordinateReferenceSystem(desurvey_crs)
        tr = QgsCoordinateTransform(layer_crs, destCrs, QgsProject.instance())
        _ = [point.transform(tr) for point in vertices]

    distance = 0

    for i in range(len(vertices) - 1):
        p1 = vertices[i]
        p2 = vertices[i + 1]

        # arc length
        d = abs(p1.distance3D(p2))
        distance += d
    return distance
