// SPDX-FileCopyrightText: Contributors to the Power Grid Model project <powergridmodel@lfenergy.org>
//
// SPDX-License-Identifier: MPL-2.0

#pragma once

// main model class

// main include
#include "calculation_parameters.hpp"
#include "calculation_preparation.hpp"
#include "main_model_fwd.hpp"
#include "math_solver/y_bus.hpp"

// common
#include "common/common.hpp"
#include "common/enum.hpp"
#include "common/exception.hpp"
#include "common/logging.hpp"
#include "common/three_phase_tensor.hpp"
#include "common/timer.hpp"

// component include
#include "auxiliary/dataset.hpp"
#include "auxiliary/meta_data.hpp"
#include "auxiliary/update.hpp"
#include "component/component.hpp"

// math model include
#include "math_solver/math_solver_dispatch.hpp"

#include "optimizer/optimizer.hpp"

// main model implementation
#include "calculation.hpp"
#include "main_core/core_utils.hpp"
#include "main_core/input.hpp"
#include "main_core/main_model_type.hpp"
#include "main_core/output.hpp"
#include "main_core/topology.hpp"
#include "main_core/update.hpp"

// stl library
#include <algorithm>
#include <array>
#include <cassert>
#include <concepts>
#include <limits>
#include <memory>
#include <ranges>
#include <span>
#include <string_view>
#include <type_traits>
#include <utility>
#include <vector>

namespace power_grid_model {
// solver output type to output type getter meta function
template <solver_output_type SolverOutputType> struct output_type_getter;
template <short_circuit_solver_output_type SolverOutputType> struct output_type_getter<SolverOutputType> {
    using type = meta_data::sc_output_getter_s;
};
template <> struct output_type_getter<SolverOutput<symmetric_t>> {
    using type = meta_data::sym_output_getter_s;
};
template <> struct output_type_getter<SolverOutput<asymmetric_t>> {
    using type = meta_data::asym_output_getter_s;
};

template <class ModelType>
    requires(main_core::is_main_model_type_v<ModelType>)
class MainModelImpl {

  private:
    // internal type traits
    using MainModelState = typename ModelType::MainModelState;

    using SequenceIdx = typename ModelType::SequenceIdx;
    using SequenceIdxRefWrappers = ModelType::SequenceIdxRefWrappers;
    using SequenceIdxView = typename ModelType::SequenceIdxView;
    using OwnedUpdateDataset = typename ModelType::OwnedUpdateDataset;
    using ComponentFlags = typename ModelType::ComponentFlags;

  public:
    using ImplType = ModelType;
    using Options = MainModelOptions;
    using MetaData = meta_data::MetaData;

    // constructor with data
    explicit MainModelImpl(double system_frequency, ConstDataset const& input_data,
                           SolverPreparationContext solver_preparation_context, Idx pos = 0)
        : system_frequency_{system_frequency},
          meta_data_{&input_data.meta_data()},
          solver_preparation_context_{std::move(solver_preparation_context)} {
        assert(input_data.get_description().dataset->name == std::string_view("input"));
        add_components(input_data, pos);
        set_construction_complete();
    }

    // constructor with only frequency
    explicit MainModelImpl(double system_frequency, meta_data::MetaData const& meta_data,
                           SolverPreparationContext solver_preparation_context)
        : system_frequency_{system_frequency},
          meta_data_{&meta_data},
          solver_preparation_context_{std::move(solver_preparation_context)} {}

    // helper function to get what components are present in the update data
    ComponentFlags get_components_to_update(ConstDataset const& update_data) const {
        return ModelType::run_functor_with_all_component_types_return_array([&update_data]<typename CompType>() {
            return (update_data.find_component(CompType::name, false) != main_core::utils::invalid_index);
        });
    }

  private:
    // template to construct components
    // using forward interators
    // different selection based on component type
    template <std::derived_from<Base> CompType, std::ranges::viewable_range Inputs>
    void add_component(Inputs&& components) {
        assert(!construction_complete_);
        main_core::add_component<CompType>(state_.components, std::forward<Inputs>(components), system_frequency_);
    }

    void add_components(ConstDataset const& input_data, Idx pos = 0) {
        auto const add_func = [this, pos, &input_data]<typename CT>() {
            if (input_data.is_columnar(CT::name)) {
                this->add_component<CT>(input_data.get_columnar_buffer_span<meta_data::input_getter_s, CT>(pos));
            } else {
                this->add_component<CT>(input_data.get_buffer_span<meta_data::input_getter_s, CT>(pos));
            }
        };
        ModelType::run_functor_with_all_component_types_return_void(add_func);
    }

    // template to update components
    // using forward interators
    // different selection based on component type
    // if sequence_idx is given, it will be used to load the object instead of using IDs via hash map.
    template <class CompType, cache_type_c CacheType, std::ranges::viewable_range Updates>
    void update_component(Updates&& updates, std::span<Idx2D const> sequence_idx) {
        constexpr auto comp_index = ModelType::template index_of_component<CompType>;

        assert(construction_complete_);
        assert(std::ranges::ssize(sequence_idx) == std::ranges::ssize(updates));

        if constexpr (CacheType::value) {
            main_core::update::update_inverse<CompType>(
                state_.components, updates, std::back_inserter(std::get<comp_index>(cached_inverse_update_)),
                sequence_idx);
        }

        UpdateChange const changed = main_core::update::update_component<CompType>(
            state_.components, std::forward<Updates>(updates),
            std::back_inserter(std::get<comp_index>(solvers_cache_status_.changed_components_indices())), sequence_idx);

        // update, get changed variable
        solvers_cache_status_.update(changed);
        if constexpr (CacheType::value) {
            cached_state_changes_ = cached_state_changes_ || changed;
        }
    }

    // entry point overload to update one row or column based component type
    template <class CompType, cache_type_c CacheType>
    void update_component_row_col(ConstDataset const& update_data, Idx pos, std::span<Idx2D const> sequence_idx) {
        assert(construction_complete_);
        assert(update_data.get_description().dataset->name == std::string_view("update"));

        if (update_data.is_columnar(CompType::name)) {
            this->update_component<CompType, CacheType>(
                update_data.get_columnar_buffer_span<meta_data::update_getter_s, CompType>(pos), sequence_idx);
        } else {
            this->update_component<CompType, CacheType>(
                update_data.get_buffer_span<meta_data::update_getter_s, CompType>(pos), sequence_idx);
        }
    }

  public:
    // overload to update all components across all scenarios
    template <cache_type_c CacheType, typename SequenceIdxMap>
        requires(std::same_as<SequenceIdxMap, SequenceIdx> || std::same_as<SequenceIdxMap, SequenceIdxView>)
    void update_components(ConstDataset const& update_data, Idx pos, SequenceIdxMap const& sequence_idx_map) {
        ModelType::run_functor_with_all_component_types_return_void(
            [this, pos, &update_data, &sequence_idx_map]<typename CT>() {
                this->update_component_row_col<CT, CacheType>(
                    update_data, pos, std::get<ModelType::template index_of_component<CT>>(sequence_idx_map));
            });
    }

    // overload to update all components in the first scenario (e.g. permanent update)
    template <cache_type_c CacheType> void update_components(ConstDataset const& update_data) {
        auto const components_to_update = get_components_to_update(update_data);
        auto const update_independence =
            main_core::update::independence::check_update_independence<ModelType>(state_.components, update_data);
        auto const sequence_idx_map = main_core::update::get_all_sequence_idx_map<ModelType>(
            state_.components, update_data, 0, components_to_update, update_independence, false);
        update_components<CacheType>(update_data, 0, sequence_idx_map);
    }

  private:
    // set complete construction
    // initialize internal arrays
    void set_construction_complete() {
        assert(!construction_complete_);
#ifndef NDEBUG
        // set construction_complete for debug assertions
        construction_complete_ = true;
#endif // !NDEBUG
        state_.components.set_construction_complete();
        state_.comp_topo =
            std::make_shared<ComponentTopology const>(main_core::construct_topology<ModelType>(state_.components));
    }

  public:
    /*
    the the sequence indexer given an input array of ID's for a given component type
    */
    void get_indexer(std::string_view component_type, ID const* id_begin, Idx size, Idx* indexer_begin) const {
        auto const get_index_func = [&components = this->state_.components, component_type, id_begin, size,
                                     indexer_begin]<typename CT>() {
            if (component_type == CT::name) {
                std::transform(id_begin, id_begin + size, indexer_begin, [&components](ID id) {
                    return main_core::get_component_idx_by_id<CT>(components, id).pos;
                });
            }
        };
        ModelType::run_functor_with_all_component_types_return_void(get_index_func);
    }

  private:
    // Entry point for main_model.hpp
    SequenceIdx get_all_sequence_idx_map(ConstDataset const& update_data) {
        auto const components_to_update = get_components_to_update(update_data);
        auto const update_independence =
            main_core::update::independence::check_update_independence<ModelType>(state_.components, update_data);
        return main_core::update::get_all_sequence_idx_map<ModelType>(state_.components, update_data, 0,
                                                                      components_to_update, update_independence, false);
    }

    template <typename CompType> void restore_component(SequenceIdxView const& sequence_idx) {
        constexpr auto component_index = ModelType::template index_of_component<CompType>;

        auto& cached_inverse_update = std::get<component_index>(cached_inverse_update_);
        auto const& component_sequence = std::get<component_index>(sequence_idx);

        if (!cached_inverse_update.empty()) {
            update_component<CompType, permanent_update_t>(cached_inverse_update, component_sequence);
            cached_inverse_update.clear();
        }
    }

  public:
    // restore the initial values of all components
    void restore_components(SequenceIdxView const& sequence_idx) {
        ModelType::run_functor_with_all_component_types_return_void(
            [this, &sequence_idx]<typename CompType>() { this->restore_component<CompType>(sequence_idx); });

        solvers_cache_status_.update(cached_state_changes_);
        // we can probably reset topo, solvers and ybus here
        cached_state_changes_ = {};
    }

  private:
    void restore_components(SequenceIdxRefWrappers const& sequence_idx) {
        ModelType::run_functor_with_all_component_types_return_void([this, &sequence_idx]<typename CompType>() {
            this->restore_component<CompType>(std::array{std::span<Idx2D const>{
                std::get<ModelType::template index_of_component<CompType>>(sequence_idx).get()}});
        });
    }

    void restore_components(ModelType::SequenceIdx const& sequence_idx) {
        ModelType::run_functor_with_all_component_types_return_void([this, &sequence_idx]<typename CompType>() {
            this->restore_component<CompType>(std::array{
                std::span<Idx2D const>{std::get<ModelType::template index_of_component<CompType>>(sequence_idx)}});
        });
    }

    template <typename MathSolverType, typename YBus, typename PrepareInputFn, typename SolveFn>
        requires std::invocable<std::remove_cvref_t<PrepareInputFn>, Idx /*n_math_solvers*/> &&
                 std::ranges::range<std::invoke_result_t<PrepareInputFn, Idx /*n_math_solvers*/>> &&
                 std::invocable<
                     std::remove_cvref_t<SolveFn>, MathSolverType&, YBus const&,
                     typename std::invoke_result_t<PrepareInputFn, Idx /*n_math_solvers*/>::const_reference> &&
                 solver_output_type<std::invoke_result_t<
                     SolveFn, MathSolverType&, YBus const&,
                     typename std::invoke_result_t<PrepareInputFn, Idx /*n_math_solvers*/>::const_reference>>
    auto calculate_(PrepareInputFn&& prepare_input, SolveFn&& solve, Logger& logger) {
        using InputType = typename std::invoke_result_t<PrepareInputFn, Idx /*n_math_solvers*/>::const_reference;
        using SolverOutputType = typename std::invoke_result_t<SolveFn, MathSolverType&, YBus const&, InputType>;
        using sym = typename SolverOutputType::sym;

        assert(construction_complete_);
        // prepare
        auto const& input = [this, &logger, prepare_input_ = std::forward<PrepareInputFn>(prepare_input)] {
            Timer const timer{logger, LogEvent::prepare};
            assert(construction_complete_);
            prepare_solvers<sym>(state_, solver_preparation_context_, solvers_cache_status_);
            assert(solvers_cache_status_.is_topology_valid());
            assert(solvers_cache_status_.template is_parameter_valid<sym>());
            return prepare_input_(get_n_math_solvers<ModelType>(state_));
        }();
        // calculate
        return [this, &logger, &input, solve_ = std::forward<SolveFn>(solve)] {
            Timer const timer{logger, LogEvent::math_calculation};
            auto& solvers = main_core::get_solvers<sym>(solver_preparation_context_.math_state);
            auto& y_bus_vec = main_core::get_y_bus<sym>(solver_preparation_context_.math_state);
            std::vector<SolverOutputType> solver_output;
            solver_output.reserve(get_n_math_solvers<ModelType>(state_));
            for (Idx i = 0; i != get_n_math_solvers<ModelType>(state_); ++i) {
                solver_output.emplace_back(solve_(solvers[i], y_bus_vec[i], input[i]));
            }
            return solver_output;
        }();
    }

    // Calculate with optimization, e.g., automatic tap changer
    template <calculation_type_tag calculation_type, symmetry_tag sym>
    auto calculate_with_optimizer(Options const& options, Logger& logger) {
        auto const get_calculator = [this, &options, &logger] {
            using Calc = Calculator<calculation_type, sym>;

            assert(options.optimizer_type == OptimizerType::no_optimization ||
                   (std::derived_from<calculation_type, power_flow_t>));

            return [this, &mutable_comp_coup = state_.comp_coup, &options,
                    &logger](MainModelState const& state, CalculationMethod calculation_method) {
                (void)state; // to avoid unused-lambda-capture when in Release build
                assert(&state == &state_);

                return calculate_<MathSolverProxy<sym>, YBus<sym>>(Calc::preparer(state, mutable_comp_coup, options),
                                                                   Calc::solver(calculation_method, options, logger),
                                                                   logger);
            };
        };

        SearchMethod const& search_method = options.optimizer_strategy == OptimizerStrategy::any
                                                ? SearchMethod::linear_search
                                                : SearchMethod::binary_search;

        return optimizer::get_optimizer<MainModelState, ConstDataset>(
                   options.optimizer_type, options.optimizer_strategy, get_calculator(),
                   [this](ConstDataset const& update_data) {
                       this->update_components<permanent_update_t>(update_data);
                   },
                   *meta_data_, search_method)
            ->optimize(state_, options.calculation_method);
    }

    // Single calculation, propagating the results to result_data
    void calculate(Options options, MutableDataset const& result_data, Logger& logger) {
        assert(construction_complete_);

        if (options.calculation_type == CalculationType::short_circuit) {
            auto const faults = state_.components.template citer<Fault>();
            auto const is_three_phase = std::ranges::all_of(
                faults, [](Fault const& fault) { return fault.get_fault_type() == FaultType::three_phase; });
            options.calculation_symmetry =
                is_three_phase ? CalculationSymmetry::symmetric : CalculationSymmetry::asymmetric;
        }

        if (options.calculation_type == CalculationType::power_flow &&
            options.calculation_method != CalculationMethod::newton_raphson &&
            state_.components.template size<VoltageRegulator>() > 0) {
            throw InvalidCalculationMethod{};
        }

        calculation_type_symmetry_func_selector(
            options.calculation_type, options.calculation_symmetry,
            []<calculation_type_tag calculation_type, symmetry_tag sym>(
                MainModelImpl& main_model_, Options const& options_, MutableDataset const& result_data_,
                Logger& logger) {
                auto const math_output = main_model_.calculate_with_optimizer<calculation_type, sym>(options_, logger);
                main_model_.output_result(math_output, result_data_, logger);
            },
            *this, options, result_data, logger);
    }

  public:
    static auto calculator(Options const& options, MainModelImpl& model, MutableDataset const& target_data,
                           bool cache_run, Logger& logger) {
        auto sub_opt = options; // copy
        sub_opt.err_tol = cache_run ? std::numeric_limits<double>::max() : options.err_tol;
        sub_opt.max_iter = cache_run ? 1 : options.max_iter;

        model.calculate(sub_opt, target_data, logger);
    }

    auto const& state() const {
        assert(construction_complete_);
        return state_;
    }
    auto const& meta_data() const {
        assert(construction_complete_);
        assert(meta_data_ != nullptr);
        return *meta_data_;
    }

    void check_no_experimental_features_used(Options const& /*options*/, ConstDataset const* batch_dataset) const {
        if (!std::ranges::all_of(
                state_.components.template citer<VoltageRegulator>(),
                [](auto const& regulator) { return is_nan(regulator.q_min()) && is_nan(regulator.q_max()); }) ||
            (batch_dataset != nullptr &&
             batch_dataset->for_each_component<meta_data::update_getter_s, VoltageRegulator>([](auto const& span) {
                 return !std::ranges::all_of(span, [](VoltageRegulatorUpdate const& regulator) {
                     return is_nan(regulator.q_min) && is_nan(regulator.q_max);
                 });
             }))) {
            throw ExperimentalFeature{"Voltage Regulator with Qmin/Qmax limits is an experimental feature"};
        }
    }

  private:
    template <solver_output_type SolverOutputType>
    void output_result(MathOutput<std::vector<SolverOutputType>> const& math_output, MutableDataset const& result_data,
                       Logger& logger) const {
        assert(!result_data.is_batch());

        auto const output_func = [this, &math_output, &result_data]<typename CT>() {
            result_data.for_each_component<typename output_type_getter<SolverOutputType>::type, CT>(
                [this, &math_output](auto const& span) {
                    if (std::empty(span)) {
                        return;
                    }
                    main_core::output_result<CT>(state_, math_output, span);
                });
        };

        Timer const t_output{logger, LogEvent::produce_output};
        ModelType::run_functor_with_all_component_types_return_void(output_func);
    }

    double system_frequency_;
    MetaData const* meta_data_;

    MainModelState state_;

    SolverPreparationContext solver_preparation_context_;

    SolversCacheStatus<ImplType> solvers_cache_status_{};

    OwnedUpdateDataset cached_inverse_update_{};
    UpdateChange cached_state_changes_{};
#ifndef NDEBUG
    // construction_complete is used for debug assertions only
    bool construction_complete_{false};
#endif // !NDEBUG
};

} // namespace power_grid_model
