import csv
import os
import time
from qgis.PyQt.QtCore import QVariant, QThread, pyqtSignal
from qgis.PyQt.QtWidgets import QFileDialog
from qgis.core import (
    QgsWkbTypes, QgsProject, QgsGeometry, QgsVectorLayer, QgsField,
    QgsFeature, QgsDistanceArea, QgsSpatialIndex, QgsPointXY
)
from qgis.utils import iface
from qgis.PyQt import QtWidgets, uic, QtCore

FORM_CLASS, _ = uic.loadUiType(os.path.join(
    os.path.dirname(__file__), 'cyanlove_calc_point_min_geometry_distince_base.ui'))


class cyanlove_calc_point_min_geometry_distince(QtWidgets.QDockWidget, FORM_CLASS):
    closingPlugin = pyqtSignal()

    def __init__(self, parent=None):
        super(cyanlove_calc_point_min_geometry_distince, self).__init__(parent)
        self.export_thread = None
        self.iface = iface
        self.canvas = iface.mapCanvas()
        self.setupUi(self)
        self.setFloating(True)
        self.property_dict = None
        self.is_running = False
        self.setWindowFlags(QtCore.Qt.Dialog)

        self.btn_calc_point_distinct_geometry.clicked.connect(self.calc_point_min_distinct_geometry)

        # 初始化距离限制控件的联动
        self.checkBoxDistanceLimit.toggled.connect(self.doubleSpinBoxDistance.setEnabled)
        self.doubleSpinBoxDistance.setEnabled(False)

    def calc_point_min_distinct_geometry(self):
        options = QFileDialog.Options()
        options |= QFileDialog.ReadOnly
        csv_file_path, _ = QFileDialog.getSaveFileName(self,
                                                       "选择文件",
                                                       "",
                                                       "CSV文件 (*.csv);;所有文件 (*)",
                                                       options=options)
        if not csv_file_path:
            return

        layer1 = self.mMapLayerComboBox.currentLayer()
        layer2 = self.mMapLayerComboBox_2.currentLayer()
        if layer1 is None or layer2 is None:
            self.label_show.setText("请选择两个图层！")
            return

        geom_type1 = layer1.geometryType()  # 0=点,1=线,2=面
        geom_type2 = layer2.geometryType()

        # 允许的组合: (点, 线), (点, 面), (线, 点), (面, 点)
        valid = False
        if geom_type1 == 0 and geom_type2 in (1, 2):
            valid = True
        elif geom_type1 == 1 and geom_type2 == 0:
            valid = True
        elif geom_type1 == 2 and geom_type2 == 0:
            valid = True

        if not valid:
            self.label_show.setText("不支持的图层组合。请使用：点+线/面，线+点，或面+点。")
            return

        max_nearest_count = self.spinBoxNearestCount.value()
        use_distance_limit = self.checkBoxDistanceLimit.isChecked()
        distance_limit_m = self.doubleSpinBoxDistance.value() if use_distance_limit else 0.0

        self.progressBar.setRange(0, 100)
        self.progressBar.setValue(0)
        self.progressBar.show()
        self.export_thread = ExportThread(layer1, layer2, csv_file_path,
                                          max_nearest_count, use_distance_limit, distance_limit_m)
        self.export_thread.progress.connect(self.update_progress)
        self.export_thread.updatelabel.connect(self.update_label)
        self.export_thread.finished.connect(self.export_finished)
        self.export_thread.start()

    def update_label(self, value):
        self.label_show.setText(value)

    def update_progress(self, value):
        self.progressBar.setValue(value)

    def export_finished(self):
        self.progressBar.setValue(100)


class ExportThread(QThread):
    progress = pyqtSignal(int)
    updatelabel = pyqtSignal(str)

    def __init__(self, layer1, layer2, csv_file_path,
                 max_nearest_count, use_distance_limit, distance_limit_m):
        super().__init__()
        self.layer1 = layer1
        self.layer2 = layer2
        self.csv_file_path = csv_file_path
        self.max_nearest_count = max_nearest_count
        self.use_distance_limit = use_distance_limit
        self.distance_limit_m = distance_limit_m

    def run(self):
        try:
            geom_type1 = self.layer1.geometryType()
            geom_type2 = self.layer2.geometryType()

            if geom_type1 == 0 and geom_type2 in (1, 2):
                self._run_point_to_line_or_polygon()
            elif geom_type1 == 1 and geom_type2 == 0:
                self._run_line_to_point()
            elif geom_type1 == 2 and geom_type2 == 0:
                self._run_polygon_to_point()
            else:
                self.updatelabel.emit(f"不支持的组合: 图层1类型={geom_type1}, 图层2类型={geom_type2}")
        except Exception as e:
            self.updatelabel.emit(f"运行错误: {str(e)}")
            import traceback
            traceback.print_exc()

    # --------------------- 辅助方法 ---------------------
    def _apply_count_limit(self, candidates):
        """
        对已排序的候选列表应用数量限制，按唯一距离值分组。
        取前 max_nearest_count 个不同的距离值，返回所有这些距离值对应的所有要素。
        """
        if not candidates:
            return []

        # 按距离分组
        distance_groups = {}
        for candidate in candidates:
            dist = candidate[1]
            if dist not in distance_groups:
                distance_groups[dist] = []
            distance_groups[dist].append(candidate)

        # 所有唯一的距离值（已排序）
        unique_distances = sorted(distance_groups.keys())

        # 如果唯一距离值数量 <= max_nearest_count，返回所有候选
        if len(unique_distances) <= self.max_nearest_count:
            return candidates

        # 只取前 max_nearest_count 个距离值
        selected_distances = unique_distances[:self.max_nearest_count]

        # 返回这些距离值对应的所有要素
        result = []
        for dist in selected_distances:
            result.extend(distance_groups[dist])
        return result

    def _get_candidate_features(self, source_geom, target_features_dict, target_index,
                                distance_area, is_point_source=True):
        """
        获取候选要素列表。
        核心策略：无论是否使用距离限制，都先按无距离限制的逻辑获取足够多的最近候选
        （确保不同距离值的数量达到 max_nearest_count 或已获取全部要素），
        然后再根据距离限制进行过滤。
        """
        if not target_features_dict:
            return []

        # 获取源点坐标
        if is_point_source:
            source_point = source_geom.asPoint()
        else:
            source_point = source_geom.centroid().asPoint()

        total_targets = len(target_features_dict)

        # 第一步：获取候选（不受距离限制，保证不同距离值数量足够）
        fetch_count = min(total_targets, max(self.max_nearest_count * 10, 100))
        nearest_ids = target_index.nearestNeighbor(source_point, fetch_count)
        candidates = []

        for fid in nearest_ids:
            feat = target_features_dict.get(fid)
            if not feat or not feat.geometry() or feat.geometry().isEmpty():
                continue
            line = source_geom.shortestLine(feat.geometry())
            if line is None or line.isEmpty():
                continue
            dist = distance_area.measureLength(line)
            pts = line.asPolyline()
            nearest_pt = pts[-1] if len(pts) >= 2 else (pts[0] if pts else source_point)
            candidates.append((fid, dist, nearest_pt, line))

        # 如果不同距离值数量不足 max_nearest_count 且还没获取全部要素，继续获取更多
        unique_distances = len(set([c[1] for c in candidates]))
        if unique_distances < self.max_nearest_count and fetch_count < total_targets:
            all_ids = target_index.nearestNeighbor(source_point, total_targets)
            existing_fids = {fid for fid, _, _, _ in candidates}
            for fid in all_ids:
                if fid in existing_fids:
                    continue
                feat = target_features_dict.get(fid)
                if not feat or not feat.geometry() or feat.geometry().isEmpty():
                    continue
                line = source_geom.shortestLine(feat.geometry())
                if line is None or line.isEmpty():
                    continue
                dist = distance_area.measureLength(line)
                pts = line.asPolyline()
                nearest_pt = pts[-1] if len(pts) >= 2 else (pts[0] if pts else source_point)
                candidates.append((fid, dist, nearest_pt, line))
                if len(set([c[1] for c in candidates])) >= self.max_nearest_count:
                    break

        # 按距离排序
        candidates.sort(key=lambda x: x[1])

        # 第二步：如果勾选了距离限制，过滤掉超出限制的候选
        if self.use_distance_limit:
            candidates = [c for c in candidates if c[1] <= self.distance_limit_m]

        return candidates

    # --------------------- 点 → 线/面 ---------------------
    def _run_point_to_line_or_polygon(self):
        layer1 = self.layer1  # 点
        layer2 = self.layer2  # 线或面

        if layer1.featureCount() == 0:
            self.updatelabel.emit("错误：点图层没有要素！")
            return
        if layer2.featureCount() == 0:
            self.updatelabel.emit("错误：线/面图层没有要素！")
            return

        fields1 = [f.name() for f in layer1.fields()]
        fields2 = [f.name() for f in layer2.fields()]
        out_fields = fields1 + fields2 + ["最短距离(米)", "最近点经度", "最近点纬度"]

        total_points = layer1.featureCount()
        is_polygon = (layer2.geometryType() == 2)

        self.updatelabel.emit(
            f"开始计算，共 {total_points} 个点，目标要素 {layer2.featureCount()} 个，最多返回 {self.max_nearest_count} 个不同距离值...")

        distance_area = QgsDistanceArea()
        distance_area.setEllipsoid('WGS84')
        distance_area.setSourceCrs(layer1.crs(), QgsProject.instance().transformContext())

        target_index = QgsSpatialIndex(layer2.getFeatures())
        target_features = {f.id(): f for f in layer2.getFeatures()}

        line_layer = QgsVectorLayer(f"LineString?crs={layer1.crs().authid()}", "最短连接线", "memory")
        line_provider = line_layer.dataProvider()
        line_provider.addAttributes([
            QgsField("source_id", QVariant.Int),
            QgsField("distance_m", QVariant.Double),
            QgsField("nearest_x", QVariant.Double),
            QgsField("nearest_y", QVariant.Double)
        ])
        line_layer.updateFields()
        line_feat_list = []

        found_sources = 0
        total_output_rows = 0

        with open(self.csv_file_path, 'w', newline='', encoding='utf-8-sig') as csv_file:
            writer = csv.writer(csv_file)
            writer.writerow(out_fields)

            for idx, point_feat in enumerate(layer1.getFeatures()):
                point_geom = point_feat.geometry()
                base_row = [str(point_feat[f]) if point_feat[f] is not None else "" for f in fields1]

                if point_geom is None or point_geom.isEmpty():
                    writer.writerow(base_row + [""] * (len(fields2) + 3))
                    self._update_progress(idx, total_points)
                    continue

                candidates = []

                # 如果是面图层，先找出所有包含该点的面（距离0）
                if is_polygon:
                    bbox = point_geom.boundingBox()
                    candidate_ids = target_index.intersects(bbox)
                    for fid in candidate_ids:
                        feat = target_features.get(fid)
                        if feat and feat.geometry() and feat.geometry().contains(point_geom):
                            line = QgsGeometry.fromPolylineXY([point_geom.asPoint(), point_geom.asPoint()])
                            candidates.append((fid, 0.0, point_geom.asPoint(), line))

                # 获取其他候选（距离>0）
                other_candidates = self._get_candidate_features(
                    point_geom, target_features, target_index,
                    distance_area, is_point_source=True
                )
                existing_fids = {fid for fid, _, _, _ in candidates}
                for fid, dist, nearest_pt, line in other_candidates:
                    if fid not in existing_fids:
                        candidates.append((fid, dist, nearest_pt, line))

                candidates.sort(key=lambda x: x[1])
                candidates = self._apply_count_limit(candidates)

                if not candidates:
                    writer.writerow(base_row + [""] * (len(fields2) + 3))
                    self._update_progress(idx, total_points)
                    continue

                found_sources += 1
                for fid, dist, nearest_pt, line in candidates:
                    feat = target_features[fid]
                    row = base_row.copy()
                    for fname in fields2:
                        val = feat[fname]
                        row.append(str(val) if val is not None else "")
                    row.append(f"{dist:.2f}")
                    row.append(f"{nearest_pt.x():.6f}" if nearest_pt else "")
                    row.append(f"{nearest_pt.y():.6f}" if nearest_pt else "")
                    writer.writerow(row)
                    total_output_rows += 1

                    if line and not line.isEmpty() and dist > 0:
                        line_feat = QgsFeature()
                        line_feat.setGeometry(line)
                        line_feat.setAttributes([point_feat.id(), dist,
                                                 nearest_pt.x() if nearest_pt else None,
                                                 nearest_pt.y() if nearest_pt else None])
                        line_feat_list.append(line_feat)

                self._update_progress(idx, total_points)
                if (idx + 1) % max(1, total_points // 10) == 0:
                    unique_dist_count = len(set([c[1] for c in candidates]))
                    self.updatelabel.emit(
                        f"已处理 {idx + 1}/{total_points} 点，"
                        f"已输出 {total_output_rows} 行，涉及源要素 {found_sources} 个，"
                        f"当前要素输出了 {len(candidates)} 个候选（{unique_dist_count} 个不同距离）"
                    )

        if line_feat_list:
            line_provider.addFeatures(line_feat_list)
            line_layer.updateExtents()
            QgsProject.instance().addMapLayer(line_layer)
            self.updatelabel.emit(f"已生成最短连接线图层，共 {len(line_feat_list)} 条线")

        self.updatelabel.emit(
            f"计算完成！共 {total_points} 点，其中 {found_sources} 个点找到了候选，"
            f"总共输出 {total_output_rows} 行（每个候选一行）。"
        )

    # --------------------- 线 → 点 ---------------------
    def _run_line_to_point(self):
        line_layer = self.layer1
        point_layer = self.layer2

        if line_layer.featureCount() == 0:
            self.updatelabel.emit("错误：线图层没有要素！")
            return
        if point_layer.featureCount() == 0:
            self.updatelabel.emit("错误：点图层没有要素！")
            return

        line_fields = [f.name() for f in line_layer.fields()]
        point_fields = [f.name() for f in point_layer.fields()]
        out_fields = line_fields + point_fields + ["最短距离(米)", "线上最近点经度", "线上最近点纬度"]

        total_lines = line_layer.featureCount()
        self.updatelabel.emit(
            f"开始计算，共 {total_lines} 条线，点要素 {point_layer.featureCount()} 个，最多返回 {self.max_nearest_count} 个不同距离值...")

        distance_area = QgsDistanceArea()
        distance_area.setEllipsoid('WGS84')
        distance_area.setSourceCrs(line_layer.crs(), QgsProject.instance().transformContext())

        point_index = QgsSpatialIndex(point_layer.getFeatures())
        point_features = {f.id(): f for f in point_layer.getFeatures()}

        conn_layer = QgsVectorLayer(f"LineString?crs={line_layer.crs().authid()}", "线到点连接线", "memory")
        conn_provider = conn_layer.dataProvider()
        conn_provider.addAttributes([
            QgsField("line_id", QVariant.Int),
            QgsField("point_id", QVariant.Int),
            QgsField("distance_m", QVariant.Double),
            QgsField("nearest_on_line_x", QVariant.Double),
            QgsField("nearest_on_line_y", QVariant.Double)
        ])
        conn_layer.updateFields()
        conn_feat_list = []

        found_sources = 0
        total_output_rows = 0

        with open(self.csv_file_path, 'w', newline='', encoding='utf-8-sig') as csv_file:
            writer = csv.writer(csv_file)
            writer.writerow(out_fields)

            for idx, line_feat in enumerate(line_layer.getFeatures()):
                line_geom = line_feat.geometry()
                base_row = [str(line_feat[f]) if line_feat[f] is not None else "" for f in line_fields]

                if line_geom is None or line_geom.isEmpty():
                    writer.writerow(base_row + [""] * (len(point_fields) + 3))
                    self._update_progress(idx, total_lines)
                    continue

                candidates = self._get_candidate_features(line_geom, point_features, point_index,
                                                          distance_area, is_point_source=False)
                candidates = self._apply_count_limit(candidates)

                if not candidates:
                    writer.writerow(base_row + [""] * (len(point_fields) + 3))
                    self._update_progress(idx, total_lines)
                    continue

                found_sources += 1
                for fid, dist, nearest_pt, line_seg in candidates:
                    point_feat = point_features[fid]
                    row = base_row.copy()
                    for fname in point_fields:
                        val = point_feat[fname]
                        row.append(str(val) if val is not None else "")
                    row.append(f"{dist:.2f}")
                    row.append(f"{nearest_pt.x():.6f}" if nearest_pt else "")
                    row.append(f"{nearest_pt.y():.6f}" if nearest_pt else "")
                    writer.writerow(row)
                    total_output_rows += 1

                    if line_seg and not line_seg.isEmpty() and dist > 0:
                        conn_feat = QgsFeature()
                        conn_feat.setGeometry(line_seg)
                        conn_feat.setAttributes([
                            line_feat.id(),
                            point_feat.id(),
                            dist,
                            nearest_pt.x() if nearest_pt else None,
                            nearest_pt.y() if nearest_pt else None
                        ])
                        conn_feat_list.append(conn_feat)

                self._update_progress(idx, total_lines)
                if (idx + 1) % max(1, total_lines // 10) == 0:
                    unique_dist_count = len(set([c[1] for c in candidates]))
                    self.updatelabel.emit(
                        f"已处理 {idx + 1}/{total_lines} 条线，"
                        f"已输出 {total_output_rows} 行，涉及源要素 {found_sources} 个，"
                        f"当前要素输出了 {len(candidates)} 个候选（{unique_dist_count} 个不同距离）"
                    )

        if conn_feat_list:
            conn_provider.addFeatures(conn_feat_list)
            conn_layer.updateExtents()
            QgsProject.instance().addMapLayer(conn_layer)
            self.updatelabel.emit(f"已生成线到点连接线图层，共 {len(conn_feat_list)} 条线")

        self.updatelabel.emit(
            f"计算完成！共 {total_lines} 条线，其中 {found_sources} 条线找到了候选，"
            f"总共输出 {total_output_rows} 行（每个候选一行）。"
        )

    # --------------------- 面 → 点 ---------------------
    def _run_polygon_to_point(self):
        polygon_layer = self.layer1
        point_layer = self.layer2

        try:
            if polygon_layer.featureCount() == 0:
                self.updatelabel.emit("错误：面图层没有要素！")
                return
            if point_layer.featureCount() == 0:
                self.updatelabel.emit("错误：点图层没有要素！")
                return

            polygon_fields = [f.name() for f in polygon_layer.fields()]
            point_fields = [f.name() for f in point_layer.fields()]
            out_fields = polygon_fields + point_fields + ["最短距离(米)", "面上最近点经度", "面上最近点纬度",
                                                          "是否在面内"]

            total_polygons = polygon_layer.featureCount()
            self.updatelabel.emit(
                f"开始计算，共 {total_polygons} 个面，点要素 {point_layer.featureCount()} 个，最多返回 {self.max_nearest_count} 个不同距离值...")

            distance_area = QgsDistanceArea()
            distance_area.setEllipsoid('WGS84')
            distance_area.setSourceCrs(polygon_layer.crs(), QgsProject.instance().transformContext())

            point_index = QgsSpatialIndex(point_layer.getFeatures())
            point_features = {f.id(): f for f in point_layer.getFeatures()}

            conn_layer = QgsVectorLayer(f"LineString?crs={polygon_layer.crs().authid()}", "面到点连接线", "memory")
            conn_provider = conn_layer.dataProvider()
            conn_provider.addAttributes([
                QgsField("polygon_id", QVariant.Int),
                QgsField("point_id", QVariant.Int),
                QgsField("distance_m", QVariant.Double),
                QgsField("nearest_on_polygon_x", QVariant.Double),
                QgsField("nearest_on_polygon_y", QVariant.Double),
                QgsField("inside", QVariant.Int)
            ])
            conn_layer.updateFields()
            conn_feat_list = []

            found_sources = 0
            total_output_rows = 0

            with open(self.csv_file_path, 'w', newline='', encoding='utf-8-sig') as csv_file:
                writer = csv.writer(csv_file)
                writer.writerow(out_fields)

                for idx, polygon_feat in enumerate(polygon_layer.getFeatures()):
                    polygon_geom = polygon_feat.geometry()
                    base_row = [str(polygon_feat[f]) if polygon_feat[f] is not None else "" for f in polygon_fields]

                    if polygon_geom is None or polygon_geom.isEmpty():
                        writer.writerow(base_row + [""] * (len(point_fields) + 4))
                        self._update_progress(idx, total_polygons)
                        continue

                    # 找出所有在面内的点（距离0）
                    inside_candidates = []
                    bbox = polygon_geom.boundingBox()
                    candidate_ids = point_index.intersects(bbox)
                    for fid in candidate_ids:
                        pt_feat = point_features.get(fid)
                        if pt_feat and pt_feat.geometry() and polygon_geom.contains(pt_feat.geometry()):
                            line = QgsGeometry.fromPolylineXY(
                                [pt_feat.geometry().asPoint(), pt_feat.geometry().asPoint()])
                            inside_candidates.append((fid, 0.0, pt_feat.geometry().asPoint(), line))

                    # 获取外部候选（距离>0）
                    other_candidates = self._get_candidate_features(
                        polygon_geom, point_features, point_index,
                        distance_area, is_point_source=False
                    )
                    existing_fids = {fid for fid, _, _, _ in inside_candidates}
                    for fid, dist, nearest_pt, line in other_candidates:
                        if fid not in existing_fids:
                            inside_candidates.append((fid, dist, nearest_pt, line))

                    inside_candidates.sort(key=lambda x: x[1])
                    candidates = self._apply_count_limit(inside_candidates)

                    if not candidates:
                        writer.writerow(base_row + [""] * (len(point_fields) + 4))
                        self._update_progress(idx, total_polygons)
                        continue

                    found_sources += 1
                    for fid, dist, nearest_pt, line in candidates:
                        point_feat = point_features[fid]
                        row = base_row.copy()
                        for fname in point_fields:
                            val = point_feat[fname]
                            row.append(str(val) if val is not None else "")
                        row.append(f"{dist:.2f}")
                        row.append(f"{nearest_pt.x():.6f}" if nearest_pt else "")
                        row.append(f"{nearest_pt.y():.6f}" if nearest_pt else "")
                        row.append("是" if dist == 0 else "否")
                        writer.writerow(row)
                        total_output_rows += 1

                        if line and not line.isEmpty() and dist > 0:
                            conn_feat = QgsFeature()
                            conn_feat.setGeometry(line)
                            conn_feat.setAttributes([
                                polygon_feat.id(),
                                point_feat.id(),
                                dist,
                                nearest_pt.x() if nearest_pt else None,
                                nearest_pt.y() if nearest_pt else None,
                                1 if dist == 0 else 0
                            ])
                            conn_feat_list.append(conn_feat)

                    self._update_progress(idx, total_polygons)
                    if (idx + 1) % max(1, total_polygons // 10) == 0:
                        unique_dist_count = len(set([c[1] for c in candidates]))
                        self.updatelabel.emit(
                            f"已处理 {idx + 1}/{total_polygons} 个面，"
                            f"已输出 {total_output_rows} 行，涉及源要素 {found_sources} 个，"
                            f"当前要素输出了 {len(candidates)} 个候选（{unique_dist_count} 个不同距离）"
                        )

            if conn_feat_list:
                conn_provider.addFeatures(conn_feat_list)
                conn_layer.updateExtents()
                QgsProject.instance().addMapLayer(conn_layer)
                self.updatelabel.emit(f"已生成面到点连接线图层，共 {len(conn_feat_list)} 条线")
            else:
                self.updatelabel.emit("未生成连接线图层（所有点都在面内或距离为0）")

            self.updatelabel.emit(
                f"计算完成！共 {total_polygons} 个面，其中 {found_sources} 个面找到了候选，"
                f"总共输出 {total_output_rows} 行（每个候选一行）。"
            )

        except Exception as e:
            self.updatelabel.emit(f"面到点计算错误: {str(e)}")
            import traceback
            traceback.print_exc()
            raise

    def _update_progress(self, current, total):
        progress_val = int((current + 1) / total * 100)
        self.progress.emit(progress_val)