/*******************************************************************************
 * Copyright (C) 2017-2026 Theodore Chang
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 ******************************************************************************/

#include "Interaction.h"

#include <Domain/DomainBase.h>
#include <Domain/Factory.hpp>
#include <Element/Element.h>
#include <Solver/Integrator/Integrator.h>
#include <Step/Step.h>

InteractionPair::InteractionPair(const shared_ptr<Element>& obj_i, const shared_ptr<Element>& obj_j)
    : object_i(obj_i)
    , object_j(obj_j) {
    effective_mass = 1. / (1. / obj_i->get(Element::Parameter::MASS) + 1. / obj_j->get(Element::Parameter::MASS));
    effective_radius = 1. / (1. / obj_i->get(Element::Parameter::RADIUS) + 1. / obj_j->get(Element::Parameter::RADIUS));
    effective_modulus = 1. / ((1. - std::pow(obj_i->get(Element::Parameter::POISSON), 2.)) / obj_i->get(Element::Parameter::ELASTIC) + (1. - std::pow(obj_j->get(Element::Parameter::POISSON), 2.)) / obj_j->get(Element::Parameter::ELASTIC));
    effective_damping = .5 * (obj_i->get(Element::Parameter::DAMPING) + obj_j->get(Element::Parameter::DAMPING));
}

vec InteractionPair::position_i() const { return object_i->get_coordinate(dimension).t() + object_i->get_trial_displacement().head(dimension); }

vec InteractionPair::position_j() const { return object_j->get_coordinate(dimension).t() + object_j->get_trial_displacement().head(dimension); }

vec InteractionPair::velocity_i() const { return object_i->get_trial_velocity().head(dimension); }

vec InteractionPair::velocity_j() const { return object_j->get_trial_velocity().head(dimension); }

const uvec& InteractionPair::dof_i() const { return object_i->get_dof_encoding(); }

const uvec& InteractionPair::dof_j() const { return object_j->get_dof_encoding(); }

double InteractionPair::initial_gap() const { return object_i->get(Element::Parameter::RADIUS) + object_j->get(Element::Parameter::RADIUS); }

void Interaction::initialize(const shared_ptr<DomainBase>& D) { domain = D; }

void Hertzian::apply(const bool full, const shared_ptr<InteractionPair>& pair) const {
    const vec position_i = pair->position_i(), position_j = pair->position_j();
    const vec chord = position_i - position_j;
    const auto chord_length = norm(chord);
    const auto compression = pair->initial_gap() - chord_length;

    if(compression <= 0.) return;

    const vec unit_chord = chord / chord_length;
    const vec velocity_rel = pair->velocity_i() - pair->velocity_j();
    const auto velocity_proj = dot(velocity_rel, unit_chord);

    const auto normal_factor = four_third * std::sqrt(pair->effective_radius * compression) * pair->effective_modulus;
    const auto damping_factor = normal_factor * pair->effective_damping;
    const auto total_force = damping_factor * velocity_proj - normal_factor * compression;

    if(total_force > 0.) return;

    const uvec &dof_i = pair->dof_i(), &dof_j = pair->dof_j();

    const auto domain_ptr = domain.lock();
    auto& factory = domain_ptr->get_factory();

    {
        const vec repulsive = total_force * unit_chord;
        auto& t_resistance = factory->modify_trial_constraint_resistance();
        std::scoped_lock resistance_lock(factory->get_trial_constraint_resistance_mutex());
        for(uword I{0}; I < chord.n_elem; ++I) {
            t_resistance(dof_i(I)) += repulsive(I);
            t_resistance(dof_j(I)) -= repulsive(I);
        }
    }

    if(!full || Integrator::Type::Explicit == domain_ptr->get_current_step()->get_integrator()->type()) return;

    {
        const mat der_unit_chord = eye(chord.n_elem, chord.n_elem) - unit_chord * unit_chord.t();
        const mat der_repulsive = total_force / chord_length * der_unit_chord + unit_chord * (damping_factor / chord_length * velocity_rel.t() * der_unit_chord + (1.5 * normal_factor - .5 * damping_factor / compression * velocity_proj) * unit_chord.t());
        auto& t_mat = factory->get_stiffness();
        std::scoped_lock lock(factory->get_stiffness_mutex());
        for(uword I{0}; I < chord.n_elem; ++I)
            for(uword J{0}; J < chord.n_elem; ++J) {
                t_mat->at(dof_i(I), dof_i(J)) += der_repulsive(I, J);
                t_mat->at(dof_j(I), dof_j(J)) += der_repulsive(I, J);
                t_mat->at(dof_i(I), dof_j(J)) -= der_repulsive(I, J);
                t_mat->at(dof_j(I), dof_i(J)) -= der_repulsive(I, J);
            }
    }

    {
        const mat der_repulsive = damping_factor * unit_chord * unit_chord.t();
        auto& t_mat = factory->get_damping();
        std::scoped_lock lock(factory->get_damping_mutex());
        for(uword I{0}; I < chord.n_elem; ++I)
            for(uword J{0}; J < chord.n_elem; ++J) {
                t_mat->at(dof_i(I), dof_i(J)) += der_repulsive(I, J);
                t_mat->at(dof_j(I), dof_j(J)) += der_repulsive(I, J);
                t_mat->at(dof_i(I), dof_j(J)) -= der_repulsive(I, J);
                t_mat->at(dof_j(I), dof_i(J)) -= der_repulsive(I, J);
            }
    }
}

FixedParticle::FixedParticle(const unsigned T, const double M, std::set<unsigned>&& P)
    : Interaction(T)
    , particles(std::move(P))
    , multiplier(M) {}

void FixedParticle::apply(const bool full, const shared_ptr<Element>& element) const {
    if(!particles.contains(element->get_tag())) return;

    const auto domain_ptr = domain.lock();

    if(full && Integrator::Type::Implicit == domain_ptr->get_current_step()->get_integrator()->type()) {
        auto& factory = domain_ptr->get_factory();

        if(auto& t_mass = factory->get_mass()) {
            const auto penalty = multiplier * element->get(Element::Parameter::MASS);
            std::scoped_lock lock(factory->get_mass_mutex());
            for(const auto I : element->get_dof_encoding()) t_mass->at(I, I) += penalty;
        }
        if(auto& t_stiffness = factory->get_stiffness()) {
            const auto penalty = multiplier * element->get(Element::Parameter::ELASTIC);
            std::scoped_lock lock(factory->get_stiffness_mutex());
            for(const auto I : element->get_dof_encoding()) t_stiffness->at(I, I) += penalty;
        }
    }

    domain_ptr->insert_constrained_dof(element->get_dof_encoding());
}
