use super::c64;
use super::time_evolution::State;
use ndarray::prelude::*;

pub struct TwoLevel {
    state: Array1<c64>,
}

impl TwoLevel {
    pub fn new(alpha: c64, beta: c64) -> TwoLevel {
        TwoLevel {
            state: array![alpha, beta],
        }
    }
}

impl State for TwoLevel {
    fn fibrate(&self) -> Vec<f64> {
        self.state
            .iter()
            .flat_map(|c| [c.re, c.im])
            .map(|v| *v)
            .collect()
    }

    fn evolve<G: super::time_evolution::MatrixGen>(mut self, t: f64) -> Self {
        self.state = self.state.dot(&G::gen(t));
        self
    }
}