/*
Functions relating to encoding optimisation.
*/

use crate::encoding::Encode;
use crate::encoding::MajoranaEncoding;

use crate::operators::{CoefficientPauliWeight, MajoranaSparse, PauliWeight};
use argmin::{
    core::{CostFunction, Error, Executor},
    solver::simulatedannealing::{Anneal, SATempFunc, SimulatedAnnealing},
};
use ndarray::ArrayView1;
use numpy::ndarray::Array1;
use rand::{distr::Uniform, prelude::*};
use rand_xoshiro::Xoshiro256PlusPlus;
use std::sync::{Arc, Mutex};

struct OptimalEnumeration {
    msparse: MajoranaSparse,
    encoding: MajoranaEncoding,
    coefficient_weighted: bool,
    rng: Arc<Mutex<Xoshiro256PlusPlus>>,
}

impl OptimalEnumeration {
    fn new(
        msparse: MajoranaSparse,
        encoding: MajoranaEncoding,
        coefficient_weighted: bool,
    ) -> Self {
        OptimalEnumeration {
            msparse,
            encoding,
            coefficient_weighted,
            rng: Arc::new(Mutex::new(Xoshiro256PlusPlus::seed_from_u64(1017))),
        }
    }
}

impl CostFunction for OptimalEnumeration {
    type Param = Array1<usize>;
    type Output = f64;

    fn cost(&self, param: &Self::Param) -> Result<Self::Output, Error> {
        let enumerated_encoding = self.encoding.apply_mode_enumeration(param.to_vec());
        let qham = enumerated_encoding.encode(&self.msparse);
        let weight = match self.coefficient_weighted {
            true => qham.coeff_pauli_weight(),
            false => qham.pauli_weight() as f64,
        };
        Ok(weight)
    }
}

impl Anneal for OptimalEnumeration {
    type Param = Array1<usize>;
    type Output = Array1<usize>;
    type Float = f64;

    fn anneal(&self, param: &Array1<usize>, temp: f64) -> Result<Array1<usize>, Error> {
        let mut next_perm = param.clone();
        let n_modes = next_perm.len();
        let mut rng = self.rng.lock().unwrap();
        let distr = Uniform::try_from(0..n_modes).unwrap();
        let temp_int = temp.floor() as u64 + 1;

        for _ in 0..temp_int {
            let pos: usize = rng.sample(distr);
            let move_distance = rng.random_range(0..temp_int) as usize % n_modes;
            let pos2: usize = if rng.random_bool(0.5) {
                (pos + move_distance) % n_modes
            } else {
                (pos + n_modes - move_distance) % n_modes
            };
            let swap_val = next_perm[[pos]];
            next_perm[[pos]] = next_perm[[pos2]];
            next_perm[[pos2]] = swap_val;
        }
        Ok(next_perm)
    }
}

pub fn anneal_enumerations(
    msparse: MajoranaSparse,
    encoding: MajoranaEncoding,
    temperature: f64,
    initial_guess: ArrayView1<usize>,
    coefficient_weighted: bool,
) -> Result<(f64, Array1<usize>), Error> {
    let operator = OptimalEnumeration::new(msparse, encoding, coefficient_weighted);

    // Define initial parameter vector

    // Set up simulated annealing solver
    // An alternative random number generator (RNG) can be provided to `new_with_rng`:
    // SimulatedAnnealing::new_with_rng(temp, Xoshiro256PlusPlus::from_entropy())?
    let solver = SimulatedAnnealing::new(temperature)?
        // Optional: Define temperature function (defaults to `SATempFunc::TemperatureFast`)
        .with_temp_func(SATempFunc::Boltzmann)
        /////////////////////////
        // Stopping criteria   //
        /////////////////////////
        // Optional: stop if there was no new best solution after 1000 iterations
        .with_stall_best(250);
    // Optional: stop if there was no accepted solution after 1000 iterations
    // .with_stall_accepted(1000);
    /////////////////////////
    // Reannealing         //
    /////////////////////////
    // Optional: Reanneal after 1000 iterations (resets temperature to initial temperature)
    // .with_reannealing_fixed(1000)
    // Optional: Reanneal after no accepted solution has been found for `iter` iterations
    // .with_reannealing_accepted(500)
    // Optional: Start reannealing after no new best solution has been found for 800 iterations
    // .with_reannealing_best(800);

    /////////////////////////
    // Run solver          //
    /////////////////////////
    let res = Executor::new(operator, solver)
        .configure(|state| {
            state
                .param(initial_guess.to_owned())
                // Optional: Set maximum number of iterations (defaults to `std::u64::MAX`)
                .max_iters(1_000)
            // Optional: Set target cost function value (defaults to `std::f64::NEG_INFINITY`)
            // .target_cost(0.0)
        })
        // Optional: Attach a observer
        // .add_observer(SlogLogger::term(), ObserverMode::Never)
        .run()?;

    let final_state = res.state();
    let best_permutation = final_state
        .best_param
        .clone()
        .expect("No best param in final anneling state.");
    Ok((final_state.best_cost, best_permutation))
}
