//! Fused per-locus ports of `simurg.genotype.gencall.call_genotypes` and `simurg.genotype.baf_lrr.baf_lrr`
//! (which port gtc2vcf `gen_std_flair`/`modilik`/`compute_score_call_prelim`/`gencall_score_map`/`make_calls`
//! and `compute_baf_lrr`, MIT, © Giulio Genovese). Every intermediate keeps the NumPy reference's type (f32 or
//! f64) and evaluation order; NumPy's `fmin`/`fmax` are `f32::min`/`max`; libm calls use CPython `math`
//! semantics. One pass per locus instead of dozens of full-array temporaries.

use numpy::{IntoPyArray, PyArray1, PyReadonlyArray1};
use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;

use crate::{math_exp, math_log, math_pow};

const M_2_SQRTPI: f64 = 1.12837916709551257390;
const M_SQRT1_2: f64 = 0.707106781186547524401;
const NEG_NAN32: f32 = f32::from_bits(0xffc0_0000);

#[inline]
fn smf(x: f32, a: f32, b: f32) -> f32 {
    let mid = (a + b) / 2.0;
    if a >= b {
        return if x >= mid { 1.0 } else { 0.0 };
    }
    let u = (x - a) / (b - a);
    let v = (x - b) / (b - a);
    if x <= a {
        0.0
    } else if a < x && x <= mid {
        2.0 * (u * u)
    } else if mid < x && x <= b {
        1.0 - 2.0 * (v * v)
    } else {
        1.0
    }
}

#[inline]
fn zmf(x: f32, a: f32, b: f32) -> f32 {
    let mid = (a + b) / 2.0;
    let u = (x - a) / (b - a);
    let v = (x - b) / (b - a);
    if x <= a {
        1.0
    } else if a < x && x <= mid {
        1.0 - 2.0 * (u * u)
    } else if mid < x && x <= b {
        2.0 * (v * v)
    } else {
        0.0
    }
}

/// Normal pdf pieces of `_pdf`: (full, peak, bad) for value x, mean mu, sd sigma.
#[inline]
fn pdf(x: f32, mu: f32, sigma: f32) -> (f64, f64, bool) {
    let z = ((x - mu) / sigma) as f64;
    let e = math_exp(-0.5 * (z * z));
    let sd = sigma as f64;
    let full = e * 0.5 * M_2_SQRTPI * M_SQRT1_2 / sd;
    let peak = (0.5 * M_2_SQRTPI * M_SQRT1_2) / sd;
    (full, peak, sigma <= 0.0)
}

#[derive(Clone, Copy)]
enum Side {
    Left,
    Right,
    Full,
}

#[inline]
fn get(p: (f64, f64, bool), x: f32, mu: f32, side: Side) -> f64 {
    let (full, peak, bad) = p;
    if bad {
        return f64::NAN;
    }
    match side {
        Side::Left => {
            if x < mu {
                peak
            } else {
                full
            }
        }
        Side::Right => {
            if x > mu {
                peak
            } else {
                full
            }
        }
        Side::Full => full,
    }
}

fn check(len: usize, n: usize, what: &str) -> PyResult<()> {
    if len != n {
        return Err(PyValueError::new_err(format!("gencall: {what} has length {len}, expected {n}")));
    }
    Ok(())
}

/// Genotypes (0 NC, 1 AA, 2 AB, 3 BB) and GenCall scores; (n, 3) cluster arrays are passed flattened (C order).
#[pyfunction]
#[allow(clippy::too_many_arguments)]
pub fn call_genotypes<'py>(
    py: Python<'py>,
    theta: PyReadonlyArray1<'py, f32>,
    r: PyReadonlyArray1<'py, f32>,
    raw_zero: PyReadonlyArray1<'py, bool>,
    present: PyReadonlyArray1<'py, bool>,
    n_cl: PyReadonlyArray1<'py, i32>,
    theta_mean: PyReadonlyArray1<'py, f32>,
    theta_dev: PyReadonlyArray1<'py, f32>,
    r_mean: PyReadonlyArray1<'py, f32>,
    r_dev: PyReadonlyArray1<'py, f32>,
    intensity_threshold: PyReadonlyArray1<'py, f32>,
    gentrain: PyReadonlyArray1<'py, f32>,
    intensity_only: PyReadonlyArray1<'py, bool>,
    cutoff: f32,
) -> PyResult<(Bound<'py, PyArray1<u8>>, Bound<'py, PyArray1<f32>>)> {
    let (th, rr, rz, pres) = (theta.as_slice()?, r.as_slice()?, raw_zero.as_slice()?, present.as_slice()?);
    let (ncl, tmv, tdv, rmv, rdv) =
        (n_cl.as_slice()?, theta_mean.as_slice()?, theta_dev.as_slice()?, r_mean.as_slice()?, r_dev.as_slice()?);
    let (ithr, gt, io) = (intensity_threshold.as_slice()?, gentrain.as_slice()?, intensity_only.as_slice()?);
    let n = th.len();
    for (l, w) in [(rr.len(), "r"), (rz.len(), "raw_zero"), (pres.len(), "present"), (ithr.len(), "threshold"),
                   (gt.len(), "gentrain"), (io.len(), "intensity_only")] {
        check(l, n, w)?;
    }
    for (l, w) in [(ncl.len(), "n"), (tmv.len(), "theta_mean"), (tdv.len(), "theta_dev"), (rmv.len(), "r_mean"),
                   (rdv.len(), "r_dev")] {
        check(l, 3 * n, w)?;
    }
    let tiny_d = f64::MIN_POSITIVE * f64::EPSILON;
    let third: f32 = 1.0 / 3.0;
    let mut geno = vec![0u8; n];
    let mut scores = vec![0f32; n];

    for i in 0..n {
        let t = th[i];
        let r = rr[i];
        let tm = [tmv[3 * i], tmv[3 * i + 1], tmv[3 * i + 2]];
        let td = [tdv[3 * i], tdv[3 * i + 1], tdv[3 * i + 2]];
        let rm = [rmv[3 * i], rmv[3 * i + 1], rmv[3 * i + 2]];
        let rd = [rdv[3 * i], rdv[3 * i + 1], rdv[3 * i + 2]];
        let nc = [ncl[3 * i], ncl[3 * i + 1], ncl[3 * i + 2]];

        // gen_std_flair
        let z1 = 0.5f32 * (tm[1] - tm[0]) / (td[1] + td[0]);
        let z2 = 0.5f32 * (tm[2] - tm[1]) / (td[2] + td[1]);
        let mz = z1.min(z2);
        let alpha = (third * mz).max(1.0);
        let beta = (third * mz).min(1.0);
        let eta = alpha * beta;
        let mut ftd = [0f32; 3];
        for k in 0..3 {
            let v = td[k] * eta;
            ftd[k] = if v < 0.02 { 0.02 } else { v };
        }
        let med = td[0].min(td[1]).max(td[0].max(td[1]).min(td[2]));
        let min_r = if 0.2f32 < med { med } else { 0.2 };
        let mut frd = [0f32; 3];
        for k in 0..3 {
            let mrk0 = rm[k] / 7.0;
            let mrk = if mrk0 < min_r { min_r } else { mrk0 };
            frd[k] = if rd[k] < mrk { mrk } else { rd[k] };
        }

        // modilik
        let d: [f64; 3] = std::array::from_fn(|k| 100.0 * ((t - tm[k]).abs() as f64) + ((r - rm[k]).abs() as f64));
        let (daa, dab, dbb) = (d[0], d[1], d[2]);
        let ok_t = !t.is_nan();
        let pt: [(f64, f64, bool); 3] = std::array::from_fn(|k| pdf(t, tm[k], ftd[k]));
        let pr: [(f64, f64, bool); 3] = std::array::from_fn(|k| pdf(r, rm[k], frd[k]));
        let pdf2 = |k: usize, st: Side, sr: Side| get(pt[k], t, tm[k], st) * get(pr[k], r, rm[k], sr);
        use Side::{Full, Left, Right};
        let mut lik = [0f64; 3];
        let mut covered = false;
        if ok_t {
            let cases: [(bool, [(usize, Side, Side); 2]); 6] = [
                (daa <= dbb && dab <= dbb && rm[0] <= rm[1], [(0, Left, Left), (1, Right, Right)]),
                (daa <= dab && dbb <= dab && rm[0] <= rm[2], [(0, Left, Left), (2, Right, Right)]),
                (dab <= daa && dbb <= daa && rm[1] <= rm[2], [(1, Left, Left), (2, Right, Right)]),
                (daa <= dbb && dab <= dbb && rm[0] > rm[1], [(0, Left, Right), (1, Right, Left)]),
                (dab <= daa && dbb <= daa && rm[1] > rm[2], [(1, Left, Right), (2, Right, Left)]),
                (daa <= dab && dbb <= dab && rm[0] > rm[2], [(0, Left, Right), (2, Right, Left)]),
            ];
            for (cond, assigns) in cases {
                if cond {
                    for (k, st, sr) in assigns {
                        lik[k] = pdf2(k, st, sr);
                    }
                    covered = true;
                }
            }
        }
        if !covered {
            for (k, l) in lik.iter_mut().enumerate() {
                *l = pdf2(k, Full, Full);
            }
        }
        for l in lik.iter_mut() {
            if l.is_nan() {
                *l = 0.0;
            }
        }
        if !ok_t {
            lik = [f64::NAN; 3];
        }

        // priors and posteriors
        let ntot = nc[0] as i64 + nc[1] as i64 + nc[2] as i64;
        let p: [f64; 3] = std::array::from_fn(|k| (((nc[k] as f32) / (ntot as f32)) as f64 + 1.0) / 4.0);
        let lp: [f64; 3] = std::array::from_fn(|k| lik[k] * p[k]);
        let ev = (lp[0] + lp[1]) + lp[2];
        let (aa, ab, bb) = (lp[0] / ev, lp[1] / ev, lp[2] / ev);
        let mut iap: u8 = if aa >= ab && aa >= bb {
            1
        } else if ab >= aa && ab >= bb {
            2
        } else if bb >= aa && bb >= ab {
            3
        } else {
            0
        };
        let mut mx = if aa > ab { aa } else { ab };
        let mut scnd = if aa > ab { ab } else { aa };
        let b_gt_mx = bb > mx;
        scnd = if b_gt_mx { mx } else if bb > scnd { bb } else { scnd };
        mx = if b_gt_mx { bb } else { mx };
        let ap_lod = math_log(mx / (tiny_d + scnd), true);
        let score_ap = smf(ap_lod as f32, 0.0, 2.0);

        // sub-scores
        let score_r1 = smf(r, 0.0, 0.1);
        let nclus = nc.iter().filter(|&&v| v > 0).count();
        let misclust: f32 = if nclus == 1 && nc[1] == 0 {
            0.7
        } else if nclus != 3 {
            0.95
        } else {
            1.0
        };
        let g = (iap as i64 - 1).clamp(0, 2) as usize;
        let (gm_t, gd_t, gm_r, gd_r) = (tm[g], ftd[g], rm[g], frd[g]);
        let st_aa = zmf(t, gm_t + 2.0 * gd_t, gm_t + 6.0 * gd_t);
        let st_ab = zmf(((t - gm_t) / gd_t).abs(), 2.0, 6.0);
        let st_bb = smf(t, gm_t - 6.0 * gd_t, gm_t - 2.0 * gd_t);
        let mut score_t = match iap {
            1 => st_aa,
            2 => st_ab,
            3 => st_bb,
            _ => 1.0,
        };
        let r2 = smf(r, 0.0, gm_r / 10.0);
        let r3 = smf(r, gm_r - 6.0 * gd_r, gm_r - 2.0 * gd_r);
        let r4 = zmf(r, gm_r + 6.0 * gd_r, gm_r + 12.0 * gd_r);
        let mut score_r = if iap > 0 { ((score_r1 * r2) * r3) * r4 } else { 0.0 };
        let mut score_n: f32 = 1.0;
        if gt[i].is_nan() {
            score_t = f32::NAN;
            score_r = f32::NAN;
        }
        if !ok_t {
            score_t = f32::NAN;
            score_r = f32::NAN;
            score_n = f32::NAN;
        }
        let mut prelim = ((((score_ap * gt[i]) * score_t) * score_r) * score_n) * misclust;
        if r < ithr[i] {
            prelim = f32::NAN;
            iap = 0;
        }

        // gencall_score_map
        let x = prelim as f64;
        let q = (x - 1.08) / 0.504;
        let gbell = 1.0 / (1.0 + math_pow(q * q, 1.71));
        let mut score = (math_pow(x, 0.35) * gbell) as f32;
        if score.is_nan() {
            score = 0.0;
        }
        if rz[i] || !pres[i] {
            score = 0.0;
            iap = 0;
        }
        geno[i] = if score < cutoff || io[i] { 0 } else { iap };
        scores[i] = score;
    }
    Ok((geno.into_pyarray(py), scores.into_pyarray(py)))
}

/// BAF and LRR (float32, canonical NaN 0xFFC00000 where undefined).
#[pyfunction]
#[allow(clippy::too_many_arguments)]
pub fn baf_lrr<'py>(
    py: Python<'py>,
    theta: PyReadonlyArray1<'py, f32>,
    r: PyReadonlyArray1<'py, f32>,
    raw_zero: PyReadonlyArray1<'py, bool>,
    present: PyReadonlyArray1<'py, bool>,
    theta_mean: PyReadonlyArray1<'py, f32>,
    r_mean: PyReadonlyArray1<'py, f32>,
    mean_r: PyReadonlyArray1<'py, f32>,
    intensity_only: PyReadonlyArray1<'py, bool>,
) -> PyResult<(Bound<'py, PyArray1<f32>>, Bound<'py, PyArray1<f32>>)> {
    let (th, rr, rz, pres) = (theta.as_slice()?, r.as_slice()?, raw_zero.as_slice()?, present.as_slice()?);
    let (tmv, rmv, mr, io) = (theta_mean.as_slice()?, r_mean.as_slice()?, mean_r.as_slice()?, intensity_only.as_slice()?);
    let n = th.len();
    for (l, w) in [(rr.len(), "r"), (rz.len(), "raw_zero"), (pres.len(), "present"), (mr.len(), "mean_r"), (io.len(), "intensity_only")] {
        check(l, n, w)?;
    }
    check(tmv.len(), 3 * n, "theta_mean")?;
    check(rmv.len(), 3 * n, "r_mean")?;
    let mut baf = vec![0f32; n];
    let mut lrr = vec![0f32; n];
    for i in 0..n {
        let (t, r) = (th[i], rr[i]);
        let (at, bt, ct) = (tmv[3 * i], tmv[3 * i + 1], tmv[3 * i + 2]);
        let (ar, br, cr) = (rmv[3 * i], rmv[3 * i + 1], rmv[3 * i + 2]);
        let (eq, lo, hi) = (t == bt, t < bt, t > bt);
        let (r_ref, mut b) = if eq {
            (br, 0.5f32)
        } else if lo {
            (ar + ((t - at) * (ar - br)) / (at - bt), ((t - at) / (bt - at)) * 0.5)
        } else if hi {
            (br + ((t - bt) * (cr - br)) / (ct - bt), 0.5 + ((t - bt) / (ct - bt)) * 0.5)
        } else {
            (f32::NAN, NEG_NAN32)
        };
        let denom = if io[i] { mr[i] } else { r_ref };
        let ratio = (r / denom) as f64;
        let mut l = math_log(ratio, false) as f32;
        if r == 0.0 {
            l = -f32::MAX;
        }
        if !(eq || lo || hi) {
            l = NEG_NAN32;
        }
        b = if b < 0.0 {
            0.0
        } else if b > 1.0 {
            1.0
        } else {
            b
        };
        if rz[i] || !pres[i] || b.is_nan() {
            b = NEG_NAN32;
        }
        if rz[i] || !pres[i] || l.is_nan() {
            l = NEG_NAN32;
        }
        baf[i] = b;
        lrr[i] = l;
    }
    Ok((baf.into_pyarray(py), lrr.into_pyarray(py)))
}