//! Simulation configuration parameters.

use anyhow::{Context, Result, bail};
use serde::{Deserialize, Serialize};
use std::{fmt::Debug, fs, ops::RangeBounds, path::Path};

/// Simulation configuration parameters.
#[derive(Debug, PartialEq, Clone, Serialize, Deserialize)]
pub struct Config {
    /// Stochastic agent-based model parameters.
    pub model: ModelParams,
    /// State initialization parameters.
    pub init: InitParams,
    /// Output format parameters.
    pub output: OutputParams,
}

/// Stochastic agent-based model parameters.
#[derive(Debug, PartialEq, Clone, Serialize, Deserialize)]
pub struct ModelParams {
    /// Number of environments.
    pub n_env: usize,
    /// Number of phenotypes.
    pub n_phe: usize,

    /// Environment transition rates (matrix `n_env x n_env`).
    pub rates_trans: Vec<Vec<f64>>,
    /// Agent birth rates (matrix `n_env x n_phe`).
    pub rates_birth: Vec<Vec<f64>>,
    /// Agent death rates (matrix `n_env x n_phe`).
    pub rates_death: Vec<Vec<f64>>,

    /// Mutation probability.
    pub prob_mut: f64,
}

/// State initialization parameters.
#[derive(Debug, PartialEq, Clone, Serialize, Deserialize)]
pub struct InitParams {
    /// Number of agents.
    pub n_agents: usize,

    /// Phenotypic strategy.
    pub strat_phe: Option<Vec<f64>>,
}

/// Output format parameters.
#[derive(Debug, PartialEq, Clone, Serialize, Deserialize)]
pub struct OutputParams {
    /// Number of steps per output file.
    pub steps_per_file: usize,

    /// Number of steps per saved observables.
    pub steps_per_save: usize,

    /// Number of histogram bins.
    pub hist_bins: usize,
}

impl Config {
    /// Load a `Config` from a TOML file.
    ///
    /// Performs validation on all parameters before returning.
    pub fn from_file<P: AsRef<Path>>(file: P) -> Result<Self> {
        let file = file.as_ref();
        let file = fs::read_to_string(file).with_context(|| format!("failed to read {file:?}"))?;

        let config: Config = toml::from_str(&file).context("failed to deserialize config")?;

        config.validate().context("failed to validate config")?;

        Ok(config)
    }

    fn validate(&self) -> Result<()> {
        let model = &self.model;
        let init = &self.init;
        let output = &self.output;

        check_num(model.n_env, 1..=16).context("invalid number of environments")?;
        check_num(model.n_phe, 1..=16).context("invalid number of phenotypes")?;

        check_mat(&model.rates_trans, (model.n_env, model.n_env))
            .context("invalid environment transition rates")?;
        check_mat(&model.rates_birth, (model.n_env, model.n_phe))
            .context("invalid agent birth rates")?;
        check_mat(&model.rates_death, (model.n_env, model.n_phe))
            .context("invalid agent death rates")?;

        check_num(model.prob_mut, 0.0..=1.0).context("invalid mutation probability")?;

        check_num(init.n_agents, 1..=16_384).context("invalid number of agents")?;

        if let Some(strat_phe) = &init.strat_phe {
            check_vec(strat_phe, model.n_phe).context("invalid phenotypic strategy")?;
        }

        check_num(output.steps_per_file, 1_024..)
            .context("invalid number of steps per output file")?;

        check_num(output.steps_per_save, 0..output.steps_per_file)
            .context("invalid number of steps per saved observables")?;

        check_num(output.hist_bins, 1..=16).context("invalid number of histogram bins")?;

        Ok(())
    }
}

fn check_num<T, R>(num: T, range: R) -> Result<()>
where
    T: PartialOrd + Debug,
    R: RangeBounds<T> + Debug,
{
    if !range.contains(&num) {
        bail!("number must be in the range {range:?}, but is {num:?}");
    }
    Ok(())
}

fn check_vec(vec: &[f64], exp_len: usize) -> Result<()> {
    let len = vec.len();
    if len != exp_len {
        bail!("vector length must be {exp_len}, but is {len}");
    }
    Ok(())
}

fn check_mat(mat: &[Vec<f64>], exp_shape: (usize, usize)) -> Result<()> {
    let exp_n_rows = exp_shape.0;
    let exp_n_cols = exp_shape.1;
    let n_rows = mat.len();
    if n_rows != exp_n_rows {
        bail!("matrix must have {exp_n_rows} rows, but has {n_rows}");
    }
    if mat.iter().any(|row| row.len() != exp_n_cols) {
        bail!("matrix must have {exp_n_cols} columns");
    }
    Ok(())
}
