"""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."""

    # 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.")

    # 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_model, has_redshift,
                         collection, single_record):
    """Extract all data from records into arrays."""
    for k, record in enumerate(records):
        if single_record:
            # For single record, assign directly (1D)
            flux[:] = record.flux
            uncertainty[:] = record.ivar
            mask[:] = record.mask
            if has_model:
                model[:] = record.model
        else:
            # For multiple records, assign to row (2D)
            flux[k, :] = record.flux
            uncertainty[k, :] = record.ivar
            mask[k, :] = record.mask
            if has_model:
                model[k, :] = record.model

        # Store redshift values if available
        if has_redshift:
            redshift.append(record.redshift)

        # For collections, each record can have its own wavelength grid
        if collection:
            spectral_axis[k, :] = record.wavelength

        # Extract all additional metadata attributes
        for attribute, value in record.items():
            if attribute not in ('flux', 'ivar', 'mask', 'model', 'redshift',
                                 'wavelength'):
                if single_record:
                    # For single record, store metadata values as scalars
                    meta[attribute] = value
                else:
                    # For multiple records, accumulate metadata values into
                    # lists (creates list on first encounter, then appends)
                    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 : :class:`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
    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.zeros((len(records), r0.wavelength.shape[0]),
                                 dtype=r0.wavelength.dtype)
    else:
        spectral_axis = r0.wavelength

    # Initialize arrays
    flux = np.zeros(flux_shape, dtype=r0.flux.dtype)
    uncertainty = np.zeros(flux_shape, dtype=r0.ivar.dtype)
    mask = np.zeros(flux_shape, dtype=r0.mask.dtype)
    model = np.zeros(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_model, has_redshift,
                         collection, single_record)

    # Convert redshift list to numpy array if exists
    if has_redshift:
        redshift = np.array(redshift)
    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('1e-17 erg cm-2 s-1 AA-1'),
        'spectral_axis': spectral_axis * u.AA,
        'uncertainty': InverseVariance(uncertainty),
        '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
        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('1e-17 erg cm-2 s-1 AA-1'),
                      spectral_axis=r.wavelength*u.AA,
                      uncertainty=InverseVariance(r.ivar),
                      redshift=redshift,
                      mask=r.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
