use crate::errors::DataError;
use async_trait::async_trait;
use ndarray::{
    s, Array1, Array2, Array3, ArrayBase, DataMut, Dimension, Ix1, Ix2, OwnedRepr, RawDataClone,
};
use npyz::{npz::NpzArchive, Deserialize};
use num_traits::{AsPrimitive, FromPrimitive, Num};
use std::io::{Cursor, Seek};
use std::{self, f64::consts::PI, fmt::Debug, io::BufReader, path::Path};
use tokio::fs::File;
use tokio::io::{self, AsyncReadExt};

pub type Radar2d<T> = RadarData2d<T, OwnedRepr<T>>;

pub trait MultiDimensionData<T>
where
    T: Num + Clone + PartialEq + PartialOrd,
{
    fn map_by_fn<F>(&mut self, f: F)
    where
        F: FnMut(&mut T);
}

#[derive(Clone)]
pub struct RadarData2d<T, Raw, X = f64, Y = f64, I = Ix1>
where
    T: Num + Clone + PartialEq + PartialOrd,
    Raw: ndarray::Data<Elem = T> + Clone + ndarray::RawDataClone,
    X: Num,
    Y: Num,
{
    pub dim1: ArrayBase<OwnedRepr<X>, I>,
    pub dim2: ArrayBase<OwnedRepr<Y>, I>,
    pub data: ArrayBase<Raw, Ix2>,
    pub fill_value: T,
    pub coord_type: CoordType,
}

impl<T, Raw, X, Y, I> Debug for RadarData2d<T, Raw, X, Y, I>
where
    T: Num + Clone + PartialEq + PartialOrd + Debug,
    Raw: ndarray::Data<Elem = T> + Clone + ndarray::RawDataClone,
    X: Num + Debug,
    Y: Num + Debug,
    I: Dimension,
{
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("RadarData2d")
            .field("dim1", &self.dim1)
            .field("dim2", &self.dim2)
            .field("data", &self.data)
            .field("fill_value", &self.fill_value)
            .field("coord_type", &self.coord_type)
            .finish()
    }
}

impl<T: Num + Clone + PartialEq + PartialOrd> Radar2d<T> {
    pub fn load(path: impl AsRef<Path>, meth: impl DataLoader<T, Self>) -> Result<Self, DataError> {
        Ok(meth.load(path)?)
    }
}

pub struct RadarData3d<T, X = f64, Y = f64, Z = f64>
where
    T: Num,
    X: Num,
    Y: Num,
    Z: Num,
{
    pub dim1: Array1<X>,
    pub dim2: Array1<Y>,
    pub dim3: Array1<Z>,
    pub data: Array3<T>,
}

impl<T, Raw> RadarData2d<T, Raw>
where
    T: Num + AsPrimitive<f64> + FromPrimitive + Clone + Debug + PartialOrd + PartialEq,
    Raw: ndarray::Data<Elem = T> + Clone + ndarray::RawDataClone,
{
    pub fn downsample(&self, output_shape: (usize, usize), meth: DownSampleMeth) -> Radar2d<T> {
        let shape = self.data.shape();
        assert!(shape[0] > output_shape.0 && shape[1] > output_shape.1);

        if shape[0] == output_shape.0 && shape[1] == output_shape.1 {
            return Radar2d {
                dim1: self.dim1.to_owned(),
                dim2: self.dim2.to_owned(),
                data: self.data.to_owned(),
                fill_value: self.fill_value.clone(),
                coord_type: self.coord_type.clone(),
            };
        }

        let scale_x = shape[1] as f64 / output_shape.1 as f64;
        let scale_y = shape[0] as f64 / output_shape.0 as f64;

        let mut output: Array2<T> = Array2::zeros(output_shape);

        let dim1 = Array1::linspace(
            *self.dim1.first().unwrap(),
            *self.dim1.last().unwrap(),
            output_shape.1,
        );

        let dim2 = Array1::linspace(
            *self.dim2.first().unwrap(),
            *self.dim2.last().unwrap(),
            output_shape.0,
        );

        output.iter_mut().enumerate().for_each(|(s, vv)| {
            let ri = s / output_shape.1;
            let ci = s % output_shape.1;
            let src_yf0 = scale_y * ri as f64;
            let src_yf1 = src_yf0 + scale_y;
            let src_y0 = src_yf0.floor() as usize;
            let src_y1 = src_yf1.floor() as usize;
            let src_xf0 = scale_x * ci as f64;
            let src_xf1 = src_xf0 + scale_x;
            let src_x0 = src_xf0.floor() as usize;
            let src_x1 = src_xf1.floor() as usize;

            match meth {
                DownSampleMeth::MEAN => {}
                DownSampleMeth::VAR | DownSampleMeth::STD => self.var(
                    src_yf0, src_yf1, src_xf0, src_xf1, src_y0, src_y1, src_x0, src_x1, vv,
                ),
            }
        });

        if let DownSampleMeth::STD = meth {
            output = output.mapv(|x| T::from_f64(f64::sqrt(x.as_())).unwrap());
        }

        return Radar2d {
            dim1,
            dim2,
            fill_value: self.fill_value.clone(),
            data: output,
            coord_type: self.coord_type.clone(),
        };
    }

    fn var(
        &self,
        src_yf0: f64,
        src_yf1: f64,
        src_xf0: f64,
        src_xf1: f64,
        src_y0: usize,
        mut src_y1: usize,
        src_x0: usize,
        mut src_x1: usize,
        row: &mut T,
    ) {
        let wy0 = 1f64 - (src_yf0 - src_y0 as f64);
        let mut wy1 = src_yf1 - src_y1 as f64;
        let wx0 = 1f64 - (src_xf0 - src_x0 as f64);
        let mut wx1 = src_xf1 - src_x1 as f64;

        if wy1 < f64::EPSILON {
            wy1 = 1f64;
            if src_y1 > src_y0 {
                src_y1 -= 1;
            }
        }

        if wx1 < f64::EPSILON {
            wx1 = 1f64;
            if src_x1 > src_x0 {
                src_x1 -= 1;
            }
        }

        let mut v_sum: f64 = 0.0;
        let mut w_sum = 0f64;
        let mut wv_sum = 0f64;
        let mut wvv_sum = 0f64;

        for src_y in src_y0..=src_y1 {
            let wy = if src_y == src_y0 {
                wy0
            } else {
                if src_y == src_y1 {
                    wy1
                } else {
                    1.0
                }
            };
            for src_x in src_x0..=src_x1 {
                let wx = if src_x == src_x0 {
                    wx0
                } else {
                    if src_x == src_x1 {
                        wx1
                    } else {
                        1.0
                    }
                };

                let v = self.data[[src_y, src_x]].as_();
                let w = wx * wy;

                v_sum += v;
                w_sum += w;
                wv_sum += w * v;
                wvv_sum += w * v * v;
            }

            if w_sum < f64::EPSILON {
                return;
            }
            let v = ((wvv_sum * w_sum - wv_sum * wv_sum) / w_sum / w_sum).clamp(-125f64, 124f64);
            *row = T::from_f64(v).unwrap();
        }
    }

    fn split(&self) -> [RadarData2d<T, ndarray::ViewRepr<&T>>; 4] {
        let middle_dim1 = self.dim1.len() / 2;
        let middle_dim2 = self.dim2.len() / 2;

        let lt = self.data.slice(s![..middle_dim1, ..middle_dim2]);
        let rt = self.data.slice(s![middle_dim1.., ..middle_dim2]);
        let lb = self.data.slice(s![..middle_dim1, middle_dim2..]);
        let rb = self.data.slice(s![middle_dim1.., middle_dim2..]);

        return [
            RadarData2d {
                dim1: self.dim1.slice(s![..middle_dim1]).to_owned(),
                dim2: self.dim2.slice(s![..middle_dim2]).to_owned(),
                data: lt,
                fill_value: self.fill_value.clone(),
                coord_type: self.coord_type,
            },
            RadarData2d {
                dim1: self.dim1.slice(s![middle_dim1..]).to_owned(),
                dim2: self.dim2.slice(s![..middle_dim2]).to_owned(),
                data: rt,
                fill_value: self.fill_value.clone(),
                coord_type: self.coord_type,
            },
            RadarData2d {
                dim1: self.dim1.slice(s![..middle_dim1]).to_owned(),
                dim2: self.dim2.slice(s![middle_dim2..]).to_owned(),
                data: lb,
                fill_value: self.fill_value.clone(),
                coord_type: self.coord_type,
            },
            RadarData2d {
                dim1: self.dim1.slice(s![middle_dim1..]).to_owned(),
                dim2: self.dim2.slice(s![middle_dim2..]).to_owned(),
                data: rb,
                fill_value: self.fill_value.clone(),
                coord_type: self.coord_type,
            },
        ];
    }
}

impl<T, Raw> MultiDimensionData<T> for RadarData2d<T, Raw>
where
    T: Num + Clone,
    T: PartialEq + PartialOrd,
    Raw: ndarray::Data<Elem = T> + Clone + RawDataClone + DataMut,
{
    fn map_by_fn<F>(&mut self, f: F)
    where
        F: FnMut(&mut T),
    {
        self.data.map_inplace(f);
    }
}

pub trait DataLoader<V, T>
where
    V: Num + Clone + PartialEq + PartialOrd,
    T: MultiDimensionData<V>,
{
    fn load<P: AsRef<Path>>(&self, path: P) -> Result<T, DataError>;
}

#[async_trait]
pub trait AsyncDataLoader<V, T>
where
    V: Num + Clone + PartialEq + PartialOrd,
    T: MultiDimensionData<V> + Send + Sync,
{
    async fn load<P: AsRef<Path> + Send>(&self, path: P) -> Result<T, DataError>;
}

pub struct Npz;
impl Npz {
    #[inline]
    fn load_1d<T: Num + Deserialize, R: std::io::Read + Seek>(
        &self,
        data: &mut NpzArchive<R>,
        name: &str,
    ) -> Result<Array1<T>, DataError> {
        // let mut data = npyz::npz::NpzArchive::open(path)?;
        let a = data.by_name(name)?.unwrap();
        let b: Vec<T> = a.into_vec().unwrap();

        Ok(Array1::from_shape_vec(b.len(), b).unwrap())
    }
    #[inline]
    fn load_2d<T: Num + Deserialize, R: std::io::Read + Seek>(
        &self,
        // path: &Path,
        data: &mut NpzArchive<R>,
        name: &str,
    ) -> Result<Array2<T>, DataError> {
        // let mut data = npyz::npz::NpzArchive::open(path)?;
        let a = data.by_name(name)?.unwrap();

        let shape = a.shape().to_vec();
        let b: Vec<T> = a.into_vec().unwrap();

        Ok(Array2::from_shape_vec((shape[0] as usize, shape[1] as usize), b).unwrap())
    }
    fn load_3d(&self, data: &mut NpzArchive<BufReader<std::fs::File>>) {}
}

impl<T> DataLoader<T, RadarData2d<T, OwnedRepr<T>>> for Npz
where
    T: Num + Clone + Deserialize + FromPrimitive,
    T: PartialEq + PartialOrd,
{
    fn load<P: AsRef<Path>>(&self, path: P) -> Result<RadarData2d<T, OwnedRepr<T>>, DataError> {
        let mut data: NpzArchive<BufReader<std::fs::File>> = npyz::npz::NpzArchive::open(path)?;
        let dim1 = self.load_1d::<f64, BufReader<std::fs::File>>(&mut data, "lon")?;
        let dim2 = self.load_1d::<f64, BufReader<std::fs::File>>(&mut data, "lat")?;
        let value = self.load_2d::<T, BufReader<std::fs::File>>(&mut data, "value")?;

        Ok(RadarData2d {
            dim1: dim1,
            dim2: dim2,
            fill_value: T::from_f64(-125.0).unwrap(),
            data: value,
            coord_type: CoordType::LatLon,
        })
    }
}

impl<T> AsyncDataLoader<T, RadarData2d<T, OwnedRepr<T>>> for Npz
where
    T: Num + Clone + Deserialize + FromPrimitive,
    T: PartialEq + PartialOrd + Send + Sync,
{
    fn load<'life0, 'async_trait, P>(
        &'life0 self,
        path: P,
    ) -> core::pin::Pin<
        Box<
            dyn core::future::Future<Output = Result<RadarData2d<T, OwnedRepr<T>>, DataError>>
                + core::marker::Send
                + 'async_trait,
        >,
    >
    where
        P: 'async_trait + AsRef<Path> + Send,
        'life0: 'async_trait,
        Self: 'async_trait,
    {
        Box::pin(async {
            let mut f = File::open(path).await?;
            let mut buffer = Vec::new();
            f.read_to_end(&mut buffer).await?;

            let mut data = npyz::npz::NpzArchive::new(Cursor::new(buffer))?;
            let dim1 = self.load_1d::<f64, Cursor<Vec<u8>>>(&mut data, "lon")?;
            let dim2 = self.load_1d::<f64, Cursor<Vec<u8>>>(&mut data, "lat")?;
            let value = self.load_2d::<T, Cursor<Vec<u8>>>(&mut data, "value")?;

            Ok(RadarData2d {
                dim1: dim1,
                dim2: dim2,
                fill_value: T::from_f64(-125.0).unwrap(),
                data: value,
                coord_type: CoordType::LatLon,
            })
        })
    }
}

// fn resample(
//     &self,
//     width_rate: f64,
//     height_rate: f64,
//     filter_len: f64,
// ) -> Result<RadarData2d<T, OwnedRepr<T>>, DataError> {
//     let width_rate = width_rate.min(1.0);
//     let height_rate = height_rate.min(1.0);
//     match self.coord_type {
//         CoordType::Polar => Err(DataError::FormatError),
//         CoordType::LatLon => {
//             let width_filtered: ArrayBase<ndarray::OwnedRepr<T>, Ix2> =
//                 Self::_resample(&self.data, width_rate, filter_len);

//             let result: ArrayBase<OwnedRepr<T>, Ix2> =
//                 Self::_resample(&width_filtered.t(), height_rate, filter_len)
//                     .t()
//                     .to_owned();

//             let new_dim1: ArrayBase<OwnedRepr<f64>, Ix1> = Self::_resample(
//                 &Array2::from_shape_vec((1, self.dim1.len()), self.dim1.to_vec()).unwrap(),
//                 width_rate,
//                 filter_len,
//             )
//             .slice(s![0, ..])
//             .to_owned();

//             let new_dim2: ArrayBase<OwnedRepr<f64>, Ix1> = Self::_resample(
//                 &Array2::from_shape_vec((1, self.dim2.len()), self.dim2.to_vec()).unwrap(),
//                 height_rate,
//                 filter_len,
//             )
//             .slice(s![0, ..])
//             .to_owned();

//             Ok(RadarData2d {
//                 dim1: new_dim1,
//                 dim2: new_dim2,
//                 data: result,
//                 coord_type: self.coord_type.to_owned(),
//             })
//         }
//     }
// }

// fn _resample<'a, V, R: ndarray::Data<Elem = V>>(
//         data: &'a ArrayBase<R, Ix2>,
//         rate: f64,
//         filter_len: f64,
//     ) -> Array2<V>
//     where
//         V: Num + Clone + AsPrimitive<f64> + FromPrimitive,
//     {
//         let ori_width = data.ncols();
//         let ori_height = data.nrows();

//         let new_width = (ori_width as f64 * rate).ceil() as usize;
//         let mut result: Array2<V> = Array2::zeros((ori_height, new_width));
//         (0..ori_height).into_iter().for_each(|height| {
//             for width in 0..new_width {
//                 let center_x = (width as f64 + 0.5) / new_width as f64 * ori_width as f64;
//                 let filter_start = center_x - filter_len / 2.0;
//                 let start_idx = (filter_start - 0.5).ceil() as usize;

//                 let mut value_sum = 0.0;
//                 let mut filter_sum = 0.0;

//                 for i in 0..filter_len as usize {
//                     let input_x = start_idx + i;
//                     let weight = windowed_sinc(
//                         (input_x as f64 + 0.5 - center_x) * rate,
//                         (input_x as f64 + 0.5 - filter_start) / filter_len,
//                     );
//                     value_sum += weight * data[[height, input_x.clamp(0, ori_width - 1)]].as_();
//                     filter_sum += weight;
//                 }

//                 result[[height, width]] = V::from_f64(value_sum / filter_sum).unwrap();
//             }
//         });
//         result
//     }

// pub struct LevelData<T: Num + Clone + PartialEq + PartialOrd>(
//     pub Quadtree<i64, RadarData2d<T, OwnedRepr<T>>>,
// );

// impl<T> LevelData<T>
// where
//     T: Num + Clone + AsPrimitive<f64> + FromPrimitive + Debug + PartialEq + PartialOrd,
// {
//     fn value(
//         level_data: Vec<RadarData2d<T, OwnedRepr<T>>>,
//         level_num: usize,
//     ) -> Vec<RadarData2d<T, OwnedRepr<T>>> {
//         if level_num == 0 {
//             return level_data;
//         }

//         let mut result: Vec<RadarData2d<T, OwnedRepr<T>>> =
//             Vec::with_capacity(level_data.len() * 4);

//         let results = level_data
//             .iter()
//             .flat_map(|v| v.split().into_iter().map(|x| x.value_to_owned()));

//         result.extend(results);

//         return Self::value(result, level_num - 1);
//     }

//     fn new(data: &RadarData2d<T, OwnedRepr<T>>, level: usize, rate: f64) -> Self {
//         // let rate = 1.0 / level as f64;
//         let resampled = data.resample(rate, rate, 2.0).unwrap();
//         let blocks = Self::value(vec![resampled], level);
//         let mut tree: Quadtree<i64, RadarData2d<T, OwnedRepr<T>>> =
//             quadtree_rs::Quadtree::new(level);

//         blocks.into_iter().for_each(|block| {
//             tree.insert(
//                 AreaBuilder::default()
//                     .anchor(quadtree_rs::point::Point {
//                         x: *block.dim1.first().unwrap() as i64,
//                         y: *block.dim2.first().unwrap() as i64,
//                     })
//                     .dimensions((block.dim1.len() as i64, block.dim2.len() as i64))
//                     .build()
//                     .unwrap(),
//                 block,
//             );
//         });

//         Self(tree)
//     }
// }

// pub fn levels<T>(data: Radar2d<T>, levels: usize) -> Vec<LevelData<T>>
// where
//     T: Num + Clone + AsPrimitive<f64> + FromPrimitive + Debug,
//     T: PartialEq + PartialOrd,
// {
//     let numerator = 1.0 / levels as f64;
//     (0..levels)
//         .into_iter()
//         .map(|level| LevelData::new(&data, level + 1, 1.0 - level as f64 * numerator))
//         .collect()
// }

#[inline]
fn windowed_sinc(x: f64, y: f64) -> f64 {
    let x = x * PI;
    let sinc = if x != 0.0 { f64::sin(x) / x } else { 1.0 };
    let window = if 0f64 <= y && y <= 1.0 {
        1.0 - (y - 0.5).abs() * 2.0
    } else {
        0f64
    };
    sinc * window
}

pub enum DownSampleMeth {
    STD,
    MEAN,
    VAR,
}

#[derive(Clone, Copy, Debug)]
pub enum CoordType {
    Polar,
    LatLon,
}
