use arrow::array::Float32Array;
use arrow::record_batch::RecordBatch;
use nalgebra::{vector, Vector3, Vector4};
use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
use std::ops::{Deref, DerefMut, Index, IndexMut};
use std::sync::Arc;
use std::{fmt::Display, fs::File};

#[cfg(feature = "rayon")]
use rayon::prelude::*;

use crate::{utils::vectors::FourMomentum, Float};

/// An event that can be used to test the implementation of an
/// [`Amplitude`](crate::amplitudes::Amplitude). This particular event contains the reaction
/// $`\gamma p \to K_S^0 K_S^0 p`$ with a polarized photon beam.
pub fn test_event() -> Event {
    use crate::utils::vectors::*;
    Event {
        p4s: vec![
            Vector4::from_momentum(&vector![0.0, 0.0, 8.747], 0.0), // beam
            Vector4::from_momentum(&vector![0.119, 0.374, 0.222], 1.007), // "proton"
            Vector4::from_momentum(&vector![-0.112, 0.293, 3.081], 0.498), // "kaon"
            Vector4::from_momentum(&vector![-0.007, -0.667, 5.446], 0.498), // "kaon"
        ],
        eps: vec![vector![0.385, 0.022, 0.000]],
        weight: 0.48,
    }
}

/// A single event in a [`Dataset`] containing all the relevant particle information.
#[derive(Debug, Clone)]
pub struct Event {
    /// A list of four-momenta for each particle.
    pub p4s: Vec<Vector4<Float>>,
    /// A list of polarization vectors for each particle.
    pub eps: Vec<Vector3<Float>>,
    /// The weight given to the event.
    pub weight: Float,
}

impl Display for Event {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        writeln!(f, "Event:")?;
        writeln!(f, "  p4s:")?;
        for p4 in &self.p4s {
            writeln!(f, "    {}", p4.to_p4_string())?;
        }
        writeln!(f, "  eps:")?;
        for eps_vec in &self.eps {
            writeln!(
                f,
                "    [{}]",
                eps_vec
                    .iter()
                    .map(|v| v.to_string())
                    .collect::<Vec<String>>()
                    .join(", ")
            )?;
        }
        writeln!(f, "  weight:")?;
        writeln!(f, "    {}", self.weight)?;
        Ok(())
    }
}

impl Event {
    /// Return a four-momentum from the sum of four-momenta at the given indices in the [`Event`].
    pub fn get_p4_sum(&self, indices: &[usize]) -> Vector4<Float> {
        indices.iter().map(|i| self.p4s[*i]).sum::<Vector4<Float>>()
    }
}

/// A collection of [`Event`]s.
#[derive(Debug, Clone)]
pub struct Dataset {
    pub(crate) events: Vec<Event>,
}

impl Index<usize> for Dataset {
    type Output = Event;

    fn index(&self, index: usize) -> &Self::Output {
        &self.events[index]
    }
}

impl IndexMut<usize> for Dataset {
    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
        &mut self.events[index]
    }
}

impl Deref for Dataset {
    type Target = Vec<Event>;

    fn deref(&self) -> &Self::Target {
        &self.events
    }
}

impl DerefMut for Dataset {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.events
    }
}

impl Dataset {
    /// The number of [`Event`]s in the [`Dataset`].
    pub fn len(&self) -> usize {
        self.events.len()
    }

    /// Checks whether or not the [`Dataset`] is empty.
    pub fn is_empty(&self) -> bool {
        self.events.is_empty()
    }

    /// Produces an iterator over the [`Event`]s in the [`Dataset`].
    pub fn iter(&self) -> std::slice::Iter<'_, Event> {
        self.events.iter()
    }

    /// Produces an parallelized iterator over the [`Event`]s in the [`Dataset`].
    #[cfg(feature = "rayon")]
    pub fn par_iter(&self) -> rayon::slice::Iter<'_, Event> {
        self.events.par_iter()
    }

    /// Extract a list of weights over each [`Event`] in the [`Dataset`].
    #[cfg(not(feature = "rayon"))]
    pub fn weights(&self) -> Vec<Float> {
        self.iter().map(|e| e.weight).collect()
    }

    /// Extract a list of weights over each [`Event`] in the [`Dataset`].
    #[cfg(feature = "rayon")]
    pub fn weights(&self) -> Vec<Float> {
        self.par_iter().map(|e| e.weight).collect()
    }

    /// Returns the sum of the weights for each [`Event`] in the [`Dataset`].
    #[cfg(not(feature = "rayon"))]
    pub fn weighted_len(&self) -> Float {
        self.iter().map(|e| e.weight).sum()
    }

    /// Returns the sum of the weights for each [`Event`] in the [`Dataset`].
    #[cfg(feature = "rayon")]
    pub fn weighted_len(&self) -> Float {
        self.par_iter().map(|e| e.weight).sum()
    }
}

/// Open a Parquet file and read the data into a [`Dataset`].
#[cfg(feature = "rayon")]
pub fn open(file_path: &str) -> Result<Arc<Dataset>, Box<dyn std::error::Error>> {
    let file = File::open(file_path)?;
    let builder = ParquetRecordBatchReaderBuilder::try_new(file)?;
    let reader = builder.build()?;
    let batches: Vec<RecordBatch> = reader.collect::<Result<Vec<_>, _>>()?;

    // First pass: Collect events without setting the index
    let events: Vec<Event> = batches
        .into_par_iter()
        .flat_map(|batch| {
            let num_rows = batch.num_rows();
            let mut local_events = Vec::with_capacity(num_rows);

            // Process each row in the batch
            for row in 0..num_rows {
                let mut p4s = Vec::new();
                let mut eps = Vec::new();

                let p4_count = batch
                    .schema()
                    .fields()
                    .iter()
                    .filter(|field| field.name().starts_with("p4_"))
                    .count()
                    / 4;
                let eps_count = batch
                    .schema()
                    .fields()
                    .iter()
                    .filter(|field| field.name().starts_with("eps_"))
                    .count()
                    / 3;

                for i in 0..p4_count {
                    let e = batch
                        .column_by_name(&format!("p4_{}_E", i))
                        .unwrap()
                        .as_any()
                        .downcast_ref::<Float32Array>()
                        .unwrap()
                        .value(row) as Float;
                    let px = batch
                        .column_by_name(&format!("p4_{}_Px", i))
                        .unwrap()
                        .as_any()
                        .downcast_ref::<Float32Array>()
                        .unwrap()
                        .value(row) as Float;
                    let py = batch
                        .column_by_name(&format!("p4_{}_Py", i))
                        .unwrap()
                        .as_any()
                        .downcast_ref::<Float32Array>()
                        .unwrap()
                        .value(row) as Float;
                    let pz = batch
                        .column_by_name(&format!("p4_{}_Pz", i))
                        .unwrap()
                        .as_any()
                        .downcast_ref::<Float32Array>()
                        .unwrap()
                        .value(row) as Float;
                    p4s.push(Vector4::new(e, px, py, pz));
                }

                // TODO: insert empty vectors if not provided
                for i in 0..eps_count {
                    let x = batch
                        .column_by_name(&format!("eps_{}_x", i))
                        .unwrap()
                        .as_any()
                        .downcast_ref::<Float32Array>()
                        .unwrap()
                        .value(row) as Float;
                    let y = batch
                        .column_by_name(&format!("eps_{}_y", i))
                        .unwrap()
                        .as_any()
                        .downcast_ref::<Float32Array>()
                        .unwrap()
                        .value(row) as Float;
                    let z = batch
                        .column_by_name(&format!("eps_{}_z", i))
                        .unwrap()
                        .as_any()
                        .downcast_ref::<Float32Array>()
                        .unwrap()
                        .value(row) as Float;
                    eps.push(Vector3::new(x, y, z));
                }

                let weight = batch
                    .column(19)
                    .as_any()
                    .downcast_ref::<Float32Array>()
                    .unwrap()
                    .value(row) as Float;

                local_events.push(Event { p4s, eps, weight });
            }
            local_events
        })
        .collect();
    Ok(Arc::new(Dataset { events }))
}

/// Open a Parquet file and read the data into a [`Dataset`].
#[cfg(not(feature = "rayon"))]
pub fn open(file_path: &str) -> Result<Arc<Dataset>, Box<dyn std::error::Error>> {
    let file = File::open(file_path)?;
    let builder = ParquetRecordBatchReaderBuilder::try_new(file)?;
    let reader = builder.build()?;
    let batches: Vec<RecordBatch> = reader.collect::<Result<Vec<_>, _>>()?;

    // First pass: Collect events without setting the index
    let events: Vec<Event> = batches
        .into_iter()
        .flat_map(|batch| {
            let num_rows = batch.num_rows();
            let mut local_events = Vec::with_capacity(num_rows);

            // Process each row in the batch
            for row in 0..num_rows {
                let mut p4s = Vec::new();
                let mut eps = Vec::new();

                let p4_count = batch
                    .schema()
                    .fields()
                    .iter()
                    .filter(|field| field.name().starts_with("p4_"))
                    .count()
                    / 4;
                let eps_count = batch
                    .schema()
                    .fields()
                    .iter()
                    .filter(|field| field.name().starts_with("eps_"))
                    .count()
                    / 3;

                for i in 0..p4_count {
                    let e = batch
                        .column_by_name(&format!("p4_{}_E", i))
                        .unwrap()
                        .as_any()
                        .downcast_ref::<Float32Array>()
                        .unwrap()
                        .value(row) as Float;
                    let px = batch
                        .column_by_name(&format!("p4_{}_Px", i))
                        .unwrap()
                        .as_any()
                        .downcast_ref::<Float32Array>()
                        .unwrap()
                        .value(row) as Float;
                    let py = batch
                        .column_by_name(&format!("p4_{}_Py", i))
                        .unwrap()
                        .as_any()
                        .downcast_ref::<Float32Array>()
                        .unwrap()
                        .value(row) as Float;
                    let pz = batch
                        .column_by_name(&format!("p4_{}_Pz", i))
                        .unwrap()
                        .as_any()
                        .downcast_ref::<Float32Array>()
                        .unwrap()
                        .value(row) as Float;
                    p4s.push(Vector4::new(e, px, py, pz));
                }

                for i in 0..eps_count {
                    let x = batch
                        .column_by_name(&format!("eps_{}_x", i))
                        .unwrap()
                        .as_any()
                        .downcast_ref::<Float32Array>()
                        .unwrap()
                        .value(row) as Float;
                    let y = batch
                        .column_by_name(&format!("eps_{}_y", i))
                        .unwrap()
                        .as_any()
                        .downcast_ref::<Float32Array>()
                        .unwrap()
                        .value(row) as Float;
                    let z = batch
                        .column_by_name(&format!("eps_{}_z", i))
                        .unwrap()
                        .as_any()
                        .downcast_ref::<Float32Array>()
                        .unwrap()
                        .value(row) as Float;
                    eps.push(Vector3::new(x, y, z));
                }

                let weight = batch
                    .column(19)
                    .as_any()
                    .downcast_ref::<Float32Array>()
                    .unwrap()
                    .value(row) as Float;

                local_events.push(Event { p4s, eps, weight });
            }
            local_events
        })
        .collect();
    Ok(Arc::new(Dataset { events }))
}
