"""Functions for converting SPARCL results to specutils objects.
"""
import warnings
import numpy as np
try:
    # specutils >= 2.0
    from specutils import Spectrum
except ImportError:
    from specutils import Spectrum1D as Spectrum
from specutils import SpectrumCollection, SpectrumList
from astropy.nddata import InverseVariance
import astropy.units as u

def _validate_records(records, r0, collection):
    """Validate that records can be converted to Spectrum.

    Parameters
    ----------
    records : list of dict
        All records to validate.
    r0 : dict
        First record, used as reference for validation.
    collection: bool
        If ``True``, attempt to convert to a
        :class:`~specutils.SpectrumCollection` instead.

    Raises
    ------
    ValueError
        If records lack 'wavelength' attribute.
        If records lack 'flux' attribute.
        If records have different data releases.
        If wavelength array lengths differ (suggests using SpectrumList).
        If wavelength pixel values differ and collection=False
        (suggests using SpectrumCollection).

    Warnings
    --------
    UserWarning
        If records come from different data releases.
    """

    # Check if the first record has wavelength data
    if 'wavelength' not in r0:
        raise ValueError("Results do not have a wavelength attribute. "
                         "Conversion is not possible.")

    if 'flux' not in r0:
        raise ValueError("Results do not have a flux attribute. "
                         "Conversion is not possible.")

    # Check if all records come from the same data release
    if not all([r.data_release == r0.data_release for r in records]):
        warnings.warn("Results are not all from the same data release, "
                      "conversion may not be possible.", UserWarning)

    # Check if all records have the same number of wavelength points
    if not all([len(r.wavelength) == len(r0.wavelength) for r in records]):
        raise ValueError("Results do not have the same wavelength solution. "
                         "Consider using .to_SpectrumList instead.")

    # If not creating a SpectrumCollection, check that wavelength values
    # are identical across all records
    if not collection and not all([(r.wavelength == r0.wavelength).all()
                                   for r in records]):
        raise ValueError("Results do not have the same wavelength pixels. "
                         "Consider using SpectrumCollection instead.")

def _extract_record_data(records, flux, uncertainty, mask, model, redshift,
                         meta, spectral_axis, has_ivar, has_mask, has_model,
                         has_redshift, collection, single_record):
    """Extract all data from records into arrays, then store any remaining
    attributes as metadata. Defines assign() and store_meta() helper
    functions before the main loop to avoid checking single_record on every
    iteration. All outputs are written in-place rather than returned.

    Parameters
    ----------
    records: list of dict
        Records containing flux, ivar, mask, wavelength, and optional
        model/redshift.
    flux : np.ndarray
        Pre-allocated array for flux values. Shape: (n_pixels,) if
        single_record, else (n_records, n_pixels).
    uncertainty : np.ndarray or None
        Pre-allocated array for inverse variance if has_ivar=True.
    mask : np.ndarray or None
        Pre-allocated array for data quality masks if has_mask=True.
    model : np.ndarray or None
        Pre-allocated array for model values if has_model=True.
    redshift : list
        Empty list to populate with redshift values.
    meta : dict
        Empty dict to populate with metadata. For a single record, values
        are stored as scalars. For multiple records, values are accumulated
        into lists.
    spectral_axis : np.ndarray
        For collections, pre-allocated 2D array to store per-record
        wavelength grids. For non-collections, this is just r0.wavelength
        (not modified).
    has_ivar : bool
        Whether records contain 'ivar' attribute.
    has_mask : bool
        Whether records contain 'mask' attribute.
    has_model : bool
        Whether records contain 'model' attribute.
    has_redshift : bool
        Whether records contain 'redshift' attribute.
    collection : bool
        If True, stores each record's wavelength array in spectral_axis.
    single_record : bool
        If True, assign() writes to a 1D array and store_meta() stores
        scalar values. If False, assign() writes to rows of a 2D array
        and store_meta() accumulates values into lists.

    Returns
    -------
    None
        All outputs are written to the input arrays/containers in-place.
    """
    # Build a list of (array, field_name) pairs for optional fields so we
    # can loop over them instead of having a separate if-block for each one
    optional_fields = []
    if has_ivar:
        optional_fields.append((uncertainty, 'ivar'))
    if has_mask:
        optional_fields.append((mask, 'mask'))
    if has_model:
        optional_fields.append((model, 'model'))

    # Pre-compute the set of fields we've already handled so the metadata
    # loop below can skip them quickly
    skip = {'flux', 'ivar', 'mask', 'model', 'wavelength'}

    for k, record in enumerate(records):
        # flux is always present so just copy it in
        if single_record:
            flux[:] = record.flux
            for arr, field in optional_fields:
                arr[:] = record[field]
        else:
            flux[k, :] = record.flux
            for arr, field in optional_fields:
                arr[k, :] = record[field]

        # Accumulate redshifts into a list for later conversion to an array
        if has_redshift:
            redshift.append(record.redshift)

        # Collections can have different wavelengths per record, so save
        # each one separately
        if collection:
            spectral_axis[k, :] = record.wavelength

        # Store any remaining attributes as metadata, skipping fields
        # that we already handled above
        for attribute, value in record.items():
            if attribute not in skip:
                if single_record:
                    meta[attribute] = value
                else:
                    meta.setdefault(attribute, []).append(value)

def to_Spectrum(results, *, collection=False):
    """Convert `results` to :class:`specutils.Spectrum`.

    Parameters
    ----------
    results : :class:`sparcl.Results.Retrieved`
        Retrieved results, or a single record from a set of results.
    collection : bool, optional
        If ``True``, attempt to convert to a
        :class:`~specutils.SpectrumCollection` instead.

    Returns
    -------
    :class:`~specutils.Spectrum` or :class:`~specutils.SpectrumCollection`
        The requested object.

    Raises
    ------
    ValueError
        If `results` can't be converted to a :class:`~specutils.Spectrum`
        object in a valid way. For example, if some of the spectra have a
        different wavelength solution.
    """
    # Prepare records
    if isinstance(results, dict):
        records = [results]
        r0 = results
    else:
        try:
            records = results.records
            if len(records) == 0:
                raise ValueError("No records found in results. Cannot "
                                 "convert empty results to Spectrum.")
            r0 = results.records[0]
        except (IndexError, AttributeError) as e:
            raise ValueError("No records found in results. Cannot "
                             "convert empty results to Spectrum.") from e

    # Validate
    _validate_records(records, r0, collection)

    # Determine which optional data components exist in records
    has_redshift = 'redshift' in r0
    has_model = 'model' in r0
    has_ivar = 'ivar' in r0
    has_mask = 'mask' in r0
    n_pixels = r0.flux.shape[0]
    single_record = len(records) == 1

    # Set flux shape based on number of records
    if single_record:
        flux_shape = (n_pixels,)
    else:
        flux_shape = (len(records), n_pixels)

    # Build spectral axis
    if collection:
        spectral_axis = np.empty((len(records), r0.wavelength.shape[0]),
                                 dtype=r0.wavelength.dtype)
    else:
        spectral_axis = r0.wavelength

    # Initialize arrays
    flux = np.empty(flux_shape, dtype=r0.flux.dtype)
    uncertainty = (np.empty(flux_shape, dtype=r0.ivar.dtype)
                   if has_ivar else None)
    mask = np.empty(flux_shape, dtype=r0.mask.dtype) if has_mask else None
    model = np.empty(flux_shape, dtype=r0.model.dtype) if has_model else None
    redshift = []
    meta = {}

    # Populate arrays by iterating through records
    _extract_record_data(records, flux, uncertainty, mask, model, redshift,
                         meta, spectral_axis, has_ivar, has_mask, has_model,
                         has_redshift, collection, single_record)

    # Convert redshift list to numpy array if exists
    if has_redshift:
        redshift = np.array(redshift)
        if single_record and len(redshift) == 1:
            # Convert to scalar if single record
            redshift = redshift[0]
    else:
        redshift = None

    # Add model to metadata if exists
    if has_model:
        meta['model'] = model

    # Prepare arguments common to both Spectrum and SpectrumCollection
    common_args = {
        'flux': flux * u.Unit('10**-17 erg cm-2 s-1 AA-1'),
        'spectral_axis': spectral_axis * u.AA,
        'uncertainty': InverseVariance(uncertainty) if has_ivar else None,
        'mask': mask,
        'meta': meta}

    if collection:
        return SpectrumCollection(**common_args)

    return Spectrum(**common_args, redshift=redshift)

def to_SpectrumList(results):
    """Convert `results` to :class:`specutils.SpectrumList`.

    Parameters
    ----------
    results : :class:`sparcl.Results.Retrieved`
        Retrieved results.

    Returns
    -------
    :class:`~specutils.SpectrumList`
        The requested object.
    """
    s = SpectrumList()
    if isinstance(results, dict):
        records = [results]
    else:
        records = results.records
    for r in records:
        if 'redshift' in r:
            redshift = r.redshift
        else:
            redshift = None
        if 'ivar' in r:
            uncertainty = InverseVariance(r.ivar)
        else:
            uncertainty = None
        if 'mask' in r:
            mask = r.mask
        else:
            mask = None
        meta = dict()
        for attribute in r:
            if attribute not in ('flux', 'wavelength', 'ivar',
                                 'redshift', 'mask'):
                meta[attribute] = r[attribute]
        s1 = Spectrum(flux=r.flux*u.Unit('10**-17 erg cm-2 s-1 AA-1'),
                      spectral_axis=r.wavelength*u.AA,
                      uncertainty=uncertainty,
                      redshift=redshift,
                      mask=mask,
                      meta=meta)
        s.append(s1)
    return s

def to_specutils(results):
    """Convert `results` to a specutils object.

    Parameters
    ----------
    results : :class:`sparcl.Results.Retrieved`
        Retrieved results.

    Returns
    -------
    :class:`~specutils.Spectrum` or :class:`~specutils.SpectrumCollection`
    or :class:`~specutils.SpectrumList`
        The most natural conversion to a specutils object.

    Raises
    ------
    ValueError
        If no valid conversion can be performed.
    """
    try:
        # Try standard Spectrum conversion first
        s = to_Spectrum(results)
    except ValueError as ve:
        # Check the error message to determine appropriate alternative
        if 'SpectrumList' in str(ve):
            # Different wavelength array lengths use SpectrumList
            s = to_SpectrumList(results)
        elif 'SpectrumCollection' in str(ve):
            # Same wavelength length but diff pixels use SpectrumCollection
            s = to_Spectrum(results, collection=True)
        else:
            raise ValueError("Could not find a valid conversion to "
                             "specutils objects!")
    return s
