// Copyright (C) 2026 Kiyotsugu Arai
// SPDX-License-Identifier: LGPL-3.0-or-later

// FactoredInteger.hpp
// Factored representation class -- stores an integer as a product of prime factors
//
// Design:
//   - Stores pairs of factor and exponent
//   - Performs multiplication and division at the factor level (without expanding)
//   - GCD/LCM computable in O(number of prime factors)
//   - expand() returns the value with all factors multiplied together
//   - Also supports an additive form

#pragma once

#include <vector>
#include <algorithm>
#include <string>
#include <sstream>
#include <cmath>
#include <cassert>

namespace sangi {

// ================================================================
// FactoredInteger<T> class
// ================================================================

template<typename T>
class FactoredInteger {
private:
    std::vector<T> factors_;      // Factors
    std::vector<int> exponents_;  // Exponents
    bool additive_;               // true: additive form (sum f_i^e_i), false: product form (prod f_i^e_i)

public:
    // ============================================================
    // Constructors
    // ============================================================

    /// Default: empty (product form = 1, additive form = 0)
    FactoredInteger() : additive_(false) {}

    /// Reserve space for the given number of factors
    explicit FactoredInteger(size_t n, bool additive = false)
        : factors_(n), exponents_(n, 1), additive_(additive)
    {
        T identity = additive ? T(0) : T(1);
        int defaultExp = additive ? 1 : 0;
        std::fill(factors_.begin(), factors_.end(), identity);
        std::fill(exponents_.begin(), exponents_.end(), defaultExp);
    }

    /// Construct from a list of (factor, exponent) pairs (product form)
    FactoredInteger(std::initializer_list<std::pair<T, int>> pairs)
        : additive_(false)
    {
        for (const auto& [f, e] : pairs) {
            factors_.push_back(f);
            exponents_.push_back(e);
        }
    }

    // ============================================================
    // Accessors
    // ============================================================

    /// Number of factors
    [[nodiscard]] size_t size() const { return factors_.size(); }

    /// Whether the container is empty
    [[nodiscard]] bool empty() const { return factors_.empty(); }

    /// Whether the form is additive
    [[nodiscard]] bool isAdditive() const { return additive_; }

    /// The i-th factor
    [[nodiscard]] const T& factor(size_t i) const { return factors_[i]; }
    [[nodiscard]] T& factor(size_t i) { return factors_[i]; }

    /// The i-th exponent
    [[nodiscard]] int exponent(size_t i) const { return exponents_[i]; }
    [[nodiscard]] int& exponent(size_t i) { return exponents_[i]; }

    /// Access a factor via operator[]
    [[nodiscard]] const T& operator[](size_t i) const { return factors_[i]; }
    [[nodiscard]] T& operator[](size_t i) { return factors_[i]; }

    // ============================================================
    // Adding / removing factors
    // ============================================================

    /// Add factor f with exponent e
    void addFactor(const T& f, int e = 1) {
        // Merge with an existing factor if it matches
        for (size_t i = 0; i < factors_.size(); ++i) {
            if (factors_[i] == f) {
                exponents_[i] += e;
                return;
            }
        }
        factors_.push_back(f);
        exponents_.push_back(e);
    }

    /// Remove factors with zero exponent
    void removeZeroExponents() {
        size_t j = 0;
        for (size_t i = 0; i < factors_.size(); ++i) {
            if (exponents_[i] != 0) {
                if (j != i) {
                    factors_[j] = std::move(factors_[i]);
                    exponents_[j] = exponents_[i];
                }
                ++j;
            }
        }
        factors_.resize(j);
        exponents_.resize(j);
    }

    /// Change the number of elements
    void resize(size_t n) {
        size_t old = factors_.size();
        factors_.resize(n);
        exponents_.resize(n);
        if (n > old) {
            T identity = additive_ ? T(0) : T(1);
            int defaultExp = additive_ ? 1 : 0;
            for (size_t i = old; i < n; ++i) {
                factors_[i] = identity;
                exponents_[i] = defaultExp;
            }
        }
    }

    // ============================================================
    // Expansion (compute the value)
    // ============================================================

    /// Return the value of multiplying (or summing) all factors together
    [[nodiscard]] T expand() const {
        if (additive_) {
            T result = T(0);
            for (size_t i = 0; i < factors_.size(); ++i) {
                T term = factors_[i];
                for (int j = 1; j < exponents_[i]; ++j)
                    term *= factors_[i];
                result += term;
            }
            return result;
        } else {
            T result = T(1);
            for (size_t i = 0; i < factors_.size(); ++i) {
                for (int j = 0; j < exponents_[i]; ++j)
                    result *= factors_[i];
            }
            return result;
        }
    }

    // ============================================================
    // Multiplication and division (factor level)
    // ============================================================

    /// Multiplication between product-form values: merge factors
    [[nodiscard]] FactoredInteger operator*(const FactoredInteger& rhs) const {
        assert(!additive_ && !rhs.additive_);
        FactoredInteger result = *this;
        for (size_t i = 0; i < rhs.factors_.size(); ++i) {
            result.addFactor(rhs.factors_[i], rhs.exponents_[i]);
        }
        return result;
    }

    /// Division between product-form values: subtract exponents
    [[nodiscard]] FactoredInteger operator/(const FactoredInteger& rhs) const {
        assert(!additive_ && !rhs.additive_);
        FactoredInteger result = *this;
        for (size_t i = 0; i < rhs.factors_.size(); ++i) {
            result.addFactor(rhs.factors_[i], -rhs.exponents_[i]);
        }
        return result;
    }

    FactoredInteger& operator*=(const FactoredInteger& rhs) {
        return *this = *this * rhs;
    }
    FactoredInteger& operator/=(const FactoredInteger& rhs) {
        return *this = *this / rhs;
    }

    // ============================================================
    // GCD / LCM (factor level -- O(number of factors))
    // ============================================================

    /// GCD between factored values: min of exponents over common factors
    [[nodiscard]] friend FactoredInteger gcd(const FactoredInteger& a,
                                              const FactoredInteger& b) {
        assert(!a.additive_ && !b.additive_);
        FactoredInteger result;
        for (size_t i = 0; i < a.factors_.size(); ++i) {
            for (size_t j = 0; j < b.factors_.size(); ++j) {
                if (a.factors_[i] == b.factors_[j]) {
                    int minExp = std::min(a.exponents_[i], b.exponents_[j]);
                    if (minExp > 0)
                        result.addFactor(a.factors_[i], minExp);
                    break;
                }
            }
        }
        return result;
    }

    /// LCM between factored values: max of exponents over all factors
    [[nodiscard]] friend FactoredInteger lcm(const FactoredInteger& a,
                                              const FactoredInteger& b) {
        assert(!a.additive_ && !b.additive_);
        FactoredInteger result = a;
        for (size_t j = 0; j < b.factors_.size(); ++j) {
            bool found = false;
            for (size_t i = 0; i < result.factors_.size(); ++i) {
                if (result.factors_[i] == b.factors_[j]) {
                    result.exponents_[i] = std::max(result.exponents_[i],
                                                     b.exponents_[j]);
                    found = true;
                    break;
                }
            }
            if (!found) {
                result.addFactor(b.factors_[j], b.exponents_[j]);
            }
        }
        return result;
    }

    // ============================================================
    // Exponentiation
    // ============================================================

    /// Multiply every exponent by n
    [[nodiscard]] FactoredInteger pow(int n) const {
        FactoredInteger result = *this;
        for (auto& e : result.exponents_)
            e *= n;
        return result;
    }

    // ============================================================
    // Comparison
    // ============================================================

    [[nodiscard]] bool operator==(const FactoredInteger& rhs) const {
        if (additive_ != rhs.additive_) return false;
        if (factors_.size() != rhs.factors_.size()) return false;
        // Check whether the factor sets are equal (order-independent)
        std::vector<bool> matched(rhs.factors_.size(), false);
        for (size_t i = 0; i < factors_.size(); ++i) {
            bool found = false;
            for (size_t j = 0; j < rhs.factors_.size(); ++j) {
                if (!matched[j] && factors_[i] == rhs.factors_[j]
                    && exponents_[i] == rhs.exponents_[j]) {
                    matched[j] = true;
                    found = true;
                    break;
                }
            }
            if (!found) return false;
        }
        return true;
    }

    [[nodiscard]] bool operator!=(const FactoredInteger& rhs) const {
        return !(*this == rhs);
    }

    // ============================================================
    // String conversion
    // ============================================================

    [[nodiscard]] std::string toString() const {
        if (factors_.empty()) {
            return additive_ ? "0" : "1";
        }

        std::ostringstream oss;
        for (size_t i = 0; i < factors_.size(); ++i) {
            if (i > 0) {
                oss << (additive_ ? " + " : " * ");
            }
            oss << "(" << factors_[i] << ")";
            if (exponents_[i] != 1) {
                oss << "^" << exponents_[i];
            }
        }
        return oss.str();
    }

    friend std::ostream& operator<<(std::ostream& os, const FactoredInteger& fi) {
        return os << fi.toString();
    }
};

// ================================================================
// Type traits
// ================================================================

template<typename T>
struct is_factored_integer : std::false_type {};

template<typename T>
struct is_factored_integer<FactoredInteger<T>> : std::true_type {};

template<typename T>
inline constexpr bool is_factored_integer_v = is_factored_integer<T>::value;

} // namespace sangi