optionstratlib_simulation/simulation/

randomwalk.rs

/******************************************************************************
   Author: Joaquín Béjar García
   Email: jb@taunais.com
   Date: 23/3/25
******************************************************************************/

use crate::simulation::WalkParams;
use crate::simulation::steps::Step;
use optionstratlib_core::model::Positive;
use optionstratlib_core::utils::Len;
use optionstratlib_pricing::error::PricingError;
use optionstratlib_pricing::pricing::Profit;
use rust_decimal::Decimal;
use std::fmt::Display;
use std::ops::{AddAssign, Index, IndexMut};

/// A struct that represents a two-dimensional random walk simulation.
///
/// `RandomWalk` stores a sequence of steps that describe a path in a two-dimensional space,
/// typically used for financial modeling, time series analysis, or statistical simulations.
/// It maintains both the steps of the random walk and a descriptive title.
///
/// # Type Parameters
///
/// * `X` - The type for x-axis values (typically representing time or sequence position),
///   which must implement `AddAssign` (allowing values to be accumulated), be convertible
///   to `Positive`, and be `Copy`.
///
/// * `Y` - The type for y-axis values (typically representing price, value, or position),
///   which must implement `AddAssign`, be convertible to `Positive`, be `Copy`, and implement
///   the `Walktypable` trait for additional functionality.
///
#[derive(Debug, Clone, Default)]
pub struct RandomWalk<X, Y>
where
    X: Copy + TryInto<Positive> + AddAssign + Display,
    Y: TryInto<Positive> + Display + Clone,
{
    /// The descriptive title of the random walk
    title: String,

    /// The collection of steps that make up the random walk path
    steps: Vec<Step<X, Y>>,
}

impl<X, Y> RandomWalk<X, Y>
where
    X: Copy + TryInto<Positive> + AddAssign + Display,
    Y: TryInto<Positive> + Display + Clone,
{
    /// Creates a new random walk instance with the given title and steps.
    ///
    /// This constructor takes a title, walk parameters, and a fallible
    /// generator function that produces the steps of the random walk
    /// from the provided parameters. Errors from the generator are
    /// propagated unchanged.
    ///
    /// # Parameters
    ///
    /// * `title` - A descriptive title for the random walk
    /// * `params` - Parameters that define the properties of the random walk
    /// * `generator` - A fallible function that generates the steps of
    ///   the random walk
    ///
    /// # Returns
    ///
    /// `Ok(RandomWalk)` on success, or `Err(E)` when the generator fails.
    ///
    /// # Errors
    ///
    /// Returns the error type produced by the supplied generator. For
    /// chain-backed generators (e.g. [`crate::simulation::generator_positive`])
    /// this is [`crate::error::SimulationError`].
    pub fn new<F, E>(title: String, params: &WalkParams<X, Y>, generator: F) -> Result<Self, E>
    where
        F: FnOnce(&WalkParams<X, Y>) -> Result<Vec<Step<X, Y>>, E>,
        X: Copy + TryInto<Positive> + AddAssign + Display,
        Y: TryInto<Positive> + Display + Clone,
    {
        let steps = generator(params)?;
        Ok(Self { title, steps })
    }

    /// Returns the title of the random walk.
    ///
    /// # Returns
    ///
    /// A string slice containing the title of the random walk.
    #[must_use]
    pub fn get_title(&self) -> &str {
        &self.title
    }

    /// Updates the title of the random walk.
    ///
    /// # Parameters
    ///
    /// * `title` - The new title to set
    pub fn set_title(&mut self, title: String) {
        self.title = title;
    }

    /// Returns a vector of references to all steps in the random walk.
    ///
    /// # Returns
    ///
    /// A vector containing references to all steps in the random walk.
    #[must_use]
    pub fn get_steps(&self) -> Vec<&Step<X, Y>> {
        self.steps.iter().collect::<Vec<&Step<X, Y>>>()
    }

    /// Returns a reference to the step at the specified index.
    ///
    /// # Parameters
    ///
    /// * `index` - The zero-based index of the step to retrieve
    ///
    /// # Returns
    ///
    /// * `Some(&Step<X, Y>)` - The step at `index`.
    /// * `None` - If `index` is out of bounds.
    ///
    /// Named `get_*`, so it follows [`slice::get`] rather than
    /// [`std::ops::Index`]: an out-of-bounds index is reported, not aborted
    /// on. Use the `walk[index]` operator when the index is an invariant of
    /// the caller and a panic is the wanted contract.
    #[must_use]
    pub fn get_step(&self, index: usize) -> Option<&Step<X, Y>> {
        self.steps.get(index)
    }

    /// Returns a mutable reference to the step at the specified index.
    ///
    /// # Parameters
    ///
    /// * `index` - The zero-based index of the step to retrieve
    ///
    /// # Returns
    ///
    /// * `Some(&mut Step<X, Y>)` - The step at `index`.
    /// * `None` - If `index` is out of bounds.
    ///
    /// Mutable counterpart of [`RandomWalk::get_step`], with the same
    /// out-of-bounds contract.
    pub fn get_step_mut(&mut self, index: usize) -> Option<&mut Step<X, Y>> {
        self.steps.get_mut(index)
    }

    /// Returns a reference to the first step in the random walk, if any.
    ///
    /// # Returns
    ///
    /// * `Some(&Step<X, Y>)` - A reference to the first step if the random walk is not empty
    /// * `None` - If the random walk has no steps
    #[must_use]
    pub fn first(&self) -> Option<&Step<X, Y>> {
        self.steps.first()
    }

    /// Returns a reference to the last step in the random walk, if any.
    ///
    /// # Returns
    ///
    /// * `Some(&Step<X, Y>)` - A reference to the last step if the random walk is not empty
    /// * `None` - If the random walk has no steps
    #[must_use]
    pub fn last(&self) -> Option<&Step<X, Y>> {
        self.steps.last()
    }
}

/// Implementation of the `Len` trait for `RandomWalk<X, Y>`.
///
/// This implementation provides methods to determine the length and emptiness
/// of a random walk by delegating to the underlying `steps` collection.
///
/// # Type Parameters
///
/// * `X` - The type for x-axis values (typically time or sequence position),
///   which must implement `AddAssign`, be convertible to `Positive`, and be `Copy`.
///
/// * `Y` - The type for y-axis values (typically price or value),
///   which must implement `AddAssign`, be convertible to `Positive`, be `Copy`,
///   and implement the `Walktypable` trait.
impl<X, Y> Len for RandomWalk<X, Y>
where
    X: Copy + TryInto<Positive> + AddAssign + Display,
    Y: TryInto<Positive> + Display + Clone,
{
    /// Returns the number of steps in the random walk.
    ///
    /// # Returns
    ///
    /// A `usize` representing the number of steps.
    fn len(&self) -> usize {
        self.steps.len()
    }

    /// Determines whether the random walk contains any steps.
    ///
    /// # Returns
    ///
    /// `true` if the random walk has no steps, `false` otherwise.
    fn is_empty(&self) -> bool {
        self.steps.is_empty()
    }
}

/// Implementation of the `Index` trait for `RandomWalk<X, Y>`.
///
/// This allows accessing the steps of a random walk using array indexing notation:
/// `walk[index]`.
///
/// # Type Parameters
///
/// * `X` - The type for x-axis values, with constraints as described above.
/// * `Y` - The type for y-axis values, with constraints as described above.
impl<X, Y> Index<usize> for RandomWalk<X, Y>
where
    X: Copy + TryInto<Positive> + AddAssign + Display,
    Y: TryInto<Positive> + Display + Clone,
{
    /// The type returned when indexing the random walk.
    type Output = Step<X, Y>;

    /// Provides read access to a specific step in the random walk by index.
    ///
    /// # Parameters
    ///
    /// * `index` - The zero-based index of the step to access.
    ///
    /// # Returns
    ///
    /// A reference to the `Step<X, Y>` at the specified index.
    ///
    /// # Panics
    ///
    /// Panics if the index is out of bounds, as indexing a `Vec` does. Use
    /// [`RandomWalk::get_step`] for the checked form, which returns `None`
    /// instead.
    #[allow(clippy::indexing_slicing)]
    fn index(&self, index: usize) -> &Self::Output {
        // `Index` / `IndexMut` follow the std contract: an out-of-bounds index
        // is a caller bug and panics, exactly like `Vec`. The checked
        // accessor is `get_step`, which returns `None` instead (#341, #788).
        &self.steps[index]
    }
}

/// Implementation of the `IndexMut` trait for `RandomWalk<X, Y>`.
///
/// This allows modifying steps in a random walk using array indexing notation:
/// `walk[index] = new_step`.
///
/// # Type Parameters
///
/// * `X` - The type for x-axis values, with constraints as described above.
/// * `Y` - The type for y-axis values, with constraints as described above.
impl<X, Y> IndexMut<usize> for RandomWalk<X, Y>
where
    X: Copy + TryInto<Positive> + AddAssign + Display,
    Y: TryInto<Positive> + Display + Clone,
{
    /// Provides mutable access to a specific step in the random walk by index.
    ///
    /// # Parameters
    ///
    /// * `index` - The zero-based index of the step to modify.
    ///
    /// # Returns
    ///
    /// A mutable reference to the `Step<X, Y>` at the specified index.
    ///
    /// # Panics
    ///
    /// Panics if the index is out of bounds, as indexing a `Vec` does. Use
    /// [`RandomWalk::get_step_mut`] for the checked form, which returns `None`
    /// instead.
    #[allow(clippy::indexing_slicing)]
    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
        // `Index` / `IndexMut` follow the std contract: an out-of-bounds index
        // is a caller bug and panics, exactly like `Vec`. The checked
        // accessor is `get_step_mut`, which returns `None` instead (#341, #788).
        &mut self.steps[index]
    }
}

impl<X, Y> Display for RandomWalk<X, Y>
where
    X: Copy + TryInto<Positive> + AddAssign + Display,
    Y: TryInto<Positive> + Display + Clone,
{
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        writeln!(f, "RandomWalk Title: {}, Steps:  ", self.title)?;
        for step in &self.steps {
            write!(f, "\t{step}")?;
        }
        Ok(())
    }
}

impl<X, Y> Profit for RandomWalk<X, Y>
where
    X: AddAssign + Copy + Display + TryInto<Positive>,
    Y: TryInto<Positive> + Display + Clone,
{
    fn calculate_profit_at(&self, _price: &Positive) -> Result<Decimal, PricingError> {
        Err(PricingError::other(
            "Profit calculation not implemented for RandomWalk",
        ))
    }
}

#[cfg(test)]
#[allow(irrefutable_let_patterns)]
mod tests_random_walk {
    use super::*;
    use optionstratlib_core::model::ExpirationDate;

    use crate::error::SimulationError;

    use crate::simulation::WalkParams;
    use crate::simulation::WalkType;
    use crate::simulation::WalkTypeAble;
    use crate::simulation::steps::Step;
    use optionstratlib_core::utils::TimeFrame;

    use num_traits::ToPrimitive;
    use rust_decimal::Decimal;

    use optionstratlib_core::pos_or_panic;
    use std::convert::Infallible;
    use std::fmt::Display;
    use std::ops::AddAssign;

    // Mock implementation of WalkTypeAble for testing
    #[derive(Clone)]
    struct TestWalker {}

    impl<X, Y> WalkTypeAble<X, Y> for TestWalker
    where
        X: Copy + TryInto<Positive> + AddAssign + Display,
        Y: TryInto<Positive> + Display + Clone,
    {
        // We'll implement the simplest possible method for testing
        fn brownian(&self, params: &WalkParams<X, Y>) -> Result<Vec<Positive>, SimulationError> {
            let mut values = Vec::new();
            let init_value: Positive = params.ystep_as_positive()?;
            values.push(init_value);

            // Generate some simple steps for test purposes
            for i in 1..params.size {
                values.push(pos_or_panic!(
                    init_value.value().to_f64().unwrap() + i as f64
                ));
            }

            Ok(values)
        }
    }

    // Helper function to create a walk parameters struct for testing
    fn create_test_params<X, Y>(
        size: usize,
        x_value: X,
        y_value: Y,
        walk_type: WalkType,
    ) -> WalkParams<X, Y>
    where
        X: Copy + TryInto<Positive> + AddAssign + Display,
        Y: TryInto<Positive> + Display + Clone,
    {
        let init_step = Step::new(
            x_value,
            TimeFrame::Day,
            ExpirationDate::Days(pos_or_panic!(30.0)),
            y_value,
        );

        WalkParams {
            size,
            init_step,
            walk_type,
            walker: Box::new(TestWalker {}),
            seed: None,
        }
    }

    // Helper function to generate test steps for a random walk
    fn generate_test_steps<X, Y>(params: &WalkParams<X, Y>) -> Result<Vec<Step<X, Y>>, Infallible>
    where
        X: Copy + TryInto<Positive> + AddAssign + Display,
        Y: TryInto<Positive> + Display + Clone,
    {
        let mut steps = Vec::new();
        steps.push(params.init_step.clone());

        let test_walker = TestWalker {};
        let values = test_walker
            .brownian(params)
            .expect("test brownian generator");

        let mut current_step = params.init_step.clone();

        // Skip the first value as it's the initial step
        for _value in values.iter().skip(1) {
            // Convert Positive back to Y type (for test we'll just use the same value)
            let new_y_value = current_step.y.value();

            // Create next step
            let next_step = current_step
                .next(new_y_value.clone())
                .expect("test step.next");
            steps.push(next_step.clone());

            current_step = next_step;
        }

        Ok(steps)
    }

    #[test]
    fn test_random_walk_creation() {
        let params = create_test_params(
            5,
            1.0,
            100.0,
            WalkType::Brownian {
                dt: Positive::ONE,
                drift: Decimal::ZERO,
                volatility: pos_or_panic!(0.2),
            },
        );

        let title = "Test Random Walk".to_string();
        let Ok(walk) = RandomWalk::new(title.clone(), &params, generate_test_steps) else {
            unreachable!()
        };

        assert_eq!(walk.get_title(), title);
        assert_eq!(walk.len(), 5);
        assert!(!walk.is_empty());
    }

    #[test]
    fn test_random_walk_empty() {
        let params = create_test_params(
            0,
            1.0,
            100.0,
            WalkType::Brownian {
                dt: Positive::ONE,
                drift: Decimal::ZERO,
                volatility: pos_or_panic!(0.2),
            },
        );

        let title = "Empty Walk".to_string();
        let Ok(walk) = RandomWalk::new(title.clone(), &params, |_| Ok::<_, Infallible>(Vec::new()))
        else {
            unreachable!()
        };

        assert_eq!(walk.get_title(), title);
        assert_eq!(walk.len(), 0);
        assert!(walk.is_empty());
        assert!(walk.first().is_none());
        assert!(walk.last().is_none());
    }

    #[test]
    fn test_random_walk_title_update() {
        let params = create_test_params(
            5,
            1.0,
            100.0,
            WalkType::Brownian {
                dt: Positive::ONE,
                drift: Decimal::ZERO,
                volatility: pos_or_panic!(0.2),
            },
        );

        let title = "Initial Title".to_string();
        let Ok(mut walk) = RandomWalk::new(title, &params, generate_test_steps) else {
            unreachable!()
        };

        let new_title = "Updated Title".to_string();
        walk.set_title(new_title.clone());

        assert_eq!(walk.get_title(), new_title);
    }

    #[test]
    fn test_random_walk_first_last() {
        let params = create_test_params(
            5,
            1.0,
            100.0,
            WalkType::Brownian {
                dt: Positive::ONE,
                drift: Decimal::ZERO,
                volatility: pos_or_panic!(0.2),
            },
        );

        let Ok(walk) = RandomWalk::new("Test Walk".to_string(), &params, generate_test_steps)
        else {
            unreachable!()
        };

        let first = walk.first().unwrap();
        let last = walk.last().unwrap();

        assert_eq!(*first.x.index(), 0);
        assert_eq!(*last.x.index(), 4); // 5 steps, zero-indexed
    }

    #[test]
    fn test_random_walk_get_steps() {
        let params = create_test_params(
            5,
            1.0,
            100.0,
            WalkType::Brownian {
                dt: Positive::ONE,
                drift: Decimal::ZERO,
                volatility: pos_or_panic!(0.2),
            },
        );

        let Ok(walk) = RandomWalk::new("Test Walk".to_string(), &params, generate_test_steps)
        else {
            unreachable!()
        };

        let steps = walk.get_steps();
        assert_eq!(steps.len(), 5);

        // Test that steps have sequential index values
        for (i, step) in steps.iter().enumerate() {
            assert_eq!(*step.x.index(), i as i32);
        }
    }

    #[test]
    fn test_random_walk_get_step() {
        let params = create_test_params(
            5,
            1.0,
            100.0,
            WalkType::Brownian {
                dt: Positive::ONE,
                drift: Decimal::ZERO,
                volatility: pos_or_panic!(0.2),
            },
        );

        let Ok(walk) = RandomWalk::new("Test Walk".to_string(), &params, generate_test_steps)
        else {
            unreachable!()
        };

        let step_0 = walk.get_step(0).unwrap();
        let step_3 = walk.get_step(3).unwrap();

        assert_eq!(*step_0.x.index(), 0);
        assert_eq!(*step_3.x.index(), 3);
    }

    /// An out-of-bounds index is reported, not aborted on. `get_step` follows
    /// `slice::get`; the panicking contract lives on the `Index` operator,
    /// which `test_random_walk_index_operator` covers.
    #[test]
    fn test_random_walk_get_step_out_of_bounds() {
        let params = create_test_params(
            5,
            1.0,
            100.0,
            WalkType::Brownian {
                dt: Positive::ONE,
                drift: Decimal::ZERO,
                volatility: pos_or_panic!(0.2),
            },
        );

        let Ok(mut walk) = RandomWalk::new("Test Walk".to_string(), &params, generate_test_steps)
        else {
            unreachable!()
        };

        assert!(walk.get_step(10).is_none());
        assert!(walk.get_step_mut(10).is_none());
    }

    #[test]
    fn test_random_walk_get_step_mut() {
        let params = create_test_params(
            5,
            1.0,
            100.0,
            WalkType::Brownian {
                dt: Positive::ONE,
                drift: Decimal::ZERO,
                volatility: pos_or_panic!(0.2),
            },
        );

        let Ok(mut walk) = RandomWalk::new("Test Walk".to_string(), &params, generate_test_steps)
        else {
            unreachable!()
        };

        // Get a mutable reference and verify initial state
        let step_2 = walk.get_step_mut(2).unwrap();
        assert_eq!(*step_2.x.index(), 2);

        // Get a new step by calling next on the current step
        let new_y_value = *step_2.y.value();
        let new_step = step_2.clone();
        *step_2 = new_step.next(new_y_value * 2.0).unwrap();

        // Verify the step was updated
        assert_eq!(*walk.get_step(2).unwrap().x.index(), 3);
    }

    #[test]
    fn test_random_walk_index_operator() {
        let params = create_test_params(
            5,
            1.0,
            100.0,
            WalkType::Brownian {
                dt: Positive::ONE,
                drift: Decimal::ZERO,
                volatility: pos_or_panic!(0.2),
            },
        );

        let Ok(walk) = RandomWalk::new("Test Walk".to_string(), &params, generate_test_steps)
        else {
            unreachable!()
        };

        // Test read access via index operator
        let step_1 = &walk[1];
        assert_eq!(*step_1.x.index(), 1);

        // Test comparison between get_step and index operator
        assert_eq!(*walk.get_step(3).unwrap().x.index(), *walk[3].x.index());
    }

    #[test]
    fn test_random_walk_index_mut_operator() {
        let params = create_test_params(
            5,
            1.0,
            100.0,
            WalkType::Brownian {
                dt: Positive::ONE,
                drift: Decimal::ZERO,
                volatility: pos_or_panic!(0.2),
            },
        );

        let Ok(mut walk) = RandomWalk::new("Test Walk".to_string(), &params, generate_test_steps)
        else {
            unreachable!()
        };

        // Get initial step via index
        let initial_index = *walk[2].x.index();

        // Modify step via index_mut operator
        let new_y_value = *walk[2].y.value();
        let new_step = walk[2].clone();
        walk[2] = new_step.next(new_y_value * 2.0).unwrap();

        // Verify the change
        assert_ne!(*walk[2].x.index(), initial_index);
    }

    #[test]
    fn test_random_walk_new_propagates_generator_error() {
        // Regression for #349: ensure the fallible generator's error is
        // surfaced unchanged to the caller.
        #[derive(Debug, PartialEq)]
        struct FakeError(&'static str);

        let params = create_test_params(
            3,
            1.0_f64,
            100.0_f64,
            WalkType::Brownian {
                dt: Positive::ONE,
                drift: Decimal::ZERO,
                volatility: pos_or_panic!(0.2),
            },
        );

        let result: Result<RandomWalk<f64, f64>, FakeError> =
            RandomWalk::new("err".to_string(), &params, |_| Err(FakeError("boom")));

        match result {
            Err(FakeError(msg)) => assert_eq!(msg, "boom"),
            Ok(_) => panic!("expected generator error to propagate"),
        }
    }

    #[test]
    fn test_random_walk_display() {
        let params = create_test_params(
            3,
            1.0,
            100.0,
            WalkType::Brownian {
                dt: Positive::ONE,
                drift: Decimal::ZERO,
                volatility: pos_or_panic!(0.2),
            },
        );

        let Ok(walk) = RandomWalk::new("Display Test".to_string(), &params, generate_test_steps)
        else {
            unreachable!()
        };

        // Test that the display output contains the title
        let display_output = format!("{walk}");
        assert!(display_output.contains("Display Test"));
    }

    #[test]
    fn test_random_walk_graph_implementation() {
        let params = create_test_params(
            5,
            1.0,
            100.0,
            WalkType::Brownian {
                dt: Positive::ONE,
                drift: Decimal::ZERO,
                volatility: pos_or_panic!(0.2),
            },
        );

        let Ok(walk) = RandomWalk::new("Graph Test".to_string(), &params, generate_test_steps)
        else {
            unreachable!()
        };

        // The title the visualization layer charts the walk under.
        assert_eq!(walk.get_title(), "Graph Test");
    }

    #[test]
    fn test_with_different_types() {
        // Test with custom types for X and Y
        #[derive(Debug, Copy, Clone, PartialEq)]
        struct TestX(f64);

        impl Display for TestX {
            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
                write!(f, "{}", self.0)
            }
        }

        impl AddAssign for TestX {
            fn add_assign(&mut self, other: Self) {
                self.0 += other.0;
            }
        }

        impl From<TestX> for Positive {
            fn from(val: TestX) -> Self {
                pos_or_panic!(val.0)
            }
        }

        #[derive(Debug, Copy, Clone, PartialEq)]
        struct TestY(f64);

        impl Display for TestY {
            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
                write!(f, "{}", self.0)
            }
        }

        impl From<TestY> for Positive {
            fn from(val: TestY) -> Self {
                pos_or_panic!(val.0)
            }
        }

        // Create params with custom types
        let params = create_test_params(
            3,
            TestX(1.0),
            TestY(100.0),
            WalkType::Brownian {
                dt: Positive::ONE,
                drift: Decimal::ZERO,
                volatility: pos_or_panic!(0.2),
            },
        );

        // Custom generator for TestX and TestY
        let generator = |params: &WalkParams<TestX, TestY>| -> Result<_, Infallible> {
            let mut steps = Vec::new();
            steps.push(params.init_step.clone());

            let mut current_step = params.init_step.clone();
            for i in 1..params.size {
                let next_step = current_step
                    .next(TestY((100.0 + i as f64) * 1.1))
                    .expect("test step.next");
                steps.push(next_step.clone());
                current_step = next_step;
            }

            Ok(steps)
        };

        let Ok(walk) = RandomWalk::new("Custom Types Test".to_string(), &params, generator) else {
            unreachable!()
        };

        assert_eq!(walk.len(), 3);
        assert_eq!(*walk[0].y.value(), TestY(100.0));
    }
}