import os
import logging
import numpy as np
import rasterio
from rasterio.warp import transform_bounds
import cv2
from scipy import ndimage

def extract_features_from_geotiff(image_path, output_folder, feature_type="trees"):
    """Improved feature extraction with contour-based tree detection."""
    try:
        logging.info(f"Extracting {feature_type} from {image_path}")
        
        with rasterio.open(image_path) as src:
            logging.info(f"Image shape: {src.shape}, bands: {src.count}")
            
            # Enhanced NDVI calculation
            if src.count >= 3:
                red = src.read(1).astype(float)
                green = src.read(2).astype(float)
                nir = src.read(4).astype(float) if src.count >= 4 else green
                
                logging.info(f"Red band range: {red.min():.3f} to {red.max():.3f}")
                logging.info(f"Green band range: {green.min():.3f} to {green.max():.3f}")
                logging.info(f"NIR band range: {nir.min():.3f} to {nir.max():.3f}")
                
                # Improved NDVI calculation
                ndvi = np.divide(nir - red, nir + red + 1e-10)
                
                # Adaptive thresholding based on image statistics
                ndvi_mean = np.mean(ndvi)
                ndvi_std = np.std(ndvi)
                threshold = max(0.05, ndvi_mean + 0.1 * ndvi_std)  # More lenient threshold
                
                logging.info(f"NDVI mean: {ndvi_mean:.3f}, std: {ndvi_std:.3f}, threshold: {threshold:.3f}")
                
                mask = ndvi > threshold
                logging.info(f"Mask pixels above threshold: {np.sum(mask)} out of {mask.size}")
            else:
                band = src.read(1)
                logging.info(f"Single band range: {band.min():.3f} to {band.max():.3f}")
                # Adaptive threshold for single-band images
                threshold = np.percentile(band, 40)  # Reduced from 70 to be more lenient
                logging.info(f"Single band threshold: {threshold:.3f}")
                mask = band > threshold
                logging.info(f"Mask pixels above threshold: {np.sum(mask)} out of {mask.size}")
            
            # Get bounds
            bounds = src.bounds
            if src.crs:
                west, south, east, north = transform_bounds(
                    src.crs, 'EPSG:4326',
                    bounds.left, bounds.bottom, bounds.right, bounds.top
                )
            else:
                west, south, east, north = -74.1, 40.6, -73.9, 40.8
        
        # Convert to uint8 for OpenCV processing
        mask_uint8 = (mask * 255).astype(np.uint8)
        logging.info(f"Mask uint8 range: {mask_uint8.min()} to {mask_uint8.max()}")
        
        # Morphological operations to clean up the mask
        kernel = np.ones((3, 3), np.uint8)
        mask_cleaned = cv2.morphologyEx(mask_uint8, cv2.MORPH_CLOSE, kernel)
        mask_cleaned = cv2.morphologyEx(mask_cleaned, cv2.MORPH_OPEN, kernel)
        
        logging.info(f"After morphological ops: {np.sum(mask_cleaned > 0)} pixels")
        
        # Find contours instead of using grid
        contours, _ = cv2.findContours(mask_cleaned, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        logging.info(f"Found {len(contours)} contours")
        
        # Filter contours by area and create features
        features = []
        height, width = mask.shape
        min_area = (height * width) * 0.0001  # Minimum 0.01% of image area (reduced from 0.1%)
        max_area = (height * width) * 0.2     # Maximum 20% of image area (increased from 10%)
        
        logging.info(f"Area filter: min={min_area:.0f}, max={max_area:.0f}")
        
        for i, contour in enumerate(contours):
            area = cv2.contourArea(contour)
            logging.info(f"Contour {i}: area={area:.0f}")
            
            # Filter by area
            if area < min_area or area > max_area:
                logging.info(f"Contour {i}: filtered out (area {area:.0f} not in range)")
                continue
            
            # Simplify contour for better polygon shape
            epsilon = 0.005 * cv2.arcLength(contour, True)
            approx = cv2.approxPolyDP(contour, epsilon, True)
            
            # Convert contour points to geographic coordinates
            polygon_coords = []
            for point in approx:
                x, y = point[0]
                
                # Convert pixel coordinates to geographic coordinates
                x_ratio = x / width
                y_ratio = y / height
                
                lon = west + x_ratio * (east - west)
                lat = north - y_ratio * (north - south)
                
                polygon_coords.append([lon, lat])
            
            # Close the polygon
            if len(polygon_coords) > 2:
                polygon_coords.append(polygon_coords[0])
                
                # Calculate confidence based on area and shape
                area_ratio = area / (height * width)
                confidence = min(0.95, 0.5 + area_ratio * 10)  # Adaptive confidence
                
                feature = {
                    "type": "Feature",
                    "id": i,
                    "properties": {
                        "feature_type": feature_type,
                        "confidence": round(confidence, 2),
                        "area_pixels": int(area),
                        "area_ratio": round(area_ratio, 4)
                    },
                    "geometry": {
                        "type": "Polygon",
                        "coordinates": [polygon_coords]
                    }
                }
                
                features.append(feature)
                logging.info(f"Contour {i}: added as feature with {len(polygon_coords)} points")
        
        logging.info(f"Extracted {len(features)} tree features")
        
        return {
            "type": "FeatureCollection",
            "features": features,
            "feature_type": feature_type
        }
        
    except Exception as e:
        logging.error(f"Error extracting features: {str(e)}")
        return {"type": "FeatureCollection", "features": []}