use std::{convert::Infallible, fmt::{Binary, Debug}, fs::File, io::{BufRead, BufReader}, marker::PhantomData, num::ParseFloatError, str::FromStr};
use rand::Rng;

fn add(a: u16, b: u16) -> u16 {
    a + b
}

#[derive(Debug, Clone, PartialEq)]
struct BinaryString {
    vec: Vec<i8>
}

struct Bounds {
    min: f64,
    max: f64,
}

impl Bounds {
    pub fn new(min: f64, max: f64) -> Self {
        Bounds {
            min,
            max
        }
    }
}

#[derive(Debug, Clone, PartialEq)]
enum BinaryStringConversionError {
    DimensionMismatch,
    NoBounds
}

impl BinaryString {
    pub fn new(vec: Vec<i8>) -> BinaryString {
        BinaryString {
            vec
        }
    }

    pub fn perturb<TRng>(self: &Self, rng: &mut TRng, p: f64) -> Self
        where TRng : Rng
    {
        BinaryString::new(
            self.into_iter()
                .map(|c| if rng.random::<f64>() <= p { 1 - *c } else { *c })
                .collect::<Vec<i8>>()
        )
    }

    fn to_real_internal<'a, T: DoubleEndedIterator<Item = &'a i8>>(vec: T, len: usize, min: f64, max: f64) -> f64
    {
        let diff = max - min;
        let len = len as i32;
        let max_represent_num = 2f64.powi(len) - 1.0;
        let represented_num = vec
            .rev()
            .enumerate()
            .map(|(bit, c)| diff * (*c as f64) * 2f64.powi(bit as i32))
            .sum::<f64>();

        min + (represented_num / max_represent_num)
    }

    pub fn to_real_single(self: &Self, min: f64, max: f64) -> f64 {
        BinaryString::to_real_internal(self.vec.iter(), self.vec.len(), min, max)
    }

    pub fn to_real(self: &Self, bounds: &Vec<Bounds>) -> Result<Vec<f64>, BinaryStringConversionError> {
        if bounds.len() == 0 {
            return Err(BinaryStringConversionError::NoBounds);
        }

        let chunk_size = self.vec.len() / bounds.len();
        if self.vec.len() % bounds.len() != 0 {
            return Err(BinaryStringConversionError::DimensionMismatch);
        }

        Ok(self.vec.chunks(chunk_size)
            .zip(bounds)
            .map(|(chunk, bound)| BinaryString::to_real_internal(chunk.iter(), chunk_size, bound.min, bound.max))
            .collect::<Vec<f64>>())
    }
}

impl FromStr for BinaryString {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let binary_vec: Vec<i8> = s
            .chars()
            // skip spaces
            .filter(|c| *c != ' ')
            // Map ones and zeros
            .map(|c| match c {
                '0' => Ok(0),
                '1' => Ok(1),
                _ => Err(format!("Invalid binary character: {}", c)),
            })
            .collect::<Result<Vec<i8>, Self::Err>>()?;

        Ok(BinaryString::new(binary_vec))
    }
}

impl<'a> IntoIterator for &'a BinaryString {
    type Item = &'a i8;
    type IntoIter = std::slice::Iter<'a, i8>;

    fn into_iter(self) -> Self::IntoIter {
        self.vec.iter()
    }
}

trait FittingFunction {
    type In;
    type Out;
    type Err;

    fn fit(self: &Self, inp: &Self::In) -> Result<Self::Out, Self::Err>;
}

// Functions
struct OneMax;
impl OneMax {
    pub fn new() -> Self {
        OneMax
    }
}
impl FittingFunction for OneMax {
    type In = BinaryString;
    type Out = i32;
    type Err = Infallible;

    fn fit(self: &Self, chromosome: &BinaryString) -> Result<i32, Infallible> {
        Ok(chromosome.into_iter()
           .map(|x| *x as i32)
           .sum())
    }
}
struct LABS;

impl LABS {
    fn new() -> Self {
        LABS
    }

    fn Ck(k: usize, S: &Vec<i32>) -> i32 {
        let D = S.len();
        S.iter()
            .take(D - k)
            .zip(S.iter().skip(k))
            .map(|(x, y)| x * y)
            .sum()
    }
}

impl FittingFunction for LABS {
    type In = BinaryString;
    type Out = i32;
    type Err = Infallible;

    fn fit(self: &Self, chromosome: &BinaryString) -> Result<i32, Infallible> {
        let S: Vec<i32> = chromosome
            .into_iter()
            .map(|c| (*c as i32) * 2 - 1)
            .collect();
        let D = S.len();

        Ok((1..=D-1)
           .map(|k| LABS::Ck(k, &S).pow(2))
           .sum())
    }
}

struct Sphere {
    offset: Vec<f64>
}

impl Sphere {
    pub fn new(offset: Vec<f64>) -> Self {
        Sphere {
            offset
        }
    }
}
impl FittingFunction for Sphere {
    type In = Vec<f64>;
    type Out = f64;
    type Err = Infallible;

    fn fit(self: &Self, chromosome: &Vec<f64>) -> Result<f64, Infallible> {
        Ok(chromosome
           .iter()
           .zip(&self.offset)
           .map(|(x, o)| (x - o).powi(2))
           .sum())
    }
}

struct Rosenbrock;

impl Rosenbrock {
    pub fn new() -> Self {
        Rosenbrock
    }
}

impl FittingFunction for Rosenbrock {
    type In = Vec<f64>;
    type Out = f64;
    type Err = Infallible;

    fn fit(self: &Self, inp: &Vec<f64>) -> Result<f64, Infallible> {
        Ok(inp.windows(2)
           .map(|xs| 100.0 * (xs[1] - xs[0].powi(2)).powi(2) + (1.0 - xs[0]).powi(2))
           .sum())
    }
}

struct BinaryFittingWrapper<TFitting> {
    bounds: Vec<Bounds>,
    fitting_function: TFitting,
}

impl<TFitting> BinaryFittingWrapper<TFitting> {
    pub fn new(fitting_function: TFitting, bounds: Vec<Bounds>) -> Self {
        BinaryFittingWrapper {
            fitting_function,
            bounds,
        }
    }
}

impl<TFitting> FittingFunction for BinaryFittingWrapper<TFitting>
where
    TFitting: FittingFunction<In = Vec<f64>, Out = f64, Err = Infallible>
{
    type In = BinaryString;
    type Out = f64;
    type Err = BinaryStringConversionError;

    fn fit(self: &Self, inp: &BinaryString) -> Result<f64, BinaryStringConversionError> {
        Ok(self.fitting_function.fit(&inp.to_real(&self.bounds)?).unwrap())
    }
}

#[derive(Debug, Clone, PartialEq)]
struct LocalSearchCandidate<T>
    where T: Clone
{
    fit: T,
    pos: BinaryString,
    cycle: usize
}

#[derive(Debug, Clone, PartialEq)]
struct LocalSearchResult<T>
    where T: Clone
{
    best_candidate: LocalSearchCandidate<T>,

    // How many cycles there were
    cycles: usize,

    // statistics
    best_candidates: Vec<LocalSearchCandidate<T>>
}

trait TerminatingCondition<T>
    where T: Clone
{
    fn should_terminate(
        self: &mut Self,
        candidate: &LocalSearchCandidate<T>,
        stats: &Vec<LocalSearchCandidate<T>>,
        cycle: usize
    ) -> bool;
}

struct EqualTerminatingCondition {
    target: BinaryString
}

impl EqualTerminatingCondition {
    pub fn new(target: BinaryString) -> Self {
        Self {
            target
        }
    }
}

impl<T: Clone> TerminatingCondition<T> for EqualTerminatingCondition {
    fn should_terminate(
        self: &mut Self,
        candidate: &LocalSearchCandidate<T>,
        _: &Vec<LocalSearchCandidate<T>>,
        _: usize
    ) -> bool {
        candidate.pos == self.target
    }
}

struct NoBetterForCyclesTerminatingCondition {
    cycles: usize
}

impl NoBetterForCyclesTerminatingCondition {
    pub fn new(cycles: usize) -> Self {
        Self {
            cycles
        }
    }
}

struct AndTerminatingConditions<'a, T: Clone> {
    terminating_conditions: Vec<&'a mut dyn TerminatingCondition<T>>
}

impl<'a, T: Clone> AndTerminatingConditions<'a, T> {
    pub fn new(terminating_conditions: Vec<&'a mut dyn TerminatingCondition<T>>) -> Self {
        Self {
            terminating_conditions
        }
    }
}

impl<'a, T: Clone> TerminatingCondition<T> for AndTerminatingConditions<'a, T> {
    fn should_terminate(
        self: &mut Self,
        candidate: &LocalSearchCandidate<T>,
        stats: &Vec<LocalSearchCandidate<T>>,
        cycle: usize
    ) -> bool {
        return self.terminating_conditions.iter_mut()
            .all(
                |cond| cond.should_terminate(candidate, stats, cycle)
            )
    }
}

impl<T: Clone> TerminatingCondition<T> for NoBetterForCyclesTerminatingCondition {
    fn should_terminate (
        self: &mut Self,
        candidate: &LocalSearchCandidate<T>,
        _: &Vec<LocalSearchCandidate<T>>,
        cycle: usize
    ) -> bool {
        (cycle - candidate.cycle) > self.cycles
    }
}

trait PerturbationOperator {
    type Chromosome;

    fn perturb(self: &mut Self, chromosome: &Self::Chromosome) -> Self::Chromosome;
}

struct BinaryStringBitPerturbation<TRng: Rng> {
    rng: TRng,
    p: f64,
}

impl BinaryStringBitPerturbation<rand::rngs::ThreadRng> {
    pub fn new(p: f64) -> Self {
        Self {
            rng: rand::rng(),
            p
        }
    }
}

impl<TRng: Rng> PerturbationOperator for BinaryStringBitPerturbation<TRng> {
    type Chromosome = BinaryString;

    fn perturb(self: &mut Self, chromosome: &Self::Chromosome) -> Self::Chromosome {
        chromosome.perturb(&mut self.rng, self.p)
    }
}

trait BetterThanOperator<T> {
    fn better_than(self: &Self, a: &T, b: &T) -> bool;
}

struct DefaultBetterThan;
impl DefaultBetterThan {
    pub fn new() -> Self {
        Self
    }
}

impl<T> BetterThanOperator<T> for DefaultBetterThan
    where T: PartialOrd
{
    fn better_than(self: &Self, a: &T, b: &T) -> bool {
        a < b
    }
}

fn local_search_first_improving<
        T, TErr, TFit, TTerminatingCondition, TPerturbationOperator, TBetterThanOperator>(
    fit: &TFit,
    terminating_condition: &mut TTerminatingCondition,
    perturbation_operator: &mut TPerturbationOperator,
    better_than_operator: &TBetterThanOperator,
    initial: &BinaryString
) -> Result<LocalSearchResult<T>, TErr>
where
    T: Clone,
    TFit: FittingFunction<In = BinaryString, Out = T, Err = TErr>,
    TTerminatingCondition: TerminatingCondition<T>,
    TPerturbationOperator: PerturbationOperator<Chromosome = BinaryString>,
    TBetterThanOperator: BetterThanOperator<T>,
{
    let mut best_candidate = LocalSearchCandidate {
        pos: initial.clone(),
        fit: fit.fit(&initial)?,
        cycle: 0
    };

    let mut stats: Vec<LocalSearchCandidate<T>> = vec![];
    let mut cycle: usize = 0;

    while !terminating_condition.should_terminate(&best_candidate, &stats, cycle) {
        let perturbed = perturbation_operator.perturb(&best_candidate.pos);
        let perturbed_fit = fit.fit(&perturbed)?;

        // Minimize
        if better_than_operator.better_than(&perturbed_fit, &best_candidate.fit) {
            best_candidate = LocalSearchCandidate {
                pos: perturbed.clone(),
                fit: perturbed_fit,
                cycle
            };

            stats.push(best_candidate.clone());
        }

        cycle += 1;
    }

    Ok(LocalSearchResult {
        best_candidate,
        best_candidates: stats,
        cycles: cycle
    })
}

#[test]
fn test_local_search_one_max() {
    let one_max = OneMax::new();
    let optimum = BinaryString::new(vec![0; 10]);

    let result = local_search_first_improving(
        &one_max,
        &mut
            AndTerminatingConditions::new(
                vec![
                    &mut EqualTerminatingCondition::new(optimum.clone()),
                    &mut NoBetterForCyclesTerminatingCondition::new(100)
                ]
            ),
        &mut BinaryStringBitPerturbation::new(0.3),
        &DefaultBetterThan::new(),
        &BinaryString::new(vec![1; 10]),
    ).unwrap();

    println!("{:?}", result);

    assert_eq!(
        result.best_candidate.fit,
        0
    );

    assert_eq!(
        result.best_candidate.pos,
        optimum
    );
}

#[test]
fn test_local_search_sphere() {
    let optimum = BinaryString::new(vec![0, 0, 1, 0, 0,
                                         0, 0, 1, 0, 0]);
    let bounds = vec![Bounds::new(0.0, 31.0), Bounds::new(0.0, 31.0)];
    let optimum_real = optimum.to_real(&bounds).unwrap();
    let sphere = Sphere::new(optimum_real);
    let sphere_wrapped = BinaryFittingWrapper::new(sphere, bounds);

    let result = local_search_first_improving(
        &sphere_wrapped,
        &mut
            AndTerminatingConditions::new(
                vec![
                    &mut EqualTerminatingCondition::new(optimum.clone()),
                    &mut NoBetterForCyclesTerminatingCondition::new(100)
                ]
            ),
        &mut BinaryStringBitPerturbation::new(0.3),
        &DefaultBetterThan::new(),
        &BinaryString::new(vec![1; 10]),
    ).unwrap();

    println!("{:?}", result);

    assert_eq!(
        result.best_candidate.fit,
        0.0
    );

    assert_eq!(
        result.best_candidate.pos,
        optimum
    );
}

#[test]
fn test_perturb() {
    let mut rng = rand::rng();

    assert_eq!(
        BinaryString::new(vec![1, 1, 0, 0])
            .perturb(&mut rng, 1.0)
            .vec,
        vec![0, 0, 1, 1]
    );

    assert_eq!(
        BinaryString::new(vec![1, 1, 0, 0])
            .perturb(&mut rng, 0.0)
            .vec,
        vec![1, 1, 0, 0]
    );
}

#[test]
fn test_binary_string_to_real_single() {
    assert_eq!(
        BinaryString::new(vec![1])
            .to_real_single(0.0, 32.0),
        32.0
    );
    assert_eq!(
        BinaryString::new(vec![1, 1])
            .to_real_single(0.0, 32.0),
        32.0
    );
    assert_eq!(
        BinaryString::new(vec![0, 1])
            .to_real_single(0.0, 32.0),
        32.0 / 3.0
    );
    assert_eq!(
        BinaryString::new(vec![0, 0])
            .to_real_single(-16.0, 16.0),
        -16.0
    );
    assert_eq!(
        BinaryString::new(vec![0, 0, 0, 0, 1])
            .to_real_single(0.0, 31.0),
        1.0
    );
    assert_eq!(
        BinaryString::new(vec![1, 1, 1, 1, 1])
            .to_real_single(0.0, 31.0),
        31.0
    );
    assert_eq!(
        BinaryString::new(vec![0, 0, 0, 1, 0])
            .to_real_single(0.0, 31.0),
        2.0
    );

    assert_eq!(
        BinaryString::new(vec![1; 512])
            .to_real_single(0.0, 31.0),
        31.0
    );
}

#[derive(Debug, Clone, PartialEq)]
struct DataArrOfReals {
    vec: Vec<f64>,
    valid: bool
}

impl FromStr for DataArrOfReals {
    type Err = <f64 as FromStr>::Err;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        // TODO: maybe better handling, as an error?
        // this would mean also reimplementing load_test_file to be able to supply
        // out the error, but only for the output...
        if !s.starts_with('[') {
            return Ok(
                DataArrOfReals {
                    valid: false,
                    vec: vec![]
                }
            )
        }

        let trimmed = &s[1..s.len()-1];

        Ok(DataArrOfReals {
            valid: true,
            vec: trimmed.split(',')
                .map(|x| x.trim().parse::<f64>())
                .collect::<Result<Vec<f64>, _>>()?
        })
    }
}

fn test_binary_string_to_real(file_name: &str, bounds: Vec<Bounds>) {
    let data = load_test_file::<i8, DataArrOfReals>(file_name);

    for test in data {
        let res = BinaryString::new(test.inp)
                .to_real(&bounds);
        if !test.out.valid {
            assert_eq!(
                res,
                Err(BinaryStringConversionError::DimensionMismatch)
            );
        }
        else {
          assert_eq!(
              res.unwrap(),
              test.out.vec
          );
        }
    }
}

#[test]
fn test_binary_string_to_real_1D_1() {
    test_binary_string_to_real(
        "tests/Bin2Real_1D_1.txt",
        vec![Bounds::new(0.0, 1.0)]
    );
}

#[test]
fn test_binary_string_to_real_1D_2() {
    test_binary_string_to_real(
        "tests/Bin2Real_1D_2.txt",
        vec![Bounds::new(0.0, 4095.0)]
    );
}

#[test]
fn test_binary_string_to_real_1D_3() {
    test_binary_string_to_real(
        "tests/Bin2Real_1D_3.txt",
        vec![Bounds::new(-5.0, 5.0)]
    );
}

#[test]
fn test_binary_string_to_real_2D_1() {
    test_binary_string_to_real(
        "tests/Bin2Real_2D_1.txt",
        vec![Bounds::new(0.0, 1.0), Bounds::new(0.0, 1.0)]
    );
}

#[test]
fn test_binary_string_to_real_2D_2() {
    test_binary_string_to_real(
        "tests/Bin2Real_2D_2.txt",
        vec![Bounds::new(0.0, 63.0), Bounds::new(-32.0, 31.0)]
    );
}

#[test]
fn test_binary_string_to_real_2D_3() {
    test_binary_string_to_real(
        "tests/Bin2Real_2D_3.txt",
        vec![Bounds::new(-5.0, 5.0), Bounds::new(0.0, 10.0)]
    );
}

#[test]
fn test_binary_string_to_real_3D_1() {
    test_binary_string_to_real(
        "tests/Bin2Real_3D_1.txt",
        vec![Bounds::new(0.0, 1.0), Bounds::new(0.0, 1.0), Bounds::new(0.0, 1.0)]
    );
}

#[test]
fn test_binary_string_to_real_3D_2() {
    test_binary_string_to_real(
        "tests/Bin2Real_3D_2.txt",
        vec![Bounds::new(0.0, 15.0), Bounds::new(-8.0, 7.0), Bounds::new(-8.0, 8.0)]
    );
}

#[test]
fn test_binary_string_to_real_4D_1() {
    test_binary_string_to_real(
        "tests/Bin2Real_4D_1.txt",
        vec![Bounds::new(0.0, 1.0), Bounds::new(0.0, 1.0), Bounds::new(0.0, 1.0), Bounds::new(0.0, 1.0)]
    );
}

#[test]
fn test_binary_string_to_real_4D_2() {
    test_binary_string_to_real(
        "tests/Bin2Real_4D_2.txt",
        vec![Bounds::new(0.0, 7.0), Bounds::new(-4.0, 3.0), Bounds::new(-4.0, 4.0), Bounds::new(-8.0, 0.0)]
    );
}

#[test]
fn test_binary_string_to_real_6D_1() {
    test_binary_string_to_real(
        "tests/Bin2Real_6D_1.txt",
        vec![Bounds::new(0.0, 1.0), Bounds::new(0.0, 1.0), Bounds::new(0.0, 1.0), Bounds::new(0.0, 1.0), Bounds::new(0.0, 1.0), Bounds::new(0.0, 1.0)]
    );
}

#[test]
fn test_binary_string_to_real_6D_2() {
    test_binary_string_to_real(
        "tests/Bin2Real_6D_2.txt",
        (0..=5).map(|x| Bounds::new(x as f64, (2 * (x + 1)) as f64)).collect::<Vec<_>>()
    );
}

#[test]
fn test_one_max() {
    let data = load_test_file::<i8, i32>("tests/onemax.txt");

    for test in data {
        assert_eq!(
            OneMax::new().fit(&BinaryString::new(test.inp)).unwrap(),
            test.out
        );
    }
}

#[test]
fn test_LABS() {
    let data = load_test_file::<i8, i32>("tests/labs.txt");

    for test in data {
        println!("Test vector {}", test.inp.iter()
            .map(|x| x.to_string())
            .collect::<Vec<String>>()
            .join(", "));
        assert_eq!(
            LABS::new().fit(&BinaryString::new(test.inp)).unwrap(),
            test.out
        )
     }
}

#[test]
fn test_sphere() {
    let data = load_test_file::<f64, f64>("tests/sphere.txt");

    for test in data {
        assert_eq!(
            Sphere::new(vec![1.0; 10]).fit(&test.inp).unwrap(),
            test.out
        )
    }
}

#[test]
fn test_rosenbrock() {
    let data = load_test_file::<f64, f64>("tests/rosenbrock.txt");

    for test in data {
        assert_eq!(
            Rosenbrock::new().fit(&test.inp).unwrap(),
            test.out
        )
    }
}

// Test infra
struct TestVector<TIn, TOut> {
    inp: Vec<TIn>,
    out: TOut
}

fn load_test_file<TIn, TOut>(file: &str) -> Vec<TestVector::<TIn, TOut>>
where
  TIn : FromStr + Debug,
  TIn::Err: Debug,
  TOut : FromStr + Debug,
  TOut::Err: Debug,
{
    let mut vectors: Vec<TestVector::<TIn, TOut>> = vec![];

    let file = File::open(file).expect("Could not read test data!");
    let reader = BufReader::new(file);

    for (_, line) in reader.lines().enumerate() {
        let line = line.expect("Could not read a line!");

        if line.starts_with('#') || line.len() == 0 {
            continue;
        }

        let (inp_str, out_str) = line.split_once(":").unwrap();

        let out: TOut = out_str.trim().parse::<TOut>().unwrap();
        let inp: Vec<TIn> = inp_str.split(' ')
            .filter(|num| num.len() > 0)
            .map(|num| num.trim().parse().unwrap())
            .collect();

        vectors.push(TestVector::<TIn, TOut> {
            inp,
            out
        });
    }

    vectors
}

fn main() {
    println!("Hello, world! {}", add(1, 2));
}