use std::fmt::Debug;

use crate::types::Max;
use crate::types::Result;
use crate::util::check_argument;

#[repr(C)]
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct Sampler<P> {
    capacity: usize,
    weights: Vec<f32>,
    points: Vec<P>,
    store_attributes: bool,
    point_attributes: Vec<P>,
    current_size: usize,
    accepted_state: (f32, usize, usize),
}

#[repr(C)]
pub struct SamplerAcceptState {
    pub(crate) eviction_occurred: bool,
    pub(crate) point_index: usize,
    pub(crate) evicted_weight: f32,
    pub(crate) point_attribute: usize,
}

impl<P: Max + Copy + std::cmp::PartialEq> Sampler<P>
where
    P: TryFrom<usize>,
    usize: From<P>,
{
    pub fn new(capacity: usize, store_attributes: bool) -> Result<Self> {
        let attrib_vec: Vec<P> = if store_attributes {
            vec![P::MAX; capacity]
        } else {
            Vec::new()
        };
        Ok(Sampler {
            store_attributes,
            capacity,
            weights: vec![0.0; capacity],
            points: vec![P::MAX; capacity],
            point_attributes: attrib_vec,
            accepted_state: (0.0, usize::MAX, usize::MAX),
            current_size: 0,
        })
    }

    pub fn get_references(&self) -> &[P] {
        &self.points[0..self.current_size]
    }

    fn swap_down(&mut self, start_index: usize, validate: bool) -> Result<()> {
        let mut current: usize = start_index;
        while 2 * current + 1 < self.current_size {
            let mut max_index: usize = 2 * current + 1;
            if 2 * current + 2 < self.current_size
                && self.weights[2 * current + 2] > self.weights[max_index]
            {
                max_index = 2 * current + 2;
            }
            if self.weights[max_index] > self.weights[current] {
                check_argument(!validate, "the heap property is not satisfied")?;
                self.swap_weights(current, max_index)?;
                current = max_index;
            } else {
                break;
            }
        }
        Ok(())
    }

    pub fn reheap(&mut self, validate: bool) -> Result<()> {
        for i in ((self.current_size + 1) / 2)..=0 {
            self.swap_down(i, validate)?;
        }
        Ok(())
    }

    fn swap_weights(&mut self, a: usize, b: usize) -> Result<()> {
        check_argument(
            self.points[a] != P::MAX && self.points[b] != P::MAX,
            "incorrect state in sampler swap",
        )?;

        let tmp: P = self.points[a];
        self.points[a] = self.points[b];
        self.points[b] = tmp;

        let tmp_weight: f32 = self.weights[a];
        self.weights[a] = self.weights[b];
        self.weights[b] = tmp_weight;

        if self.store_attributes {
            let tmp_attrib: P = self.point_attributes[a];
            self.point_attributes[a] = self.point_attributes[b];
            self.point_attributes[b] = tmp_attrib;
        }
        Ok(())
    }

    pub fn add_point(&mut self, point_index: usize) -> Result<()>
    where
        <P as TryFrom<usize>>::Error: Debug,
    {
        if point_index != usize::MAX {
            check_argument(self.current_size < self.capacity.into(), "sampler full")?;
            check_argument(
                self.accepted_state.1 != usize::MAX,
                "this method should only be called after a successful call to accept_sample(long)",
            )?;

            self.weights[self.current_size] = self.accepted_state.0;
            self.points[self.current_size] = point_index.try_into().unwrap();
            // note, not self.accepted_state.1, even though we want that to not
            // P::MAX This corresponds to the change in the index value via
            // duplicates in the trees
            if self.store_attributes {
                self.point_attributes[self.current_size] = if self.accepted_state.2 != usize::MAX {
                    self.accepted_state.2.try_into().unwrap()
                } else {
                    P::MAX
                }
            };

            let mut current = self.current_size;
            self.current_size += 1;

            while current > 0 {
                let tmp = (current - 1) / 2;
                if self.weights[tmp] < self.weights[current] {
                    self.swap_weights(current, tmp)?;
                    current = tmp;
                } else {
                    break;
                }
            }
            // resetting the state
            self.accepted_state = (0.0, usize::MAX, usize::MAX);
        }
        Ok(())
    }

    pub fn accept_point(
        &mut self,
        initial: bool,
        weight: f32,
        point_index: usize,
        attribute: usize,
    ) -> Result<SamplerAcceptState> {
        let mut return_val = (true, usize::MAX, weight, usize::MAX);
        if initial || (weight < self.weights[0]) {
            self.accepted_state = (weight, point_index, attribute);
            if !initial {
                let partial = self.evict_max()?;
                return_val = (true, partial.0, partial.1, partial.2);
            }
        } else {
            return_val.0 = false;
        }
        Ok(SamplerAcceptState {
            eviction_occurred: return_val.0,
            point_index: return_val.1,
            evicted_weight: return_val.2,
            point_attribute: return_val.3,
        })
    }

    /**
     * evicts the maximum weight point from the sampler. can be used repeatedly to
     * change the size of the sampler and associated tree
     */

    pub fn evict_max(&mut self) -> Result<(usize, f32, usize)> {
        let evicted_attribute_index: usize = if self.store_attributes {
            self.point_attributes[0].into()
        } else {
            usize::MAX
        };

        let evicted_point = (
            self.points[0].into(),
            self.weights[0],
            evicted_attribute_index,
        );
        self.current_size -= 1;
        let current: usize = self.current_size.into();
        self.weights[0] = self.weights[current];
        self.points[0] = self.points[current];
        if self.store_attributes {
            self.point_attributes[0] = self.point_attributes[current];
        }
        self.swap_down(0, false)?;

        Ok(evicted_point)
    }

    pub fn is_full(&self) -> bool {
        self.current_size == self.capacity
    }

    pub fn get_fill_fraction(&self) -> f64 {
        if self.current_size == self.capacity {
            return 1.0;
        }
        let fill_fraction: f64 = self.current_size as f64 / self.capacity as f64;
        fill_fraction
    }

    pub fn get_size(&self) -> usize {
        (self.weights.len()) * std::mem::size_of::<f32>()
            + (self.points.len()) * std::mem::size_of::<P>()
            + std::mem::size_of::<Sampler<P>>()
    }
}
