import os
import numpy as np
import pandas as pd
import ezdxf
from pathlib import Path
from scipy.spatial import KDTree

def fit_and_solve_closed_polynomial(pts: np.ndarray) -> dict:
    """
    Fits 5 points into a 2D quadric polynomial curve: Ax^2 + Bxy + Cy^2 + Dx + Ey + F = 0
    Returns solved algebraic coefficients and roots for the dataset.
    """
    if len(pts) < 5:
        # Pad with zeros if fewer than 5 points exist
        return {"A": 0.0, "B": 0.0, "C": 0.0, "D": 0.0, "E": 0.0, "F": 0.0, "root_x": 0.0, "root_y": 0.0}

    x = pts[:, 0]
    y = pts[:, 1]

    # Form design matrix for implicit conic curve fitting: Ax^2 + Bxy + Cy^2 + Dx + Ey = 1 (F = -1)
    M = np.column_stack([x**2, x * y, y**2, x, y])
    rhs = np.ones(len(pts))

    try:
        # Solve algebraic coefficients via least squares SVD
        coeffs, _, _, _ = np.linalg.lstsq(M, rhs, rcond=None)
        A, B, C, D, E = coeffs
        F = -1.0
    except Exception:
        A, B, C, D, E, F = 0.0, 0.0, 0.0, 0.0, 0.0, 0.0

    # Solve critical manifold center point (root solutions where derivatives = 0)
    # dP/dx = 2Ax + By + D = 0 ; dP/dy = Bx + 2Cy + E = 0
    det = 4 * A * C - B**2
    if abs(det) > 1e-6:
        root_x = (B * E - 2 * C * D) / det
        root_y = (B * D - 2 * A * E) / det
    else:
        root_x, root_y = np.mean(x), np.mean(y)

    return {
        "poly_A": float(A),
        "poly_B": float(B),
        "poly_C": float(C),
        "poly_D": float(D),
        "poly_E": float(E),
        "poly_F": float(F),
        "root_x": float(root_x),
        "root_y": float(root_y)
    }

def process_mesh(input_csv_path: str, dxf_dir_path: str, max_faces=500, precision=2) -> pd.DataFrame:
    """
    Reads DXF files, interpolates closed curves from 5 nearest neighbor points,
    solves quadric polynomials, and returns the solved polynomial dataset.
    """
    print(f"    -> Processing Polynomial Geometry directly from DXF: {dxf_dir_path}")
    
    dxf_dir = Path(dxf_dir_path)
    if not dxf_dir.exists():
        raise FileNotFoundError(f"DXF directory not found at {dxf_dir}")

    all_points = []
    dxf_files = sorted(list(dxf_dir.glob("*.dxf")))

    if len(dxf_files) > 60:
        dxf_files = dxf_files[::2] 

    for dxf_file in dxf_files:
        try:
            frame_id = int(''.join(filter(str.isdigit, dxf_file.stem)) or 0)
            doc = ezdxf.readfile(dxf_file)
            msp = doc.modelspace()
            
            entities = msp.query('LINE')
            sampling_rate = 2 if len(entities) > 500 else 1
            
            for i, entity in enumerate(entities):
                if i % sampling_rate != 0:
                    continue
                all_points.append({'t': frame_id, 'x': entity.dxf.start.x, 'y': entity.dxf.start.y, 'z': entity.dxf.start.z})
                all_points.append({'t': frame_id, 'x': entity.dxf.end.x, 'y': entity.dxf.end.y, 'z': entity.dxf.end.z})
        except Exception:
            continue

    if not all_points:
        raise ValueError("No valid point geometry vectors could be parsed from DXF layers.")

    df_vertices = pd.DataFrame(all_points)
    solved_dataset_rows = []

    for frame_id, group in df_vertices.groupby('t'):
        if len(group) < 5: 
            continue
        
        pts_2d = group[['x', 'y']].drop_duplicates().values
        if len(pts_2d) < 5:
            continue

        # Construct KD-Tree to locate 5 closest points for interpolated closed curves
        tree = KDTree(pts_2d)
        visited = set()

        for idx, pt in enumerate(pts_2d):
            if idx in visited or len(solved_dataset_rows) >= max_faces:
                continue

            # Query 5 nearest neighbors to form interpolated closed boundary loop
            distances, indices = tree.query(pt, k=min(5, len(pts_2d)))
            cluster_pts = pts_2d[indices]
            
            # Solve quadric polynomial equations for the closed boundary
            poly_sol = fit_and_solve_closed_polynomial(cluster_pts)
            poly_sol['frame_id'] = frame_id
            poly_sol['mean_z'] = float(group['z'].mean())
            
            solved_dataset_rows.append(poly_sol)

    # Compile solved polynomial representations directly into pandas dataset
    dataset_df = pd.DataFrame(solved_dataset_rows)
    print(f"    ✅ Generated {len(dataset_df)} solved polynomial curves for dataset.")
    return dataset_df