pokestat 0.1.1

guess a pokemon's IVs and EVs from in-game information
Documentation
use std::ops::{Index, IndexMut};
use std::fmt;

mod db;

pub use db::get_base_stats;

/// identifies an individual stat
// pokemon showdown is very particular that status should be capitalized exactly
// like this.
#[derive(Debug, Copy, Clone, PartialEq)]
#[repr(u8)]
pub enum Stat {
	HP = 0,
	Atk,
	Def,
	SpA,
	SpD,
	Spe,
}

static STAT_LIST: [Stat; 6] = { use Stat::*; [HP, Atk, Def, SpA, SpD, Spe] };

/// a group of stats. (hp, atk, def, sp. atk, sp. def, speed)
///
/// can represent base stats, EVs, IVs, or total stats.
#[derive(Default, Debug, PartialEq, Clone)]
pub struct Stats(pub [u16; 6]);

impl Index<usize> for Stats {
    type Output = u16;

    fn index(&self, index: usize) -> &Self::Output {
		&self.0[index]
	}
}

impl IndexMut<usize> for Stats {
    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
		&mut self.0[index]
	}
}

impl Index<Stat> for Stats {
    type Output = u16;

    fn index(&self, index: Stat) -> &Self::Output {
		&self.0[index as usize]
	}
}

impl IndexMut<Stat> for Stats {
    fn index_mut(&mut self, index: Stat) -> &mut Self::Output {
		&mut self.0[index as usize]
	}
}

impl Stats {
	pub fn total(&self) -> u16 {
		self.0.iter().sum()
	}
}

static NATURE_TABLE: [[&str; 5]; 5] = [
	[ "Hardy", 	"Lonely", 	"Adamant", 	"Naughty", 	"Brave" ],
	[ "Bold", 	"Docile", 	"Impish", 	"Lax", 	"Relaxed" ],
	[ "Modest", 	"Mild", 	"Bashful", 	"Rash", 	"Quiet" ],
	[ "Calm", 	"Gentle", 	"Careful", 	"Quirky", 	"Sassy" ],
	[ "Timid", 	"Hasty", 	"Jolly", 	"Naive", 	"Serious" ],
];

/// +/-
#[derive(Debug, PartialEq, Clone, Copy)]
pub struct Nature(pub Stat, pub Stat);

#[derive(Debug, thiserror::Error)]
#[error("unknown nature '{0}'")]
pub struct NatureError(String);

impl std::str::FromStr for Nature {
    type Err = NatureError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
		// iterate over all stats besides hp
		for plus in &STAT_LIST[1..] {
			for minus in &STAT_LIST[1..] {
				if s.eq_ignore_ascii_case(
					NATURE_TABLE[(*plus as usize) - 1][(*minus as usize) - 1])
				{
					return Ok(Nature(*plus, *minus));
				}
			}
		}
		Err(NatureError(s.to_string()))
	}
}

impl fmt::Display for Nature {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
		let &Nature(plus, minus) = self;
		write!(f, "{}", NATURE_TABLE[(plus as usize) -1][(minus as usize) -1])
	}
}

#[derive(Debug, thiserror::Error)]
pub enum ConvError {
	#[error("unable to find pokemon named {0}")]
	NoPokemon(String),
}

#[derive(Debug, PartialEq, Clone)]
pub struct Spread {
	pub base: Stats,
	pub iv: Stats,
	pub ev: Stats,
	pub nature: Nature,
	pub level: u8,
}

impl fmt::Display for Spread {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
		// uses the very particular format used by pokemon showdown
        writeln!(f, "Level: {}", self.level)?;
		writeln!(f, "{} Nature", self.nature)?;
		let mut ev_strs = vec![];
		let mut iv_strs = vec![];
		for st in STAT_LIST {
			if self.iv[st] != 31 {
				iv_strs.push(format!("{} {:?}", self.iv[st], st));
			}
			if self.ev[st] != 0 {
				ev_strs.push(format!("{} {:?}", self.ev[st], st));
			}
		}
		if !ev_strs.is_empty() {
			writeln!(f, "EVs: {}", ev_strs.join(" / "))?;
		}
		if !iv_strs.is_empty() {
			writeln!(f, "IVs: {}", iv_strs.join(" / "))?;
		}
		Ok(())
    }
}

impl Spread {
	pub fn from_summary(pokemon: &str, level: u8, nature: Nature, stats: &Stats)
						-> Result<Spread, ConvError> {
		let Some(base) = get_base_stats(pokemon)
		else { return Err(ConvError::NoPokemon(pokemon.to_string())) };
		let mut spread = Spread{
			base, nature, level,
			iv: Stats::default(),
			ev: Stats::default(),
		};
		spread.guess_eviv(stats);
		Ok(spread)
	}

	pub fn calculate(&self) -> Stats {
		let mut r = Stats::default();
		use Stat::*;
		//r[Hp] = 2*self.base[Hp]+self.iv[Hp]+self.ev[Hp]/
		for i in 0..6 {
			r[i] = (2*self.base[i]+self.iv[i]+self.ev[i]/4)*(self.level as u16)/100 + 5
		}
		r[HP] += (self.level as u16) + 5;
		let Nature(plus, minus) = self.nature;
		if plus != minus {
			r[plus] *= 11;
			r[minus] *= 9;
			r[plus] /= 10;
			r[minus] /= 10;
		}
		// handle shedinja
		if self.base[HP] == 1 {
			r[HP] = 1;
		}
		return r;
	}

	fn calculate_stat(&self, stat: Stat) -> u16 {
		// this should be optimized in the future,
		// it calculates all the stats, not just the one we want.
		self.calculate()[stat]
	}

	fn guess_eviv(&mut self, target: &Stats) {
		self.iv = Stats([31, 31, 31, 31, 31, 31]);
		//while self.calculate() != *target {
			// first, try to find EVs assuming max IVs
			for i in STAT_LIST {
				while self.calculate_stat(i) < target[i] {
					self.ev[i] += 4;
				}
			}
			// next, lower IVs appropriatly
			for i in STAT_LIST {
				while self.calculate_stat(i) > target[i] {
					if self.iv[i] == 0 {
						dbg!(i, target, self);
						panic!("0 IVs isn't low enough?!");
					}
					self.iv[i] -= 1;
				}
			}
		//	dbg!(self.calculate(), target);
		//}
	}
}

#[cfg(test)]
mod tests {
	use super::*;
	#[test]
	fn calc_stats() {
		let spread = Spread{
			base: Stats([60, 45, 70, 45, 90, 95]),
			iv: Stats([31, 31, 31, 31, 31, 31]),
			ev: Stats([252, 0, 236, 0, 20, 0]),
			nature: Nature(Stat::Def, Stat::SpA),
			level: 50,
		};
		let stats = spread.calculate();
		assert_eq!(stats, Stats([167, 65, 132, 58, 113, 115]));
	}
	#[test]
	fn calc_eviv() {
		let target_spread = Spread{
			base: Stats([60, 45, 70, 45, 90, 95]),
			iv: Stats([31, 31, 31, 31, 31, 31]),
			ev: Stats([252, 0, 236, 0, 20, 0]),
			nature: Nature(Stat::Def, Stat::SpA),
			level: 50,
		};
		let mut input_spread = Spread{
			iv: Stats::default(),
			ev: Stats::default(),
			.. target_spread.clone()
		};
		let target_stats = Stats([167, 65, 132, 58, 113, 115]);
		input_spread.guess_eviv(&target_stats);
		assert_eq!(input_spread, target_spread);
		assert_eq!(input_spread.calculate(), target_spread.calculate());
	}

	#[test]
	fn calc_eviv_low() {
		let target_spread = Spread{
			base: Stats([60, 45, 70, 45, 90, 95]),
			iv: Stats([31, 1, 31, 31, 31, 1]),
			ev: Stats([252, 0, 236, 0, 20, 0]),
			nature: Nature(Stat::Def, Stat::SpA),
			level: 50,
		};
		let mut input_spread = Spread{
			iv: Stats::default(),
			ev: Stats::default(),
			.. target_spread.clone()
		};
		//panic!("{:?}", &target_spread.calculate());
		let target_stats =
			Stats([167, 50, 132, 58, 113, 100]);
		input_spread.guess_eviv(&target_stats);
		assert_eq!(input_spread, target_spread);
		assert_eq!(input_spread.calculate(), target_spread.calculate());
	}
	#[test]
	fn summary() {
		let spread1 = Spread::from_summary(
			"Farigiraf", 50, "Modest".parse().unwrap(),
			&Stats([220, 85, 115, 144, 104, 81])).unwrap();
		assert_eq!(spread1.iv, Stats([31, 1, 31, 31, 31, 31]));
		let spread2 = Spread::from_summary(
			"shedinja", 100, "naive".parse().unwrap(),
			&Stats([1, 279, 126, 97, 86, 196])).unwrap();
		assert_eq!(spread2.ev, Stats([0, 252, 0, 4, 0, 252]));
	}

	#[test]
	fn base_stat_db() {
		let mew_stats = get_base_stats("mew").unwrap();
		assert_eq!(mew_stats, Stats([100, 100, 100, 100, 100, 100]));
	}
}