AareDAQ/scripts/camera_stat_thread.py

import json
import threading
import time
from typing import Any

import cv2
import numpy as np
import zmq
from aarecommon.config.logger import setup_logger

logger = setup_logger("aareDAQ")


class ImageStatsReceiver:
    def __init__(self, zmq_url: str):
        """Initialize ZeroMQ receiver for image statistics calculation"""
        context = zmq.Context()
        self._socket = context.socket(zmq.SUB)
        self._socket.setsockopt(zmq.SUBSCRIBE, b"")
        self._socket.setsockopt(zmq.RCVTIMEO, 1000)  # 1 second timeout
        self._socket.setsockopt(zmq.LINGER, 0)  # Don't linger on close

        try:
            self._socket.connect(zmq_url)
            print(f"Connected to ZeroMQ socket: {zmq_url}")
        except Exception as e:
            print(f"Failed to connect to {zmq_url}: {e}")
            raise

        self.running = False
        self.latest_stats = None
        self.stats_lock = threading.Lock()
        self.message_count = 0
        self.connection_attempts = 0
        self.last_message_time = None
        self.connection_attempts = 0
        self.last_message_time = None

    def calculate_projections(self, image: np.ndarray) -> dict[str, Any]:
        """Calculate X and Y projections of the image"""
        # Convert to grayscale if color image
        if len(image.shape) == 3:
            gray_image = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
        else:
            gray_image = image.copy()

        # Calculate projections (sum along each axis)
        x_projection = np.mean(gray_image, axis=0)  # Sum along Y axis -> X profile
        y_projection = np.mean(gray_image, axis=1)  # Sum along X axis -> Y profile

        # Also calculate max projections (useful for finding peaks)
        x_projection_max = np.max(gray_image, axis=0)
        y_projection_max = np.max(gray_image, axis=1)

        # Find peak positions and values
        x_peak_idx = np.argmax(x_projection)
        y_peak_idx = np.argmax(y_projection)
        x_peak_value = x_projection[x_peak_idx]
        y_peak_value = y_projection[y_peak_idx]

        # Calculate centroid positions (intensity-weighted center)
        x_coords = np.arange(len(x_projection))
        y_coords = np.arange(len(y_projection))

        x_centroid = (
            np.sum(x_coords * x_projection) / np.sum(x_projection)
            if np.sum(x_projection) > 0
            else len(x_projection) / 2
        )
        y_centroid = (
            np.sum(y_coords * y_projection) / np.sum(y_projection)
            if np.sum(y_projection) > 0
            else len(y_projection) / 2
        )

        # Calculate FWHM (Full Width at Half Maximum) approximation
        def calculate_fwhm(profile):
            peak_val = np.max(profile)
            half_max = peak_val / 2
            indices = np.where(profile >= half_max)[0]
            if len(indices) > 0:
                return indices[-1] - indices[0]
            return 0

        x_fwhm = calculate_fwhm(x_projection)
        y_fwhm = calculate_fwhm(y_projection)

        return {
            "x_projection": x_projection.tolist(),
            "y_projection": y_projection.tolist(),
            "x_projection_max": x_projection_max.tolist(),
            "y_projection_max": y_projection_max.tolist(),
            "x_peak_position": int(x_peak_idx),
            "y_peak_position": int(y_peak_idx),
            "x_peak_value": float(x_peak_value),
            "y_peak_value": float(y_peak_value),
            "x_centroid": float(x_centroid),
            "y_centroid": float(y_centroid),
            "x_fwhm": float(x_fwhm),
            "y_fwhm": float(y_fwhm),
            "x_coords": list(range(len(x_projection))),
            "y_coords": list(range(len(y_projection))),
        }

    def calculate_radial_integration(self, image: np.ndarray, num_bins: int = 50) -> dict[str, Any]:
        """Calculate radial integration of the image"""
        # Convert to grayscale if color image
        if len(image.shape) == 3:
            gray_image = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
        else:
            gray_image = image.copy()

        # Get image center
        center_y, center_x = np.array(gray_image.shape) // 2

        # Create coordinate grids
        y_coords, x_coords = np.ogrid[: gray_image.shape[0], : gray_image.shape[1]]

        # Calculate distance from center for each pixel
        distances = np.sqrt((x_coords - center_x) ** 2 + (y_coords - center_y) ** 2)

        # Maximum possible radius (to corner of image)
        max_radius = np.sqrt((gray_image.shape[0] / 2) ** 2 + (gray_image.shape[1] / 2) ** 2)

        # Create radial bins
        r_bins = np.linspace(0, max_radius, num_bins + 1)
        r_centers = (r_bins[:-1] + r_bins[1:]) / 2

        # Calculate radial profile
        radial_profile = []
        radial_std = []
        pixel_counts = []

        for i in range(num_bins):
            # Create mask for current radial bin
            mask = (distances >= r_bins[i]) & (distances < r_bins[i + 1])

            if np.any(mask):
                # Get pixel values in this radial bin
                pixel_values = gray_image[mask]
                radial_profile.append(float(np.mean(pixel_values)))
                radial_std.append(float(np.std(pixel_values)))
                pixel_counts.append(int(np.sum(mask)))
            else:
                radial_profile.append(0.0)
                radial_std.append(0.0)
                pixel_counts.append(0)

        return {
            "r_centers": r_centers.tolist(),
            "radial_profile": radial_profile,
            "radial_std": radial_std,
            "pixel_counts": pixel_counts,
            "max_radius": float(max_radius),
            "center": [int(center_x), int(center_y)],
            "num_bins": num_bins,
        }

    def calculate_image_stats(self, image: np.ndarray) -> dict[str, Any]:
        """Calculate comprehensive statistics for an image"""
        image_float = image.astype(np.float64)

        stats = {
            "timestamp": time.time(),
            "shape": image.shape,
            "dtype": str(image.dtype),
            "mean": float(np.mean(image_float)),
            "std": float(np.std(image_float)),
            "median": float(np.median(image_float)),
            "min": float(np.min(image_float)),
            "max": float(np.max(image_float)),
            "message_count": self.message_count,
        }

        # Calculate per-channel stats if color image
        if len(image.shape) == 3 and image.shape[2] > 1:
            for channel in range(image.shape[2]):
                channel_data = image_float[:, :, channel]
                stats[f"mean_ch{channel}"] = float(np.mean(channel_data))
                stats[f"std_ch{channel}"] = float(np.std(channel_data))
                stats[f"median_ch{channel}"] = float(np.median(channel_data))

        radial_data = self.calculate_radial_integration(image)
        stats["radial"] = radial_data

        # projection_data = self.calculate_projections(image)
        # stats['projections'] = projection_data

        return stats

    def process_message(self):
        """Process a single ZeroMQ message"""
        try:
            r = self._socket.recv_multipart(zmq.NOBLOCK)
            if len(r) != 2:
                return None

            meta, data = r
            header = json.loads(meta)
            header_shape = header["shape"]
            self.message_count += 1

            if header["type"] == "uint8" and len(header_shape) == 2:
                # Bayer image
                bayer_image = np.frombuffer(data, dtype=np.uint8)
                bayer_image = bayer_image.reshape(header_shape)

                # Convert to RGB for main stats
                rgb_image = cv2.cvtColor(bayer_image, cv2.COLOR_BAYER_GB2RGB)
                rgb_image = rgb_image[:, ::-1, :].copy()

                stats = self.calculate_image_stats(rgb_image)
                stats["image_type"] = "rgb_from_bayer"

                with self.stats_lock:
                    self.latest_stats = stats

                return stats

            elif header["type"] == "uint8" and len(header_shape) == 3:
                # RGB image
                rgb_image = np.frombuffer(data, dtype=np.uint8)
                rgb_image = rgb_image.reshape(header_shape)

                stats = self.calculate_image_stats(rgb_image)
                stats["image_type"] = "rgb_direct"

                with self.stats_lock:
                    self.latest_stats = stats

                return stats

        except zmq.Again:
            return None
        except Exception:
            logger.exception("Error processing ZMQ message")
            return None

    def receiver_thread(self):
        """Background thread for receiving messages"""
        while self.running:
            self.process_message()
            time.sleep(0.001)  # Small sleep to prevent busy waiting

    def print_stats_thread(self):
        """Background thread for printing stats every second"""
        while self.running:
            time.sleep(5.0)  # Print once per second

            with self.stats_lock:
                if self.latest_stats:
                    self.print_formatted_stats(self.latest_stats)
                else:
                    print(f"[{time.strftime('%H:%M:%S')}] Waiting for image data...")

    def print_formatted_stats(self, stats: dict[str, Any]):
        """Print formatted statistics in a compact terminal format"""
        timestamp = time.strftime("%H:%M:%S", time.localtime(stats["timestamp"]))

        # Compact one-line format
        print(
            f"[{timestamp}] #{stats['message_count']:4d} | "
            f"Shape: {stats['shape']} | "
            f"Mean: {stats['mean']:6.1f} | "
            f"Std: {stats['std']:6.1f} | "
            f"Median: {stats['median']:6.1f} | "
            f"Range: [{stats['min']:3.0f}-{stats['max']:3.0f}]",
            flush=True,
        )

        if "radial" in stats:
            radial = stats["radial"]
            center_intensity = radial["radial_profile"][0] if radial["radial_profile"] else 0
            edge_intensity = radial["radial_profile"][-1] if radial["radial_profile"] else 0
            peak_radius_idx = np.argmax(radial["radial_profile"]) if radial["radial_profile"] else 0
            peak_radius = radial["r_centers"][peak_radius_idx] if radial["r_centers"] else 0

            print(
                f"{'':21} Radial: Center={center_intensity:.1f} | "
                f"Edge={edge_intensity:.1f} | "
                f"Peak@r={peak_radius:.1f} | "
                f"Center=({radial['center'][0]},{radial['center'][1]})",
                flush=True,
            )

        # Print projection summary
        if "projections" in stats:
            proj = stats["projections"]
            print(
                f"{'':21} X-Profile: Peak@{proj['x_peak_position']}({proj['x_peak_value']:.1f}) | "
                f"Centroid={proj['x_centroid']:.1f} | FWHM={proj['x_fwhm']:.1f}",
                flush=True,
            )
            print(
                f"{'':21} Y-Profile: Peak@{proj['y_peak_position']}({proj['y_peak_value']:.1f}) | "
                f"Centroid={proj['y_centroid']:.1f} | FWHM={proj['y_fwhm']:.1f}",
                flush=True,
            )

        # Optional: Print per-channel stats if available
        if any(k.startswith("mean_ch") for k in stats):
            channels = []
            i = 0
            while f"mean_ch{i}" in stats:
                channels.append(f"Ch{i}({stats[f'mean_ch{i}']:.1f})")
                i += 1
            if channels:
                print(f"{'':21} Channels: {' '.join(channels)}", flush=True)

    def get_radial_profile_summary(self):
        """Get a summary of the current radial profile"""
        with self.stats_lock:
            if self.latest_stats and "radial" in self.latest_stats:
                radial = self.latest_stats["radial"]
                return {
                    "r_centers": radial["r_centers"],
                    "radial_profile": radial["radial_profile"],
                    "center_intensity": radial["radial_profile"][0]
                    if radial["radial_profile"]
                    else 0,
                    "edge_intensity": radial["radial_profile"][-1]
                    if radial["radial_profile"]
                    else 0,
                    "max_intensity_radius": radial["r_centers"][np.argmax(radial["radial_profile"])]
                    if radial["radial_profile"]
                    else 0,
                    "max_intensity_value": max(radial["radial_profile"])
                    if radial["radial_profile"]
                    else 0,
                }
        return None

    def save_radial_profile_to_file(self, filename: str | None = None):
        """Save the current radial profile to a file"""
        if filename is None:
            filename = f"radial_profile_{int(time.time())}.txt"

        with self.stats_lock:
            if self.latest_stats and "radial" in self.latest_stats:
                radial = self.latest_stats["radial"]

                with open(filename, "w") as f:
                    f.write("# Radial Integration Profile\n")
                    f.write(f"# Timestamp: {time.ctime(self.latest_stats['timestamp'])}\n")
                    f.write(f"# Image shape: {self.latest_stats['shape']}\n")
                    f.write(f"# Center: {radial['center']}\n")
                    f.write("# Radius(pixels)\tMean_Intensity\tStd_Intensity\tPixel_Count\n")

                    f.writelines(
                        f"{radial['r_centers'][i]:.2f}\t"
                        f"{radial['radial_profile'][i]:.2f}\t"
                        f"{radial['radial_std'][i]:.2f}\t"
                        f"{radial['pixel_counts'][i]}\n"
                        for i in range(len(radial["r_centers"]))
                    )

                print(f"Radial profile saved to {filename}")
                return filename

        print("No radial profile data available")
        return None

    def start(self):
        """Start the background threads"""
        self.running = True

        # Start receiver thread
        self.receiver_thread_obj = threading.Thread(target=self.receiver_thread, daemon=True)
        self.receiver_thread_obj.start()

        # Start stats printing thread
        self.print_thread_obj = threading.Thread(target=self.print_stats_thread, daemon=True)
        self.print_thread_obj.start()

        print("Image stats receiver started - printing every 1 second")

    def stop(self):
        """Stop the receiver and all threads"""
        self.running = False
        if self._socket:
            self._socket.close()
        print("Image stats receiver stopped")


# Global stats receiver instance
stats_receiver = None


def start_image_stats_receiver(zmq_url: str = "tcp://localhost:5555"):
    """Start the image stats receiver in background"""
    global stats_receiver

    if stats_receiver is None or not stats_receiver.running:
        try:
            stats_receiver = ImageStatsReceiver(zmq_url)
            stats_receiver.start()
            print(f"Started image statistics monitoring on {zmq_url}")
        except Exception:
            logger.exception("Failed to start image stats receiver")
            stats_receiver = None


def stop_image_stats_receiver():
    """Stop the image stats receiver"""
    global stats_receiver

    if stats_receiver and stats_receiver.running:
        stats_receiver.stop()
        stats_receiver = None


def get_latest_image_stats():
    """Get the latest image statistics"""
    if stats_receiver and stats_receiver.latest_stats:
        with stats_receiver.stats_lock:
            return stats_receiver.latest_stats.copy()
    return None


def get_radial_profile():
    """Get just the radial profile data"""
    if stats_receiver:
        return stats_receiver.get_radial_profile_summary()
    return None


def save_current_radial_profile(filename: str | None = None):
    """Save the current radial profile to a file"""
    if stats_receiver:
        return stats_receiver.save_radial_profile_to_file(filename)
    return None