import os
import gradio as gr
import torch
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
import requests
import io
import matplotlib.colors as mcolors
import cv2
from io import BytesIO
import urllib.request
import tempfile
import rasterio
from rasterio.plot import reshape_as_image
import warnings
warnings.filterwarnings("ignore")

# Try to import segmentation_models_pytorch
try:
    import segmentation_models_pytorch as smp
    smp_available = True
    print("Successfully imported segmentation_models_pytorch")
except ImportError:
    smp_available = False
    print("Warning: segmentation_models_pytorch not available, will try to install it")
    import subprocess
    try:
        subprocess.check_call([
            "pip", "install", "segmentation-models-pytorch"
        ])
        import segmentation_models_pytorch as smp
        smp_available = True
        print("Successfully installed and imported segmentation_models_pytorch")
    except:
        print("Failed to install segmentation_models_pytorch")

# Try to import albumentations if needed for preprocessing
try:
    import albumentations as A
    albumentations_available = True
    print("Successfully imported albumentations")
except ImportError:
    albumentations_available = False
    print("Warning: albumentations not available, will use OpenCV for transforms")

# Set device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# Initialize the model
if smp_available:
    # Define the DeepLabV3+ model using smp
    model = smp.DeepLabV3Plus(
        encoder_name="resnet34",  # Using ResNet34 backbone as in your training
        encoder_weights=None,     # We'll load your custom weights
        in_channels=3,            # RGB input
        classes=1,                # Binary segmentation
    )
else:
    # Fallback to a simple model that won't actually work but allows the UI to load
    print("Warning: Using a placeholder model that won't produce correct predictions.")
    from torch import nn
    class PlaceholderModel(nn.Module):
        def __init__(self):
            super().__init__()
            self.conv = nn.Conv2d(3, 1, 3, padding=1)
        def forward(self, x):
            return self.conv(x)
    model = PlaceholderModel()

# Download model weights from HuggingFace
MODEL_REPO = "dcrey7/wetlands_segmentation_deeplabsv3plus"
MODEL_FILENAME = "DeepLabV3plus_best_model.pth"

def download_model_weights():
    """Download model weights from HuggingFace repository"""
    try:
        os.makedirs('weights', exist_ok=True)
        local_path = os.path.join('weights', MODEL_FILENAME)
        
        # Check if weights are already downloaded
        if os.path.exists(local_path):
            print(f"Model weights already downloaded at {local_path}")
            return local_path
        
        # Download weights
        print(f"Downloading model weights from {MODEL_REPO}...")
        url = f"https://huggingface.co/{MODEL_REPO}/resolve/main/{MODEL_FILENAME}"
        urllib.request.urlretrieve(url, local_path)
        print(f"Model weights downloaded to {local_path}")
        return local_path
    except Exception as e:
        print(f"Error downloading model weights: {e}")
        return None

# Load the model weights
weights_path = download_model_weights()
if weights_path:
    try:
        # Try to load with strict=False to allow for some parameter mismatches
        state_dict = torch.load(weights_path, map_location=device)
        # Check if we need to modify the state dict keys
        if all(key.startswith('encoder.') or key.startswith('decoder.') for key in list(state_dict.keys())[:5]):
            print("Model weights use encoder/decoder format, loading directly")
            model.load_state_dict(state_dict, strict=False)
        else:
            print("Attempting to adapt state dict to match model architecture")
            # This is a placeholder for state dict adaptation if needed
            model.load_state_dict(state_dict, strict=False)
        print("Model weights loaded successfully")
    except Exception as e:
        print(f"Error loading model weights: {e}")
else:
    print("No weights available. Model will not produce valid predictions.")

model.to(device)
model.eval()

def read_tiff_image(tiff_path):
    """
    Read a TIFF image using rasterio, focusing on RGB bands
    This matches your training data loading approach
    """
    try:
        # Read the image using rasterio (get RGB channels)
        with rasterio.open(tiff_path) as src:
            # Check if we have enough bands
            if src.count >= 3:
                red = src.read(1)
                green = src.read(2)
                blue = src.read(3)
                
                # Stack to create RGB image
                image = np.dstack((red, green, blue)).astype(np.float32)
                
                # Normalize to [0, 1]
                if image.max() > 0:
                    image = image / image.max()
                
                return image
            else:
                # If less than 3 bands, handle accordingly
                print(f"Warning: TIFF file has only {src.count} bands, RGB expected")
                bands = [src.read(i+1) for i in range(src.count)]
                # If only one band, duplicate to create RGB
                if len(bands) == 1:
                    image = np.dstack((bands[0], bands[0], bands[0]))
                else:
                    # Use available bands and pad with zeros if needed
                    while len(bands) < 3:
                        bands.append(np.zeros_like(bands[0]))
                    image = np.dstack(bands[:3])  # Use first 3 bands
                
                # Normalize
                if image.max() > 0:
                    image = image / image.max()
                
                return image
    except Exception as e:
        print(f"Error reading TIFF file: {e}")
        return None

def read_tiff_mask(mask_path):
    """
    Read a TIFF mask using rasterio
    This matches your training data loading approach
    """
    try:
        # Read mask
        with rasterio.open(mask_path) as src:
            mask = src.read(1).astype(np.uint8)
        return mask
    except Exception as e:
        print(f"Error reading mask file: {e}")
        return None

def preprocess_image(image, target_size=(128, 128)):
    """
    Preprocess an image for inference
    """
    # If image is a TIFF path, read it
    if isinstance(image, str) and (image.lower().endswith('.tif') or image.lower().endswith('.tiff')):
        image = read_tiff_image(image)
        if image is None:
            return None, None
    # Convert to numpy array if PIL image
    elif isinstance(image, Image.Image):
        image = np.array(image)
    
    # Ensure RGB format
    if len(image.shape) == 2:  # Grayscale
        image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
    elif image.shape[2] == 4:  # RGBA
        image = image[:, :, :3]
    
    # Make a copy for display
    display_image = (image * 255).astype(np.uint8) if image.max() <= 1.0 else image.copy()
    
    # Resize image to the target size
    if albumentations_available:
        # Use albumentations to match training preprocessing
        aug = A.Compose([
            A.PadIfNeeded(min_height=target_size[0], min_width=target_size[1], 
                         border_mode=cv2.BORDER_CONSTANT, value=0),
            A.CenterCrop(height=target_size[0], width=target_size[1])
        ])
        augmented = aug(image=image)
        image_resized = augmented['image']
    else:
        # Fallback to OpenCV
        image_resized = cv2.resize(image, target_size, interpolation=cv2.INTER_LINEAR)
    
    # Convert to tensor [C, H, W]
    image_tensor = torch.from_numpy(image_resized.transpose(2, 0, 1)).float().unsqueeze(0)
    
    return image_tensor, display_image

def preprocess_uploaded_file(file_obj, target_size=(128, 128)):
    """
    Process an uploaded file (could be image or TIFF)
    """
    try:
        # Save to a temporary file
        with tempfile.NamedTemporaryFile(delete=False, suffix='.tif') as temp_file:
            temp_path = temp_file.name
            # Write the content to the temp file
            if hasattr(file_obj, 'read'):
                # If it's a file-like object
                temp_file.write(file_obj.read())
            elif isinstance(file_obj, (str, bytes)):
                # If it's a string path or bytes
                if isinstance(file_obj, str):
                    with open(file_obj, 'rb') as f:
                        temp_file.write(f.read())
                else:
                    temp_file.write(file_obj)
        
        # Check if it's a TIFF file
        if temp_path.lower().endswith(('.tif', '.tiff')):
            # Read as TIFF
            image = read_tiff_image(temp_path)
            if image is None:
                os.unlink(temp_path)
                return None, None
            
            # Make a copy for display
            display_image = (image * 255).astype(np.uint8) if image.max() <= 1.0 else image.copy()
            
            # Resize/preprocess
            if albumentations_available:
                aug = A.Compose([
                    A.PadIfNeeded(min_height=target_size[0], min_width=target_size[1], 
                                 border_mode=cv2.BORDER_CONSTANT, value=0),
                    A.CenterCrop(height=target_size[0], width=target_size[1])
                ])
                augmented = aug(image=image)
                image_resized = augmented['image']
            else:
                image_resized = cv2.resize(image, target_size, interpolation=cv2.INTER_LINEAR)
            
            # Convert to tensor
            image_tensor = torch.from_numpy(image_resized.transpose(2, 0, 1)).float().unsqueeze(0)
        else:
            # Try to open as a regular image
            try:
                pil_image = Image.open(temp_path)
                image_tensor, display_image = preprocess_image(pil_image, target_size)
            except Exception as e:
                print(f"Error opening as regular image: {e}")
                os.unlink(temp_path)
                return None, None
        
        # Clean up
        os.unlink(temp_path)
        return image_tensor, display_image
    
    except Exception as e:
        print(f"Error processing uploaded file: {e}")
        import traceback
        traceback.print_exc()
        return None, None

def preprocess_mask(mask_input, target_size=(128, 128)):
    """
    Preprocess a ground truth mask
    """
    try:
        # If mask is a file path or file-like object, process it
        if isinstance(mask_input, str) or hasattr(mask_input, 'read'):
            # Save to temp file if it's a file-like object
            if hasattr(mask_input, 'read'):
                with tempfile.NamedTemporaryFile(delete=False, suffix='.tif') as temp_file:
                    temp_path = temp_file.name
                    temp_file.write(mask_input.read())
                mask_path = temp_path
            else:
                mask_path = mask_input
            
            # Check if it's a TIFF
            if mask_path.lower().endswith(('.tif', '.tiff')):
                mask = read_tiff_mask(mask_path)
                if mask is None:
                    if hasattr(mask_input, 'read'):
                        os.unlink(temp_path)
                    return None
            else:
                # Try as regular image
                try:
                    mask = np.array(Image.open(mask_path))
                    if len(mask.shape) == 3:
                        mask = cv2.cvtColor(mask, cv2.COLOR_RGB2GRAY)
                except Exception as e:
                    print(f"Error opening mask as image: {e}")
                    if hasattr(mask_input, 'read'):
                        os.unlink(temp_path)
                    return None
            
            # Clean up if temp file
            if hasattr(mask_input, 'read'):
                os.unlink(temp_path)
        
        # Convert to numpy array if PIL image
        elif isinstance(mask_input, Image.Image):
            mask = np.array(mask_input)
            if len(mask.shape) == 3:
                mask = cv2.cvtColor(mask, cv2.COLOR_RGB2GRAY)
        
        # Already numpy array
        else:
            mask = mask_input
            if len(mask.shape) == 3:
                mask = cv2.cvtColor(mask, cv2.COLOR_RGB2GRAY)
        
        # Resize mask to target size
        if albumentations_available:
            aug = A.Compose([
                A.PadIfNeeded(min_height=target_size[0], min_width=target_size[1], 
                             border_mode=cv2.BORDER_CONSTANT, value=0),
                A.CenterCrop(height=target_size[0], width=target_size[1])
            ])
            augmented = aug(image=mask)
            mask_resized = augmented['image']
        else:
            mask_resized = cv2.resize(mask, target_size, interpolation=cv2.INTER_NEAREST)
        
        # Binarize the mask (0: background, 1: wetland)
        mask_binary = (mask_resized > 127).astype(np.uint8)
        
        return mask_binary
    
    except Exception as e:
        print(f"Error preprocessing mask: {e}")
        import traceback
        traceback.print_exc()
        return None

def predict_segmentation(image_tensor):
    """
    Run inference on the model
    """
    try:
        image_tensor = image_tensor.to(device)
        
        with torch.no_grad():
            output = model(image_tensor)
            
            # Handle different model output formats
            if isinstance(output, dict):
                output = output['out']
            if output.shape[1] > 1:  # Multi-class output
                pred = torch.argmax(output, dim=1).squeeze(0).cpu().numpy()
            else:  # Binary output (from smp models)
                pred = (torch.sigmoid(output) > 0.5).squeeze().cpu().numpy().astype(np.uint8)
                
        return pred
    except Exception as e:
        print(f"Error during prediction: {e}")
        import traceback
        traceback.print_exc()
        return None

def calculate_metrics(pred_mask, gt_mask):
    """
    Calculate evaluation metrics between prediction and ground truth
    """
    # Ensure binary masks
    pred_binary = (pred_mask > 0).astype(np.uint8)
    gt_binary = (gt_mask > 0).astype(np.uint8)
    
    # Calculate intersection and union
    intersection = np.logical_and(pred_binary, gt_binary).sum()
    union = np.logical_or(pred_binary, gt_binary).sum()
    
    # Calculate IoU
    iou = intersection / union if union > 0 else 0
    
    # Calculate precision and recall
    true_positive = intersection
    false_positive = pred_binary.sum() - true_positive
    false_negative = gt_binary.sum() - true_positive
    
    precision = true_positive / (true_positive + false_positive) if (true_positive + false_positive) > 0 else 0
    recall = true_positive / (true_positive + false_negative) if (true_positive + false_negative) > 0 else 0
    
    # Calculate F1 score
    f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0
    
    metrics = {
        "IoU": float(iou),
        "Precision": float(precision),
        "Recall": float(recall),
        "F1 Score": float(f1)
    }
    
    return metrics

def process_images(input_image=None, input_tiff=None, gt_mask=None):
    """
    Process input images and generate predictions
    """
    try:
        # Check if we have input
        if input_image is None and input_tiff is None:
            return None, "Please upload an image or TIFF file."
        
        # Process the input image
        if input_tiff is not None:
            # Process uploaded TIFF file
            image_tensor, display_image = preprocess_uploaded_file(input_tiff)
        else:
            # Process regular image
            image_tensor, display_image = preprocess_image(input_image)
        
        if image_tensor is None or display_image is None:
            return None, "Failed to process the input image."
        
        # Get prediction
        pred_mask = predict_segmentation(image_tensor)
        if pred_mask is None:
            return None, "Failed to generate prediction."
        
        # Process ground truth mask if provided
        gt_mask_processed = None
        metrics_text = ""
        
        if gt_mask is not None:
            gt_mask_processed = preprocess_mask(gt_mask)
            
            if gt_mask_processed is not None:
                metrics = calculate_metrics(pred_mask, gt_mask_processed)
                metrics_text = "\n".join([f"{k}: {v:.4f}" for k, v in metrics.items()])
        
        # Create visualization
        fig = plt.figure(figsize=(12, 6))
        
        if gt_mask_processed is not None:
            # Show original, ground truth, and prediction
            plt.subplot(1, 3, 1)
            plt.imshow(display_image)
            plt.title("Input Image")
            plt.axis('off')
            
            plt.subplot(1, 3, 2)
            plt.imshow(gt_mask_processed, cmap='binary')
            plt.title("Ground Truth")
            plt.axis('off')
            
            plt.subplot(1, 3, 3)
            plt.imshow(pred_mask, cmap='binary')
            plt.title("Prediction")
            plt.axis('off')
        else:
            # Show original and prediction
            plt.subplot(1, 2, 1)
            plt.imshow(display_image)
            plt.title("Input Image")
            plt.axis('off')
            
            plt.subplot(1, 2, 2)
            plt.imshow(pred_mask, cmap='binary')
            plt.title("Predicted Wetlands")
            plt.axis('off')
        
        # Calculate wetland percentage
        wetland_percentage = np.mean(pred_mask) * 100
        
        # Add metrics info
        result_text = f"Wetland Coverage: {wetland_percentage:.2f}%"
        if metrics_text:
            result_text += f"\n\nEvaluation Metrics:\n{metrics_text}"
        
        # Convert figure to image
        buf = BytesIO()
        plt.tight_layout()
        plt.savefig(buf, format='png', dpi=150)
        buf.seek(0)
        result_image = Image.open(buf)
        plt.close(fig)
        
        return result_image, result_text
        
    except Exception as e:
        print(f"Error in processing: {e}")
        import traceback
        traceback.print_exc()
        return None, f"Error: {str(e)}"

# Create Gradio interface
with gr.Blocks(title="Wetlands Segmentation from Satellite Imagery") as demo:
    gr.Markdown("# Wetlands Segmentation from Satellite Imagery")
    gr.Markdown("Upload a satellite image or TIFF file to identify wetland areas. Optionally, you can also upload a ground truth mask for evaluation.")
    
    with gr.Row():
        with gr.Column():
            # Input options
            gr.Markdown("### Input")
            with gr.Tab("Upload Image"):
                input_image = gr.Image(label="Upload Satellite Image", type="numpy")
            
            with gr.Tab("Upload TIFF"):
                input_tiff = gr.File(label="Upload TIFF File", file_types=[".tif", ".tiff"])
            
            gt_mask = gr.Image(label="Ground Truth Mask (Optional)", type="numpy")
            
            process_btn = gr.Button("Analyze Image", variant="primary")
        
        with gr.Column():
            # Output
            gr.Markdown("### Results")
            output_image = gr.Image(label="Segmentation Results", type="pil")
            output_text = gr.Textbox(label="Statistics", lines=6)
    
    # Information about the model
    gr.Markdown("### About this model")
    gr.Markdown("""
    This application uses a DeepLabv3+ model trained to segment wetland areas in satellite imagery.
    
    **Model Details:**
    - Architecture: DeepLabv3+ with ResNet-34
    - Input: RGB satellite imagery
    - Output: Binary segmentation mask (Wetland vs Background)
    - Resolution: 128×128 pixels
    
    **Tips for best results:**
    - The model works best with RGB satellite imagery
    - For optimal results, use images with similar characteristics to those used in training
    - The model focuses on identifying wetland regions in natural landscapes
    
    **Repository:** [dcrey7/wetlands_segmentation_deeplabsv3plus](https://huggingface.co/dcrey7/wetlands_segmentation_deeplabsv3plus)
    """)
    
    # Set up event handlers
    process_btn.click(
        fn=process_images,
        inputs=[input_image, input_tiff, gt_mask],
        outputs=[output_image, output_text]
    )

# Launch the app
demo.launch()