//! Persistent store of all calculated image descriptions.
//! Can be queried to find identical or similar images.

use std::collections::{HashMap, HashSet};
use crate::descriptors::Descriptor;
use crate::mutators::get_all_mutators;
use std::fs;
use std::path::Path;
use std::fmt;
use image::DynamicImage;
use log::{info, debug, error};
use bktree::*;

#[derive(Debug)]
pub enum SaveError {
	Serialization,
	File,
}

/// Struct for calculating Recall-Precision per mutation
pub struct PRStats {
	true_positives: [u64; 64],
	false_positives: [u64; 64],
	false_negatives: [u64; 64],
	pub tag: String,
	pub name: String
}

impl PRStats {
	pub fn precision(&self, threshold: usize) -> f64 {
		let t = match threshold {
			0..=64 => threshold,
			_ => std::cmp::max(0, std::cmp::min(64, threshold))
		};
		let tpos = self.true_positives[t] as f64;
		let fpos = self.false_positives[t] as f64;
		tpos / (tpos + fpos)
	}
	pub fn recall(&self, threshold: usize) -> f64 {
		let t = match threshold {
			0..=64 => threshold,
			_ => std::cmp::max(0, std::cmp::min(64, threshold))
		};
		let tpos = self.true_positives[t] as f64;
		let fneg = self.false_negatives[t] as f64;
		tpos / (tpos + fneg)
	}
	pub fn pr(&self, threshold: usize) -> (f64, f64) {
		let t = match threshold {
			0..=64 => threshold,
			_ => std::cmp::max(0, std::cmp::min(64, threshold))
		};
		let tpos = self.true_positives[t] as f64;
		let fpos = self.false_positives[t] as f64;
		let fneg = self.false_negatives[t] as f64;
		let p = tpos / (tpos + fpos);
		let r = tpos / (tpos + fneg);
		(r, p)
	}
	pub fn pr_curve(&self, threshold: usize) -> Vec<(f64, f64)> {
		let mut curve = Vec::new();
		for i in 0..threshold {
			curve.push(self.pr(i));
		}
		curve
	}

}

/// Uses a hashmap to map descriptors to buckets of files.
/// Also keeps a BK-tree for quick distance ranking
pub struct DescriptorStore {
	pub descriptor: Box<dyn Descriptor>,
	/// Main hashmap that maps descriptors to buckets of filenames
	map: HashMap<u64, Vec<String>>, 

	/// Place to load and store the map on file
	save_location: std::path::PathBuf, 

	/// BK-tree for fast nearest neighbour
	bktree: BkTree<u64>, 

	/// Cache for seen files to skip on initial load
	seen: Option<HashSet<String>>, 
}

impl DescriptorStore {
	/// Makes a new empty `DescriptorStore` with default settings
	pub fn new(descriptor: Box<dyn Descriptor>) -> Self {
		let map: HashMap<u64, Vec<String>> = HashMap::new();
		//let seen: HashSet<String> = HashSet::new();
		//let seen = None;
		DescriptorStore {
			map,
			descriptor,
			save_location: std::path::PathBuf::from("store.messagepack"),
			bktree: BkTree::new(hamming_distance),
			seen: None,
		}
	}

	/// Sets the file location and loads data from file (if available)
	pub fn with_file<P: AsRef<Path>>(mut self, path: P) -> Self {
		self.save_location = std::path::PathBuf::from(path.as_ref());
		let map_file = fs::read(&self.save_location);
		if let Ok(f) = map_file {
			self.map = rmp_serde::from_slice(&f).unwrap();
			for i in self.map.keys() {
				self.bktree.insert(*i);
			}
			let mut seen: HashSet<String> = HashSet::new();
			for bucket in self.map.values() {
				for element in bucket {
					if !element.starts_with("mut.") {
						seen.insert(element.clone());
					}
				}
			}
			self.seen = Some(seen);
		} else {
			info!("{} not found, starting from empty store.", self.save_location.display());
		};
		self
	}

	/// Stores non-mutated images in a hashset for quick membership checks
	pub fn see(&mut self, value: String) {
		if !value.starts_with("mut.") {
			match &mut self.seen {
				Some(set) => {
					set.insert(value);
				},
				None => {
					let mut seen: HashSet<String> = HashSet::new();
					seen.insert(value);
					self.seen = Some(seen);
				}
			}
		}
	}

	pub fn save(&self) -> Result<(), SaveError> {
		let serialized: Vec<u8> = match rmp_serde::to_vec(&self.map) {
			Ok(value) => value,
			Err(_e) => return Err(SaveError::Serialization),
		};
		match fs::write(&self.save_location, serialized) {
			Ok(()) => Ok(()),
			Err(_e) => Err(SaveError::File),
		}
	}

	/// Returns true iff the store already contains the key
	pub fn contains(&self, key: u64) -> bool {
		self.map.contains_key(&key)
	}

	/// Returns true iff the store already contains the value in some bucket. 
	/// Will use a hashmap cache if available
	pub fn has_value(&self, value: &String) -> bool {
		match &self.seen {
			Some(set) => {
				set.contains(value)
			},
			None => {
				for bucket in self.map.values() {
					if bucket.contains(value) {
						return true;
					}
				}
				false
			}
		}
	}

	pub fn get(&self, value: &String) -> Option<u64> {
		for (key, bucket) in self.map.iter() {
			if bucket.contains(value) {
				return Some(*key);
			}
		}
		None
	}

	/// Inserts a single value into the store
	pub fn insert(&mut self, key: &u64, value: String) {
		let bucket = match self.map.get(key) {
			Some(b) => {
				let mut n = b.clone();
				if !n.contains(&value) {
					n.push(value.clone());
				}
				n
			},
			None => vec![value.clone()],
		};
		self.map.insert(*key, bucket);
		self.see(value);
		self.bktree.insert(*key);
	}

	/// Calculates all descriptions with a given descriptor 
	pub fn insert_directory<U: AsRef<Path>>(&mut self, dir: U) {
		for node in fs::read_dir(dir).unwrap() {
			let file = node.expect("Error walking directory");
			let name = match file.file_name().into_string() {
				Ok(v) => v,
				Err(e) => {
					error!("Error reading {}: {:?}:", file.path().to_string_lossy(), e);
					continue
				}
			};
			if !self.has_value(&name) {
				debug!("{} already known, skipping.", name);
			} else {
				info!("Processing {}", name);
				let img = match image::open(file.path()) {
					Ok(v) => v,
					Err(e) => {
						error!("Failed to process {}: {}", name, e);
						continue
					}
				};
				let phash = self.descriptor.describe(&img);
				if self.contains(phash) {
					println!("{} duplicate of {:?}", name, self.map.get(&phash));
				}
				self.insert(&phash, name);
			}
			self.save().expect("error");
		}
		// We have done our bulk loading, so we can unset our hashset cache:
		//self.seen = None;
	}

	pub fn store(&mut self, img: &DynamicImage, name: String) {
		if !self.has_value(&name) {
			let phash = self.descriptor.describe(img);
			self.insert(&phash, name);
		}
	}

	/// Nearest neighbours
	pub fn nn(&self, from: u64, max_distance: usize) -> Vec<(&u64, isize)> {
		let neighbours = self.bktree.find(from, max_distance.try_into().unwrap());
		//neighbours.sort_by(|a, b| a.1.cmp(&b.1));
		neighbours
	}


	/// Returns a flat vector with all the filenames that are within a given distance
	pub fn nn_flat_results(&self, from: u64, max_distance: usize) -> Vec<String> {
		let nn = self.nn(from, max_distance);
		let mut results: Vec<String> = Vec::new();
		for (key, _) in nn {
			// We know the key exists, so getting the result should never be none
			let mut bucket = self.map.get(key).unwrap().clone();
			results.append(&mut bucket);
		}
		results
	}

	pub fn print_most_dups(&self) {
		let mut most = 0;
		for bucket in self.map.values() {
			let len = bucket.len();
			if len > most {
				most = len;
			}
		}

		let mut output = String::new();
		for (key, bucket) in self.map.iter() {
			if bucket.len() == most {
				let key_entry = format!("{key:064b}:\n\t\t");
				output.push_str(&key_entry);
				for value in bucket {
					let value_entry = format!("{value}\n");
					output.push_str(&value_entry);
				}
				output.push('\n');
			}
		}
		print!("DescriptorStore:\n{}Total: {}", output, self.map.len())
	}

	pub fn get_stats(&self, mut max_threshold: usize) -> Vec<PRStats> {
		max_threshold = std::cmp::min(64, max_threshold);
		let mut mutator_stats = Vec::new();
		for mutator in get_all_mutators() {
			mutator_stats.push(
				PRStats {
					true_positives: [0; 64],
					false_positives: [0; 64],
					false_negatives: [0; 64],
					tag: "mut".to_string() + &mutator.tag(),
					name: mutator.info(),
				}
			);
		}

		// Seen should only contain a list of base images
		let set = self.seen.clone().unwrap();
		for image in set.iter() {
			let phash = self.get(image).unwrap();
			debug!("Checking {}, with phash: {}", image, phash);
			for threshold in 0..max_threshold {
				debug!("Threshold: {}", threshold);
				//Assume miss, therefore a false negative
				//Undo the miss when there is a true positive
				for mutator in &mut mutator_stats {
					mutator.false_negatives[threshold] += 1;
				}

				for found in self.nn_flat_results(phash, threshold) {
					if found.ends_with(&format!(".{}", image)) { // True positive
						for mutator in &mut mutator_stats {
							if found.starts_with(&mutator.tag) {
								debug!("{} is hit for {}", found, mutator.name);
								mutator.true_positives[threshold] += 1;
								mutator.false_negatives[threshold] -= 1;
							}
						}
					} else if found.eq(image) {
						continue;
					} else { // False positive!
						// Mutated misses only count for the mutator
						// Unmutated misses count for everyone
						if found.starts_with("mut") {
							for mutator in &mut mutator_stats {
								if found.starts_with(&mutator.tag) {
									debug!("{} is false positive for {}", found, mutator.name);
									mutator.false_positives[threshold] += 1;
								}
							}
						} else {
							for mutator in &mut mutator_stats {
								mutator.false_positives[threshold] += 1;
							}
						}
					}
				}
			}
		}
		mutator_stats
	}

	pub fn print_stats(&self) {
		let mut mutator_stats = self.get_stats(64);
		for mutator in &mut mutator_stats {
			println!("{}:", mutator.name);
			for threshold in 0..64 {
				let tpos = mutator.true_positives[threshold] as f64;
				let fpos = mutator.false_positives[threshold] as f64;
				let fneg = mutator.false_negatives[threshold] as f64;
				let p = tpos / (tpos + fpos);
				let r = tpos / (tpos + fneg);
				let f1 = 2.0*tpos / (2.0*tpos + fpos + fneg);
				debug!("Precision: {}, Recall: {}, F1 score: {}", p, r, f1);
				debug!("Tp: {}, Fp: {}, Fn: {}", tpos, fpos, fneg);
				println!("{}, {}" , p, r)
			}
		}
	}
}

impl fmt::Display for DescriptorStore {
	fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
		let mut output = String::new();
		for (key, bucket) in self.map.iter() {
			let key_entry = format!("{key:064b}:\n");
			output.push_str(&key_entry);
			for value in bucket {
				let value_entry = format!("\t{value}\n");
				output.push_str(&value_entry);
			}
			output.push('\n');
		}
		write!(f, "DescriptorStore:\n{}Total: {}", output, self.map.len())
	}
}