// SPDX-License-Identifier: LGPL-3.0-or-later
// Author: Kristian Lytje

#include <data/symmetry/CyclicSymmetry.h>
#include <math/MatrixUtils.h>

#include <cassert>

using namespace ausaxs;
using namespace ausaxs::symmetry;

bool sanity_checks(const CyclicSymmetry& s) {
    assert(1 <= s._repetitions && "Zero or negative repeats does not make sense.");
    assert(s._initial_relation.translation.dot(s._repeat_relation.translation) == 0 && "The translation vectors must be orthogonal.");
    assert([&]() -> bool {
        if (1e-6 < std::abs(s._repeat_relation.angle)) { return 1e-6 < s._repeat_relation.axis.magnitude(); }
        return true;
    }() && "Rotation angle is non-zero but rotation axis is zero vector.");
    assert([&]() -> bool {
        auto normed_axis = s._repeat_relation.axis / s._repeat_relation.axis.magnitude();
        auto R_r = matrix::rotation_matrix<double>(normed_axis, s._repeat_relation.angle);
        return std::abs(s._repeat_relation.angle) < 1e-6 || R_r*s._repeat_relation.translation == s._repeat_relation.translation;
    }() && "The translation vector must lie in the invariant space of the rotation matrix.");
    return true;
}

CyclicSymmetry::CyclicSymmetry() = default;
CyclicSymmetry::CyclicSymmetry(_Relation initial_relation, _Repeat repeat_relation, int repetitions) {
    this->_initial_relation = initial_relation;
    this->_repeat_relation  = repeat_relation;
    this->_repetitions      = repetitions;
    assert(sanity_checks(*this));
}

CyclicSymmetry::CyclicSymmetry(Vector3<double> offset, Vector3<double> repeat_translation, Vector3<double> repeat_axis, double repeat_rotation, int repetitions) {
    this->_initial_relation = _Relation{offset};
    this->_repeat_relation  = _Repeat{repeat_translation, repeat_axis, repeat_rotation};
    this->_repetitions      = repetitions;
    assert(sanity_checks(*this));
}

ISymmetry& CyclicSymmetry::add(observer_ptr<const ISymmetry> other) {
    auto cast = dynamic_cast<const CyclicSymmetry*>(other);
    assert(cast != nullptr && "Can only add Symmetry with another Symmetry.");
    this->_initial_relation.translation += cast->_initial_relation.translation;
    this->_repeat_relation.axis += cast->_repeat_relation.axis;
    this->_repeat_relation.angle += cast->_repeat_relation.angle;
    return *this;
}

std::unique_ptr<ISymmetry> CyclicSymmetry::clone() const { return std::make_unique<CyclicSymmetry>(*this); }

std::function<ausaxs::Vector3<double>(ausaxs::Vector3<double>)> CyclicSymmetry::get_transform(const Vector3<double>& cm, int rep) const {
    Matrix<double>  R_final;
    Vector3<double> T_final;

    {
        auto t_i = _initial_relation.translation;
        auto t_r = _repeat_relation.translation;
        auto normed_axis = _repeat_relation.axis / _repeat_relation.axis.magnitude();
        assert(sanity_checks(*this));
        assert((
            1e-6 < t_r.magnitude() || 1e-6 < std::abs(_repeat_relation.angle)
            || (rep == 1 && 1e-6 < t_i.magnitude())
        ) && "The symmetry is trivial; no transformation to apply.");

        // The symmetric ensemble is generated by:
        //   For copy k: v_k = R_r^k*(v - cm) + R_r^k*t_i - t_i + sum_{j=1}^{k} R_r^{j-1}*t_r + cm
        //
        // Accumulated iteratively:
        //   base_T = cm - r_r*cm + (r_r*t_i + t_r - t_i)   (k=1 translation)
        //   R_k    = r_r^k
        //   T_k    = r_r * T_{k-1} + base_T,   T_1 = base_T
        auto R_r = matrix::rotation_matrix<double>(normed_axis, _repeat_relation.angle);
        auto base_T = cm - R_r * cm + (R_r * t_i + t_r - t_i);
        R_final = R_r;
        T_final = base_T;
        for (int i = 1; i < rep; ++i) {
            T_final = R_r*T_final + base_T;
            R_final = R_r*R_final;
        }
    }

    return [R_final=std::move(R_final), T_final=std::move(T_final)](Vector3<double> v) {
        return R_final*v + T_final;
    };
}

unsigned int CyclicSymmetry::repetitions() const {return _repetitions;}
std::span<double> CyclicSymmetry::span_translation() {return std::span<double>(_initial_relation.translation.begin(), _initial_relation.translation.end());}
std::span<double> CyclicSymmetry::span_rotation() {return std::span<double>(_repeat_relation.axis.begin(), _repeat_relation.axis.begin());}

bool CyclicSymmetry::is_closed() const {
    if (_repeat_relation.translation.magnitude() != 0) { return false; }
    if (std::abs(_repeat_relation.angle) < 1e-9) { return true; }

    // (repetitions+1) steps must accumulate to a multiple of 2pi
    double total = (_repetitions + 1)*_repeat_relation.angle;
    constexpr int cmp = 628; // 100*2pi rounded to nearest integer
    int rounded = static_cast<int>(std::round(100 * std::abs(total)));
    return (rounded % cmp) == 0;
}
