"""Functions for the spectroscopic mode of the Lux ETC"""

import numpy as np
from astropy import units as u
from specutils import Spectrum1D
from specutils.manipulation import FluxConservingResampler
from astropy.io import ascii
from astropy.nddata import StdDevUncertainty
from specutils import Spectrum
from specutils.manipulation import gaussian_smooth
from specutils import SpectralRegion
from specutils.manipulation import extract_region
# Standard library
from dataclasses import dataclass

# Local
from . import config
from . import hardware
from . import nirchannel
from . import uvchannel
from . import optchannel

from .optchannel import OptChannel
from .uvchannel import UVChannel
from .nirchannel import NIRChannel

optchannel = OptChannel()
uvchannel = UVChannel()
nirchannel = NIRChannel()

def luxreturnwave():
    return config.WAV

def luxreturnaeffoptspec():
    print('hello')
    yo = optchannel.Aeff_optspec
    return yo
    #return optchannel.OptChannel.Aeff_optspec()



def luxspecformat(inspec_path,mode="opt"):
    """for an input spectrum, will spit out a spectrum rebinned and 
    smoothed to Lux dispersion and resolution """

    #input spectrum should be an ascii file where first two columns are ang, flux (erg/cm**2/A)
    
    data = ascii.read(inspec_path)

    # Extract columns and assign units
    wavelength = data['col1'] * u.AA
    flux = data['col2'] * u.erg / (u.cm**2 * u.s * u.AA)

    #create spectrum object, find wmin and wmax
    spec = Spectrum(spectral_axis=wavelength, flux=flux)
    w_min = spec.spectral_axis.min()
    w_max = spec.spectral_axis.max()

    #initialize resampler
    resampler = FluxConservingResampler()
    
    #which mode are we in, assign wavelength range
    if mode=="opt":
        lux_wmin = 4000. * u.AA
        lux_wmax = 8000. * u.AA
        #Create new grid from w_min to w_max with 10A steps
        #Assuming R~100 at 5000A.  
        #10A/pix gives us 5 pixel per resolution element
        pixscale = 10. #10A per pixel
        sampling_factor = 5.
        new_grid = np.arange(w_min.value, w_max.value, pixscale) * u.AA

        #this returns a new spectrum object on the 10A (pixscale) grid
        spec_rebinned = resampler(spec, new_grid)
        
        # Check the results
        #print(f"Original bins: {len(spec.spectral_axis)}")
        #print(f"New bins: {len(spec_rebinned.spectral_axis)}")

        dw = np.diff(spec_rebinned.spectral_axis).mean()

        #find sigma and smooth spectrum
        sigma_angstrom = pixscale * sampling_factor * u.AA
        sigma_pixels = (sigma_angstrom / dw).decompose().value
        smoothed_spec = gaussian_smooth(spec_rebinned, stddev=sigma_pixels)

        #cut out relevant wavelength range
        region = SpectralRegion(lux_wmin, lux_wmax)

        # Extract the sub-spectrum from your original object
        sub_smoothedspectrum = extract_region(smoothed_spec, region)

    elif mode=="uv":
        lux_wmin = 2000. * u.AA
        lux_wmax = 4000. * u.AA
        #Create new grid from w_min to w_max with 10A steps
        #Assuming R~100 at 3000A.  
        #6A/pix gives us 5 pixel per resolution element
        pixscale = 6. #6A per pixel
        sampling_factor = 5.
        
        new_grid = np.arange(w_min.value, w_max.value, pixscale) * u.AA

        #this returns a new spectrum object on the 6A (pixscale) grid
        spec_rebinned = resampler(spec, new_grid)
        dw = np.diff(spec_rebinned.spectral_axis).mean()

         #find sigma and smooth spectrum
        sigma_angstrom = pixscale * sampling_factor * u.AA
        sigma_pixels = (sigma_angstrom / dw).decompose().value
        smoothed_spec = gaussian_smooth(spec_rebinned, stddev=sigma_pixels)

        #cut out relevant wavelength range
        region = SpectralRegion(lux_wmin, lux_wmax)

        # Extract the sub-spectrum from your original object
        sub_smoothedspectrum = extract_region(smoothed_spec, region)
       

    
    elif mode=="nir":
        raise ValueError("NIR spectroscopic mode is not implemented yet")

    else:
        raise ValueError("No such mode in set: %s" % mode)

    return sub_smoothedspectrum


def luxspec_cnts(luxspec_in, mode='opt'):

    #this function will take an input spectrum in lux resolution 
    #and pixel scale (from luxspecformat) and multiply it by 
    #the effective area.
    
    #input spectroscopic units should be u.erg / (u.cm**2 * u.s * u.AA)
    #this is not being checked right now

    #input should be the rebinned, smoothed spectrum made by luxspecformat

    
    # Define the target unit
    target_unit = u.photon / (u.cm**2 * u.s * u.AA)

    # We use the spectral_density equivalency, passing the spectral axis 
    # so the code knows how to calculate energy-per-photon at each point.
    spec_photons = luxspec_in.with_flux_unit(
        target_unit, 
        equivalencies=u.spectral_density(luxspec_in.spectral_axis)
    )

    resampler = FluxConservingResampler()
    
    #which mode are we in
    if mode=="opt":
        #get the optical effective area into specutils format
        Aeff_optspec = Spectrum(flux=optchannel.Aeff_optspec,
                     spectral_axis=config.WAV)

        # Match the units of the spectral axis before resampling
        Aeff_optspec_converted = \
            Aeff_optspec.with_spectral_axis_unit(spec_photons.spectral_axis.unit)

        
        # Resample Aeff_optspec to match the grid of luxspec_in
        Aeff_optspec_matched = \
            resampler(Aeff_optspec_converted,spec_photons.spectral_axis)

        src_cts = spec_photons * Aeff_optspec_matched
    
    elif mode=="uv":

        #get the optical effective area into specutils format
        Aeff_uvspec = Spectrum(flux=uvchannel.Aeff_uvspec,
                     spectral_axis=config.WAV)

        # Match the units of the spectral axis before resampling
        Aeff_uvspec_converted = \
            Aeff_uvspec.with_spectral_axis_unit(spec_photons.spectral_axis.unit)

        # Resample Aeff_uvspec to match the grid of luxspec_in
        Aeff_uvspec_matched = \
            resampler(Aeff_uvspec_converted,spec_photons.spectral_axis)

        src_cts = spec_photons * Aeff_uvspec_matched
   
    elif mode=="nir":
        raise ValueError("NIR spectroscopic mode is not implemented yet")

    else:
        raise ValueError("No such mode in set: %s" % mode)


    return src_cts

def rebin_SNR_spectrum(wavelengths, snr_per_pixel, factor=5):
    # Determine the number of new bins
    n_bins = len(wavelengths) // factor

    # Trim the arrays to be exactly divisible by the factor
    w_trimmed = wavelengths[:n_bins * factor]
    s_trimmed = snr_per_pixel[:n_bins * factor]

    # Average the wavelengths within each bin
    new_wavelength = w_trimmed.reshape(-1, factor).mean(axis=1)

    # Calculate the new SNR: Average SNR in the bin * sqrt(factor)
    new_snr = s_trimmed.reshape(-1, factor).mean(axis=1) * np.sqrt(factor)

    return new_wavelength, new_snr

def add_lux_noise(luxspec,snrperpix):

    #calculate per-pixel sigma
    #sigma(lambda) = flux(lambda) / SNR(lambda)
    flux_vals = luxspec.flux.value
    sigma_array = flux_vals / snrperpix.value

    # Generate noise scaled by the local sigma
    # We multiply a standard normal distribution by our sigma array
    noise = np.random.normal(0, 1, size=flux_vals.shape) * sigma_array

    # Create the new degraded spectrum
    degraded_flux = (flux_vals + noise) * luxspec.flux.unit
    uncertainty = StdDevUncertainty(sigma_array * luxspec.flux.unit)

    degraded_spectrum = Spectrum(
        spectral_axis=luxspec.spectral_axis, 
        flux=degraded_flux 
        #uncertainty=uncertainty
    )

    return degraded_spectrum

def luxspec_snr(inspec_path,texp,mode='opt'):

    exptime = texp * u.s

    Nframes = exptime / (config.TFRAME * u.s)

    #print(Nframes)
    #this is the number of pixels we are summing in spatial direction
    #think about implementing this in a smarter way later.
    Nspatial = 5. 

    #this is for converting from SNR/pixel to SNR/resolution element
    #needs to be implemented better, but is how we have set up luxspecformat
    Nspectral = 5. 

    gain = 1. * u.electron / u.photon

        
    #read in input spectrum and put into lux format
    luxspec = luxspecformat(inspec_path,mode=mode)

    #calculate spec counts per second and multiply by effective area.
    luxspec_counts = luxspec_cnts(luxspec, mode=mode)
    
    #which mode are we in
    if mode=="opt":
    #read in input spectrum and put into lux format
        #this really needs to be defined more universally, is in luxspecformat
        pixscale_opt = 10. * u.AA
        
        
        
        dark_noise = optchannel.background_dark_rate * texp * u.s
        dark_noise_withspatial = dark_noise * Nspatial 

        
        
        phot_elec = luxspec_counts * (exptime) * (pixscale_opt) * gain

        #calculate SNR per pixel
        SNRperpix = phot_elec.flux / \
            np.sqrt(phot_elec.flux + dark_noise_withspatial + Nframes * Nspatial * optchannel.background_read_noise**2 / u.electron ) 


        #calculate SNR per resolution element and get the wavelength array too
        SNRperres_wave, SNRperres = \
            rebin_SNR_spectrum(luxspec.spectral_axis, SNRperpix, factor=5)

        #calculate output spectrum with noise
        observed_spectrum = add_lux_noise(luxspec,SNRperpix)
    
    elif mode=="uv":

        #this really needs to be defined more universally, is in luxspecformat
        pixscale_uv = 6. * u.AA

        #Nframes = exptime / LETC.UVChannel.integration_time

        dark_noise = uvchannel.background_dark_rate * texp * u.s
        dark_noise_withspatial = dark_noise * Nspatial 

        phot_elec = luxspec_counts * (exptime) * (pixscale_uv) * gain

        #calculate SNR per pixel
        SNRperpix = phot_elec.flux / \
            np.sqrt(phot_elec.flux + dark_noise_withspatial + Nframes * Nspatial * uvchannel.background_read_noise**2 / u.electron ) 


        #calculate SNR per resolution element and get the wavelength array too
        SNRperres_wave, SNRperres = \
            rebin_SNR_spectrum(luxspec.spectral_axis, SNRperpix, factor=5)

        #calculate output spectrum with noise
        observed_spectrum = add_lux_noise(luxspec,SNRperpix)
    
    elif mode=="nir":
        raise ValueError("NIR spectroscopic mode is not implemented yet")

    else:
        raise ValueError("No such mode in set: %s" % mode)


    #1. return the spectrum rebinned and smoothed to Lux resolution.
    return luxspec, observed_spectrum, SNRperpix, SNRperres, SNRperres_wave
