import { Sequence } from './SeqNode'
import { initializeTable, lastColumnOf, indicesOf, lastRowOf, initializePointerTable, indicesInTable, fillByIndex } from './Table'
import { AffineGapAlignment, AlignmentPair, AlignmentResult, AlignType, DirectedPosition, DynamicTable, GapsParams, GapType, Graph, isAffineGapAlignment, LinearGapAlignment, Path, PathMap, Position, SimilarityFunction, TableName, Triple, TripleTable } from './types'

/**
 * Generic sequence aligner
 * 
 * @param alignType: string, defining type of alignment. There are four alignment types:
 *                           - `global` (i.e. Needlman-Wunch or more precisely a variant of David Sankoff algorithm
 *                                            see  Sankoff, D. (1972). "Matching sequences under deletion-insertion constraints". 
 *                                                 Proceedings of the National Academy of Sciences of the United States of America. 69 (1): 4–6
 *                                       )
 *                           - `local` (i.e. Smith-Waterman algorithm)
 *                           - `semi-global` alignment
 *                           - `end-gap-free` alignment
 *  @todo to implement overlap alignment such as the following pattern:
 *            +++++++******--------
 *            -------******========
 *            - is the gap and +,*,= are the symbols/characters
 *        which is roughly performed by assigning 0 to the first column and take the max score in the last row of the dynamic table
 *                           
 *  @description After reading multiple articles and implementations, I came to a conclusion that there were inconsistencies
 *               in many of the implementations and/or the formulation of the dynamic programming recursion -- especially
 *               for the affine gap penalty case. Interestingly -- after figuring out that -- I found a paper discussing/reporting
 *               the same frustration I witnessed. 
 *                       
 *                       ` Are all global alignment algorithms and implementations correct? <http://www.biorxiv.org/content/biorxiv/early/2015/11/12/031500.full.pdf>`__
 *               
 *               This article should be THE reference paper for debugging/verifying alignment implementation -- especially for affine gap penalty. 
 */
export class Aligner<LabelType> {
    alignType: AlignType
    seq1: Sequence<LabelType> = []
    seq2: Sequence<LabelType> = []
    alignment?: LinearGapAlignment | AffineGapAlignment
    similarityFunction: SimilarityFunction

    constructor(alignType: AlignType) {
        this.alignType = alignType
        this.similarityFunction = SimilarityFunction.Euclidean
    }

    setAlignType(alignType: AlignType) {
        this.alignType = alignType
    }

    setSimilarityFunction(similarityFunction: SimilarityFunction) {
        this.similarityFunction = similarityFunction
    }

    align(seq1: Sequence<LabelType>, seq2: Sequence<LabelType>, gapsParams: GapsParams) {
        this.seq1 = seq1
        this.seq2 = seq2
        let alignType = this.alignType
        // parse the gaps -- assume that gaps are negative numbers
        const gapOpen = gapsParams.gapOpen
        const gapExt = gapsParams.gapExt

        if(!gapExt) { //  case of linear gap applied
            // backtrack matrix
            let Vp: TripleTable = initializePointerTable(seq1.length, seq2.length)
 
            // initialize V matrix -- holding the best alignment score for seq1[1:i] and seq2[1:j]
            let V: DynamicTable = initializeTable(seq1.length+1, seq2.length+1, 0.0)
 
            if(alignType === AlignType.Global) {
                // fill the top row  with the linear gap
                V[0] = [0, ...fillByIndex(seq2.length, (i: number) => ((i+1) * gapOpen))]
                // fill the first left column with the linear gap
                for (let i=1; i<seq1.length+1; i++) {
                    V[i][0] = gapOpen * i
                }
            }
            else if (alignType == AlignType.SemiGlobal) {
                // in this implementation, seq_2 is always considered the pattern sequence (i.e. generally the pattern sequence is the shorter sequence)
                // we penalize gaps in the pattern sequence (seq_2) but not in the long sequence (seq_1)
                // fill the top row with the linear gap
                V[0] = [0, ...fillByIndex(seq2.length, (i: number) => ((i+1) * gapOpen))]
            }

            for (let i=1; i<seq1.length+1; i++) {
                for (let j=1; j<seq2.length+1; j++) {
                    // compute score
                    let node_seq1 = seq1[i-1]
                    let node_seq2 = seq2[j-1]
                    let sim_score = node_seq1.computeSimilarity(node_seq2, /***kwargs*/)
                    let case_1 = V[i][j-1] + gapOpen // align node_j from seq_2 to a gap
                    let case_2 = V[i-1][j-1] + sim_score // align node_i from seq_1 to node_j from seq_2
                    let case_3 = V[i-1][j] + gapOpen // align node_i from seq_1 to a gap
                    let cases_vec: Triple = [case_1, case_2, case_3]
                    this.fillDynamicAndPointerTables(cases_vec, V, Vp, [i,j], /* main_dtable*/true)
                }
            }
            this.alignment = {
                V,
                Vp,
                end_i: seq1.length-1,
                end_j: seq2.length-1
            }
        }
        else { // case of affine gap applied
            // backtrack matrices
            let Vp: TripleTable = initializePointerTable(seq1.length, seq2.length)
            let Fp: TripleTable = initializePointerTable(seq1.length, seq2.length)
            let Ep: TripleTable = initializePointerTable(seq1.length, seq2.length)
            // initialize V matrix -- holding the best alignment score for seq1[1:i] and seq2[1:j]
            let V: DynamicTable = initializeTable(seq1.length+1, seq2.length+1, 0.0)
            // best alignment score where node_i in seq_1 is aligned to a space
            let F: DynamicTable = initializeTable(seq1.length+1, seq2.length+1, -Infinity)
            // best alignment score where node_j in seq_2 is aligned to a space
            let E: DynamicTable = initializeTable(seq1.length+1, seq2.length+1, -Infinity)
              
            if(alignType === AlignType.Global) {
                // the affine gap function used
                // g(k) = -gap_open-gap_ext*k where k is the gap length

                // fill the top row  with the affine gap
                V[0] = [0, ...fillByIndex(seq2.length, (i: number) => ((i+1) * gapExt + gapOpen))]

                // fill the first left column with the affine gap
                for (let i=1; i<seq1.length+1; i++) {
                    V[i][0] = gapOpen + gapExt * i
                }
            }
            else if(alignType === AlignType.SemiGlobal) {
                // seq_2 is always considered the pattern sequence (i.e. pattern sequence should be the shorter sequence)
                // we penalize gaps in the pattern sequence (seq_2) but not in the long sequence (seq_1)
                // fill the top row with the linear gap
                V[0] = [0, ...fillByIndex(seq2.length, (i: number) => ((i+1) * gapExt + gapOpen))]
            }
            for (let i=1; i<seq1.length+1; i++) {
                for (let j=1; j<seq2.length+1; j++) {
                    // compute score
                    let node_seq1 = seq1[i-1]
                    let node_seq2 = seq2[j-1]
                    let sim_score = node_seq1.computeSimilarity(node_seq2 /* kwargs */)

                    let e_vec: Triple =
                                [E[i][j-1] + gapExt,
                                 V[i][j-1] + gapOpen + gapExt, 
                                 F[i][j-1] + gapOpen + gapExt]
                    this.fillDynamicAndPointerTables(e_vec, E, Ep, [i,j])
                    
                    let f_vec: Triple =
                                [E[i-1][j] + gapOpen + gapExt,
                                 V[i-1][j] + gapOpen + gapExt, 
                                 F[i-1][j] + gapExt]
                    this.fillDynamicAndPointerTables(f_vec, F, Fp, [i,j])

                    let v_vec: Triple =
                                [E[i-1][j-1] + sim_score,
                                 V[i-1][j-1] + sim_score, 
                                 F[i-1][j-1] + sim_score] 
                    this.fillDynamicAndPointerTables(v_vec, V, Vp, [i,j], /*main_dtable =*/ true)
                }
            }
            this.alignment = {
                V, Vp, E, Ep, F, Fp,
                end_i: seq1.length,
                end_j: seq2.length
            }
        }
    }

    // todo: herauslösen
    protected fillDynamicAndPointerTables(vec: Triple,
                                            dynamic_table: DynamicTable,
                                            pointer_table: TripleTable,
                                            pos: Position, main_dtable = false) {
        let direc_flags: Triple = [0, 0, 0]
        let vec_max = Math.max(...vec)
        if (main_dtable && this.alignType === AlignType.Local) { // case of V matrix and local alignment is required
            vec_max = Math.max(vec_max, 0)
        }
        const index = vec.findIndex(element => element === vec_max)
        direc_flags[index] = 1
        const [i, j] = pos
        pointer_table[i][j] = direc_flags
        dynamic_table[i][j] = vec_max
    }

    public retrieveAlignments(numPaths: number = 1): AlignmentResult<LabelType> {
        if (!this.alignment) {
            throw new Error("align() must be performed before retrieving alignment results")
        }

        if (isAffineGapAlignment(this.alignment)) {
            return this.retrieveAlignmentsAffine(this.alignment, numPaths)
        }
        return this.retrieveAlignmentsLinear(this.alignment, numPaths)
    }

    protected retrieveAlignmentsLinear(linearGapAlignment: LinearGapAlignment,
                                     num_paths: number = 1): AlignmentResult<LabelType> {
        const {
            V, Vp, end_i, end_j
        } = linearGapAlignment

        const align_type = this.alignType
        // get the starting nodes with max align score corresponding to the specified alignment type
        let root_nodes: Position[] = []
        if (align_type === AlignType.Global) {
            root_nodes.push([end_i, end_j])
        }
        else if (align_type === AlignType.SemiGlobal) {
            // maximum in the last column
            const lastCol = lastColumnOf(V)
            const maxLastCol = Math.max(...lastCol)
            const rowIndices = indicesOf(lastCol, maxLastCol)
            for (const rowIndex of rowIndices) {
                root_nodes.push([rowIndex, end_j])
            }
        }
        else if(align_type == AlignType.EndGapFree) {
            const lastRow = lastRowOf(V)
            const maxLastRow = Math.max(...lastRow)

            const lastCol = lastColumnOf(V)
            const maxLastCol = Math.max(...lastCol)

            if(maxLastRow === maxLastCol) {
                for (const indexCol of indicesOf(lastRow, maxLastRow)) {
                    root_nodes.push([end_i, indexCol])
                }
                for (const indexRow of indicesOf(lastCol, maxLastCol)) {
                    root_nodes.push([indexRow, end_j])
                }
            }
            else if(maxLastRow > maxLastCol) {
                for (const indexCol of indicesOf(lastRow, maxLastRow)) {
                    root_nodes.push([end_i, indexCol])
                }
            }
            else if(maxLastRow < maxLastCol) {
                for (const indexRow of indicesOf(lastCol, maxLastCol)) {
                    root_nodes.push([indexRow, end_j])
                }
            }
        }
        else if(align_type == AlignType.Local) {
            // coordinates of the maximum in V
            const max_matrix = Math.max(...V.flat())
            for (const coord of indicesInTable(V, max_matrix)) {
                root_nodes.push(coord)
            }
        }

        const [x, y] = root_nodes[0]
        const score = V[x][y]
        if (num_paths === 0) return {
            score,
            paths: []
        }
        
        let alignmentPaths = []
        // building the graph
        for (const root_node of root_nodes) {
            // building a graph from the nodes (i.e. coordinates (i,j))
            let coord_graph: Graph<Position> = new Map<Position, Position[]>()
            let terminal_nodes: Position[] = []
            this.build_coord_graph_linear(root_node, V, Vp, coord_graph, terminal_nodes)

            // prepare for traversing the graph using depth-first 
            let all_paths: PathMap<LabelType> = this.traversePath(root_node, coord_graph, GapType.LINEAR)
            alignmentPaths.push(this.buildAlignments(all_paths, terminal_nodes))
            // delete all_paths 

            if(alignmentPaths.length >= num_paths) {
                return {
                    score,
                    paths: alignmentPaths
                }
            }
        }
        return {
            score,
            paths: alignmentPaths
        }
    }
    
    private build_coord_graph_linear(curr_node: Position, V: DynamicTable, Vp: TripleTable, coord_graph: Graph<Position>, terminal_nodes: Position[]) {
        // determine if we have reached a terminal node
        let [i, j] = curr_node
        let conds: Map<AlignType, boolean> = new Map(
            [[AlignType.Global,     curr_node === [0,0]],
             [AlignType.Local,      V[i][j] === 0],
             [AlignType.SemiGlobal, curr_node === [0,0]],
             [AlignType.EndGapFree, curr_node === [0,0]]])

        if(conds.get(this.alignType)) {
            coord_graph.set(curr_node, [])
            terminal_nodes.push(curr_node)
            return
        }

        // explore the children of parent node
        if(!coord_graph.has(curr_node)) {
            // get children nodes
            const children: Position[] = []
            for (const [position, direction] of Vp[i][j].entries()) {
                if (direction) {
                    switch (position) {
                        case 0: children.push([i, j-1]); break;     // horizontal direction (left)
                        case 1: children.push([i-1, j-1]); break;   // diagonal direction
                        case 2: children.push([i-1, j]); break;     // vertical direction (up)
                    }
                }
            }

            coord_graph.set(curr_node, children)
            for (const child of children) {
                this.build_coord_graph_linear(child, V, Vp, coord_graph, terminal_nodes)
            }
        }
    }

    private retrieveAlignmentsAffine(affineGapAlignment: AffineGapAlignment, num_paths: number=1): AlignmentResult<LabelType> {
        const {
            V, Vp, 
            E, Ep, 
            F, Fp,
            end_i,
            end_j
        } = affineGapAlignment

        // get the starting nodes with max align score corresponding to the specified alignment type
        let root_nodes: DirectedPosition[] = []
        let score: number = 0
        if(this.alignType === AlignType.Global) {
            const vec: Triple = [
                Math.max(E[end_i][end_j]), 
                Math.max(V[end_i][end_j]), 
                Math.max(F[end_i][end_j]), 
            ]
            const v_max = Math.max(...vec)

            let dtable_name: TableName = TableName.None
            for (const indexPosition of indicesOf(vec, v_max)) {
                switch (indexPosition) {
                    case 0: 
                        dtable_name = TableName.E
                        score = E[end_i][end_j]
                        break;
                    case 1:
                        dtable_name = TableName.V
                        score = V[end_i][end_j]
                        break;
                    case 2: 
                        dtable_name = TableName.F
                        score = F[end_i][end_j]
                        break;
                }
                root_nodes.push([[end_i, end_j], dtable_name])
            }
        }
        else if(this.alignType === AlignType.SemiGlobal) {
            let lastCol = V.map((value: number[]) => value[value.length-1])
            let maxLastCol = Math.max(...lastCol)
            for (const indexRow of indicesOf(lastCol, maxLastCol)) {
                root_nodes.push([[indexRow, end_j], TableName.V])
            }
            let [row, column] = root_nodes[0][0]
            score = V[row][column]
        }
        else if(this.alignType === AlignType.EndGapFree) {
            const maxLastRow = Math.max(...lastRowOf(V))
            const maxLastCol = Math.max(...lastColumnOf(V))
            if(maxLastRow == maxLastCol) {
                for (const indexColumn of indicesOf(lastRowOf(V), maxLastRow)) {
                    root_nodes.push([[end_i, indexColumn], TableName.V])
                } 
                for (const indexRow of indicesOf(lastColumnOf(V), maxLastCol)) {
                    root_nodes.push([[indexRow, end_j], TableName.V])
                }
            }
            else if(maxLastRow > maxLastCol) {
                for (const indexColumn of indicesOf(lastRowOf(V), maxLastRow)) {
                    root_nodes.push([[end_i, indexColumn], TableName.V])
                }
            }
            else if(maxLastCol < maxLastCol) {
                for (const indexRow of indicesOf(lastColumnOf(V), maxLastCol)) {
                    root_nodes.push([[indexRow, end_j], TableName.V])
                }
            }
            const [row, column] = root_nodes[0][0]
            score = V[row][column]
        }
        else if(this.alignType === AlignType.Local) {
            const max_matrix = Math.max(...V.flat())
            for (const coord of indicesInTable(V, max_matrix)) {
                root_nodes.push([coord, TableName.V])
            }
            const [row, column] = root_nodes[0][0]
            score = V[row][column]
        }

        if(num_paths === 0) return {
            score, paths: []
        }

        // building a graph from the nodes (i.e. coordinates (i,j))
        let alignmentPaths: Path<LabelType>[] = []
        for (const root_node of root_nodes) {
            let coord_graph: Graph<DirectedPosition> = new Map<DirectedPosition, DirectedPosition[]>()
            let terminal_nodes: DirectedPosition[] = []
            this.build_coord_graph_affine(root_node, V, Vp, Ep, Fp, coord_graph, terminal_nodes)

            // prepare for traversing the graph using depth-first 
            let all_paths = this.traversePath(root_node, coord_graph, GapType.AFFINE)
            let alignmentPath = this.buildAlignments(all_paths, terminal_nodes)
            alignmentPaths.push(alignmentPath)
            //del all_paths
            if(alignmentPaths.length>=num_paths) {
                return {
                    score,
                    paths: alignmentPaths
                }
            }
        }
        return {
            score,
            paths: alignmentPaths
        }
    }

    
    private build_coord_graph_affine(curr_node: DirectedPosition,
                                     V: DynamicTable,
                                     Vp: TripleTable,
                                     Ep: TripleTable,
                                     Fp: TripleTable,
                                     coord_graph: Graph<DirectedPosition>,
                                     terminal_nodes: DirectedPosition[]) {
        // determine if we have reached a terminal node
        let [curr_pos, pointer] = curr_node
        let [i, j] = curr_pos

        // TODO this is stupid
        let conds: Map<AlignType, boolean> = new Map(
            [[AlignType.Global,     curr_pos[0] === 0 && curr_pos[1] === 0],
             [AlignType.Local,      V[i][j] === 0],
             [AlignType.SemiGlobal, curr_pos[0] === 0 && curr_pos[1] === 0],
             [AlignType.EndGapFree, curr_pos[0] === 0 && curr_pos[1] === 0]])
        
        if(conds.get(this.alignType)) {
            coord_graph.set(curr_node, [])
            terminal_nodes.push(curr_node)
            return
        }

        if (coord_graph.has(curr_node)) return

        // explore the children of parent node
        // get children nodes
        let new_pos: Position = [0,0]
        let children: DirectedPosition[] = []
        let pointer_table: TripleTable = []
        if(pointer === TableName.V) {
            pointer_table = Vp
            if(i===0 && j!==0) {
                new_pos = [i, j-1]
            }
            else if(i!==0 && j==0) {
                new_pos = [i-1, j]
            }
            else {
                new_pos = [i-1, j-1]
            }
        }
        else if(pointer === TableName.E) {
            pointer_table = Ep
            if(i===0 && j!==0) {
                new_pos = [i, j-1]
            }
            else if(i!==0 && j===0) {
                new_pos = [i-1, j]
            }
            else {
                new_pos = [i, j-1]
            }
        }
        else if(pointer == TableName.F) {
            pointer_table = Fp
            if(i==0 && j!=0) {
                new_pos = [i, j-1]
            }
            else if(i!=0 && j==0) {
                new_pos = [i-1, j]
            }
            else {
                new_pos = [i-1, j]
            }
        }

        for (const [position, direction] of pointer_table[i][j].entries()) {
            if(direction === 0) continue 

            switch (position) {
                case 0: children.push([new_pos, TableName.E]); break;       // horizontal (left)
                case 1: children.push([new_pos, TableName.V]); break;       // diagonal
                case 2: children.push([new_pos, TableName.F]); break;       // vertical

            }
        }
        coord_graph.set(curr_node, children)
        for (const child of children) {
            this.build_coord_graph_affine(child, V, Vp, Ep, Fp, coord_graph, terminal_nodes)
        }
    }

       
    private buildAlignments(paths: PathMap<LabelType>, terminal_nodes: DirectedPosition[] | Position[]): Path<LabelType> {
        let alignments_path = []
        for (const node of terminal_nodes) {
            const path = paths.get(node)

            if (!path || !path.length) {
                console.log('failed getting node', node, 'from paths')
                return []
            }
            
            for (let i=0; i<path.length; i++) {
                let segment: AlignmentPair<LabelType> = path[i]
                //segment.reverse() // TODO why?
                alignments_path.push(segment)
            }
        }
        return alignments_path
    }

    private traversePath(root_node: Position | DirectedPosition,
                          coord_graph: Graph<Position | DirectedPosition>,
                          gap_type: GapType = GapType.LINEAR): PathMap<LabelType> {
        type T = typeof root_node

        // breadth first search
        let q: T[] = [] // start a queue 
        q.push(root_node) // add the root node
        let q_parents: Map<T, T>[] = [] // queue for the parents
        let visited_path = new Map<[T, T], boolean>()
        let track_paths: PathMap<LabelType> = new Map<Position | DirectedPosition, Path<LabelType>>() // track paths
        while (q.length !== 0) {
            let curr_node = q.pop()!
            if (!curr_node) {
                console.log('current node is undefined')
            }

            let curr_parent = q_parents.pop()

            if (curr_parent && curr_parent.has(curr_node)) {
                // to protect against root node
                let parent = curr_parent.get(curr_node)
                if(parent && !visited_path.has([curr_node, parent])) {
                    let segment: AlignmentPair<LabelType> = this.getAlignmentPair(parent, curr_node, gap_type)
                    if(track_paths.has(parent)) {
                        //for (const path of track_paths.get(parent)!) {
                            const path = track_paths.get(parent)!
                            let tmp = path
                            tmp.push(segment)
                            if(track_paths.has(curr_node)) track_paths.get(curr_node)!.push(...tmp)
                            else track_paths.set(curr_node, tmp)
                        //}
                    }
                    else {
                        track_paths.set(curr_node, [segment])
                    }
                    
                    visited_path.set([curr_node, parent], true)
                }
            } 
            if (coord_graph.has(curr_node)) {
                for (const child of coord_graph.get(curr_node)!) {
                    q.push(child)
                    q_parents.push(new Map([[child, curr_node]]))
                }
            }
        }
        return track_paths
    }

    private getAlignmentPair(parent_node: DirectedPosition | Position,
                                      child_node: DirectedPosition | Position,
                                      gapType: GapType = GapType.LINEAR): AlignmentPair<LabelType> {
        const [p_i, p_j] = (gapType === GapType.LINEAR) ? (parent_node as Position) : (parent_node as DirectedPosition)[0]
        const [c_i, c_j] = (gapType === GapType.LINEAR) ? (child_node as Position) : (child_node as DirectedPosition)[0]
        
        let align_repr: [LabelType, LabelType] = ['-' as unknown as LabelType, '-' as unknown as LabelType]
        if((p_i-1 === c_i) && (p_j-1 == c_j)) // diagonal case
            align_repr = [this.seq1[p_i-1].label, this.seq2[p_j-1].label]
        else if((p_i-1 == c_i) && (p_j == c_j)) // vertical case
            align_repr = [this.seq1[p_i-1].label, '-' as unknown as LabelType]
        else if((p_i == c_i) && (p_j-1 == c_j)) // horizontal case
            align_repr = ['-' as unknown as LabelType, this.seq2[p_j-1].label]
        return align_repr 
    }
}