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

// IntFactorTable.cpp
// Implementation of the prime factor table

#include "math/core/mp/Int/IntFactorTable.hpp"
#include "math/core/mp/Int/IntOps.hpp"
#include "math/core/mp/Int/IntPrime.hpp"
#include "math/core/mp/Int/IntGCD.hpp"
#include <cmath>
#include <fstream>
#include <algorithm>

namespace sangi {

// Constructor
IntFactorTable::IntFactorTable(int max_value)
    : max_value_(max_value)
{
    if (max_value_ <= 0) {
        max_value_ = 0;
        return;
    }

    // Store odd numbers only (even numbers are fixed with prime factor 2)
    // The index of an odd number n is n/2
    size_t table_size = (max_value_ + 1) / 2;
    table_.resize(table_size, 0);

    // Try to load from file
    if (max_value_ >= 100000) {  // Use a file only for large tables
        if (loadFromFile()) {
            return;  // Load succeeded
        }
    }

    // Load failed or file not used → build the table
    buildTable();

    // Save to file (speeds up the next startup)
    if (max_value_ >= 100000) {
        saveToFile();
    }
}

// Build the table with the sieve of Eratosthenes (largest prime factor version)
void IntFactorTable::buildTable() {
    // Store the largest prime factor of an odd number n in table[n/2]
    // 0 means "prime" or "unprocessed"

    if (max_value_ < 1) return;

    // The prime factor of 1 is 1 (special case)
    if (max_value_ >= 1) {
        table_[oddIndex(1)] = 1;
    }

    // Sieve up to sqrt(max_value)
    int sqrt_max = static_cast<int>(std::sqrt(static_cast<double>(max_value_)));

    // For every odd prime p, write p into its multiples
    // Since smaller primes are processed first, the last value written is the largest prime factor
    for (int p = 3; p <= sqrt_max; p += 2) {
        // Check whether p is prime (prime if table[p] == 0)
        if (table_[oddIndex(p)] == 0) {
            // Record the largest prime factor p for the multiples of p
            // Start from the first odd composite multiple of p (skip p itself since it is prime, start from p*2)
            int start = 2 * p;  // The first composite multiple of p (even)
            if ((start & 1) == 0) {
                start += p;  // Make it odd (if p*2 is even, use p*3)
            }

            // Step through the multiples of p by 2*p (odd numbers only)
            for (int m = start; m <= max_value_; m += 2 * p) {
                // ★ Characteristic of LPF: overwrite unconditionally with no test ★
                table_[oddIndex(m)] = static_cast<uint32_t>(p);
            }
        }
    }

    // Process the multiples of primes p larger than sqrt_max
    // For these, p^2 > max_value, so only the small multiples of p fall within range
    for (int p = sqrt_max + (sqrt_max % 2 == 0 ? 1 : 2); p <= max_value_; p += 2) {
        // Check whether p is prime
        if (table_[oddIndex(p)] == 0) {
            // Start from the first odd composite multiple of p
            // Use int64_t: avoids overflow when 2*p exceeds INT_MAX
            int64_t start = 2LL * p;
            if ((start & 1) == 0) {
                start += p;  // Make it odd
            }
            if (start > max_value_) continue;  // Skip if out of range

            // Step through the multiples of p by 2*p (odd numbers only)
            for (int64_t m = start; m <= max_value_; m += 2LL * p) {
                table_[oddIndex(static_cast<int>(m))] = static_cast<uint32_t>(p);
            }
        }
    }
}

// Get the largest prime factor
int IntFactorTable::operator[](int n) const {
    // Out of range
    if (n < 0 || n > max_value_) {
        return 0;
    }

    // Special cases
    if (n == 0 || n == 1) {
        return 0;  // No prime factor
    }

    // 2 is prime
    if (n == 2) {
        return 0;
    }

    // Even (other than 2)
    if ((n & 1) == 0) {
        // For an even number, compare the largest prime factor of the odd part with 2 and return the larger
        int odd_part = n;
        while ((odd_part & 1) == 0) {
            odd_part >>= 1;
        }
        if (odd_part == 1) {
            return 2;  // A power of 2
        }
        // Get the largest prime factor of the odd part
        int odd_lpf = table_[oddIndex(odd_part)];
        if (odd_lpf == 0) {
            // odd_part is prime
            // Compare the prime factor of the odd part with 2
            return (odd_part > 2) ? odd_part : 2;
        }
        // Compare the largest prime factor recorded in the table with 2
        return (odd_lpf > 2) ? odd_lpf : 2;
    }

    // Odd case
    int lpf = table_[oddIndex(n)];
    if (lpf == 0) {
        // n is prime
        return 0;
    }
    return lpf;
}

// Prime factorization
std::vector<std::pair<Int, int>> IntFactorTable::factorize(const Int& n) const {
    std::vector<std::pair<Int, int>> factors;


    // Check for special states
    if (n.isNaN() || n.isInfinite()) {
        return factors;  // Empty
    }

    // Zero case
    if (n.isZero()) {
        factors.push_back({Int::Zero(), 1});
        return factors;
    }

    // For a negative number, factorize the absolute value
    Int abs_n = abs(n);

    // One case
    if (abs_n.isOne()) {
        return factors;  // Empty
    }

    Int remaining = abs_n;

    // Divide out as many factors of 2 as possible (number of trailing zero bits of LSB = exponent of 2)
    size_t exp_2 = remaining.countTrailingZeros();
    if (exp_2 > 0) {
        factors.push_back({Int(2), static_cast<int>(exp_2)});
        remaining = remaining >> static_cast<int>(exp_2);
    }

    // Stop once remaining becomes 1
    if (remaining.isOne()) {
        return factors;
    }

    // Factorize using the table (while remaining is at most max_value)
    bool finished_in_table = false;
    while (remaining <= Int(max_value_) && !remaining.isOne()) {
        int n_int = remaining.toInt();
        int lpf = (*this)[n_int];

        if (lpf == 0) {
            // remaining is prime
            factors.push_back({remaining, 1});
            finished_in_table = true;
            break;
        }

        // Divide out as many factors of lpf as possible
        Int p(lpf);
        int exp = 0;
        while (true) {
            Int quotient = remaining / p;
            Int remainder_div = remaining % p;

            if (!remainder_div.isZero()) {
                break;
            }

            remaining = quotient;
            exp++;
        }

        if (exp > 0) {
            factors.push_back({p, exp});
        }
    }

    // If remaining > max_value, extract small prime factors via one-limb batch GCD
    if (!remaining.isOne() && !finished_in_table) {
        IntPrime::extractSmallPrimeFactors(remaining, factors);

        // If remaining > 1 even after dividing out small primes, it is a prime or a large composite
        if (!remaining.isOne()) {
            factors.push_back({remaining, 1});
        }
    }

    // Sort the prime factors in ascending order
    std::sort(factors.begin(), factors.end(),
              [](const std::pair<Int, int>& a, const std::pair<Int, int>& b) {
                  return a.first < b.first;
              });

    return factors;
}

// Save to file
bool IntFactorTable::saveToFile(const std::string& filename) const {
    std::ofstream ofs(filename, std::ios::binary);
    if (!ofs) {
        return false;
    }

    // Write max_value
    ofs.write(reinterpret_cast<const char*>(&max_value_), sizeof(max_value_));

    // Write the size of the table
    uint32_t table_size = static_cast<uint32_t>(table_.size());
    ofs.write(reinterpret_cast<const char*>(&table_size), sizeof(table_size));

    // Write the contents of the table
    ofs.write(reinterpret_cast<const char*>(table_.data()), table_size * sizeof(uint32_t));

    return ofs.good();
}

// Load from file
bool IntFactorTable::loadFromFile(const std::string& filename) {
    std::ifstream ifs(filename, std::ios::binary);
    if (!ifs) {
        return false;
    }

    // Read max_value
    int file_max_value;
    ifs.read(reinterpret_cast<char*>(&file_max_value), sizeof(file_max_value));
    if (!ifs || file_max_value < max_value_) {
        return false;  // The table in the file is too small
    }

    // Read the size of the table
    uint32_t file_table_size;
    ifs.read(reinterpret_cast<char*>(&file_table_size), sizeof(file_table_size));
    if (!ifs || file_table_size < static_cast<uint32_t>(table_.size())) {
        return false;  // Insufficient size
    }

    // Read the contents of the table
    ifs.read(reinterpret_cast<char*>(table_.data()), table_.size() * sizeof(uint32_t));
    if (!ifs) {
        return false;
    }

    return true;
}

} // namespace sangi