import rasterio
import numpy as np
import cv2
from shapely.geometry import Polygon, mapping
import geojson
from rasterio.features import shapes
from rasterio.transform import Affine

def read_geotiff(file_path):
    with rasterio.open(file_path) as src:
        img = src.read([1, 2, 3])  # Read RGB channels
        img = np.stack(img, axis=2)  # Stack channels to create an RGB image
        transform = src.transform  # Get the transformation matrix
        crs = src.crs  # Get the coordinate reference system
        meta = src.meta  # Get metadata
    return img, transform, crs, meta

def detect_oil_palm_plantations(img):
    # Convert the image to HSV color space
    hsv_img = cv2.cvtColor(img, cv2.COLOR_RGB2HSV)
    
    # Define the range for green color in HSV space
    lower_green = np.array([30, 40, 40])
    upper_green = np.array([80, 255, 255])
    
    # Create a binary mask where green colors are white and the rest are black
    mask = cv2.inRange(hsv_img, lower_green, upper_green)
    
    # Find contours in the mask
    contours, _ = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    
    # Approximate contours to polygons
    polygons = [cv2.approxPolyDP(contour, 0.01 * cv2.arcLength(contour, True), True) for contour in contours]
    
    # Convert contours to polygons
    polygons = [Polygon(polygon[:, 0, :]) for polygon in polygons if len(polygon) > 2]
    
    return polygons, mask

def save_geojson(polygons, transform, crs, output_path):
    features = []
    for polygon in polygons:
        # Transform polygon coordinates from image to geographical coordinates
        transformed_polygon = [rasterio.transform.xy(transform, point[1], point[0]) for point in polygon.exterior.coords]
        geojson_polygon = Polygon(transformed_polygon)
        features.append(geojson.Feature(geometry=geojson_polygon))
    
    feature_collection = geojson.FeatureCollection(features)
    
    with open(output_path, 'w') as f:
        geojson.dump(feature_collection, f)

def draw_polygons_on_image(img, polygons):
    # Draw polygons on the image
    for polygon in polygons:
        points = np.array([polygon.exterior.coords], dtype=np.int32)
        cv2.polylines(img, [points], isClosed=True, color=(255, 0, 0), thickness=2)
    return img

def save_new_tif(file_path, img, meta):
    meta.update({"count": 3})  # Ensure that we save an RGB image
    with rasterio.open(file_path, 'w', **meta) as dst:
        dst.write(img[:, :, 0], 1)  # Red channel
        dst.write(img[:, :, 1], 2)  # Green channel
        dst.write(img[:, :, 2], 3)  # Blue channel

def detect_and_save_oil_palm_plantations(file_path, geojson_output_path, new_tif_output_path):
    img, transform, crs, meta = read_geotiff(file_path)
    polygons, mask = detect_oil_palm_plantations(img)
    save_geojson(polygons, transform, crs, geojson_output_path)
    img_with_polygons = draw_polygons_on_image(img, polygons)
    save_new_tif(new_tif_output_path, img_with_polygons, meta)

# Example usage:
file_path = 'Downloads/satellite_new8.tif'
geojson_output_path = 'plantations.geojson'
new_tif_output_path = 'image_with_polygons.tif'
detect_and_save_oil_palm_plantations(file_path, geojson_output_path, new_tif_output_path)