import pymupdf
import cv2
import numpy as np
import xml.etree.ElementTree as ET
from shapely.geometry import Polygon
import geopandas as gpd
from paddleocr import PaddleOCR
from osgeo import gdal, osr
import re

# 1. Extract and Render Raster from PDF
def extract_raster_from_pdf(pdf_path, dpi=300):
    doc = pymupdf.open(pdf_path)
    page = doc.load_page(0)
    pix = page.get_pixmap(dpi=dpi)
    img = np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.height, pix.width, pix.n)
    if pix.n >= 3:
        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
    return img

# 2. Parse KML Directly & Auto-Fix Unclosed Polygons
def load_kml_data(kml_path):
    tree = ET.parse(kml_path)
    root = tree.getroot()
    
    for elem in root.iter():
        if '}' in elem.tag:
            elem.tag = elem.tag.split('}', 1)[1]
            
    records = []
    
    for placemark in root.iter('Placemark'):
        survey_raw = ""
        for simple_data in placemark.iter('SimpleData'):
            name_attr = simple_data.attrib.get('name', '')
            if name_attr in ['Drawings', 'Survey Number']:
                survey_raw = simple_data.text or ""
                break
                
        numbers = re.findall(r'\d+', survey_raw)
        if not numbers:
            continue
            
        coords_elem = placemark.find('.//coordinates')
        if coords_elem is not None and coords_elem.text:
            raw_text = coords_elem.text.strip()
            pts = []
            for item in raw_text.split():
                parts = item.split(',')
                if len(parts) >= 2:
                    pts.append((float(parts[0]), float(parts[1])))
                    
            if len(pts) >= 3:
                if pts[0] != pts[-1]:
                    pts.append(pts[0])
                    
                poly = Polygon(pts)
                if poly.is_valid and not poly.is_empty:
                    centroid = poly.centroid
                    for num in numbers:
                        records.append({
                            'daag_no': num,
                            'geo_x': centroid.x,
                            'geo_y': centroid.y
                        })

    return gpd.GeoDataFrame(records)

# 3. Detect Daag Numbers via OCR
def extract_ocr_points(img, ocr_engine):
    results = ocr_engine.ocr(img)
    detected_points = []
    if not results or not results[0]:
        return detected_points
        
    for line in results[0]:
        text, conf = line[1]
        clean_numbers = re.findall(r'\d+', text)
        if clean_numbers and conf > 0.5:
            bbox = line[0]
            cx = sum([p[0] for p in bbox]) / 4.0
            cy = sum([p[1] for p in bbox]) / 4.0
            for num in clean_numbers:
                detected_points.append({
                    'daag_no': num,
                    'pixel_x': cx,
                    'pixel_y': cy
                })
    return detected_points

# 4. Generate GCPs via Daag Matching + RANSAC Filter
def match_gcps(kml_df, ocr_points):
    matched_pairs = []
    for pt in ocr_points:
        daag = pt['daag_no']
        kml_match = kml_df[kml_df['daag_no'] == daag]
        if len(kml_match) == 1:
            row = kml_match.iloc[0]
            matched_pairs.append({
                'geo_x': row['geo_x'],
                'geo_y': row['geo_y'],
                'pixel_x': pt['pixel_x'],
                'pixel_y': pt['pixel_y']
            })
            
    if len(matched_pairs) < 4:
        raise ValueError(f"Only {len(matched_pairs)} unique GCP matches found. Minimum 4 required.")
        
    src_pts = np.float32([[p['pixel_x'], p['pixel_y']] for p in matched_pairs])
    dst_pts = np.float32([[p['geo_x'], p['geo_y']] for p in matched_pairs])
    
    _, inliers = cv2.estimateAffinePartial2D(src_pts, dst_pts, method=cv2.RANSAC, ransacReprojThreshold=50.0)
    
    gdal_gcps = []
    for idx, pair in enumerate(matched_pairs):
        if inliers is not None and inliers[idx]:
            gdal_gcps.append(
                gdal.GCP(pair['geo_x'], pair['geo_y'], 0, pair['pixel_x'], pair['pixel_y'])
            )
    return gdal_gcps

# 5. Georeference and Warp
def warp_raster_with_gcps(img, gcps, output_tif):
    temp_img = "temp_raster.png"
    cv2.imwrite(temp_img, img)
    
    src_ds = gdal.Open(temp_img, gdal.GA_Update)
    srs = osr.SpatialReference()
    srs.ImportFromEPSG(4326)
    src_ds.SetGCPs(gcps, srs.ExportToWkt())
    src_ds = None
    
    gdal.Warp(
        output_tif,
        temp_img,
        dstSRS="EPSG:3857",
        tps=True,
        resampleAlg=gdal.GRA_Cubic
    )
    print(f"Success! Output generated: {output_tif}")

if __name__ == "__main__":
    PDF_FILE = "Sheet 0 - Mapদাসপল্সাময়ূরেশ্বর-2.pdf"
    KML_FILE = "1787306740331.kml"
    OUTPUT_TIF = "daspalsa_mouza_aligned.tif"
    
    print("1. Loading PDF and parsing KML...")
    img = extract_raster_from_pdf(PDF_FILE, dpi=300)
    kml_df = load_kml_data(KML_FILE)
    print(f"Loaded {len(kml_df)} plot coordinate references from KML.")
    
    print("2. Initializing PaddleOCR...")
    ocr = PaddleOCR(use_textline_orientation=True, lang='en', device='cpu')
    
    print("3. Detecting plot numbers from map sheet...")
    ocr_points = extract_ocr_points(img, ocr)
    
    print("4. Calculating GCPs...")
    gcps = match_gcps(kml_df, ocr_points)
    print(f"Verified {len(gcps)} control points.")
    
    print("5. Generating georeferenced raster...")
    warp_raster_with_gcps(img, gcps, OUTPUT_TIF)