#pragma once

#include <fesa/core/ids.hpp>

#include <vector>

namespace fesa::analysis {

struct IterationState {
    double time = 0.0;
    int increment = 0;
    int iteration = 0;
};

class AnalysisState {
public:
    explicit AnalysisState(int total_dof_count);

    const std::vector<double>& displacement() const;
    const std::vector<double>& velocity() const;
    const std::vector<double>& acceleration() const;
    const std::vector<double>& temperature() const;
    const std::vector<double>& external_force() const;
    const std::vector<double>& internal_force() const;
    const std::vector<double>& residual() const;

    void set_displacement(std::vector<double> values);
    void set_external_force(std::vector<double> values);
    void set_internal_force(std::vector<double> values);
    void update_residual();

    IterationState& iteration_state();
    const IterationState& iteration_state() const;

    void set_element_state(core::ElementId element_id, std::vector<double> state);
    const std::vector<double>* element_state(core::ElementId element_id) const;

private:
    static std::vector<double> vector_of(int size);
    static void assign_same_size(std::vector<double>& target, std::vector<double> values);

    std::vector<double> displacement_;
    std::vector<double> velocity_;
    std::vector<double> acceleration_;
    std::vector<double> temperature_;
    std::vector<double> external_force_;
    std::vector<double> internal_force_;
    std::vector<double> residual_;
    IterationState iteration_state_;
    std::vector<std::pair<core::ElementId, std::vector<double>>> element_states_;
};

} // namespace fesa::analysis