use timsrust_core::utils::vec::{argsort, group_and_sum};

use crate::{
    FrameReaderError, TdfFrameReader,
    file_readers::sql_reader::{
        ReadableSqlTable, SqlReader, SqlReaderError,
        pasef_frame_msms::SqlPasefFrameMsMs,
    },
};

use super::raw_spectra::{
    RawSpectrum, RawSpectrumReaderError, RawSpectrumReaderTrait,
};

#[derive(Debug)]
pub(crate) struct DDARawSpectrumReader {
    order: Vec<usize>,
    offsets: Vec<usize>,
    pasef_frames: Vec<SqlPasefFrameMsMs>,
    frame_reader: TdfFrameReader,
}

impl DDARawSpectrumReader {
    pub(crate) fn new(
        tdf_sql_reader: &SqlReader,
        frame_reader: TdfFrameReader,
    ) -> Result<Self, DDARawSpectrumReaderError> {
        let pasef_frames = SqlPasefFrameMsMs::from_sql_reader(tdf_sql_reader)?;
        let pasef_precursors =
            &pasef_frames.iter().map(|x| x.precursor).collect::<Vec<_>>();
        let order = argsort(pasef_precursors);
        let max_precursor = pasef_precursors
            .iter()
            .max()
            .expect("SqlReader cannot return empty vecs, so there is always a max precursor index");
        let mut offsets = Vec::with_capacity(max_precursor + 1);
        offsets.push(0);
        for (offset, &index) in order.iter().enumerate().take(order.len() - 1) {
            let second_index: usize = order[offset + 1];
            if pasef_precursors[index] != pasef_precursors[second_index] {
                offsets.push(offset + 1)
            }
        }
        offsets.push(order.len());
        let reader = Self {
            order,
            offsets,
            pasef_frames,
            frame_reader,
        };
        Ok(reader)
    }

    pub(crate) fn iterate_over_pasef_frames(
        &self,
        index: usize,
    ) -> impl Iterator<Item = &SqlPasefFrameMsMs> {
        let start: usize = self.offsets[index];
        let end: usize = self.offsets[index + 1];
        self.order[start..end]
            .iter()
            .map(|&x| &self.pasef_frames[x])
    }

    fn _get(
        &self,
        index: usize,
    ) -> Result<RawSpectrum, DDARawSpectrumReaderError> {
        let mut collision_energy = 0.0;
        let mut isolation_mz = 0.0;
        let mut isolation_width = 0.0;
        let mut tof_indices: Vec<u32> = vec![];
        let mut intensities: Vec<u32> = vec![];
        for pasef_frame in self.iterate_over_pasef_frames(index) {
            collision_energy = pasef_frame.collision_energy;
            isolation_mz = pasef_frame.isolation_mz;
            isolation_width = pasef_frame.isolation_width;
            let frame_index: usize = pasef_frame.frame;
            let frame = self
                .frame_reader
                .get_frame(frame_index)
                .map_err(FrameReaderError::from)?;
            if frame.is_empty() {
                continue;
            }
            let scan_start: usize = pasef_frame.scan_start;
            let scan_end: usize = pasef_frame.scan_end;
            let scan_offsets = frame.ions().scan_offsets();
            let offset_start: usize = scan_offsets[scan_start];
            let offset_end: usize = scan_offsets[scan_end];
            let tof_selection = &frame.ions().tof_indices()
                [offset_start..offset_end]
                .iter()
                .map(|&x| x.into())
                .collect::<Vec<_>>();
            let intensity_selection = &frame.ions().intensities()
                [offset_start..offset_end]
                .iter()
                .map(|&x| x.into())
                .collect::<Vec<_>>();
            tof_indices.extend(tof_selection);
            intensities.extend(intensity_selection);
        }
        let (raw_tof_indices, raw_intensities) = group_and_sum(
            tof_indices,
            intensities.iter().map(|x| *x as u64).collect(),
        );
        let raw_spectrum = RawSpectrum {
            tof_indices: raw_tof_indices,
            intensities: raw_intensities,
            index,
            collision_energy,
            isolation_mz,
            isolation_width,
        };
        Ok(raw_spectrum)
    }
}

impl RawSpectrumReaderTrait for DDARawSpectrumReader {
    fn get(&self, index: usize) -> Result<RawSpectrum, RawSpectrumReaderError> {
        Ok(self._get(index)?)
    }

    fn len(&self) -> usize {
        self.offsets.len() - 1
    }
}

#[derive(Debug, thiserror::Error)]
pub(crate) enum DDARawSpectrumReaderError {
    #[error("{0}")]
    SqlReaderError(#[from] SqlReaderError),
    #[error("{0}")]
    FrameReaderError(#[from] FrameReaderError),
}
