//! Pollard's kangaroo (lambda) method for d ∈ [lo, hi], in the van Oorschot–Wiener parallel form
//! with tame/wild herds and distinguished points (Pollard 2000 for the power-of-two jump set).
//!
//! Tame kangaroo c starts at position t_c = lo + ⌊w/2⌋ + c·s (point t_c·G); wild kangaroo c starts
//! at offset o_c = c·s (point Q + o_c·G); s = ⌈mean jump / (N/2)⌉ with N kangaroos in total.
//! Jump set J = {2^0 G, …, 2^{k−1} G}, j = h(x) ∈ [0, k), k chosen so that (2^k − 1)/k ≈ N·√w/4
//! (which is √w/2 for the classical single tame/wild pair). Every kangaroo keeps its position mod n;
//! a tame/wild distinguished-point collision with positions (t, o) gives d = t − o (same y) or
//! d = −t − o (opposite y). Same-herd collisions reposition the newcomer by a random jump.

use crate::curve::{CombTable, Counters, Pt};
use crate::instance::{is_distinguished, isqrt, partition, Instance};
use crate::rho::{RhoConfig, WalkStats};
use crate::stats::FxBuild;
use hashbrown::HashMap;
use rand::{Rng, SeedableRng};
use rand_chacha::ChaCha20Rng;
use serde::{Deserialize, Serialize};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Mutex;
use std::time::Instant;

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct KangarooConfig {
    pub lo: u64,
    pub hi: u64,
    pub threads: usize,
    /// kangaroos per thread (0 = automatic, even number, half tame / half wild)
    pub walks: usize,
    /// 0 = automatic
    pub dp_bits: u32,
    pub seed: u64,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct KangarooResult {
    pub d: Option<u64>,
    pub stats: WalkStats,
    pub dp_count: u64,
    pub dp_bytes: u64,
    pub table_entries: u64,
    pub threads: usize,
    pub kangaroos_per_thread: usize,
    pub dp_bits: u32,
    pub jump_count: u32,
    pub mean_jump: u64,
    pub interval_width: u64,
    pub wall_s: f64,
}

#[derive(Clone, Copy)]
struct Roo {
    p: Pt,
    dist: u64,
    tame: bool,
    /// total distance jumped since the last (re)start; a kangaroo that travels much further than
    /// the interval width has lost its herd and is restarted at a fresh random position
    travel: u128,
}

#[derive(Clone, Copy)]
struct DpEntry {
    y: u64,
    dist: u64,
    tame: bool,
}

pub fn auto_kangaroos(w: u64) -> usize {
    let bits = 64 - w.leading_zeros();
    let k = 1usize << (bits / 4).saturating_sub(2).min(6);
    k.clamp(2, 64).next_multiple_of(2)
}

/// Pick k with (2^k − 1)/k closest to `target`.
pub fn choose_jump_count(target: u64) -> u32 {
    let mut best = 1u32;
    let mut best_err = u128::MAX;
    for k in 1..=62u32 {
        let mean = (((1u128 << k) - 1) / k as u128) as u128;
        let err = mean.abs_diff(target as u128);
        if err < best_err {
            best_err = err;
            best = k;
        }
    }
    best
}

pub fn solve_kangaroo(inst: &Instance, cfg: &KangarooConfig) -> KangarooResult {
    let t0 = Instant::now();
    let n = inst.n.n;
    assert!(cfg.lo <= cfg.hi, "lo must be ≤ hi");
    let w = cfg.hi - cfg.lo + 1;
    let wbits = 64 - w.leading_zeros();
    let threads = cfg.threads.max(1);
    let per_thread = if cfg.walks == 0 { auto_kangaroos(w) } else { cfg.walks.max(2).next_multiple_of(2) };
    let total_roos = (threads * per_thread) as u64;
    let dp_bits = if cfg.dp_bits == 0 { RhoConfig::default_dp_bits(wbits) } else { cfg.dp_bits };
    let dp_mask = (1u64 << dp_bits.min(63)) - 1;
    let sqrt_w = isqrt(w as u128) as u64;
    let target_mean = (total_roos * sqrt_w / 4).max(1);
    let k = choose_jump_count(target_mean);
    let mean_jump = (((1u128 << k) - 1) / k as u128) as u64;
    let spacing = (mean_jump / (total_roos / 2)).max(1);

    let mut setup = Counters::default();
    let comb_g = CombTable::new(&inst.curve, &inst.g, &mut setup);
    let jumps: Vec<Pt> = (0..k).map(|i| comb_g.mul(&inst.curve, (1u64 << i) % n, &mut setup)).collect();
    let mut total = WalkStats { fruitless_by_length: vec![], ..Default::default() };
    total.ops.merge(&setup);

    if inst.q.inf {
        return KangarooResult { d: Some(0), stats: total, dp_count: 0, dp_bytes: 0, table_entries: 0, threads, kangaroos_per_thread: per_thread, dp_bits, jump_count: k, mean_jump, interval_width: w, wall_s: t0.elapsed().as_secs_f64() };
    }

    let expected_dps = 2.0 * (w as f64).sqrt() / (1u64 << dp_bits.min(63)) as f64;
    let table: Mutex<HashMap<u64, DpEntry, FxBuild>> = Mutex::new(HashMap::with_capacity_and_hasher((expected_dps * 3.0) as usize + 1024, Default::default()));
    let found: Mutex<Option<u64>> = Mutex::new(None);
    let stop = AtomicBool::new(false);
    let stats: Mutex<Vec<WalkStats>> = Mutex::new(Vec::new());
    let mid = cfg.lo + w / 2;
    let jumps = &jumps;
    let comb_g = &comb_g;

    std::thread::scope(|s| {
        for t in 0..threads {
            let table = &table;
            let found = &found;
            let stop = &stop;
            let stats = &stats;
            s.spawn(move || {
                let cpu0 = crate::stats::thread_cpu_s();
                let mut seed = [0u8; 32];
                seed[..8].copy_from_slice(&cfg.seed.to_le_bytes());
                seed[8..16].copy_from_slice(&(t as u64).to_le_bytes());
                seed[16..24].copy_from_slice(b"kangaroo");
                let mut rng = ChaCha20Rng::from_seed(seed);
                let mut st = WalkStats { fruitless_by_length: vec![], ..Default::default() };
                let f = &inst.curve.f;
                let curve = &inst.curve;
                let half = per_thread / 2;
                let mut roos: Vec<Roo> = Vec::with_capacity(per_thread);
                let mut solved: Option<u64> = None;
                let start = |st: &mut WalkStats, tame: bool, dist: u64| -> Roo {
                    let base = comb_g.mul(curve, dist % n, &mut st.ops);
                    let p = if tame { base } else { curve.add(&inst.q, &base, &mut st.ops) };
                    Roo { p, dist: dist % n, tame, travel: 0 }
                };
                // random fresh start for a lost / merged kangaroo
                let fresh = |st: &mut WalkStats, rng: &mut ChaCha20Rng, tame: bool| -> Roo {
                    let off = rng.random_range(0..w);
                    let dist = if tame { (cfg.lo as u128 + off as u128) % n as u128 } else { off as u128 % n as u128 } as u64;
                    st.restarts += 1;
                    start(st, tame, dist)
                };
                let max_travel: u128 = 3 * w as u128 + 4 * (mean_jump as u128) * (1u128 << dp_bits.min(40));
                for c in 0..per_thread {
                    let gc = (t * half + c % half) as u64;
                    let tame = c < half;
                    let dist = if tame { (mid as u128 + gc as u128 * spacing as u128) % n as u128 } else { (gc as u128 * spacing as u128) % n as u128 } as u64;
                    let r = start(&mut st, tame, dist);
                    if r.p.inf {
                        // wild: Q = −dist·G ⇒ d = −dist ; tame: dist ≡ 0 — harmless, re-seed
                        if !tame {
                            let d = inst.n.neg(r.dist);
                            if inst.verify(d, &mut st.ops) {
                                solved = Some(d);
                            }
                        }
                        let fr = fresh(&mut st, &mut rng, tame);
                        roos.push(fr);
                    } else {
                        roos.push(r);
                    }
                }
                let mut dens = vec![0u64; per_thread];
                let mut scratch = vec![0u64; per_thread];
                let mut js = vec![0u32; per_thread];
                let mut special = vec![false; per_thread];
                let mut dps: Vec<(u64, DpEntry)> = Vec::new();
                'main: while solved.is_none() && !stop.load(Ordering::Relaxed) {
                    for i in 0..per_thread {
                        let j = partition(roos[i].p.x, k as u64);
                        js[i] = j;
                        let jp = &jumps[j as usize];
                        let sp = roos[i].p.inf || roos[i].p.x == jp.x;
                        special[i] = sp;
                        dens[i] = if sp { f.one } else { f.sub(jp.x, roos[i].p.x) };
                    }
                    let (m, inv) = f.batch_inv(&mut dens[..per_thread], &mut scratch[..per_thread]);
                    st.ops.mul += m;
                    st.ops.inv += inv;
                    for i in 0..per_thread {
                        let j = js[i] as usize;
                        let jp = jumps[j];
                        let jd = (1u64 << j) % n;
                        let r = &mut roos[i];
                        if special[i] {
                            // P = ±2^j G
                            if !r.tame {
                                // Q + o·G = ±2^j·G ⇒ d = ±2^j − o
                                for d in [inst.n.sub(jd, r.dist), inst.n.sub(inst.n.neg(jd), r.dist)] {
                                    if inst.verify(d, &mut st.ops) {
                                        solved = Some(d);
                                        break 'main;
                                    }
                                }
                            }
                            let np = curve.add(&r.p, &jp, &mut st.ops);
                            st.iterations += 1;
                            if np.inf {
                                *r = fresh(&mut st, &mut rng, r.tame);
                            } else {
                                r.p = np;
                                r.dist = inst.n.add(r.dist, jd);
                                r.travel += 1u128 << j;
                            }
                        } else {
                            let lam = f.mul(f.sub(jp.y, r.p.y), dens[i]);
                            r.p = curve.finish_add(&r.p, jp.x, lam);
                            r.dist = inst.n.add(r.dist, jd);
                            r.travel += 1u128 << j;
                            st.ops.add += 1;
                            st.ops.mul += 3;
                            st.iterations += 1;
                        }
                        if r.travel > max_travel {
                            st.abandoned_walks += 1;
                            *r = fresh(&mut st, &mut rng, r.tame);
                            continue;
                        }
                        if is_distinguished(r.p.x, dp_mask) {
                            st.dp_found += 1;
                            dps.push((r.p.x, DpEntry { y: r.p.y, dist: r.dist, tame: r.tame }));
                        }
                    }
                    if !dps.is_empty() {
                        let mut tab = table.lock().unwrap();
                        for (x, e) in dps.drain(..) {
                            match tab.get(&x) {
                                None => {
                                    tab.insert(x, e);
                                }
                                Some(old) => {
                                    let old = *old;
                                    if old.tame == e.tame {
                                        if old.dist != e.dist {
                                            // merged trails inside one herd: reposition the newcomer
                                            st.useless_collisions += 1;
                                            if let Some(r) = roos.iter_mut().find(|r| r.p.x == x && r.tame == e.tame) {
                                                *r = fresh(&mut st, &mut rng, r.tame);
                                            }
                                        }
                                        continue;
                                    }
                                    let (tame_d, wild_d) = if old.tame { (old.dist, e.dist) } else { (e.dist, old.dist) };
                                    let cand = if old.y == e.y { inst.n.sub(tame_d, wild_d) } else { inst.n.sub(inst.n.neg(tame_d), wild_d) };
                                    if inst.verify(cand, &mut st.ops) {
                                        solved = Some(cand);
                                        break;
                                    }
                                    st.useless_collisions += 1;
                                }
                            }
                        }
                    }
                }
                if let Some(d) = solved {
                    let mut fd = found.lock().unwrap();
                    if fd.is_none() {
                        *fd = Some(d);
                    }
                    stop.store(true, Ordering::Relaxed);
                }
                st.cpu_s = crate::stats::thread_cpu_s() - cpu0;
                stats.lock().unwrap().push(st);
            });
        }
    });
    for st in stats.lock().unwrap().iter() {
        total.merge(st);
    }
    let tab = table.lock().unwrap();
    let entries = tab.len() as u64;
    let bytes = tab.capacity() as u64 * (std::mem::size_of::<(u64, DpEntry)>() as u64 + 1);
    let d = *found.lock().unwrap();
    KangarooResult { d, stats: total, dp_count: entries, dp_bytes: bytes, table_entries: entries, threads, kangaroos_per_thread: per_thread, dp_bits, jump_count: k, mean_jump, interval_width: w, wall_s: t0.elapsed().as_secs_f64() }
}