import os
import logging
import numpy as np
import rasterio
from rasterio.warp import transform_bounds

def extract_features_from_geotiff(image_path, output_folder, feature_type="trees"):
    """Simple feature extraction for HF Spaces."""
    try:
        logging.info(f"Extracting {feature_type} from {image_path}")
        
        with rasterio.open(image_path) as src:
            # Simple 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
                
                ndvi = np.divide(nir - red, nir + red + 1e-10)
                mask = ndvi > 0.2
            else:
                band = src.read(1)
                mask = band > np.percentile(band, 60)
            
            # 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
        
        # Create simple features
        features = []
        height, width = mask.shape
        grid_size = max(10, min(height, width) // 50)
        
        feature_id = 0
        for y in range(0, height, grid_size):
            for x in range(0, width, grid_size):
                cell = mask[y:y+grid_size, x:x+grid_size]
                if np.sum(cell) > grid_size * grid_size * 0.3:
                    
                    x_ratio = x / width
                    y_ratio = y / height
                    
                    lon1 = west + x_ratio * (east - west)
                    lat1 = north - y_ratio * (north - south)
                    
                    x2_ratio = min((x + grid_size) / width, 1.0)
                    y2_ratio = min((y + grid_size) / height, 1.0)
                    
                    lon2 = west + x2_ratio * (east - west)
                    lat2 = north - y2_ratio * (north - south)
                    
                    polygon_coords = [
                        [lon1, lat1], [lon2, lat1], [lon2, lat2], [lon1, lat2], [lon1, lat1]
                    ]
                    
                    feature = {
                        "type": "Feature",
                        "id": feature_id,
                        "properties": {
                            "feature_type": feature_type,
                            "confidence": 0.8
                        },
                        "geometry": {
                            "type": "Polygon",
                            "coordinates": [polygon_coords]
                        }
                    }
                    
                    features.append(feature)
                    feature_id += 1
        
        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": []}