"""Two prime number races, followed exactly at every integer up to N.

pi(x; q, a) counts the primes up to x that are congruent to a modulo q. For the race modulo 4
this follows the lead d(x) = pi(x; 4, 3) - pi(x; 4, 1), and for the race modulo 3 the lead
e(x) = pi(x; 3, 2) - pi(x; 3, 1), at every integer x from 1 to N, and counts the integers at
which each side leads, the stretches in which the usual loser leads, and, weighting each x by
1/x, the share of the race the usual loser leads. Writes results/R1.json.

The primes come from a sieve of Eratosthenes over the odd numbers; everything but the weighted
shares is an exact integer count. The weights 1/x are summed exactly, as harmonic numbers
H(b) - H(a - 1) over each stretch, with math.fsum, and the weighted shares then divide two such
sums, so they carry only floating-point rounding.
"""

import itertools
import json
import math
import pathlib
from fractions import Fraction

N = 10**8


def odd_primes(limit):
    """The odd primes up to limit, in order, from a sieve over the odd numbers 3, 5, 7, ..."""
    size = (limit - 1) // 2  # index i stands for 2 * i + 3
    sieve = bytearray([1]) * size
    for i in range((math.isqrt(limit) - 1) // 2):
        if sieve[i]:
            p = 2 * i + 3
            start = (p * p - 3) // 2
            sieve[start::p] = bytes(len(range(start, size, p)))
    return (2 * i + 3 for i in itertools.compress(range(size), sieve))


class Race:
    """The lead of one residue class over another, as the integers x run from 1 to N."""

    def __init__(self, name, modulus, winner, loser):
        self.name, self.modulus, self.winner, self.loser = name, modulus, winner, loser
        self.lead = 0  # pi(x; q, winner) - pi(x; q, loser)
        self.counts = {"winner": 0, "tie": 0, "loser": 0}  # integers x at which each side leads
        self.stretches = []  # [first x, last x] of each stretch in which the loser leads
        self.ties = []  # [first x, last x] of each stretch of ties
        self.weighted = {"winner": [], "tie": [], "loser": []}  # stretches, for the weighted shares
        self.primes = {"winner": 0, "loser": 0}

    def step(self, prime):
        residue = prime % self.modulus
        if residue == self.winner:
            self.lead += 1
            self.primes["winner"] += 1
        elif residue == self.loser:
            self.lead -= 1
            self.primes["loser"] += 1

    def hold(self, first, last):
        """The lead holds for every integer from first to last."""
        side = "winner" if self.lead > 0 else "loser" if self.lead < 0 else "tie"
        self.counts[side] += last - first + 1
        spans = self.weighted[side]
        if spans and spans[-1][1] == first - 1:
            spans[-1][1] = last
        else:
            spans.append([first, last])
        for which, kept in (("loser", self.stretches), ("tie", self.ties)):
            if side == which:
                if kept and kept[-1][1] == first - 1:
                    kept[-1][1] = last
                else:
                    kept.append([first, last])


def harmonic(m):
    """H(m) = 1 + 1/2 + ... + 1/m, to double precision: directly for small m, else asymptotically."""
    if m < 1000:
        return math.fsum(1 / k for k in range(1, m + 1))
    m2 = m * m
    # ln m + gamma + 1/(2m) - 1/(12m^2) + 1/(120m^4) - 1/(252m^6), error below 1/(240m^8).
    return math.fsum([math.log(m), 0.5772156649015329, 1 / (2 * m), -1 / (12 * m2), 1 / (120 * m2 * m2), -1 / (252 * m2 * m2 * m2)])


def weight(spans):
    """The sum of 1/x over the integers in the given stretches."""
    return math.fsum(harmonic(last) - harmonic(first - 1) for first, last in spans)


def main():
    races = [Race("mod4", 4, 3, 1), Race("mod3", 3, 2, 1)]
    # Before 2 no prime has been counted; 2 is in neither race's classes.
    previous = 1
    for prime in itertools.chain([2], odd_primes(N)):
        for race in races:
            if prime > previous:
                race.hold(previous, prime - 1)
            race.step(prime)
        previous = prime
    for race in races:
        race.hold(previous, N)

    total = weight([[1, N]])
    results = {"n": N}
    for race in races:
        stretches = race.stretches
        longest = max(stretches, key=lambda s: s[1] - s[0], default=None)
        results[race.name] = {
            "winner_class": race.winner,
            "loser_class": race.loser,
            "primes_in_winner_class": race.primes["winner"],
            "primes_in_loser_class": race.primes["loser"],
            "lead_at_n": race.lead,
            "integers_winner_leads": race.counts["winner"],
            "integers_tied": race.counts["tie"],
            "integers_loser_leads": race.counts["loser"],
            "share_loser_leads": round(race.counts["loser"] / N, 9),
            "loser_stretches": len(stretches),
            "first_loser_lead": stretches[0][0] if stretches else None,
            "last_loser_lead": stretches[-1][1] if stretches else None,
            "longest_loser_stretch": {"first": longest[0], "last": longest[1], "length": longest[1] - longest[0] + 1} if longest else None,
            "loser_stretch_starts": [first for first, _ in stretches][:20],
            "last_tie": race.ties[-1][1] if race.ties else None,
            # Each x weighted by 1/x: the share of the race up to N the usual loser leads, the
            # complement of the logarithmic density of x at which the usual winner leads, once
            # ties, which have none, are set aside.
            "log_weighted_share_loser_leads": round(weight(race.weighted["loser"]) / total, 9),
        }
    out = pathlib.Path(__file__).resolve().parent.parent / "results" / "R1.json"
    out.write_text(json.dumps(results, indent=2) + "\n")
    print(json.dumps(results, indent=2))


if __name__ == "__main__":
    main()