use crate::types::*;
use crate::utils;
use crate::utils::backtrace_utils::*;

use rayon::prelude::*;

const NULL_OFFSET: isize = -10; // TODO: use Option<Offset>

// TODO: move to config
const MIN_WAVEFRONT_LENGTH: i32 = 10;
const MAX_DISTANCE_THRESHOLD: usize = 50;

fn reduce_wavefront_offsets(
    wavefront: &mut WaveFront,
    min_distance: usize,
    qlen: usize,
    tlen: usize,
    alignment_k: usize,
) -> bool {
    // TODO: move this into a fn
    let compute_distance = |offset: Offset, diagonal: Diagonal| -> usize {
        let v = utils::compute_v(offset, diagonal, alignment_k);
        let h = utils::compute_h(offset, diagonal, alignment_k);
        let left_v = qlen - v;
        let left_h = tlen - h;

        max!(left_v, left_h)
    };

    // Reduce from bottom
    let baz = if alignment_k == 0 {
        0
    } else {
        (alignment_k - 1) as i32
    };
    let top_limit = min!(baz, wavefront.get_hi());
    for diagonal in wavefront.get_lo()..top_limit {
        let k: usize = utils::compute_k(qlen, diagonal);
        let offset: Offset = match wavefront.get_offset_at_diagonal(k) {
            None => continue,
            Some(o) => o,
        };

        let distance = compute_distance(offset, diagonal);
        if (distance - min_distance) <= MAX_DISTANCE_THRESHOLD {
            break;
        }
        let lo = wavefront.get_lo();
        wavefront.set_lo(lo + 1);
    }

    // Reduce from top
    let baz = (alignment_k + 1) as i32;
    let bottom_limit = min!(baz, wavefront.get_lo());
    for diagonal in wavefront.get_lo()..bottom_limit {
        let k: usize = utils::compute_k(qlen, diagonal);
        let offset: Offset = match wavefront.get_offset_at_diagonal(k) {
            None => continue,
            Some(o) => o,
        };

        let distance = compute_distance(offset, diagonal);
        if (distance - min_distance) <= MAX_DISTANCE_THRESHOLD {
            break;
        }
        let hi = wavefront.get_hi();
        wavefront.set_hi(hi - 1);
    }

    let hi = wavefront.get_hi();
    let lo = wavefront.get_lo();

    // check hi/lo range
    if lo > hi {
        // let wavefront_set: &mut WaveFrontSet = wavefronts.get_wavefront_set_mut(score).unwrap();
        // wavefront_set.anull(WaveFrontTypes::M);
        true
    } else {
        false
    }
}

pub fn wf_reduce(wavefronts: &mut WaveFronts, score: Score) {
    let alignment_k = wavefronts.central_diagonal();
    let qlen: usize = wavefronts.qlen();
    let tlen: usize = wavefronts.tlen();

    // let wf_set = wavefronts.get_wavefront_set(score).unwrap();

    // let mm =  wf_set.m_mut();
    // let ii =  wf_set.m();
    // let dd =  wf_set.m();

    let mut min_distance: usize = max!(qlen, tlen);

    let m_lo: Diagonal;
    let m_hi: Diagonal;

    {
        // Fetch m-wavefront
        let maybe_m_wavefront: Option<&WaveFront> = wavefronts
            .get_wavefront_set(score)
            .and_then(|wf: &WaveFrontSet| wf.m());

        let m_wavefront: &WaveFront = match maybe_m_wavefront {
            Some(x) => x,
            None => return,
        };

        if (m_wavefront.get_hi() - m_wavefront.get_lo() + 1) < MIN_WAVEFRONT_LENGTH {
            return;
        }

        // TODO: move this into a fn
        let compute_distance = |offset: Offset, diagonal: Diagonal| -> usize {
            let v = utils::compute_v(offset, diagonal, alignment_k);
            let h = utils::compute_h(offset, diagonal, alignment_k);
            let left_v = qlen - v;
            let left_h = tlen - h;

            max!(left_v, left_h)
        };

        // Compute min distance
        for diagonal in m_wavefront.get_lo()..=m_wavefront.get_hi() {
            let k: usize = wavefronts.compute_k(diagonal);
            let offset: Offset = match m_wavefront.get_offset_at_diagonal(k) {
                None => continue,
                Some(o) => o,
            };
            let distance: usize = compute_distance(offset, diagonal);
            min_distance = min!(min_distance, distance);
        }

        m_lo = m_wavefront.get_lo();
        m_hi = m_wavefront.get_hi();
    }

    {
        // Reduce the m-wavefront
        let m_wavefront: &mut WaveFront = wavefronts
            .wavefronts
            .get_mut(score as usize)
            .and_then(|x: &mut Option<Box<WaveFrontSet>>| x.as_deref_mut())
            .and_then(|x: &mut WaveFrontSet| x.m_mut())
            .unwrap();
        // let m_wavefront: &mut WaveFront = wavefront_set.m_mut().unwrap();
        let nullify = reduce_wavefront_offsets(m_wavefront, min_distance, qlen, tlen, alignment_k);

        if nullify {
            let wavefront_set: &mut WaveFrontSet = wavefronts.get_wavefront_set_mut(score).unwrap();
            wavefront_set.anull(WaveFrontTypes::M);
        }
    }

    {
        // Reduce the i-wavefront
        let wavefront_set: &mut WaveFrontSet = wavefronts
            .wavefronts
            .get_mut(score as usize)
            .and_then(|x: &mut Option<Box<WaveFrontSet>>| x.as_deref_mut())
            .unwrap();
        let maybe_i_wavefront: Option<&mut WaveFront> = wavefront_set.i_mut();

        if maybe_i_wavefront.is_some() {
            let i_wavefront: &mut WaveFront = maybe_i_wavefront.unwrap();

            if m_lo > i_wavefront.get_lo() {
                i_wavefront.set_lo(m_lo);
            }
            if m_hi < i_wavefront.get_hi() {
                i_wavefront.set_hi(m_hi);
            }
            if i_wavefront.get_lo() > i_wavefront.get_hi() {
                wavefront_set.anull(WaveFrontTypes::I);
            }
        }
    }

    {
        // Reduce the d-wavefront
        let wavefront_set: &mut WaveFrontSet = wavefronts
            .wavefronts
            .get_mut(score as usize)
            .and_then(|x: &mut Option<Box<WaveFrontSet>>| x.as_deref_mut())
            .unwrap();
        let maybe_d_wavefront: Option<&mut WaveFront> = wavefront_set.d_mut();

        if maybe_d_wavefront.is_some() {
            let d_wavefront: &mut WaveFront = maybe_d_wavefront.unwrap();

            if m_lo > d_wavefront.get_lo() {
                d_wavefront.set_lo(m_lo);
            }
            if m_hi < d_wavefront.get_hi() {
                d_wavefront.set_hi(m_hi);
            }
            if d_wavefront.get_lo() > d_wavefront.get_hi() {
                wavefront_set.anull(WaveFrontTypes::I);
            }
        }
    }
}

pub fn wf_extend<F, G>(
    wf: &mut WaveFront,
    central_diagonal: usize,
    compute_k: &mut F,
    match_lambda: &mut G,
) where
    F: FnMut(Diagonal) -> Option<usize>,
    G: FnMut(&mut usize, &mut usize) -> bool,
{
    if false {
        eprintln!("[wf_extend]");
        eprintln!("\tdiagonal \t k \t offset \t (v, h)");
    }

    for diagonal in wf.get_lo()..=wf.get_hi() {
        let k = compute_k(diagonal);
        if k.is_none() {
            continue;
        };

        let k = k.unwrap();
        let offset = wf.get_offset_at_diagonal(k).unwrap();

        let mut v: usize = utils::compute_v(offset, diagonal, central_diagonal);
        let mut h: usize = utils::compute_h(offset, diagonal, central_diagonal);

        if false {
            eprintln!(
                "\t{} \t\t {} \t {} \t\t ({}, {})",
                diagonal, k, offset, v, h
            );
        }

        let mut previous_v_position = v;

        while match_lambda(&mut v, &mut h) {
            let matches = (v - previous_v_position) as u64;
            previous_v_position = v;

            wf.set_offset_at_diagonal_with_closure(k, |x: &mut Offset| *x += matches);
        }
    }
}

pub fn wf_backtrace<F>(
    wavefronts: &WaveFronts,
    score: Score,
    backtrace_lambda: &mut F,
) -> Result<String, String>
where
    F: FnMut((i32, i32), (i32, i32)),
{
    if false {
        eprintln!("[wfa::wf_backtrace]");
    }

    let mut cigar = String::new();

    let central_diagonal: usize = wavefronts.central_diagonal();
    let alignment_k = central_diagonal as i32;
    let mut k = central_diagonal as i32;

    // compute starting location
    let diagonal: usize = wavefronts.compute_k(alignment_k);

    let maybe_offset = wavefronts
        .get_wavefront_set(score)
        .and_then(|wf_set: &WaveFrontSet| {
            // eprintln!("{}", wf_set.m().unwrap());
            wf_set.m_get_offset_at_diagonal(diagonal)
        });

    let offset: u64 = match maybe_offset {
        Some(offset) => offset,
        _ => {
            return Err(format!(
                "could not find offset for score {} at diagonal {}",
                score, k
            ))
        }
    };

    if false {
        eprintln!("\t score={} offset={}", score, offset);
    }

    let mut v: usize = wavefronts.compute_v(offset, k);
    let mut h: usize = wavefronts.compute_h(offset, k);

    let mut backtrace_op = Operation::MatchMismatch;

    // TODO: move to config
    let x: isize = 1;
    let o: isize = 1;
    let e: isize = 1;

    let mut score = score as isize;
    let mut offset = offset as isize;

    if false {
        eprintln!("\t v={} h={} score={} offset={}", v, h, score, offset);
    }

    while v > 0 && h > 0 && score > 0 {
        if false {
            eprintln!("\tbacktrace_op = {:?}", backtrace_op);
        }

        // compute scores
        let gap_open_score = score - o - e;
        let gap_extend_score = score - e;
        let mismatch_score = score - x;

        let del_ext: isize = if backtrace_op == Operation::Insertion {
            NULL_OFFSET
        } else {
            backtrace_deletion_extend_offset(wavefronts, gap_extend_score, k)
        };

        let del_open: isize = if backtrace_op == Operation::Insertion {
            NULL_OFFSET
        } else {
            backtrace_deletion_open_offset(wavefronts, gap_open_score, k)
        };

        let ins_ext: isize = if backtrace_op == Operation::Deletion {
            NULL_OFFSET
        } else {
            backtrace_insertion_extend_offset(wavefronts, gap_extend_score, k)
        };

        let ins_open: isize = if backtrace_op == Operation::Deletion {
            NULL_OFFSET
        } else {
            backtrace_insertion_open_offset(wavefronts, gap_open_score, k)
        };

        let misms: isize = if backtrace_op != Operation::MatchMismatch {
            NULL_OFFSET
        } else {
            backtrace_mismatch_offset(wavefronts, mismatch_score, k)
        };

        // Compute maximum offset
        let max_all: isize = max!(del_ext, del_open, ins_ext, ins_open, misms);

        if false {
            eprintln!(
                "\t[wfa::wf_backtrace while loop] \
                 score={} offset={} k={} \
                 gap_open_score={} gap_extend_score={} mismatch_score={} \
                 max_all={} del_ext={} del_open={} ins_ext={} ins_open={} mims={} \
                 backtrace_op={:?}",
                score,
                offset,
                k,
                gap_open_score,
                gap_extend_score,
                mismatch_score,
                max_all,
                del_ext,
                del_open,
                ins_ext,
                ins_open,
                misms,
                backtrace_op
            );
        }

        // Traceback Matches
        if backtrace_op == Operation::MatchMismatch && offset >= max_all {
            let num_matches = (offset - max_all) as usize;
            backtrace_matches_check(
                &mut offset,
                &mut cigar,
                num_matches,
                k,
                central_diagonal,
                backtrace_lambda,
            );
            offset = max_all;
        }

        if max_all == del_ext {
            // Add Deletion
            cigar.push('D');
            // Update state
            score = gap_extend_score;
            k += 1;
            backtrace_op = Operation::Deletion;
        } else if max_all == del_open {
            // Add Deletion
            cigar.push('D');
            // Update state
            score = gap_open_score;
            k += 1;
            backtrace_op = Operation::MatchMismatch;
        } else if max_all == ins_ext {
            // Add Insertion
            cigar.push('I');
            // Update state
            score = gap_extend_score;
            k -= 1;
            offset -= 1;
            backtrace_op = Operation::Insertion;
        } else if max_all == ins_open {
            // Add Insertion
            cigar.push('I');
            // Update state
            score = gap_open_score;
            k -= 1;
            offset -= 1;
            backtrace_op = Operation::MatchMismatch;
        } else if max_all == misms {
            // Add Mismatch
            cigar.push('X');

            // Update state
            score = mismatch_score;
            offset -= 1;
        } else {
            return Err(String::from(
                "Backtrace error: No link found during backtrace",
            ));
        }

        // Update coordinates
        if false {
            eprintln!(
                "\t[wfa::wf_backtrace while loop] bad stuff. offset is {}",
                offset
            );
        }
        v = wavefronts.compute_v(offset as u64, k);
        h = wavefronts.compute_h(offset as u64, k);
    }

    if score == 0 {
        // backtrace matches check
        let num_matches = offset as usize;
        backtrace_matches_check(
            &mut offset,
            &mut cigar,
            num_matches,
            k,
            central_diagonal,
            backtrace_lambda,
        );
    } else {
        while v > 0 {
            cigar.push('D');
            v -= 1;
        }

        while h > 0 {
            cigar.push('I');
            h -= 1;
        }
    }

    Ok(run_length_encode(&cigar[..], true))
}

pub fn wf_next(wavefronts: &mut WaveFronts, score: Score, _penalties: &Penalties) {
    if false {
        eprintln!("[wf_next]");
    }

    let s: isize = score as isize;

    let x: isize = 1;
    let o: isize = 1;
    let e: isize = 1;

    // TODO: how to use these penalties and not run out of memory
    // let x = penalties.mismatch;
    // let o = penalties.gap_open;
    // let e = penalties.gap_extend;

    let ms = s - x;
    let go = s - o - e;
    let ge = s - e;

    // Compute bounds
    let compute_bounds =
        |s: isize, bound: WaveFrontBounds, wavefront_type: &WaveFrontTypes| -> Diagonal {
            let default_diagonal = wavefronts.central_diagonal();
            if s < 0 {
                return default_diagonal as i32;
            };

            wavefronts
                .get_wavefront_set(s as u64)
                .and_then(|wavefront_set: &WaveFrontSet| match wavefront_type {
                    WaveFrontTypes::I => wavefront_set.i(),
                    WaveFrontTypes::M => wavefront_set.m(),
                    WaveFrontTypes::D => wavefront_set.d(),
                })
                .map(|wf: &WaveFront| match bound {
                    WaveFrontBounds::Hi => wf.get_hi(),
                    WaveFrontBounds::Lo => wf.get_lo(),
                })
                .unwrap_or(default_diagonal as i32)
        };

    let wf_types = vec![
        WaveFrontTypes::M,
        WaveFrontTypes::M,
        WaveFrontTypes::I,
        WaveFrontTypes::D,
    ];
    let scores = vec![ms, go, ge, ge];
    let baz = wf_types.par_iter().zip(scores.par_iter());

    let mut hi: Diagonal = baz
        .map(|(wf_type, s): (&WaveFrontTypes, &isize)| {
            compute_bounds(*s, WaveFrontBounds::Hi, wf_type)
        })
        .max()
        .unwrap()
        + 1;

    let baz = wf_types.par_iter().zip(scores.par_iter());
    let mut lo: Diagonal = baz
        .map(|(wf_type, s): (&WaveFrontTypes, &isize)| {
            compute_bounds(*s, WaveFrontBounds::Lo, wf_type)
        })
        .max()
        .unwrap()
        - 1;

    if false {
        eprintln!("\t{} {} {}", hi, lo, score);
    }

    // Compute wavefronts

    if lo < wavefronts.min_k() {
        lo = wavefronts.min_k();
    };

    if hi > wavefronts.max_k() {
        hi = wavefronts.max_k();
    }

    wavefronts.add_wavefront_set(score, hi, lo);

    for k in lo..=hi {
        if k < wavefronts.min_k() || k > wavefronts.max_k() {
            continue;
        }

        let bar = |s: isize, k: i32, wf_type: WaveFrontTypes| -> Offset {
            let default_offset = 0;
            if s < 0 || k < 0 {
                return default_offset;
            };

            let k = wavefronts.compute_k(k);

            wavefronts
                .get_wavefront_set(s as u64)
                .and_then(|wavefront_set: &WaveFrontSet| match wf_type {
                    WaveFrontTypes::I => wavefront_set.i(),
                    WaveFrontTypes::M => wavefront_set.m(),
                    WaveFrontTypes::D => wavefront_set.d(),
                })
                .and_then(|wf: &WaveFront| wf.get_offset_at_diagonal(k))
                .unwrap_or(default_offset)
        };

        let i = *vec![
            bar(go, k - 1, WaveFrontTypes::M),
            bar(ge, k - 1, WaveFrontTypes::I),
        ]
        .iter()
        .max()
        .unwrap()
            + 1;
        let d = *vec![
            bar(go, k + 1, WaveFrontTypes::M),
            bar(ge, k + 1, WaveFrontTypes::D),
        ]
        .iter()
        .max()
        .unwrap();
        let m = *vec![bar(ms, k, WaveFrontTypes::M) + 1, i, d]
            .iter()
            .max()
            .unwrap();

        let k = wavefronts.compute_k(k);

        let wf: &mut WaveFrontSet = wavefronts.get_wavefront_set_mut(score).unwrap();
        wf.set_wavefront(k, i, WaveFrontTypes::I);
        wf.set_wavefront(k, d, WaveFrontTypes::D);
        wf.set_wavefront(k, m, WaveFrontTypes::M);
    }
}
