/-
  PrimeLut.lean — Prime Look-Up Table reader

  Two backends:
    1. Small embedded LUT (first 168 primes, always available)
    2. Binary LUT reader for external/precomputed-math-data/data/primes.bin

  API matches Research Stack Semantics.PrimeLut for drop-in compatibility:
    primeAt n      — nth prime (0-indexed)
    nthPrime n     — nth prime (1-indexed)
    isPrimeInLut x — membership check (embedded LUT)
    primeFloor x   — largest prime ≤ x
    primeCeil x    — smallest prime ≥ x
    pi n           — prime-counting function π(n)
-/

namespace SilverSight.PrimeLut

-- ── Embedded small LUT ─────────────────────────────────────────

/-- First 168 primes (all primes < 1000). Always available. -/
def firstPrimes : Array UInt64 := #[
  2, 3, 5, 7, 11, 13, 17, 19, 23, 29,
  31, 37, 41, 43, 47, 53, 59, 61, 67, 71,
  73, 79, 83, 89, 97, 101, 103, 107, 109, 113,
  127, 131, 137, 139, 149, 151, 157, 163, 167, 173,
  179, 181, 191, 193, 197, 199, 211, 223, 227, 229,
  233, 239, 241, 251, 257, 263, 269, 271, 277, 281,
  283, 293, 307, 311, 313, 317, 331, 337, 347, 349,
  353, 359, 367, 373, 379, 383, 389, 397, 401, 409,
  419, 421, 431, 433, 439, 443, 449, 457, 461, 463,
  467, 479, 487, 491, 499, 503, 509, 521, 523, 541,
  547, 557, 563, 569, 571, 577, 587, 593, 599, 601,
  607, 613, 617, 619, 631, 641, 643, 647, 653, 659,
  661, 673, 677, 683, 691, 701, 709, 719, 727, 733,
  739, 743, 751, 757, 761, 769, 773, 787, 797, 809,
  811, 821, 823, 827, 829, 839, 853, 857, 859, 863,
  877, 881, 883, 887, 907, 911, 919, 929, 937, 941,
  947, 953, 967, 971, 977, 983, 991, 997
]

/-- Count of primes in the embedded LUT. -/
def firstPrimesCount : UInt64 := 168

/-- Safe array lookup. -/
def arrayGet? (xs : Array α) (n : Nat) : Option α :=
  if h : n < xs.size then some (xs[n]'h) else none

/-- Get nth prime (0-indexed) from the embedded LUT. -/
def primeAt (n : UInt64) : Option UInt64 :=
  if n < firstPrimesCount then arrayGet? firstPrimes n.toNat else none

/-- Get nth prime (1-indexed, mathematical convention). -/
def nthPrime (n : UInt64) : Option UInt64 :=
  if n = 0 then none else primeAt (n - 1)

/-- Check membership in the embedded LUT. -/
def isPrimeInLut (x : UInt64) : Bool :=
  firstPrimes.contains x

/-- Largest prime ≤ x in the embedded LUT. -/
def primeFloor (x : UInt64) : Option UInt64 :=
  firstPrimes.toList.foldl (fun best p => if p ≤ x then some p else best) none

/-- Smallest prime ≥ x in the embedded LUT. -/
def primeCeil (x : UInt64) : Option UInt64 :=
  if x ≤ 2 then some 2 else
  firstPrimes.toList.foldl (fun best p => if p ≥ x then
    match best with | none => some p | some q => some (min p q)
  else best) none

/-- Number of primes ≤ n in the embedded LUT. -/
def pi (n : UInt64) : UInt64 :=
  (firstPrimes.toList.filter (λ p => p ≤ n)).length.toUInt64

-- ── Subsets (mirrors PrimeLut.lean) ────────────────────────────

/-- Twin prime pairs (p, p+2) from the embedded LUT. -/
def twinPrimes : Array (UInt64 × UInt64) := #[
  (3, 5), (5, 7), (11, 13), (17, 19), (29, 31),
  (41, 43), (59, 61), (71, 73), (101, 103), (107, 109),
  (137, 139), (149, 151), (179, 181), (191, 193), (197, 199),
  (227, 229), (239, 241), (269, 271), (281, 283), (311, 313),
  (347, 349), (419, 421), (431, 433), (461, 463), (521, 523),
  (569, 571), (599, 601), (617, 619), (641, 643), (659, 661),
  (809, 811), (821, 823), (827, 829), (857, 859), (881, 883)
]

/-- Safe primes (p where (p-1)/2 is also prime) from the embedded LUT. -/
def safePrimes : Array UInt64 := #[
  5, 7, 11, 23, 47, 59, 83, 107, 167, 179,
  227, 263, 347, 359, 383, 467, 479, 503, 563, 587,
  719, 839, 863, 887, 983
]

-- ── Binary LUT reader ──────────────────────────────────────────
/--
  Binary prime LUT reader backed by external/precomputed-math-data/data/.

  Binary format:
    primes.bin   — raw 4-byte LE u32 primes, sorted ascending
    primes.index — every 10000th prime: u32 value + u64 byte offset (12 bytes/entry)

  Both files are read entirely into memory at construction via IO.FS.readBinFile.
    primes.bin   — 401 MB (105,097,565 primes × 4 bytes)
    primes.index — 126 KB (10,510 entries × 12 bytes)
-/
structure BinaryLUT where
  binBytes   : ByteArray     -- entire primes.bin
  indexBytes : ByteArray     -- entire primes.index
  primeCount : UInt64         -- total primes (105,097,565 for full LUT)

namespace BinaryLUT

/-- Decode a u32 LE from a ByteArray at the given byte offset. -/
def readU32 (bytes : ByteArray) (offset : UInt64) : UInt64 :=
  (bytes.get! offset.toNat).toUInt64 |||
  ((bytes.get! (offset.toNat + 1)).toUInt64 <<< 8) |||
  ((bytes.get! (offset.toNat + 2)).toUInt64 <<< 16) |||
  ((bytes.get! (offset.toNat + 3)).toUInt64 <<< 24)

/-- Decode a u64 LE from a ByteArray at the given byte offset. -/
def readU64 (bytes : ByteArray) (offset : UInt64) : UInt64 :=
  (bytes.get! offset.toNat).toUInt64 |||
  ((bytes.get! (offset.toNat + 1)).toUInt64 <<< 8) |||
  ((bytes.get! (offset.toNat + 2)).toUInt64 <<< 16) |||
  ((bytes.get! (offset.toNat + 3)).toUInt64 <<< 24) |||
  ((bytes.get! (offset.toNat + 4)).toUInt64 <<< 32) |||
  ((bytes.get! (offset.toNat + 5)).toUInt64 <<< 40) |||
  ((bytes.get! (offset.toNat + 6)).toUInt64 <<< 48) |||
  ((bytes.get! (offset.toNat + 7)).toUInt64 <<< 56)

/-- Read index entry at position i (0-indexed). Returns (prime_value, byte_offset). -/
def readIndexEntry (lut : BinaryLUT) (i : UInt64) : UInt64 × UInt64 :=
  let off := i * 12
  (readU32 lut.indexBytes off, readU64 lut.indexBytes (off + 4))

/-- Load from the default external directory. Returns none if files are missing. -/
def openLUT (dataDir : System.FilePath := "external/precomputed-math-data/data") : IO (Option BinaryLUT) := do
  let binPath := dataDir / "primes.bin"
  let idxPath := dataDir / "primes.index"
  try
    let binBytes ← IO.FS.readBinFile binPath
    let indexBytes ← IO.FS.readBinFile idxPath
    let primeCount := (binBytes.size.toUInt64) / 4
    pure (some { binBytes, indexBytes, primeCount })
  catch _ =>
    pure none

/-- Get prime at 0-indexed position. -/
def get (lut : BinaryLUT) (idx : UInt64) : UInt64 :=
  readU32 lut.binBytes (idx * 4)

/-- Number of index entries. -/
def indexEntryCount (lut : BinaryLUT) : UInt64 :=
  lut.indexBytes.size.toUInt64 / 12

/-- Find the index entry index for n via binary search (using Nat indices).
    Returns entryIdx such that index_prime[entryIdx] ≤ n < next_index_prime. -/
def findIndexEntry (lut : BinaryLUT) (n : UInt64) : Option UInt64 :=
  if lut.indexBytes.size < 12 then none else
  let count := lut.indexBytes.size / 12
  let rec go (lo hi : Nat) : Option UInt64 :=
    if hi ≤ lo then
      if lo = 0 then none else some (lo.toUInt64 - 1)
    else
      let mid := (lo + hi) / 2
      let (p, _) := readIndexEntry lut mid.toUInt64
      if p ≤ n then go (mid + 1) hi else go lo mid
  termination_by hi - lo
  go 0 count

/-- Prime counting function π(n): number of primes ≤ n.

    Uses binary search on the index, then within the segment.
    Falls back to the embedded small LUT when the binary files are not found. -/
def pi (lut : BinaryLUT) (n : UInt64) : UInt64 :=
  if n < 2 then 0
  else
    match findIndexEntry lut n with
    | none => 0
    | some entryIdx =>
      let baseIdx := entryIdx * 10000
      let (_, offset) := readIndexEntry lut entryIdx
      let remaining := if lut.primeCount > baseIdx then lut.primeCount - baseIdx else 0
      let segSize := min 10000 remaining
      let rec segSearch (lo hi : Nat) : UInt64 :=
        if hi ≤ lo then baseIdx + lo.toUInt64
        else
          let mid := (lo + hi) / 2
          let p := readU32 lut.binBytes (offset + mid.toUInt64 * 4)
          if p ≤ n then segSearch (mid + 1) hi else segSearch lo mid
      termination_by hi - lo
      segSearch 0 segSize.toNat

/-- Check if n is prime. -/
def isPrime (lut : BinaryLUT) (n : UInt64) : Bool :=
  if n < 2 then false else
  let idx := pi lut (n - 1)
  if idx ≥ lut.primeCount then false else
  get lut idx = n

/-- Largest prime ≤ n. -/
def floor (lut : BinaryLUT) (n : UInt64) : Option UInt64 :=
  if n < 2 then none else
  let idx := pi lut n
  if idx = 0 then none else
  some (get lut (idx - 1))

/-- Smallest prime ≥ n. -/
def ceil (lut : BinaryLUT) (n : UInt64) : Option UInt64 :=
  if n ≤ 2 then some 2 else
  let idx := pi lut (n - 1)
  if idx ≥ lut.primeCount then none else
  some (get lut idx)

end BinaryLUT

-- ── #eval witnesses ────────────────────────────────────────────

#eval primeAt 0     -- some 2
#eval primeAt 10    -- some 31
#eval nthPrime 1    -- some 2
#eval nthPrime 26   -- some 101
#eval isPrimeInLut 997   -- true
#eval isPrimeInLut 1000  -- false
#eval primeFloor 100     -- some 97
#eval primeCeil 100      -- some 101
#eval pi 1000            -- 168

end SilverSight.PrimeLut
