use super::{CatEncs, CatTriples, CatType, Cats};
use crate::cats::LockedCats;
use crate::solution_mapping::{BaseCatState, EagerSolutionMappings};
use crate::{
    BaseRDFNodeType, LANG_STRING_LANG_FIELD, LANG_STRING_VALUE_FIELD, OBJECT_COL_NAME,
    SUBJECT_COL_NAME,
};
use nohash_hasher::NoHashHasher;
use polars::datatypes::{AnyValue, PlSmallStr};
use polars::error::{ErrString, PolarsError, PolarsResult};
use polars::frame::column::ScalarColumn;
use polars::frame::DataFrame;
use polars::prelude::{as_struct, col, Column, DataType, IntoColumn, IntoLazy, LazyFrame, Series};
use rayon::iter::{IntoParallelIterator, ParallelIterator};
use std::collections::HashMap;
use std::hash::BuildHasherDefault;
use std::sync::Arc;

#[derive(Debug, Clone)]
pub struct CatReEnc {
    pub cat_map: Arc<HashMap<u32, u32, BuildHasherDefault<NoHashHasher<u32>>>>,
}

impl CatReEnc {
    pub fn re_encode(
        self,
        mut lf: LazyFrame,
        c: &str,
        is_multi: bool,
        bt: BaseRDFNodeType,
        forget_others: bool,
    ) -> LazyFrame {
        if bt.is_lang_string() {
            let move_bt1 = bt.clone();
            let move_bt2 = bt.clone();
            let self_clone = self.clone();
            lf = lf.with_column(
                col(c)
                    .struct_()
                    .with_fields(vec![
                        col(c)
                            .struct_()
                            .field_by_name(LANG_STRING_VALUE_FIELD)
                            .map(
                                move |x| {
                                    self_clone.clone().re_encode_column(
                                        x,
                                        move_bt1.clone(),
                                        forget_others,
                                    )
                                },
                                |_, f| Ok(f.clone()),
                            )
                            .alias(LANG_STRING_VALUE_FIELD),
                        col(c)
                            .struct_()
                            .field_by_name(LANG_STRING_LANG_FIELD)
                            .map(
                                move |x| {
                                    self.clone().re_encode_column(
                                        x,
                                        move_bt2.clone(),
                                        forget_others,
                                    )
                                },
                                |_, f| Ok(f.clone()),
                            )
                            .alias(LANG_STRING_LANG_FIELD),
                    ])
                    .alias(c),
            );
        } else if is_multi {
            let move_bt = bt.clone();
            lf = lf.with_column(
                col(c)
                    .struct_()
                    .with_fields(vec![col(c)
                        .struct_()
                        .field_by_name(&bt.field_col_name())
                        .map(
                            move |x| {
                                self.clone()
                                    .re_encode_column(x, move_bt.clone(), forget_others)
                            },
                            |_, f| Ok(f.clone()),
                        )
                        .alias(&bt.field_col_name())])
                    .alias(c),
            );
        } else {
            lf = lf.with_column(
                col(c)
                    .map(
                        move |x| self.clone().re_encode_column(x, bt.clone(), forget_others),
                        |_, f| Ok(f.clone()),
                    )
                    .alias(c),
            );
        }
        lf
    }

    pub fn re_encode_column(
        self,
        c: Column,
        t: BaseRDFNodeType,
        forget_others: bool,
    ) -> PolarsResult<Column> {
        if t.is_lang_string() {
            let value_col = c
                .struct_()
                .unwrap()
                .field_by_name(LANG_STRING_VALUE_FIELD)
                .unwrap();
            let mut value_renc = self
                .clone()
                .re_encode_single_column(value_col.into_column(), forget_others)?;
            value_renc.rename(PlSmallStr::from_str(LANG_STRING_VALUE_FIELD));
            let lang_col = c
                .struct_()
                .unwrap()
                .field_by_name(LANG_STRING_LANG_FIELD)
                .unwrap();
            let mut lang_renc =
                self.re_encode_single_column(lang_col.into_column(), forget_others)?;
            lang_renc.rename(PlSmallStr::from_str(LANG_STRING_LANG_FIELD));
            let mut df = DataFrame::new(c.len(), vec![value_renc, lang_renc]).unwrap();
            df = df
                .lazy()
                .with_column(
                    as_struct(vec![
                        col(LANG_STRING_VALUE_FIELD),
                        col(LANG_STRING_LANG_FIELD),
                    ])
                    .alias(c.name().as_str()),
                )
                .select([col(c.name().as_str())])
                .collect()?;
            Ok(df.drop_in_place(c.name().as_str())?)
        } else {
            self.re_encode_single_column(c, forget_others)
        }
    }

    pub fn re_encode_single_column(self, c: Column, forget_others: bool) -> PolarsResult<Column> {
        let name = c.name().clone();
        match c {
            Column::Series(c) => {
                let uch = c.u32().unwrap();
                let mut v = Vec::with_capacity(uch.len());
                for u in uch {
                    let u = if let Some(u) = u {
                        if let Some(remap_u) = self.cat_map.get(&u) {
                            Some(*remap_u)
                        } else if !forget_others {
                            Some(u)
                        } else {
                            None
                        }
                    } else {
                        None
                    };
                    v.push(u);
                }
                let mut s = Series::from_iter(v);
                s.rename(name);
                Ok(s.into_column())
            }
            Column::Scalar(sc) => {
                let len = sc.len();
                let av = match sc.scalar().as_any_value() {
                    AnyValue::Null => AnyValue::Null,
                    AnyValue::UInt32(u) => {
                        if let Some(remap_u) = self.cat_map.get(&u) {
                            AnyValue::UInt32(*remap_u)
                        } else if !forget_others {
                            AnyValue::UInt32(u)
                        } else {
                            AnyValue::Null
                        }
                    }
                    _ => {
                        return Err(PolarsError::SchemaMismatch(ErrString::new_static(
                            "Expected UInt32 for cat, got other type",
                        )))
                    }
                };
                let ser = Series::from_any_values_and_dtype(name, &[av], &DataType::UInt32, true)?;
                Ok(ScalarColumn::from_single_value_series(ser, len).into())
            }
        }
    }
}

impl CatEncs {
    pub fn inner_join_re_enc(&self, other: &CatEncs) -> Vec<(u32, u32)> {
        self.maps.inner_join_re_enc(&other.maps)
    }
}

impl Cats {
    //Re enc is for right hand side
    pub fn join(left: LockedCats, right: LockedCats) -> CatReEnc {
        let left = left.read().unwrap();
        let right = right.read().unwrap();
        let rencs: Vec<_> = left
            .cat_map
            .iter()
            .map(|(t, left_enc)| {
                if let Some(right_enc) = right.cat_map.get(t) {
                    let renc = left_enc.inner_join_re_enc(right_enc);
                    Some(renc)
                } else {
                    None
                }
            })
            .filter(|x| x.is_some())
            .map(|x| x.unwrap())
            .collect();
        let renc_map: HashMap<_, _, BuildHasherDefault<NoHashHasher<u32>>> =
            rencs.into_iter().flatten().map(|x| x).collect();
        let cat_re_enc = CatReEnc {
            cat_map: Arc::new(renc_map),
        };
        cat_re_enc
    }

    pub fn merge(
        &mut self,
        other_cats: Vec<LockedCats>,
    ) -> HashMap<String, HashMap<CatType, CatReEnc>> {
        let mut map = HashMap::new();
        for c in other_cats {
            let c = c.read().unwrap();
            let mut other_map = HashMap::new();
            for (t, other_enc) in c.cat_map.iter() {
                let mut counter = self.get_counter(t);
                if let Some(enc) = self.cat_map.get_mut(t) {
                    let re_enc = enc.maps.merge(&other_enc.maps, &mut counter);
                    other_map.insert(t.clone(), re_enc);
                } else {
                    let mut enc = CatEncs::new_empty(
                        self.path.as_ref().map(|x| x.as_ref()),
                        &t.as_base_rdf_node_type(),
                    );
                    let re_enc = enc.maps.merge(&other_enc.maps, &mut counter);
                    self.cat_map.insert(t.clone(), enc);
                    other_map.insert(t.clone(), re_enc);
                }
                self.set_counter(counter, t);
            }
            if !other_map.is_empty() {
                map.insert(c.uuid.clone(), other_map);
            }
        }
        map
    }
}

pub fn re_encode(
    encoded: Vec<CatTriples>,
    re_enc_map: HashMap<String, HashMap<CatType, CatReEnc>>,
) -> Vec<CatTriples> {
    let re_encoded: Vec<_> = encoded
        .into_par_iter()
        .map(|mut x| {
            let mut enc = x.encoded_triples;
            let subj_encs = if let Some(s_uuid) = &enc.subject_local_cat_uuid {
                if let Some(s_map) = re_enc_map.get(s_uuid) {
                    if let Some(s) = &enc.subject {
                        s_map.get(s)
                    } else {
                        None
                    }
                } else {
                    None
                }
            } else {
                None
            };
            if let Some(subj_encs) = subj_encs {
                let mut renc_col = subj_encs
                    .clone()
                    .re_encode_column(
                        enc.df.column(SUBJECT_COL_NAME).unwrap().clone(),
                        x.subject_type.clone(),
                        false,
                    )
                    .unwrap();
                renc_col.rename(PlSmallStr::from_str(SUBJECT_COL_NAME));
                enc.df.with_column(renc_col).unwrap();
            }
            let obj_encs = if let Some(o_uuid) = &enc.object_local_cat_uuid {
                if let Some(o_map) = re_enc_map.get(o_uuid) {
                    if let Some(o) = &enc.object {
                        o_map.get(o)
                    } else {
                        None
                    }
                } else {
                    None
                }
            } else {
                None
            };
            if let Some(obj_encs) = obj_encs {
                let mut renc_col = obj_encs
                    .clone()
                    .re_encode_column(
                        enc.df.column(OBJECT_COL_NAME).unwrap().clone(),
                        x.object_type.clone(),
                        false,
                    )
                    .unwrap();
                renc_col.rename(PlSmallStr::from_str(OBJECT_COL_NAME));
                enc.df.with_column(renc_col).unwrap();
            }

            x.encoded_triples = enc;
            x
        })
        .collect();
    re_encoded
}

pub fn reencode_solution_mappings(
    sm: EagerSolutionMappings,
    reencs: &HashMap<String, HashMap<BaseRDFNodeType, CatReEnc>>,
) -> EagerSolutionMappings {
    let EagerSolutionMappings {
        mappings,
        mut rdf_node_types,
    } = sm;
    let mut mappings = mappings.lazy();
    for (c, t) in &rdf_node_types {
        let mut reenc_exprs = vec![];
        for (bt, bs) in &t.map {
            if let BaseCatState::CategoricalNative(Some(local)) = bs {
                let local = local.read().unwrap();
                if let Some(rmap) = reencs.get(&local.uuid) {
                    let r = rmap.get(bt).unwrap();
                    let e = if t.is_multi() {
                        col(c).struct_().field_by_name(&bt.field_col_name())
                    } else {
                        col(c)
                    };
                    let r_cloned = r.clone();
                    let bt = bt.clone();
                    reenc_exprs.push(e.map(
                        move |c| r_cloned.clone().re_encode_column(c, bt.clone(), false),
                        |_, y| Ok(y.clone()),
                    ));
                }
            }
        }
        if !reenc_exprs.is_empty() {
            if t.is_multi() {
                mappings = mappings.with_column(col(c).struct_().with_fields(reenc_exprs).alias(c));
            } else {
                mappings = mappings.with_column(reenc_exprs.pop().unwrap().alias(c))
            }
        }
    }
    for v in rdf_node_types.values_mut() {
        for v2 in v.map.values_mut() {
            if matches!(v2, BaseCatState::CategoricalNative(..)) {
                *v2 = BaseCatState::CategoricalNative(None);
            }
        }
    }
    let sm = EagerSolutionMappings::new(mappings.collect().unwrap(), rdf_node_types);
    sm
}
