pub mod cluster;
pub mod eval;
pub mod liblinear;
pub mod train;

use crate::mat_util::*;
use crate::{Index, IndexValueVec};
use hashbrown::HashMap;
use itertools::Itertools;
use log::{info, warn};
use ordered_float::NotNan;
use rayon::prelude::*;
use serde::{Deserialize, Serialize};
use std::cmp::Reverse;
use std::io;
use std::mem::swap;
use std::time;

/// Model training hyper-parameters.
pub type TrainHyperParam = train::HyperParam;

#[derive(Eq, PartialEq, Clone, Copy, Debug, Serialize, Deserialize)]
struct Settings {
    n_features: usize,
    classifier_loss_type: liblinear::LossType,
}

/// A Omikuji model, which contains a forest of trees.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct Model {
    trees: Vec<TreeNode>,
    settings: Settings,
}

static MODEL_SETTINGS_FILE_NAME: &str = "settings.json";
static TREE_FILE_NAME_PREFIX: &str = "tree";

impl Model {
    /// Returns a ranked list of predictions for the given input example.
    ///
    /// # Arguments
    ///
    /// * `feature_vec` - An input vector for prediction, assumed to be ordered by indices and have
    /// no duplicate or out-of-range indices
    /// * `beam_size` - Beam size for beam search.
    pub fn predict(&self, feature_vec: &[(Index, f32)], beam_size: usize) -> IndexValueVec {
        let feature_vec = self.prepare_feature_vec(feature_vec);
        let mut label_to_total_score = HashMap::<Index, f32>::new();
        let tree_predictions: Vec<_> = self
            .trees
            .iter()
            .map(|tree| tree.predict(self.settings.classifier_loss_type, &feature_vec, beam_size))
            .collect();
        for label_score_pairs in tree_predictions {
            for (label, score) in label_score_pairs {
                let total_score = label_to_total_score.entry(label).or_insert(0.);
                *total_score += score;
            }
        }

        let mut label_score_pairs = label_to_total_score
            .iter()
            .map(|(&label, &total_score)| (label, total_score / self.trees.len() as f32))
            .collect_vec();
        label_score_pairs.sort_unstable_by_key(|&(_, score)| Reverse(NotNan::new(score).unwrap()));
        label_score_pairs
    }

    /// The expected dimension of feature vectors.
    pub fn n_features(&self) -> usize {
        self.settings.n_features
    }

    /// The number of trees in the forest model.
    pub fn n_trees(&self) -> usize {
        self.trees.len()
    }

    /// Prepare the feature vector in both dense and sparse forms to make prediction more efficient.
    fn prepare_feature_vec(&self, sparse_vec: &[(Index, f32)]) -> SparseVec {
        let norm = sparse_vec
            .iter()
            .map(|(_, v)| v.powi(2))
            .sum::<f32>()
            .sqrt();

        let (mut indices, mut data): (Vec<_>, Vec<_>) = sparse_vec
            .iter()
            .cloned()
            .map(|(i, v)| (i, v / norm))
            .unzip();

        indices.push(self.settings.n_features as Index);
        data.push(1.);

        SparseVec::new(self.settings.n_features + 1, indices, data)
    }

    /// Serialize model into the directory with the given path.
    pub fn save<P: AsRef<std::path::Path>>(&self, dir_path: P) -> io::Result<()> {
        info!("Saving model...");
        let start_t = time::Instant::now();

        let dir_path = dir_path.as_ref();
        if !dir_path.exists() {
            std::fs::create_dir_all(dir_path)?;
        } else if !dir_path.is_dir() {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                "file with the given name already exists",
            ));
        }

        let settings_path = dir_path.join(MODEL_SETTINGS_FILE_NAME);
        if settings_path.exists() {
            let reader = std::io::BufReader::new(std::fs::File::open(settings_path)?);
            let existing_settings = serde_json::from_reader(reader)?;
            if self.settings != existing_settings {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidInput,
                    "a model with different settings is already saved in the given directory",
                ));
            } else {
                info!(
                    "A model is already saved at {}; trees will be added to the existing model",
                    dir_path.display(),
                );
            }
        } else {
            let writer = std::io::BufWriter::new(std::fs::File::create(settings_path)?);
            serde_json::to_writer_pretty(writer, &self.settings).map_err(|e| {
                io::Error::new(
                    io::ErrorKind::Other,
                    format!("Unable to serialize settings: {}", e),
                )
            })?;
        }

        let index_to_tree_path =
            |index: usize| dir_path.join(format!("{}{}.cbor", TREE_FILE_NAME_PREFIX, index));
        let mut curr_index = 0usize;
        for tree in &self.trees {
            let mut tree_path = index_to_tree_path(curr_index);
            while tree_path.exists() {
                info!(
                    "File {} already exists, skipping to keep it",
                    tree_path.display()
                );
                curr_index += 1;
                tree_path = index_to_tree_path(curr_index);
            }

            info!("Saving tree to {}", tree_path.display());
            let writer = std::io::BufWriter::new(std::fs::File::create(tree_path)?);
            serde_cbor::to_writer(writer, tree).map_err(|e| {
                io::Error::new(
                    io::ErrorKind::Other,
                    format!("Unable to serialize tree: {}", e),
                )
            })?;
            curr_index += 1;
        }

        info!(
            "Model saved; it took {:.2}s",
            start_t.elapsed().as_secs_f32()
        );
        Ok(())
    }

    /// Deserialize model from the given directory.
    pub fn load<P: AsRef<std::path::Path>>(dir_path: P) -> io::Result<Self> {
        let start_t = time::Instant::now();

        let dir_path = dir_path.as_ref();
        info!("Loading model from {}...", dir_path.display());

        let settings = {
            let settings_path = dir_path.join(MODEL_SETTINGS_FILE_NAME);
            info!("Loading model settings from {}...", settings_path.display());
            let reader = std::io::BufReader::new(std::fs::File::open(settings_path)?);
            serde_json::from_reader(reader)?
        };
        info!("Loaded model settings {:?}...", settings);

        let mut trees = Vec::<TreeNode>::new();
        for entry in dir_path.read_dir()? {
            let entry = entry?;

            let has_tree_file_name = {
                let file_name = entry.file_name();
                let file_name_str = file_name.to_string_lossy();
                file_name_str.starts_with(TREE_FILE_NAME_PREFIX) && file_name_str.ends_with(".cbor")
            };
            if !has_tree_file_name {
                continue;
            }

            let tree_path = entry.path();
            info!("Loading tree from {}...", tree_path.display());
            let reader = std::io::BufReader::new(std::fs::File::open(tree_path.as_path())?);
            let tree: TreeNode = serde_cbor::from_reader(reader).map_err(|e| {
                io::Error::new(
                    io::ErrorKind::InvalidData,
                    format!(
                        "Unable to deserialize tree from {} with error: {}",
                        tree_path.display(),
                        e
                    ),
                )
            })?;
            if !tree.is_valid(settings) {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidData,
                    format!("Tree loaded from {} is invalid", tree_path.display()),
                ));
            }
            trees.push(tree);
        }

        if !trees.is_empty() {
            info!(
                "Loaded model with {} trees; it took {:.2}s",
                trees.len(),
                start_t.elapsed().as_secs_f32()
            );
        } else {
            warn!(
                "Failed to load any trees from model directory {}; returning an empty model",
                dir_path.display()
            )
        }
        Ok(Self { trees, settings })
    }

    /// Densify model weights to speed up prediction at the cost of more memory usage.
    pub fn densify_weights(&mut self, max_sparse_density: f32) {
        info!("Densifying model weights...");
        let start_t = time::Instant::now();

        self.trees
            .par_iter_mut()
            .for_each(|tree| tree.densify_weights(max_sparse_density));

        info!(
            "Model weights densified; it took {:.2}s",
            start_t.elapsed().as_secs_f32()
        );
    }
}

#[derive(Clone, Debug, Serialize, Deserialize)]
enum TreeNode {
    Branch {
        weights: WeightMat,
        children: Vec<TreeNode>,
    },
    Leaf {
        weights: WeightMat,
        labels: Vec<Index>,
    },
}

impl TreeNode {
    fn is_valid(&self, settings: Settings) -> bool {
        match self {
            TreeNode::Branch {
                ref weights,
                ref children,
            } => {
                weights.shape() == (settings.n_features + 1, children.len())
                    && children.iter().all(|c| c.is_valid(settings))
            }
            TreeNode::Leaf {
                ref weights,
                ref labels,
            } => weights.shape() == (settings.n_features + 1, labels.len()),
        }
    }

    fn is_leaf(&self) -> bool {
        matches!(self, TreeNode::Leaf { .. })
    }

    fn densify_weights(&mut self, max_sparse_density: f32) {
        fn densify(weights: &mut WeightMat, max_sparse_density: f32) {
            if !weights.is_dense() && weights.density() > max_sparse_density {
                weights.densify();
            }
        }

        match self {
            TreeNode::Branch {
                ref mut weights,
                ref mut children,
            } => {
                densify(weights, max_sparse_density);
                children
                    .par_iter_mut()
                    .for_each(|child| child.densify_weights(max_sparse_density));
            }
            TreeNode::Leaf {
                ref mut weights, ..
            } => {
                densify(weights, max_sparse_density);
            }
        }
    }

    fn predict(
        &self,
        classifier_loss_type: liblinear::LossType,
        feature_vec: &SparseVec,
        beam_size: usize,
    ) -> IndexValueVec {
        assert!(beam_size > 0);
        let mut curr_level = Vec::<(&TreeNode, f32)>::with_capacity(beam_size * 2);
        let mut next_level = Vec::<(&TreeNode, f32)>::with_capacity(beam_size * 2);

        curr_level.push((self, 0.));

        // Iterate until only leaves are left
        while curr_level.iter().any(|(node, _)| !node.is_leaf()) {
            assert!(!curr_level.is_empty());
            next_level.clear();
            for &(node, node_score) in &curr_level {
                match node {
                    TreeNode::Branch { weights, children } => {
                        let mut child_scores =
                            liblinear::predict(weights, classifier_loss_type, feature_vec);
                        child_scores += node_score;
                        next_level
                            .extend(children.iter().zip_eq(child_scores.into_iter().cloned()));
                    }
                    TreeNode::Leaf { .. } => {
                        next_level.push((node, node_score));
                    }
                }
            }

            swap(&mut curr_level, &mut next_level);
            if curr_level.len() > beam_size {
                pdqselect::select_by_key(curr_level.as_mut_slice(), beam_size, |&(_, score)| {
                    Reverse(NotNan::new(score).unwrap())
                });
                curr_level.truncate(beam_size);
            }
        }

        curr_level
            .iter()
            .flat_map(|&(leaf, leaf_score)| match leaf {
                TreeNode::Leaf { weights, labels } => {
                    let mut label_scores =
                        liblinear::predict(weights, classifier_loss_type, feature_vec);
                    label_scores.mapv_inplace(|v| (v + leaf_score).exp());

                    let mut label_score_pairs = labels
                        .iter()
                        .cloned()
                        .zip_eq(label_scores.into_iter().cloned())
                        .collect_vec();
                    pdqselect::select_by_key(
                        label_score_pairs.as_mut_slice(),
                        beam_size,
                        |&(_, score)| Reverse(NotNan::new(score).unwrap()),
                    );
                    label_score_pairs.truncate(beam_size);
                    label_score_pairs
                }
                _ => unreachable!(),
            })
            .collect_vec()
    }
}
