#For PyCharm
from geosoft import gxpy

gxc = gxpy.gx.GXpy()  # Initialize Geosoft context first

# -----------------------------
# New/Updated Imports for Optimization
# -----------------------------
import geosoft.gxpy.gdb as gxdb
import geopandas as gpd
from shapely.geometry import Point, LineString
import pandas as pd
import numpy as np
from itertools import combinations
import math
import os
import time

# Suppress GeoPandas/Shapely warnings
import warnings

warnings.filterwarnings("ignore", category=UserWarning)


# -----------------------------
# Helper functions (Kept from original)
# -----------------------------
def flatten(arr):
    """Flatten nested lists/arrays from Geosoft VV channels"""
    out = []
    for item in arr:
        if item is None:
            continue
        if isinstance(item, (list, np.ndarray)):
            out.extend(flatten(item))
        else:
            try:
                out.append(float(item))
            except:
                pass
    return np.array(out, dtype=float)


def are_parallel(line1, line2, tolerance_degrees=30.0):
    """Checks if two LineString objects are approximately parallel."""
    coords1 = line1.coords
    coords2 = line2.coords

    if len(coords1) < 2 or len(coords2) < 2: return False

    dx1 = coords1[-1][0] - coords1[0][0]
    dy1 = coords1[-1][1] - coords1[0][1]
    dx2 = coords2[-1][0] - coords2[0][0]
    dy2 = coords2[-1][1] - coords2[0][1]

    if (dx1 == 0 and dy1 == 0) or (dx2 == 0 and dy2 == 0): return False

    angle1 = math.atan2(dy1, dx1)
    angle2 = math.atan2(dy2, dx2)
    diff_rad = abs(angle1 - angle2) % math.pi
    if diff_rad > math.pi / 2:
        diff_rad = math.pi - diff_rad
    diff_deg = math.degrees(diff_rad)

    return diff_deg <= tolerance_degrees


def clean_points(pts):
    """Removes NaN points and sequential duplicate points."""
    cleaned = []
    last = None
    for p in pts:
        if not (math.isnan(p.x) or math.isnan(p.y)):
            if last is None or (p.x != last.x or p.y != last.y):
                cleaned.append(p)
                last = p
    return cleaned


# -----------------------------
# User parameters
# -----------------------------
t_start = time.time()
gdb_path = r"D:\P110-2025\Database\P304-2024_Split_A1.gdb"
check_column = "Signal"
alt_column = "Alt"
buffer_distance = 0.5
parallel_tolerance_deg = 10.0

print(f"Opening GDB: {gdb_path}")

# -----------------------------
# Read GDB lines/channels (No change needed here)
# -----------------------------
records = []
with gxdb.Geosoft_gdb.open(gdb_path) as gdb:
    lines = gdb.list_lines()
    print(f"Found {len(lines)} lines.")
    for idx, line in enumerate(lines, start=1):
        gdb.current_line = line
        chs = gdb.list_channels()
        if not all(c in chs for c in ["X", "Y", check_column, alt_column]):
            print(f"[{idx}/{len(lines)}] Skipped line {line} (missing channels)")
            continue
        X = flatten(gdb.read_channel(line, "X"))
        Y = flatten(gdb.read_channel(line, "Y"))
        S = flatten(gdb.read_channel(line, check_column))
        Alt = flatten(gdb.read_channel(line, alt_column))
        n = min(len(X), len(Y), len(S), len(Alt))
        for i in range(n):
            records.append({
                "row_id": len(records),  # Add unique row_id for stable merge later
                "Line": line,
                "X": X[i],
                "Y": Y[i],
                check_column: S[i],
                alt_column: Alt[i]
            })

df = pd.DataFrame(records)

# -----------------------------
# BUILD GEODATAFRAMES (The Optimization Base)
# -----------------------------

# 1. Create GeoDataFrame of Points
geometry_points = [Point(x, y) for x, y in zip(df["X"], df["Y"])]
gdf_points = gpd.GeoDataFrame(
    df.assign(geometry=geometry_points),
    geometry='geometry',
    crs="EPSG:4326"  # Using a dummy CRS
)

# 2. Build Line Geometries
line_geoms = {}
for line, group in gdf_points.groupby("Line"):
    pts = clean_points(list(group.sort_values("row_id").geometry))
    if len(pts) >= 2:
        try:
            ls = LineString(pts)
            if ls.is_valid:
                line_geoms[line] = ls
        except:
            continue

# 3. Create GeoDataFrame of Lines (for spatial indexing)
gdf_lines = gpd.GeoDataFrame(
    {'Line': line_geoms.keys(), 'geometry': line_geoms.values()},
    geometry='geometry',
    crs=gdf_points.crs
).reset_index(drop=True)

# -----------------------------
# OPTIMIZED INTERSECTION DETECTION
# -----------------------------
intersection_data = []
sindex = gdf_lines.sindex  # Build R-tree index

print(f"\nComputing line intersections using Spatial Index on {len(gdf_lines)} lines...")
# Iterate over lines and query the index for potential neighbors (bounding box overlap)
for i in range(len(gdf_lines)):
    line_i = gdf_lines.loc[i]
    possible_matches_index = list(sindex.intersection(line_i.geometry.bounds))

    for j in possible_matches_index:
        # Check each pair once (j > i) and skip self-check
        if i < j:
            line_j = gdf_lines.loc[j]

            # Check parallel lines
            if are_parallel(line_i.geometry, line_j.geometry, tolerance_degrees=parallel_tolerance_deg):
                continue

            # Check for intersection
            if line_i.geometry.intersects(line_j.geometry):
                inter = line_i.geometry.intersection(line_j.geometry)
                if not inter.is_empty:
                    buf = inter.buffer(buffer_distance)
                    intersection_data.append({
                        'line_id1': line_i['Line'],
                        'line_id2': line_j['Line'],
                        'intersecting_lines_pair': f"{line_i['Line']}_{line_j['Line']}",
                        'geometry': buf
                    })

if not intersection_data:
    print("No line intersections found. Skipping mark/writeback.")
    exit()  # Exit early if no work is needed

# GeoDataFrame of all intersection buffers
gdf_buffers = gpd.GeoDataFrame(
    intersection_data,
    geometry='geometry',
    crs=gdf_points.crs
)
print(f"Identified {len(gdf_buffers)} unique intersection buffers.")

# -----------------------------
# OPTIMIZED POINT MARKING (SPATIAL JOIN)
# -----------------------------

# Vectorized Spatial Join: find all points contained within any buffer
joined_points = gpd.sjoin(
    gdf_points,
    gdf_buffers,
    how="inner",
    predicate="within"  # Check if the point is within the buffer polygon
)

# Get the unique original row_ids and the corresponding intersecting_lines_pair
intersecting_rows = (
    joined_points[['row_id', 'intersecting_lines_pair']]
    .drop_duplicates(subset=['row_id'])
    .set_index('row_id')
    .rename(columns={'intersecting_lines_pair': 'intersecting_lines'})
)

# Merge results back to the original DataFrame
df['intersection'] = 0
df['intersecting_lines'] = None
df.loc[intersecting_rows.index, 'intersection'] = 1
df.loc[intersecting_rows.index, 'intersecting_lines'] = intersecting_rows['intersecting_lines'].values

# -----------------------------
# Highest mean Signal (No change needed, using optimized df)
# -----------------------------
print("Calculating mean signal and altitude...")
sig_means = df.groupby(["intersecting_lines", "Line"])[check_column].mean().reset_index()
sig_means["min_sig"] = sig_means.groupby("intersecting_lines")[check_column].transform("min")
sig_means["is_highest_mean_Signal"] = (sig_means[check_column] == sig_means["min_sig"]).astype(int)

df = df.merge(
    sig_means[["intersecting_lines", "Line", "is_highest_mean_Signal"]],
    on=["intersecting_lines", "Line"],
    how="left"
).fillna({"is_highest_mean_Signal": 0})

# -----------------------------
# Lowest mean Alt (No change needed, using optimized df)
# -----------------------------
alt_means = df.groupby(["intersecting_lines", "Line"])[alt_column].mean().reset_index()
alt_means["max_alt"] = alt_means.groupby("intersecting_lines")[alt_column].transform("max")
alt_means["is_lowest_mean_Alt"] = (alt_means[alt_column] == alt_means["max_alt"]).astype(int)

df = df.merge(
    alt_means[["intersecting_lines", "Line", "is_lowest_mean_Alt"]],
    on=["intersecting_lines", "Line"],
    how="left"
).fillna({"is_lowest_mean_Alt": 0})

# -----------------------------
# Write back to GDB (Final Correction)
# -----------------------------
t_start_write = time.time()
print("\nWriting channels back to GDB...")

# Define only the final NUMERIC channels to write back
NUMERIC_CHANNELS_TO_WRITE = [
    "intersection",
    "is_highest_mean_Signal",
    "is_lowest_mean_Alt"
]

# --- FIX APPLIED HERE: Removed the 'mode' argument ---
with gxdb.Geosoft_gdb.open(gdb_path) as gdb:
    # 1. DELETE AND CREATE ALL NUMERIC CHANNELS ONCE
    for ch in NUMERIC_CHANNELS_TO_WRITE:
        if ch in gdb.list_channels():
            # This requires write access
            gdb.delete_channel(ch)

        # This requires write access and ensures the channel exists
        gdb.new_channel(ch, dtype=np.float32)

        # 2. WRITE DATA LINE BY LINE
    for line in gdb.list_lines():
        gdb.current_line = line
        chunk = df[df["Line"] == line]
        if chunk.empty: continue

        # Write data for all three numeric columns
        gdb.write_channel(line, "intersection", chunk["intersection"].astype(float).tolist())
        gdb.write_channel(line, "is_highest_mean_Signal", chunk["is_highest_mean_Signal"].astype(float).tolist())
        gdb.write_channel(line, "is_lowest_mean_Alt", chunk["is_lowest_mean_Alt"].astype(float).tolist())

t_end_write = time.time()
print("\n✔ Updated GDB successfully!")
print(f"Channels written: {', '.join(NUMERIC_CHANNELS_TO_WRITE)}")
print(f"Writeback time: {t_end_write - t_start_write:.2f} seconds.")