"""
Module for generating aligned matrix PDF plots for high-resolution spectroscopic data.

This module provides functionality to create multi-page PDF reports showing
synthetic aligned matrices for different spectral orders.
"""

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.backends.backend_pdf import PdfPages
from typing import Optional, Union
from pathlib import Path


def generate_aligned_matrix_pdf(
    synthetic_plus_error_model: np.ndarray,
    wavelength_matrix: Optional[np.ndarray],
    output_path: Union[str, Path],
    target_name: str = "Target",
    invert_x: bool = False,
    original_aligned_data: Optional[np.ndarray] = None,
) -> None:
    """
    Generate a multi-page PDF with aligned matrix plots, one page per spectral order.

    Each page contains two subplots:
    - Upper subplot: Aligned matrix visualization for the current order showing
    flux variations across wavelength points and image numbers
    - Lower subplot: Original aligned data visualization (if provided), otherwise blank

    The function creates a comprehensive visualization of synthetic spectroscopic
    data with error models, useful for quality assessment and data analysis.

    Parameters
    ----------
    synthetic_plus_error_model : np.ndarray
        3D numpy array of shape (wavelength_points, orders, images) containing
        synthetic flux data with error model applied
    wavelength_matrix : np.ndarray or None
        3D numpy array of shape (wavelength_points, orders, images) containing
        wavelength calibration values. If None, wavelength information won't
        be displayed on plots
    output_path : str or Path
        File path where the PDF report will be saved. Should include .pdf extension
    target_name : str, optional
        Name of the astronomical target for plot titles and headers, by default "Target"
    invert_x: bool, optional
        If True, the wavelength axis will be inverted (from right to left) in the plots, by default False.
    original_aligned_data : np.ndarray or None, optional
        3D numpy array of shape (wavelength_points, orders, images) containing
        original aligned data. If provided, will be plotted in the lower subplot.

    Returns
    -------
    None
        The function saves the PDF file to the specified path

    Raises
    ------
    ValueError
        If synthetic_plus_error_model is not a 3D array
    IOError
        If the output path cannot be written to

    """
    # Input validation
    if synthetic_plus_error_model.ndim != 3:
        raise ValueError(f"Expected 3D array, got {synthetic_plus_error_model.ndim}D array")
    
    print(f"=== Generating Aligned Matrix PDF ===")
    print(f"Matrix shape: {synthetic_plus_error_model.shape}")
    print(f"Output path: {output_path}")
    
    # Extract dimensions for clarity
    n_wavelength_points, n_orders, n_images = synthetic_plus_error_model.shape
    
    # Create multi-page PDF document
    with PdfPages(output_path) as pdf:
        # Generate one page per spectral order
        for order in range(n_orders):
            print(f"  Creating page for order {order+1}/{n_orders}")
            
            # Create figure with 2 vertically stacked subplots
            fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 10))
            
            # Extract data for current order: (wavelength_points, images)
            order_data = synthetic_plus_error_model[:, order, :]
            
            # ===== Upper subplot: Aligned matrix visualization =====
            # Create 2D image plot showing flux variation across wavelength and image number
            im = ax1.imshow(
                np.transpose(order_data),
                aspect='auto',           # Automatic aspect ratio scaling
                cmap='gnuplot2',         # Color map optimized for scientific data
                origin='lower',          # Place origin at bottom-left
                interpolation='nearest'  # No interpolation for data integrity
            )
            
            # Configure plot labels and title
            ax1.set_title(f'{target_name} - Order {order+1}/{n_orders} - Synthetic Aligned Matrix', 
                         fontsize=14, fontweight='bold')
            ax1.set_ylabel('Image Number', fontsize=12)
            ax1.set_xlabel('Wavelength [nm]', fontsize=12)
            
            # Add colorbar with flux units
            cbar = plt.colorbar(im, ax=ax1, label='Flux (arbitrary units)')
            cbar.ax.tick_params(labelsize=10)
            
            # Add wavelength range information if wavelength matrix is provided
            if wavelength_matrix is not None:
                # Calculate wavelength range for current order
                wl_min = np.min(wavelength_matrix[:, order, :])
                wl_max = np.max(wavelength_matrix[:, order, :])
                tick_positions = np.linspace(0, n_wavelength_points - 1, 5)
                if invert_x:
                    tick_labels = np.linspace(wl_max, wl_min, 5)
                else:
                    tick_labels = np.linspace(wl_min, wl_max, 5)
                # Add wavelength range annotation
                ax1.text(0.02, 0.98, f'λ: {wl_min:.2f} - {wl_max:.2f} nm', 
                        transform=ax1.transAxes, verticalalignment='top',
                        bbox=dict(boxstyle='round', facecolor='white', alpha=0.8),
                        fontsize=10, fontweight='bold')
                ax1.set_xticks(tick_positions)
                ax1.set_xticklabels([f'{x:.1f}' for x in tick_labels])
            
            # ===== Lower subplot: Original aligned data or placeholder =====
            if original_aligned_data is not None:
                # Extract original data for current order: (wavelength_points, images)
                original_order_data = original_aligned_data[:, order, :]

                # Create 2D image plot showing original flux variation
                im2 = ax2.imshow(
                    np.transpose(original_order_data),
                    aspect='auto',
                    cmap='gnuplot2',
                    origin='lower',
                    interpolation='nearest'
                )

                # Configure plot labels and title
                ax2.set_title(f'{target_name} - Order {order+1}/{n_orders} - Original Aligned Matrix',
                             fontsize=14, fontweight='bold')
                ax2.set_ylabel('Image Number', fontsize=12)
                ax2.set_xlabel('Wavelength [nm]', fontsize=12)

                # Add colorbar with flux units
                cbar2 = plt.colorbar(im2, ax=ax2, label='Flux (arbitrary units)')
                cbar2.ax.tick_params(labelsize=10)

                # Add wavelength range information if wavelength matrix is provided
                if wavelength_matrix is not None:
                    ax2.set_xticks(tick_positions)
                    ax2.set_xticklabels([f'{x:.1f}' for x in tick_labels])
                    # Add wavelength range annotation
                    ax2.text(0.02, 0.98, f'λ: {wl_min:.2f} - {wl_max:.2f} nm',
                            transform=ax2.transAxes, verticalalignment='top',
                            bbox=dict(boxstyle='round', facecolor='white', alpha=0.8),
                            fontsize=10, fontweight='bold')
            else:
                # Placeholder when no original data is provided
                ax2.set_title('Reserved for Future Use', fontsize=14, fontweight='bold')
                ax2.set_xlabel('Placeholder X-axis', fontsize=12)
                ax2.set_ylabel('Placeholder Y-axis', fontsize=12)

                # Add placeholder text with professional styling
                ax2.text(0.5, 0.5, 'Plot area reserved\nfor future implementation',
                        transform=ax2.transAxes, ha='center', va='center',
                        fontsize=14, style='italic', alpha=0.6,
                        bbox=dict(boxstyle='round', facecolor='lightgray', alpha=0.3))

                # Add grid for visual reference
                ax2.grid(True, alpha=0.3, linestyle='--')
            
            # Optimize layout and save current page
            plt.tight_layout(pad=2.0)  # Add padding between subplots
            pdf.savefig(fig, bbox_inches='tight', dpi=150)  # High DPI for quality
            plt.close(fig)  # Free memory by closing figure
    
    print(f"=== Aligned Matrix PDF Generated: {n_orders} pages ===")
    print(f"PDF saved successfully at: {output_path}")