use arcis_compiler::{
    traits::GetBit,
    utils::{
        field::BaseField,
        number::Number,
        packing::{DataSize, PackLocation},
        used_field::UsedField,
    },
    ArcisFloatValue,
    EvalValue,
};
use num_traits::cast::ToPrimitive;

#[doc(hidden)]
pub trait ArcisType {
    fn n_values() -> usize;
    fn gen_input(values: &mut Vec<EvalValue>) -> Self;
    fn from_values(values: &[EvalValue]) -> Self;
    fn handle_outputs(&self, outputs: &mut Vec<EvalValue>);

    fn from_mut_values(values: &mut &[EvalValue]) -> Self
    where
        Self: Sized,
    {
        let result = Self::from_values(values);
        *values = &values[Self::n_values()..];
        result
    }

    fn is_similar(&self, other: &Self) -> bool;
    fn n_bools() -> usize;
    fn from_bools(bools: &[bool]) -> Self;
    fn from_mut_bools(bools: &mut &[bool]) -> Self
    where
        Self: Sized,
    {
        let result = Self::from_bools(bools);
        *bools = &bools[Self::n_bools()..];
        result
    }
    fn data_size(acc: &mut Vec<DataSize>);
    fn pack(&self, locations: &mut &[PackLocation], containers: &mut [BaseField]);
    fn unpack(locations: &mut &[PackLocation], containers: &[BaseField]) -> Self;
    fn to_values(&self) -> Vec<BaseField> {
        let mut outputs = Vec::new();
        self.handle_outputs(&mut outputs);
        outputs
            .into_iter()
            .map(EvalValue::to_signed_number)
            .map(BaseField::from)
            .collect()
    }
}

macro_rules! impl_arcis_type_for_integer {
    ($t:ty, $to_method:ident) => {
        impl ArcisType for $t {
            fn n_values() -> usize {
                1
            }
            fn gen_input(values: &mut Vec<EvalValue>) -> Self {
                use rand::Rng;
                let mut rng = rand::thread_rng();
                let num: Self = rng.r#gen();
                let number: Number = num.into();
                values.push(EvalValue::arcis(number.into()));
                num
            }
            fn from_values(values: &[EvalValue]) -> Self {
                let num = values[0].to_signed_number();
                num.$to_method().unwrap_or_else(|| panic!("An overflow or underflow happened ({:?}) but was not detected by the rust code. Make sure that overflow-checks are enabled, even for libraries.", values[0]))
            }
            fn handle_outputs(&self, outputs: &mut Vec<EvalValue>) {
                let number: Number = self.into();
                outputs.push(EvalValue::arcis(number.into()));
            }
            fn is_similar(&self, other: &Self) -> bool {
                self == other
            }
            fn n_bools() -> usize {
                <$t>::BITS as usize
            }
            fn from_bools(bools: &[bool]) -> Self {
                let mut res = 0;
                for (i, bit) in bools.iter().take(Self::n_bools()).enumerate() {
                    res += (*bit as Self) << i;
                }
                res
            }
            fn data_size(acc: &mut Vec<DataSize>) {
                acc.push(DataSize::PowerOfTwo(Self::n_bools().ilog2() as u8))
            }
            fn pack(&self, locations: &mut &[PackLocation], containers: &mut [BaseField]) {
                let location = locations[0];
                *locations = &locations[1..];
                let val = self.abs_diff(Self::MIN); // useful for signed integers
                let val = Number::from(val);
                let val = BaseField::from(val);
                assert!(val < BaseField::power_of_two(Self::BITS as usize));
                containers[location.index] += val * BaseField::power_of_two(location.bit_offset as usize);
            }
            fn unpack(locations: &mut &[PackLocation], containers: &[BaseField]) -> Self {
                let location = locations[0];
                *locations = &locations[1..];
                let val = containers[location.index].unsigned_euclidean_division(BaseField::power_of_two(location.bit_offset as usize)).to_unsigned_number();
                ((val % Number::power_of_two(Self::BITS as usize)) + Number::from(Self::MIN)).$to_method().expect("This won't happen.")
            }
        }
    };
}

impl_arcis_type_for_integer!(u8, to_u8);
impl_arcis_type_for_integer!(u16, to_u16);
impl_arcis_type_for_integer!(u32, to_u32);
impl_arcis_type_for_integer!(u64, to_u64);
impl_arcis_type_for_integer!(u128, to_u128);
impl_arcis_type_for_integer!(usize, to_usize);
impl_arcis_type_for_integer!(i8, to_i8);
impl_arcis_type_for_integer!(i16, to_i16);
impl_arcis_type_for_integer!(i32, to_i32);
impl_arcis_type_for_integer!(i64, to_i64);
impl_arcis_type_for_integer!(i128, to_i128);
impl_arcis_type_for_integer!(isize, to_isize);

impl ArcisType for bool {
    fn n_values() -> usize {
        1
    }

    fn gen_input(values: &mut Vec<EvalValue>) -> Self {
        use rand::Rng;
        let mut rng = rand::thread_rng();
        let res = rng.r#gen::<bool>();
        values.push(EvalValue::Bit(res));
        res
    }

    fn from_values(values: &[EvalValue]) -> Self {
        values[0].to_signed_number().to_u8().unwrap() == 1
    }

    fn handle_outputs(&self, outputs: &mut Vec<EvalValue>) {
        outputs.push(EvalValue::Bit(*self));
    }

    fn is_similar(&self, other: &Self) -> bool {
        self == other
    }

    fn n_bools() -> usize {
        1
    }

    fn from_bools(bools: &[bool]) -> Self {
        bools[0]
    }

    fn data_size(acc: &mut Vec<DataSize>) {
        acc.push(DataSize::PowerOfTwo(0))
    }
    fn pack(&self, locations: &mut &[PackLocation], containers: &mut [BaseField]) {
        let location = locations[0];
        *locations = &locations[1..];
        let val = BaseField::from(*self);
        containers[location.index] += val * BaseField::power_of_two(location.bit_offset as usize);
    }

    fn unpack(locations: &mut &[PackLocation], containers: &[BaseField]) -> Self {
        let location = locations[0];
        *locations = &locations[1..];
        containers[location.index].get_bit(location.bit_offset as usize, false)
    }
}

impl<T: ArcisType, const N: usize> ArcisType for [T; N] {
    fn n_values() -> usize {
        N * T::n_values()
    }

    fn gen_input(values: &mut Vec<EvalValue>) -> Self {
        core::array::from_fn(|_| T::gen_input(values))
    }

    fn from_values(mut values: &[EvalValue]) -> Self {
        core::array::from_fn(|_| T::from_mut_values(&mut values))
    }

    fn handle_outputs(&self, outputs: &mut Vec<EvalValue>) {
        self.iter().for_each(|t| t.handle_outputs(outputs));
    }

    fn is_similar(&self, other: &Self) -> bool {
        for i in 0..N {
            if !self[i].is_similar(&other[i]) {
                return false;
            }
        }
        true
    }

    fn n_bools() -> usize {
        N * T::n_bools()
    }

    fn from_bools(mut bools: &[bool]) -> Self {
        core::array::from_fn(|_| T::from_mut_bools(&mut bools))
    }

    fn data_size(acc: &mut Vec<DataSize>) {
        for _ in 0..N {
            T::data_size(acc);
        }
    }

    fn pack(&self, locations: &mut &[PackLocation], containers: &mut [BaseField]) {
        for item in self {
            item.pack(locations, containers);
        }
    }

    fn unpack(locations: &mut &[PackLocation], containers: &[BaseField]) -> Self {
        core::array::from_fn(|_| T::unpack(locations, containers))
    }
}
impl<T: ArcisType> ArcisType for &[T] {
    fn n_values() -> usize {
        panic!("Cannot use slices as inputs.")
    }

    fn gen_input(_values: &mut Vec<EvalValue>) -> Self {
        panic!("Cannot use slices as inputs.")
    }

    fn from_values(_values: &[EvalValue]) -> Self {
        panic!("Cannot use slices as inputs.")
    }

    fn handle_outputs(&self, outputs: &mut Vec<EvalValue>) {
        self.iter().for_each(|t| t.handle_outputs(outputs));
    }

    fn is_similar(&self, other: &Self) -> bool {
        if self.len() != other.len() {
            return false;
        }
        for (s, o) in self.iter().zip(*other) {
            if !s.is_similar(o) {
                return false;
            }
        }
        true
    }

    fn n_bools() -> usize {
        panic!("Cannot generate slices.")
    }

    fn from_bools(_bools: &[bool]) -> Self {
        panic!("Cannot generate slices.")
    }

    fn data_size(_acc: &mut Vec<DataSize>) {
        panic!("Cannot size slices.")
    }
    fn pack(&self, _locations: &mut &[PackLocation], _containers: &mut [BaseField]) {
        panic!("Cannot pack slices.")
    }

    fn unpack(_locations: &mut &[PackLocation], _containers: &[BaseField]) -> Self {
        panic!("Cannot unpack into slices.")
    }
}
macro_rules! impl_tuples {
    ($($name:ident $idx:ident $other_idx:ident),*) => {
        impl<$($name),*> ArcisType for ($($name,)*)
        where
            $($name: ArcisType),*
        {
            fn n_values() -> usize {
                [$($name::n_values(),)*].into_iter().sum()
            }
            fn gen_input(values: &mut Vec<EvalValue>) -> Self {
                ($($name::gen_input(values),)*)
            }
            fn from_values(mut values: &[EvalValue]) -> Self {
                ($($name::from_mut_values(&mut values),)*)
            }
            fn handle_outputs(&self, outputs: &mut Vec<EvalValue>) {
                let ($($idx,)*) = self;
                $($idx.handle_outputs(outputs);)*
            }
            fn is_similar(&self, other: &Self) -> bool {
                let ($($idx,)*) = self;
                let ($($other_idx,)*) = other;
                $($idx.is_similar($other_idx))&&*
            }
            fn n_bools() -> usize {
                [$($name::n_bools(),)*].into_iter().sum()
            }
            fn from_bools(mut bools: &[bool]) -> Self {
                ($($name::from_mut_bools(&mut bools),)*)
            }
            fn data_size(acc: &mut Vec<DataSize>) {
                $($name::data_size(acc);)*
            }
            fn pack(&self, locations: &mut &[PackLocation], containers: &mut [BaseField]) {
                let ($($idx,)*) = self;
                $($idx.pack(locations, containers);)*
            }
            fn unpack(locations: &mut &[PackLocation], containers: &[BaseField]) -> Self {
                ($($name::unpack(locations, containers),)*)
            }
        }
    };
}

macro_rules! impl_all_tuples {
    ( $first_name:ident $first_idx:ident $first_other_idx:ident $(,$name:ident $idx:ident $other_idx:ident)*) => {
        impl_tuples!($first_name $first_idx $first_other_idx $(,$name $idx $other_idx)*);
        impl_all_tuples!($($name $idx $other_idx),*);
    };
    () => {
    }
}

// 12 just like the rust core library
impl_all_tuples!(A a other_a, B b other_b, C c other_c, D d other_d, E e other_e, F f other_f, G g other_g, H h other_h, I i other_i, J j other_j, K k other_k, L l other_l);

impl ArcisType for () {
    fn n_values() -> usize {
        0
    }

    fn gen_input(_values: &mut Vec<EvalValue>) -> Self {}

    fn from_values(_values: &[EvalValue]) -> Self {}

    fn handle_outputs(&self, _outputs: &mut Vec<EvalValue>) {}

    fn is_similar(&self, _other: &Self) -> bool {
        true
    }

    fn n_bools() -> usize {
        0
    }

    fn from_bools(_bools: &[bool]) -> Self {}

    fn data_size(_acc: &mut Vec<DataSize>) {}

    fn pack(&self, _locations: &mut &[PackLocation], _containers: &mut [BaseField]) {}

    fn unpack(_locations: &mut &[PackLocation], _containers: &[BaseField]) -> Self {}
}

impl<T: ArcisType> ArcisType for &T {
    fn n_values() -> usize {
        T::n_values()
    }

    fn gen_input(values: &mut Vec<EvalValue>) -> Self {
        <&mut T>::gen_input(values)
    }

    fn from_values(values: &[EvalValue]) -> Self {
        <&mut T>::from_values(values)
    }

    fn handle_outputs(&self, outputs: &mut Vec<EvalValue>) {
        T::handle_outputs(*self, outputs);
    }

    fn is_similar(&self, other: &Self) -> bool {
        (*self).is_similar(*other)
    }

    fn n_bools() -> usize {
        T::n_bools()
    }

    fn from_bools(bools: &[bool]) -> Self {
        <&mut T>::from_bools(bools)
    }

    fn data_size(acc: &mut Vec<DataSize>) {
        T::data_size(acc);
    }

    fn pack(&self, locations: &mut &[PackLocation], containers: &mut [BaseField]) {
        (**self).pack(locations, containers);
    }

    fn unpack(locations: &mut &[PackLocation], containers: &[BaseField]) -> Self {
        <&mut T>::unpack(locations, containers)
    }
}

impl<T: ArcisType> ArcisType for &mut T {
    fn n_values() -> usize {
        T::n_values()
    }

    fn gen_input(values: &mut Vec<EvalValue>) -> Self {
        let a = T::gen_input(values);
        let b = Box::new(a);
        Box::leak(b)
    }

    fn from_values(values: &[EvalValue]) -> Self {
        let a = T::from_values(values);
        let b = Box::new(a);
        Box::leak(b)
    }

    fn handle_outputs(&self, outputs: &mut Vec<EvalValue>) {
        T::handle_outputs(*self, outputs);
    }

    fn is_similar(&self, other: &Self) -> bool {
        (**self).is_similar(&**other)
    }

    fn n_bools() -> usize {
        T::n_bools()
    }

    fn from_bools(bools: &[bool]) -> Self {
        let a = T::from_bools(bools);
        let b = Box::new(a);
        Box::leak(b)
    }

    fn data_size(acc: &mut Vec<DataSize>) {
        T::data_size(acc)
    }

    fn pack(&self, locations: &mut &[PackLocation], containers: &mut [BaseField]) {
        (**self).pack(locations, containers);
    }

    fn unpack(locations: &mut &[PackLocation], containers: &[BaseField]) -> Self {
        let a = T::unpack(locations, containers);
        let b = Box::new(a);
        Box::leak(b)
    }
}

impl ArcisType for f64 {
    fn n_values() -> usize {
        1
    }

    fn gen_input(values: &mut Vec<EvalValue>) -> Self {
        let a = i128::gen_input(values);
        ArcisFloatValue::i128_to_f64(a)
    }

    fn from_values(values: &[EvalValue]) -> Self {
        ArcisFloatValue::number_to_f64(values[0].to_signed_number())
    }

    fn handle_outputs(&self, outputs: &mut Vec<EvalValue>) {
        outputs.push(EvalValue::Base(
            ArcisFloatValue::f64_to_number(*self).into(),
        ));
    }

    fn is_similar(&self, other: &Self) -> bool {
        // less than 0.5 difference or 12.5% difference
        self - other
            < 0.5f64
                .max(0.125f64 * self.abs())
                .max(0.125f64 * other.abs())
    }

    fn n_bools() -> usize {
        panic!("Cannot generate floats.")
    }

    fn from_bools(_bools: &[bool]) -> Self {
        panic!("Cannot generate floats.")
    }

    fn data_size(acc: &mut Vec<DataSize>) {
        acc.push(DataSize::Full);
    }

    fn pack(&self, locations: &mut &[PackLocation], containers: &mut [BaseField]) {
        BaseField::from(ArcisFloatValue::f64_to_number(*self)).pack(locations, containers);
    }

    fn unpack(locations: &mut &[PackLocation], containers: &[BaseField]) -> Self {
        let val = BaseField::unpack(locations, containers);
        ArcisFloatValue::number_to_f64(val.to_signed_number())
    }
}

impl ArcisType for f32 {
    fn n_values() -> usize {
        1
    }

    fn gen_input(values: &mut Vec<EvalValue>) -> Self {
        let a = i128::gen_input(values);
        ArcisFloatValue::i128_to_f32(a)
    }

    fn from_values(values: &[EvalValue]) -> Self {
        ArcisFloatValue::number_to_f64(values[0].to_signed_number()) as f32
    }

    fn handle_outputs(&self, outputs: &mut Vec<EvalValue>) {
        outputs.push(EvalValue::Base(
            ArcisFloatValue::f32_to_number(*self).into(),
        ));
    }

    fn is_similar(&self, other: &Self) -> bool {
        // less than 0.5 difference or 12.5% difference
        self - other
            < 0.5f32
                .max(0.125f32 * self.abs())
                .max(0.125f32 * other.abs())
    }

    fn n_bools() -> usize {
        panic!("Cannot generate floats.")
    }

    fn from_bools(_bools: &[bool]) -> Self {
        panic!("Cannot generate floats.")
    }

    fn data_size(acc: &mut Vec<DataSize>) {
        acc.push(DataSize::Full);
    }

    fn pack(&self, locations: &mut &[PackLocation], containers: &mut [BaseField]) {
        BaseField::from(ArcisFloatValue::f32_to_number(*self)).pack(locations, containers);
    }

    fn unpack(locations: &mut &[PackLocation], containers: &[BaseField]) -> Self {
        let val = BaseField::unpack(locations, containers);
        ArcisFloatValue::number_to_f64(val.to_signed_number()) as f32
    }
}

fn cell_with<T: ArcisType, R>(cell: &std::cell::Cell<T>, f: impl FnOnce(&T) -> R) -> R {
    let temp = {
        let mut temp = Vec::new();
        T::gen_input(&mut temp)
    };
    let a = cell.replace(temp);
    let res = f(&a);
    cell.replace(a);
    res
}

impl<T: ArcisType> ArcisType for std::cell::Cell<T> {
    fn n_values() -> usize {
        T::n_values()
    }

    fn gen_input(values: &mut Vec<EvalValue>) -> Self {
        let a = T::gen_input(values);
        std::cell::Cell::new(a)
    }

    fn from_values(values: &[EvalValue]) -> Self {
        let a = T::from_values(values);
        std::cell::Cell::new(a)
    }

    fn handle_outputs(&self, outputs: &mut Vec<EvalValue>) {
        cell_with(self, |a| a.handle_outputs(outputs));
    }

    fn is_similar(&self, other: &Self) -> bool {
        cell_with(self, |a| cell_with(other, |b| a.is_similar(b)))
    }

    fn n_bools() -> usize {
        T::n_bools()
    }

    fn from_bools(bools: &[bool]) -> Self {
        let a = T::from_bools(bools);
        std::cell::Cell::new(a)
    }

    fn data_size(acc: &mut Vec<DataSize>) {
        T::data_size(acc);
    }

    fn pack(&self, locations: &mut &[PackLocation], containers: &mut [BaseField]) {
        cell_with(self, |a| a.pack(locations, containers))
    }

    fn unpack(locations: &mut &[PackLocation], containers: &[BaseField]) -> Self {
        let a = T::unpack(locations, containers);
        std::cell::Cell::new(a)
    }
}

impl<T: ArcisType> ArcisType for Box<T> {
    fn n_values() -> usize {
        T::n_values()
    }

    fn gen_input(values: &mut Vec<EvalValue>) -> Self {
        let a = T::gen_input(values);
        Box::new(a)
    }

    fn from_values(values: &[EvalValue]) -> Self {
        let a = T::from_values(values);
        Box::new(a)
    }

    fn handle_outputs(&self, outputs: &mut Vec<EvalValue>) {
        (**self).handle_outputs(outputs);
    }

    fn is_similar(&self, other: &Self) -> bool {
        (**self).is_similar(other)
    }

    fn n_bools() -> usize {
        T::n_bools()
    }

    fn from_bools(bools: &[bool]) -> Self {
        let a = T::from_bools(bools);
        Box::new(a)
    }

    fn data_size(acc: &mut Vec<DataSize>) {
        T::data_size(acc);
    }

    fn pack(&self, locations: &mut &[PackLocation], containers: &mut [BaseField]) {
        (**self).pack(locations, containers);
    }

    fn unpack(locations: &mut &[PackLocation], containers: &[BaseField]) -> Self {
        let a = T::unpack(locations, containers);
        Box::new(a)
    }
}
impl<T: ArcisType> ArcisType for Box<[T]> {
    fn n_values() -> usize {
        panic!("Cannot use `Box<[T]>` as input.")
    }

    fn gen_input(_values: &mut Vec<EvalValue>) -> Self {
        panic!("Cannot use `Box<[T]>` as input.")
    }

    fn from_values(_values: &[EvalValue]) -> Self {
        panic!("Cannot use `Box<[T]>` as input.")
    }

    fn handle_outputs(&self, outputs: &mut Vec<EvalValue>) {
        for a in self {
            a.handle_outputs(outputs);
        }
    }

    fn is_similar(&self, other: &Self) -> bool {
        (&**self).is_similar(&&**other)
    }

    fn n_bools() -> usize {
        panic!("Cannot randomly generate `Box<[T]>`.")
    }

    fn from_bools(_bools: &[bool]) -> Self {
        panic!("Cannot randomly generate `Box<[T]>`.")
    }

    fn data_size(_acc: &mut Vec<DataSize>) {
        panic!("Cannot pack `Box<[T]>`.")
    }

    fn pack(&self, _locations: &mut &[PackLocation], _containers: &mut [BaseField]) {
        panic!("Cannot pack `Box<[T]>`.")
    }

    fn unpack(_locations: &mut &[PackLocation], _containers: &[BaseField]) -> Self {
        panic!("Cannot pack `Box<[T]>`.")
    }
}

impl<T: ArcisType> ArcisType for Option<T> {
    fn n_values() -> usize {
        panic!("Cannot use `Option<T>` as input nor output.")
    }

    fn gen_input(_values: &mut Vec<EvalValue>) -> Self {
        panic!("Cannot use `Option<T>` as input nor output.")
    }

    fn from_values(_values: &[EvalValue]) -> Self {
        panic!("Cannot use `Option<T>` as input nor output.")
    }

    fn handle_outputs(&self, _outputs: &mut Vec<EvalValue>) {
        panic!("Cannot use `Option<T>` as input nor output.")
    }

    fn is_similar(&self, _other: &Self) -> bool {
        panic!("Cannot use `Option<T>` as input nor output.")
    }

    fn n_bools() -> usize {
        panic!("Cannot use `Option<T>` as input nor output.")
    }

    fn from_bools(_bools: &[bool]) -> Self {
        panic!("Cannot use `Option<T>` as input nor output.")
    }

    fn data_size(_acc: &mut Vec<DataSize>) {
        panic!("Cannot use `Option<T>` as input nor output.")
    }

    fn pack(&self, _locations: &mut &[PackLocation], _containers: &mut [BaseField]) {
        panic!("Cannot use `Option<T>` as input nor output.")
    }

    fn unpack(_locations: &mut &[PackLocation], _containers: &[BaseField]) -> Self {
        panic!("Cannot use `Option<T>` as input nor output.")
    }
}
