import os
from pathlib import Path
import numpy as np
from astropy.io import fits
from scipy.interpolate import splrep, splev

from GUIBRUSHR.GUI.Input_Output_Panels.Input_Panels.TabPanels.FrameGenerationHRData.plot import \
    generate_aligned_matrix_pdf
from GUIBRUSHR.General_Constants.Classes.Night import Night
from GUIBRUSHR.General_Constants.FunctionsAndConstants.Constant_Variables import ConstantVariables
from GUIBRUSHR.General_Constants.FunctionsAndConstants.general_functions import (
    create_path_night,
    read_fits, read_depth_fits,
    convolve_resolution, rv_planet_and_star
)


def shuffle_HR(target_obj, path_target, rad_mode, instrument_obj, string_nights,
                 model_CC, order_sel, tell_rm_method, eccentricity=False):
    """
    Simulate high-resolution spectroscopic data for a given target.

    This function generates simulated high-resolution spectroscopic data by:
    1. Reading and processing model data
    2. Convolving the model with instrument resolution
    3. Generating error matrices
    4. Applying orbital motion effects
    5. Creating simulated spectra
    6. Saving the results

    Args:
        target_obj: Target object containing orbital parameters
        path_target: Path to target data directory
        rad_mode: Radiation mode (e.g., 'Transmission')
        instrument_obj: Instrument object containing specifications
        string_nights: Comma-separated string of observation nights
        model_CC: Cross-correlation model name
        order_sel: Selected order
        tell_rm_method: Telluric removal method
        eccentricity: Whether to use eccentric orbit calculations

    Returns:
        None. Results are saved to disk in a '_sim' directory.
    """
    print("=== Starting Simulation HR Data Processing ===")
    
    # Create necessary paths
    list_nights = string_nights.split(',')
    print(f"Processing {len(list_nights)} nights: {list_nights}")
    
    path_instrument = create_path_night(path_target, target_obj.name, rad_mode, instrument_obj.name)
    path_model_CC = str(Path(path_target, target_obj.name, ConstantVariables.CC_FOLDER, rad_mode, model_CC))
    print(f"Model path: {path_model_CC}")

    # Read and process model data
    print("Reading model data...")
    transit_depth = read_depth_fits(path_model_CC, rad_mode)
    wavelength = read_fits(str(Path(path_model_CC, "wavelength.fits")))
    print(f"Model wavelength range: {np.min(wavelength):.2f} - {np.max(wavelength):.2f} nm")
    print(f"Transit depth range: {np.min(transit_depth):.6f} - {np.max(transit_depth):.6f}")

    # Convolve model with instrument resolution
    print(f"Convolving model with instrument resolution (HWHM: {instrument_obj.hwhm_km_s} km/s)...")
    wl_convolved_resolution, model_convolved_resolution = convolve_resolution(
        wavelength, transit_depth, instrument_obj.hwhm_km_s, 0, 0.1
    )
    csscaled = splrep(wl_convolved_resolution, model_convolved_resolution)
    print("Model convolution and spline creation completed")

    # Process each night
    for night in list_nights:
        print(f"\n--- Processing night: {night} ---")
        
        # Setup paths for the current night
        path_night = str(Path(path_instrument, night))
        path_order = Path(path_instrument, night, "orders_selection", order_sel)
        path_order_tell = path_order / tell_rm_method
        path_order_test = str(Path(path_order_tell, "test"))
        path_good_order = str(path_order / f"good_order_{order_sel}.pkl")

        # Initialize night object
        print("Initializing night object...")
        night_obj = Night(
            target_obj.limit_phase_t14, path_instrument,
            night, rad_mode, path_good_order, path_order_tell, path_order_test,
            table_output_file=None, CC=False, simulation=True, synthetic=False, retrieval_standardize="From pkl"
        )
        print(f"Night initialized: {night_obj.n_tot_phases} images, {night_obj.n_good_orders} orders, {night_obj.n_pixels} pixels")

        # Generate error matrix using Gaussian distribution
        print("Generating error matrix...")
        gaussian_matrix = np.random.normal(loc=0.0, scale=1.0, size=night_obj.sigma_spectra_full.shape)
        error_matrix = gaussian_matrix * night_obj.sigma_spectra_full
        print(f"Error matrix shape: {error_matrix.shape}")

        # Calculate mean aligned data
        print("Processing aligned data...")
        data_aligned = night_obj.aligned
        mean_data_aligned = np.mean(data_aligned, axis=2)
        mean_data_aligned = np.expand_dims(mean_data_aligned, axis=2)
        mean_data_aligned = np.tile(mean_data_aligned, (1, 1, night_obj.n_tot_phases))
        print(f"Aligned data shape: {data_aligned.shape}")
        print(f"Mean flux range: {np.min(mean_data_aligned):.2e} to {np.max(mean_data_aligned):.2e}")

        vtot, _, _ = rv_planet_and_star(
            eccentricity, target_obj, 0, 0, 0, night_obj, "Synthetic",
            0, None, 0
        )

        # Generate spectrum
        print("Generating spectrum...")
        spectrum = np.zeros_like(night_obj.lambdas)
        for h in range(night_obj.n_good_orders):
            if h % 10 == 0:  # Print progress every 10 orders
                print(f"  Processing order {h+1}/{night_obj.n_good_orders}")
            
            dl_planet = night_obj.lambdas[:, h, :] * vtot / ConstantVariables.CLIGHT
            wlen_shifted_planet = dl_planet + night_obj.lambdas[:, h, :]
            if rad_mode == 'Transmission':
                spectrum[:, h, :] = splev(wlen_shifted_planet, csscaled, der=0)
            # TODO: Implement stellar spectrum calculation for other radiation modes
        print(f"Spectrum generation completed. Signal range: {np.min(spectrum):.6f} to {np.max(spectrum):.6f}")

        # Shuffle and process data
        print("Shuffling and processing data...")
        indices = np.arange(night_obj.n_tot_phases)
        np.random.shuffle(indices)
        data_aligned_new = np.array(data_aligned[:, :, indices])
        print(f"Data shuffled with {len(indices)} images")

        # Apply spectrum to in-transit data
        print("Applying spectrum to in-transit data...")
        idx = np.ix_(range(night_obj.n_pixels), night_obj.good_order, night_obj.intransit)
        print(f"In-transit images: {len(night_obj.intransit)}/{night_obj.n_tot_phases}")
        data_aligned_new[idx] = data_aligned_new[idx] * spectrum
        print("Spectrum applied to in-transit data")

        # Create simulation directory and copy data
        print("Saving results...")
        path_night_sim = f"{path_night}sim"
        os.system(f"mkdir -p {path_night_sim}")
        path_night += os.sep + "*"
        os.system(f"cp -r {path_night} {path_night_sim}")

        # Save results
        hdu = fits.PrimaryHDU(np.transpose(data_aligned_new, (1, 2, 0)))
        hdul = fits.HDUList([hdu])
        hdul.writeto(str(Path(path_night_sim, 'aligned.fits')), overwrite=True)

        hdu = fits.PrimaryHDU(np.transpose(spectrum, (1, 2, 0)))
        hdul = fits.HDUList([hdu])
        hdul.writeto(str(Path(path_night_sim, 'spectrum_simulation.fits')), overwrite=True)
        
        print(f"Night {night} completed. Results saved to: {path_night_sim}")

    generate_aligned_matrix_pdf(
        data_aligned_new, night_obj.complete_wavelength,
        Path(path_night_sim, "aligned_matrix.pdf"), target_obj.name,
        original_aligned_data=night_obj.aligned
    )
    
    print("=== Simulation HR Data Processing Completed ===\n")





