sorghum_pipeline/output/manager.py

"""
Minimal output manager for demo (saves only 7 required images).
"""

import os
import numpy as np
import cv2
import matplotlib
if os.environ.get('MPLBACKEND') is None:
    matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.cm as cm
from pathlib import Path
from typing import Dict, Any
import logging

logger = logging.getLogger(__name__)


class OutputManager:
    """Minimal output manager for demo."""
    
    def __init__(self, output_folder: str, settings: Any):
        """Initialize output manager."""
        self.output_folder = Path(output_folder)
        self.settings = settings
        try:
            self.minimal_demo: bool = bool(int(os.environ.get('MINIMAL_DEMO', '0')))
        except Exception:
            self.minimal_demo = False
        self.output_folder.mkdir(parents=True, exist_ok=True)

    def create_output_directories(self) -> None:
        """Create output directories."""
        self.output_folder.mkdir(parents=True, exist_ok=True)
    
    def save_plant_results(self, plant_key: str, plant_data: Dict[str, Any]) -> None:
        """Save minimal demo outputs only."""
        if not self.minimal_demo:
            logger.warning("OutputManager configured for minimal demo only")
            return
        
        self._save_minimal_demo_outputs(plant_data)

    def _save_minimal_demo_outputs(self, plant_data: Dict[str, Any]) -> None:
        """Save only the 7 required images."""
        results_dir = self.output_folder / 'results'
        veg_dir = self.output_folder / 'Vegetation_indices_images'
        tex_dir = self.output_folder / 'texture_output'
        results_dir.mkdir(parents=True, exist_ok=True)
        veg_dir.mkdir(parents=True, exist_ok=True)
        tex_dir.mkdir(parents=True, exist_ok=True)

        # 1. Mask
        try:
            mask = plant_data.get('mask')
            if isinstance(mask, np.ndarray):
                titled = self._add_title_banner(mask, 'Mask')
                cv2.imwrite(str(results_dir / 'mask.png'), titled)
        except Exception as e:
            logger.error(f"Failed to save mask: {e}")

        # 2. Overlay
        try:
            base_image = plant_data.get('composite')
            mask = plant_data.get('mask')
            if isinstance(base_image, np.ndarray) and isinstance(mask, np.ndarray):
                overlay = self._create_overlay(base_image, mask)
                # Convert BGR→RGB for correct viewing in standard image viewers
                overlay_rgb = cv2.cvtColor(overlay, cv2.COLOR_BGR2RGB)
                titled = self._add_title_banner(overlay_rgb, 'Segmentation Overlay')
                cv2.imwrite(str(results_dir / 'overlay.png'), titled)
        except Exception as e:
            logger.error(f"Failed to save overlay: {e}")

        # 2b. Composite (input to segmentation)
        try:
            base_image = plant_data.get('composite')
            if isinstance(base_image, np.ndarray):
                # Ensure uint8
                if base_image.dtype != np.uint8:
                    base_image = self._normalize_to_uint8(base_image.astype(np.float64))
                # Convert BGR→RGB for human viewing
                comp_rgb = cv2.cvtColor(base_image, cv2.COLOR_BGR2RGB)
                titled = self._add_title_banner(comp_rgb, 'Composite (Segmentation Input)')
                cv2.imwrite(str(results_dir / 'composite.png'), titled)
        except Exception as e:
            logger.error(f"Failed to save composite: {e}")

        # 3-5. Vegetation indices (NDVI, GNDVI, SAVI)
        try:
            veg = plant_data.get('vegetation_indices', {})
            for name in ['NDVI', 'GNDVI', 'SAVI']:
                data = veg.get(name, {})
                values = data.get('values') if isinstance(data, dict) else None
                if isinstance(values, np.ndarray) and values.size > 0:
                    try:
                        # Colormap with colorbar similar to src: use matplotlib savefig
                        cmap = cm.RdYlGn
                        # Value ranges
                        if name in ['NDVI', 'GNDVI']:
                            vmin, vmax = (-1, 1)
                        else:
                            vmin, vmax = (0, 1)

                        masked = np.ma.masked_invalid(values.astype(np.float64))
                        fig, ax = plt.subplots(figsize=(5, 5))
                        ax.set_axis_off()
                        ax.set_facecolor('white')
                        im = ax.imshow(masked, cmap=cmap, vmin=vmin, vmax=vmax)
                        ax.set_title(f"{name}", fontsize=12, fontweight='bold', pad=8)
                        # add colorbar
                        cbar = fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
                        cbar.ax.tick_params(labelsize=8)
                        plt.tight_layout()
                        plt.savefig(veg_dir / f"{name.lower()}.png", dpi=120, bbox_inches='tight')
                        plt.close(fig)
                    except Exception as e:
                        logger.error(f"Failed to save {name}: {e}")
        except Exception as e:
            logger.error(f"Failed to save vegetation indices: {e}")

        # 6. Texture features: ONLY LBP on green band
        try:
            tex = plant_data.get('texture_features', {})
            green_band = tex.get('green', {})
            feats = green_band.get('features', {})

            lbp = feats.get('lbp')
            if isinstance(lbp, np.ndarray) and lbp.size > 0:
                try:
                    img = lbp.astype(np.float64)
                    fig, ax = plt.subplots(figsize=(5, 5))
                    ax.set_axis_off()
                    ax.set_facecolor('white')
                    im = ax.imshow(img, cmap='gray', vmin=0, vmax=255)
                    ax.set_title('Texture: LBP (Green Band)', fontsize=12, fontweight='bold', pad=8)
                    cbar = fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
                    cbar.ax.tick_params(labelsize=8)
                    plt.tight_layout()
                    plt.savefig(tex_dir / 'lbp_green.png', dpi=120, bbox_inches='tight')
                    plt.close(fig)
                except Exception as e:
                    logger.error(f"Failed to save LBP with colorbar: {e}")
        except Exception as e:
            logger.error(f"Failed to save texture: {e}")

        # 9. Morphology size analysis
        try:
            morph = plant_data.get('morphology_features', {})
            images = morph.get('images', {})
            size_img = images.get('size_analysis')
            if isinstance(size_img, np.ndarray) and size_img.size > 0:
                titled = self._add_title_banner(size_img, 'Morphology Size')
                cv2.imwrite(str(results_dir / 'size.size_analysis.png'), titled)
            else:
                # Fallback: synthesize a simple size analysis from the mask if available
                mask_for_size = plant_data.get('mask')
                base_img_for_size = plant_data.get('composite')
                if isinstance(mask_for_size, np.ndarray) and mask_for_size.size > 0:
                    synthesized = self._create_size_analysis_from_mask(mask_for_size, base_img_for_size)
                    titled = self._add_title_banner(synthesized, 'Morphology Size')
                    cv2.imwrite(str(results_dir / 'size.size_analysis.png'), titled)
        except Exception as e:
            logger.error(f"Failed to save size analysis: {e}")
    
    def _create_overlay(self, image: np.ndarray, mask: np.ndarray) -> np.ndarray:
        """Create green overlay on brightened composite, following src pipeline style."""
        if mask is None:
            return image
        if mask.shape[:2] != image.shape[:2]:
            mask = cv2.resize(mask.astype(np.uint8), (image.shape[1], image.shape[0]),
                              interpolation=cv2.INTER_NEAREST)
        binary = (mask.astype(np.int32) > 0).astype(np.uint8) * 255
        base = image
        if base.dtype != np.uint8:
            base = self._normalize_to_uint8(base.astype(np.float64))
        bright = cv2.convertScaleAbs(base, alpha=1.2, beta=15)
        green_overlay = bright.copy()
        green_overlay[binary == 255] = (0, 255, 0)
        blended = cv2.addWeighted(bright, 1.0, green_overlay, 0.5, 0)
        return blended
    
    def _normalize_to_uint8(self, arr: np.ndarray) -> np.ndarray:
        """Normalize to uint8."""
        arr = np.nan_to_num(arr, nan=0.0, posinf=0.0, neginf=0.0)
        ptp = np.ptp(arr)
        if ptp > 0:
            normalized = (arr - arr.min()) / (ptp + 1e-6) * 255
        else:
            normalized = np.zeros_like(arr)
        return np.clip(normalized, 0, 255).astype(np.uint8)

    def _add_title_banner(self, image: np.ndarray, title: str) -> np.ndarray:
        """Add a top banner with centered title text to an image using OpenCV.
        Supports grayscale or color images; returns a BGR image.
        """
        if image is None or image.size == 0:
            return image

        # Ensure 3-channel BGR for drawing
        if image.ndim == 2:
            base_bgr = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
        elif image.ndim == 3 and image.shape[2] == 3:
            base_bgr = image.copy()
        elif image.ndim == 3 and image.shape[2] == 4:
            base_bgr = cv2.cvtColor(image, cv2.COLOR_BGRA2BGR)
        else:
            # Fallback: normalize to uint8 then convert to BGR
            norm = self._normalize_to_uint8(image.astype(np.float64))
            base_bgr = cv2.cvtColor(norm, cv2.COLOR_GRAY2BGR)

        h, w = base_bgr.shape[:2]
        banner_height = max(30, int(0.08 * h))
        banner = np.full((banner_height, w, 3), (245, 245, 245), dtype=np.uint8)

        # Compose banner + image
        composed = np.vstack([banner, base_bgr])

        # Put centered title text
        font = cv2.FONT_HERSHEY_SIMPLEX
        font_scale = max(0.5, min(1.0, w / 800.0))
        thickness = 1
        text = str(title)
        (tw, th), baseline = cv2.getTextSize(text, font, font_scale, thickness)
        x = max(5, (w - tw) // 2)
        y = (banner_height + th) // 2
        # Slight shadow for readability
        cv2.putText(composed, text, (x+1, y+1), font, font_scale, (0, 0, 0), thickness+1, cv2.LINE_AA)
        cv2.putText(composed, text, (x, y), font, font_scale, (0, 80, 0), thickness+1, cv2.LINE_AA)

        return composed

    def _create_size_analysis_from_mask(self, mask: np.ndarray, base_image: Any = None) -> np.ndarray:
        """Create a simple size analysis visualization from a binary mask.
        Draws contours and prints pixel area. If base_image is provided, overlays on it; otherwise uses a white canvas.
        """
        if mask is None or mask.size == 0:
            return np.zeros((1, 1, 3), dtype=np.uint8)

        # Prepare base image
        if isinstance(base_image, np.ndarray) and base_image.size > 0:
            img = base_image
            if img.dtype != np.uint8:
                img = self._normalize_to_uint8(img.astype(np.float64))
            if img.ndim == 2:
                base_bgr = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
            elif img.ndim == 3 and img.shape[2] == 3:
                base_bgr = img.copy()
            elif img.ndim == 3 and img.shape[2] == 4:
                base_bgr = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)
            else:
                norm = self._normalize_to_uint8(img.astype(np.float64))
                base_bgr = cv2.cvtColor(norm, cv2.COLOR_GRAY2BGR)
        else:
            h, w = mask.shape[:2]
            base_bgr = np.full((h, w, 3), 255, dtype=np.uint8)

        # Ensure binary mask
        if mask.ndim == 3 and mask.shape[2] == 3:
            gray = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
        else:
            gray = mask.astype(np.uint8)
        _, bin_mask = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY)

        # Contours and area
        contours, _ = cv2.findContours(bin_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        cv2.drawContours(base_bgr, contours, -1, (0, 0, 255), 1)
        area_px = int(cv2.countNonZero(bin_mask))

        # Bounding box for the largest contour
        if contours:
            largest = max(contours, key=cv2.contourArea)
            x, y, w, h = cv2.boundingRect(largest)
            cv2.rectangle(base_bgr, (x, y), (x + w, y + h), (255, 0, 0), 1)

        # Put area text
        cv2.putText(base_bgr, f"Area: {area_px} px", (10, 24), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2, cv2.LINE_AA)

        return base_bgr