/*
* Copyright (C) 2020-2026 MEmilio
*
* Authors: Julia Bicker
*
* Contact: Martin J. Kuehn <Martin.Kuehn@DLR.de>
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
*     http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/

#include "abm/household.h"
#include "abm/model.h"
#include "abm/infection_state.h"
#include "abm/location_type.h"
#include "abm/time.h"
#include "abm/person_id.h"
#include "graph_abm/graph_abm_mobility.h"
#include "graph_abm/graph_abmodel.h"
#include "memilio/io/history.h"
#include "memilio/mobility/graph.h"
#include "memilio/utils/abstract_parameter_distribution.h"
#include "memilio/utils/parameter_distributions.h"
#include <cstddef>
#include <cstdint>
#include <map>
#include <tuple>
#include <utility>
#include <vector>

//Logger
struct Logger : mio::LogAlways {
    /**
    * A vector of tuples with the Location information i.e. each tuple contains the following information:
    * - The LocationId (including the model id)
    * - The total number of Persons at the location
    * - A map containing the number of Persons per InfectionState at the location
    */
    using Type = std::vector<std::tuple<int, mio::abm::LocationType, mio::abm::LocationId, size_t,
                                        std::map<mio::abm::InfectionState, size_t>>>;
    static Type log(const mio::abm::Simulation<mio::GraphABModel>& sim)
    {
        Type location_information{};
        location_information.reserve(size_t(mio::abm::LocationType::Count));
        auto t = sim.get_time();
        for (auto&& loc : sim.get_model().get_locations()) {
            std::map<mio::abm::InfectionState, size_t> persons_per_infection_state;
            for (size_t i = 0; i < static_cast<size_t>(mio::abm::InfectionState::Count); ++i) {
                auto inf_state = mio::abm::InfectionState(i);
                persons_per_infection_state.insert(
                    {inf_state, sim.get_model().get_subpopulation(loc.get_id(), t, inf_state)});
            }
            location_information.push_back(std::make_tuple(loc.get_model_id(), loc.get_type(), loc.get_id(),
                                                           sim.get_model().get_number_persons(loc.get_id()),
                                                           persons_per_infection_state));
        }
        return location_information;
    }
};

int main()
{
    mio::set_log_level(mio::LogLevel::warn);
    // This is an example with three age groups representing children, adults and seniors.
    size_t num_age_groups         = 3;
    const auto age_group_children = mio::AgeGroup(0);
    const auto age_group_adults   = mio::AgeGroup(1);
    const auto age_group_seniors  = mio::AgeGroup(2);

    auto model1 = mio::GraphABModel(num_age_groups, 0);

    //Set infection parameters
    model1.parameters.get<mio::abm::TimeExposedToNoSymptoms>()           = mio::ParameterDistributionConstant(4.);
    model1.parameters.get<mio::abm::TimeInfectedNoSymptomsToSymptoms>()  = mio::ParameterDistributionConstant(2.);
    model1.parameters.get<mio::abm::TimeInfectedNoSymptomsToRecovered>() = mio::ParameterDistributionConstant(4.);
    model1.parameters.get<mio::abm::TimeInfectedSymptomsToRecovered>()   = mio::ParameterDistributionConstant(5.);
    model1.parameters.get<mio::abm::TimeInfectedSymptomsToSevere>()      = mio::ParameterDistributionConstant(6.);
    model1.parameters.get<mio::abm::TimeInfectedSevereToRecovered>()     = mio::ParameterDistributionConstant(8.);
    model1.parameters.get<mio::abm::TimeInfectedSevereToCritical>()      = mio::ParameterDistributionConstant(7.);
    model1.parameters.get<mio::abm::TimeInfectedCriticalToRecovered>()   = mio::ParameterDistributionConstant(10.);
    model1.parameters.get<mio::abm::TimeInfectedCriticalToDead>()        = mio::ParameterDistributionConstant(11.);

    //Age group 0 goes to school and age group 1 goes to work
    model1.parameters.get<mio::abm::AgeGroupGotoSchool>()[age_group_children] = true;
    model1.parameters.get<mio::abm::AgeGroupGotoWork>()[age_group_adults]     = true;

    //Household members can be child, parent or senior
    auto child = mio::abm::HouseholdMember(num_age_groups);
    child.set_age_weight(age_group_children, 1);
    auto parent = mio::abm::HouseholdMember(num_age_groups);
    parent.set_age_weight(age_group_adults, 1);
    auto adult = mio::abm::HouseholdMember(num_age_groups);
    adult.set_age_weight(age_group_adults, 1);
    adult.set_age_weight(age_group_seniors, 1);

    //Single-Person households
    auto single_hh = mio::abm::Household();
    single_hh.add_members(adult, 1);

    //Two-Adult household
    auto two_adult_hh = mio::abm::Household();
    two_adult_hh.add_members(adult, 2);

    //Single-Parent household
    auto single_parent_hh = mio::abm::Household();
    single_parent_hh.add_members(child, 1);
    single_parent_hh.add_members(parent, 1);

    //Family household
    auto family_hh = mio::abm::Household();
    family_hh.add_members(child, 1);
    family_hh.add_members(parent, 2);

    //Household groups for model 1
    auto single_hh_group_m1 = mio::abm::HouseholdGroup();
    single_hh_group_m1.add_households(single_hh, 5);
    auto two_adult_hh_group_m1 = mio::abm::HouseholdGroup();
    two_adult_hh_group_m1.add_households(two_adult_hh, 3);
    auto single_parent_hh_group_m1 = mio::abm::HouseholdGroup();
    single_parent_hh_group_m1.add_households(single_parent_hh, 5);
    auto family_hh_group_m1 = mio::abm::HouseholdGroup();
    family_hh_group_m1.add_households(family_hh, 10);
    add_household_group_to_model(model1, single_hh_group_m1);
    add_household_group_to_model(model1, two_adult_hh_group_m1);
    add_household_group_to_model(model1, single_hh_group_m1);
    add_household_group_to_model(model1, family_hh_group_m1);

    auto model2 = mio::GraphABModel(num_age_groups, 1);

    //Set infection parameters
    model2.parameters.get<mio::abm::TimeExposedToNoSymptoms>()           = mio::ParameterDistributionConstant(4.);
    model2.parameters.get<mio::abm::TimeInfectedNoSymptomsToSymptoms>()  = mio::ParameterDistributionConstant(2.);
    model2.parameters.get<mio::abm::TimeInfectedNoSymptomsToRecovered>() = mio::ParameterDistributionConstant(4.);
    model2.parameters.get<mio::abm::TimeInfectedSymptomsToRecovered>()   = mio::ParameterDistributionConstant(5.);
    model2.parameters.get<mio::abm::TimeInfectedSymptomsToSevere>()      = mio::ParameterDistributionConstant(6.);
    model2.parameters.get<mio::abm::TimeInfectedSevereToRecovered>()     = mio::ParameterDistributionConstant(8.);
    model2.parameters.get<mio::abm::TimeInfectedSevereToCritical>()      = mio::ParameterDistributionConstant(7.);
    model2.parameters.get<mio::abm::TimeInfectedCriticalToRecovered>()   = mio::ParameterDistributionConstant(10.);
    model2.parameters.get<mio::abm::TimeInfectedCriticalToDead>()        = mio::ParameterDistributionConstant(11.);

    //Age group 0 goes to school and age group 1 goes to work
    model2.parameters.get<mio::abm::AgeGroupGotoSchool>()[age_group_children] = true;
    model2.parameters.get<mio::abm::AgeGroupGotoWork>()[age_group_adults]     = true;

    //Household groups for model 2
    auto single_hh_group_m2 = mio::abm::HouseholdGroup();
    single_hh_group_m2.add_households(single_hh, 6);
    auto two_adult_hh_group_m2 = mio::abm::HouseholdGroup();
    two_adult_hh_group_m2.add_households(two_adult_hh, 2);
    auto single_parent_hh_group_m2 = mio::abm::HouseholdGroup();
    single_parent_hh_group_m2.add_households(single_parent_hh, 10);
    auto family_hh_group_m2 = mio::abm::HouseholdGroup();
    family_hh_group_m2.add_households(family_hh, 11);
    add_household_group_to_model(model2, single_hh_group_m2);
    add_household_group_to_model(model2, two_adult_hh_group_m2);
    add_household_group_to_model(model2, single_hh_group_m2);
    add_household_group_to_model(model2, family_hh_group_m2);

    //Create locations for both models
    //model 1
    auto event_m1 = model1.add_location(mio::abm::LocationType::SocialEvent);
    model1.get_location(event_m1).get_infection_parameters().set<mio::abm::MaximumContacts>(10);
    auto hospital_m1 = model1.add_location(mio::abm::LocationType::Hospital);
    model1.get_location(hospital_m1).get_infection_parameters().set<mio::abm::MaximumContacts>(10);
    auto icu_m1 = model1.add_location(mio::abm::LocationType::ICU);
    model1.get_location(icu_m1).get_infection_parameters().set<mio::abm::MaximumContacts>(5);
    auto shop_m1 = model1.add_location(mio::abm::LocationType::BasicsShop);
    model1.get_location(shop_m1).get_infection_parameters().set<mio::abm::MaximumContacts>(20);
    auto school_m1 = model1.add_location(mio::abm::LocationType::School);
    model1.get_location(school_m1).get_infection_parameters().set<mio::abm::MaximumContacts>(20);
    auto work_m1 = model1.add_location(mio::abm::LocationType::Work);
    model1.get_location(work_m1).get_infection_parameters().set<mio::abm::MaximumContacts>(10);
    //model 2
    auto event_m2 = model2.add_location(mio::abm::LocationType::SocialEvent);
    model2.get_location(event_m2).get_infection_parameters().set<mio::abm::MaximumContacts>(10);
    auto hospital_m2 = model2.add_location(mio::abm::LocationType::Hospital);
    model2.get_location(hospital_m2).get_infection_parameters().set<mio::abm::MaximumContacts>(10);
    auto icu_m2 = model2.add_location(mio::abm::LocationType::ICU);
    model2.get_location(icu_m2).get_infection_parameters().set<mio::abm::MaximumContacts>(5);
    auto shop_m2 = model2.add_location(mio::abm::LocationType::BasicsShop);
    model2.get_location(shop_m2).get_infection_parameters().set<mio::abm::MaximumContacts>(20);
    auto school_m2 = model2.add_location(mio::abm::LocationType::School);
    model2.get_location(school_m2).get_infection_parameters().set<mio::abm::MaximumContacts>(20);
    auto work_m2 = model2.add_location(mio::abm::LocationType::Work);
    model2.get_location(work_m2).get_infection_parameters().set<mio::abm::MaximumContacts>(10);

    auto start_date = mio::abm::TimePoint(0);
    auto end_date   = mio::abm::TimePoint(0) + mio::abm::days(30);

    //Assign infection states and locations to persons from model 1
    std::vector<double> infection_distribution_m1{0.5, 0.3, 0.05, 0.05, 0.05, 0.05, 0.0, 0.0};
    for (auto& person : model1.get_persons()) {
        mio::abm::InfectionState infection_state = mio::abm::InfectionState(
            mio::DiscreteDistribution<size_t>::get_instance()(model1.get_rng(), infection_distribution_m1));
        auto rng = mio::abm::PersonalRandomNumberGenerator(person);
        if (infection_state != mio::abm::InfectionState::Susceptible) {
            person.add_new_infection(mio::abm::Infection(rng, mio::abm::VirusVariant::Wildtype, person.get_age(),
                                                         model1.parameters, start_date, infection_state));
        }
        person.set_assigned_location(mio::abm::LocationType::SocialEvent, event_m1, model1.get_id());
        person.set_assigned_location(mio::abm::LocationType::BasicsShop, shop_m1, model1.get_id());
        person.set_assigned_location(mio::abm::LocationType::Hospital, hospital_m1, model1.get_id());
        person.set_assigned_location(mio::abm::LocationType::ICU, icu_m1, model1.get_id());
        if (person.get_age() == age_group_children) {
            person.set_assigned_location(mio::abm::LocationType::School, school_m1, model1.get_id());
        }
        if (person.get_age() == age_group_adults) {
            //10% of adults in model 1 work in model 2
            size_t work_model = mio::DiscreteDistribution<size_t>::get_instance()(mio::thread_local_rng(),
                                                                                  std::vector<double>{0.9, 0.1});
            if (work_model == 1) { //person works in other model
                person.set_assigned_location(mio::abm::LocationType::Work, work_m2, model2.get_id());
            }
            else { //person works in same model
                person.set_assigned_location(mio::abm::LocationType::Work, work_m1, model1.get_id());
            }
        }
    }

    //Assign infection states and locations to persons from model 2
    std::vector<double> infection_distribution_m2{0.7, 0.1, 0.05, 0.05, 0.1, 0.0, 0.0, 0.0};
    for (auto& person : model2.get_persons()) {
        mio::abm::InfectionState infection_state = mio::abm::InfectionState(
            mio::DiscreteDistribution<size_t>::get_instance()(model2.get_rng(), infection_distribution_m2));
        auto rng = mio::abm::PersonalRandomNumberGenerator(person);
        if (infection_state != mio::abm::InfectionState::Susceptible) {
            person.add_new_infection(mio::abm::Infection(rng, mio::abm::VirusVariant::Wildtype, person.get_age(),
                                                         model2.parameters, start_date, infection_state));
        }
        person.set_assigned_location(mio::abm::LocationType::SocialEvent, event_m2, model2.get_id());
        person.set_assigned_location(mio::abm::LocationType::BasicsShop, shop_m2, model2.get_id());
        person.set_assigned_location(mio::abm::LocationType::Hospital, hospital_m2, model2.get_id());
        person.set_assigned_location(mio::abm::LocationType::ICU, icu_m2, model2.get_id());
        if (person.get_age() == age_group_children) {
            person.set_assigned_location(mio::abm::LocationType::School, school_m2, model2.get_id());
        }
        if (person.get_age() == age_group_adults) {
            //20% of adults in model 2 work in model 1
            size_t work_model = mio::DiscreteDistribution<size_t>::get_instance()(mio::thread_local_rng(),
                                                                                  std::vector<double>{0.2, 0.8});
            if (work_model == 1) { //person works in same model
                person.set_assigned_location(mio::abm::LocationType::Work, work_m2, model2.get_id());
            }
            else { //person works in other model
                person.set_assigned_location(mio::abm::LocationType::Work, work_m1, model1.get_id());
            }
        }
    }

    using HistoryType = mio::History<mio::DataWriterToMemory, Logger>;
    mio::Graph<mio::ABMSimulationNode<HistoryType>, mio::ABMMobilityEdge<HistoryType>> graph;
    graph.add_node(model1.get_id(), HistoryType{}, start_date, std::move(model1));
    graph.add_node(model2.get_id(), HistoryType{}, start_date, std::move(model2));
    graph.add_edge(model1.get_id(), model2.get_id());
    graph.add_edge(model2.get_id(), model1.get_id());

    auto exchange_time_span = mio::abm::hours(12);
    auto sim                = mio::make_abm_graph_sim<HistoryType>(start_date, exchange_time_span, std::move(graph));
    sim.advance(end_date);

    return 0;
}