import json
import cv2
import numpy as np
import zmq
import time
import threading
from typing import Dict, Any


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 as e:
            print(f"Error processing ZMQ message: {e}")
            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.keys()):
            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):
        """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")

                    for i in range(len(radial['r_centers'])):
                        f.write(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")

                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(f"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 as e:
            print(f"Failed to start image stats receiver: {e}")
            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"""
    global stats_receiver

    if stats_receiver and stats_receiver.latest_stats:
        with stats_receiver.stats_lock:
            return stats_receiver.latest_stats.copy()
    return None


def get_latest_image_stats():
    """Get the latest image statistics including radial profile"""
    global stats_receiver

    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"""
    global stats_receiver

    if stats_receiver:
        return stats_receiver.get_radial_profile_summary()
    return None


def save_current_radial_profile(filename: str = None):
    """Save the current radial profile to a file"""
    global stats_receiver

    if stats_receiver:
        return stats_receiver.save_radial_profile_to_file(filename)
    return None

# # Auto-start when daq.py is imported/run
# if __name__ == "__main__":
#     # If running daq.py directly
#     start_image_stats_receiver()
#
#     try:
#         # Your existing daq.py code continues here
#         print("DAQ system running with image statistics monitoring...")
#
#         # Keep the program running
#         while True:
#             time.sleep(1)
#
#             # Optional: Access latest stats in your main code
#             latest = get_latest_image_stats()
#             if latest and latest['message_count'] % 10 == 0:  # Every 10th message
#                 print(f"Main thread sees: Mean={latest['mean']:.1f}")
#
#     except KeyboardInterrupt:
#         print("\nShutting down...")
#     finally:
#         stop_image_stats_receiver()
#
# else:
#     # If daq.py is imported as a module, auto-start the receiver
#     start_image_stats_receiver()