#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# IMPORTS
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

from src.thirdparty.types import *;
from src.thirdparty.maths import *;

from src.models.hirschberg import *;

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# EXPORTS
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

__all__ = [
    'compute_cost_matrix',
    'update_cost_matrix',
];

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# METHODS cost matrix + optimal paths
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

def compute_cost_matrix(
    X: str,
    Y: str,
) -> Tuple[np.ndarray, np.ndarray]: # Tuple[NDArray[(Any, Any), int], NDArray[(Any, Any), Directions]]:
    '''
    Berechnet Hirschberg-Costs-Matrix (ohne Rekursion).

    Annahmen:
    - X[0] = gap
    - Y[0] = gap
    '''
    m = len(X); # display vertically
    n = len(Y); # display horizontally
    Costs = np.full(shape=(m, n), dtype=int, fill_value=0);
    Moves = np.full(shape=(m, n), dtype=Directions, fill_value=Directions.UNSET);

    # zuerst 0. Spalte und 0. Zeile ausfüllen:
    for i, x in list(enumerate(X))[1:]:
        update_cost_matrix(Costs, Moves, x, '', i, 0);

    for j, y in list(enumerate(Y))[1:]:
        update_cost_matrix(Costs, Moves, '', y, 0, j);

    # jetzt alle »inneren« Werte bestimmen:
    for i, x in list(enumerate(X))[1:]:
        for j, y in list(enumerate(Y))[1:]:
            update_cost_matrix(Costs, Moves, x, y, i, j);
    return Costs, Moves;

def update_cost_matrix(
    Costs: np.ndarray, # NDArray[(Any, Any), int],
    Moves: np.ndarray, # NDArray[(Any, Any), Directions],
    x: str,
    y: str,
    i: int,
    j: int,
):
    '''
    Schrittweise Funktion zur Aktualisierung vom Eintrag `(i,j)` in der Kostenmatrix.

    Annahme:
    - alle »Vorgänger« von `(i,j)` in der Matrix sind bereits optimiert.

    @inputs
    - `Costs` - bisher berechnete Kostenmatrix
    - `Moves` - bisher berechnete optimale Schritte
    - `i`, `x` - Position und Wert in String `X` (»vertical« dargestellt)
    - `j`, `y` - Position und Wert in String `Y` (»horizontal« dargestellt)
    '''

    # nichts zu tun, wenn (i, j) == (0, 0):
    if i == 0 and j == 0:
        Costs[0, 0] = 0;
        return;

    ################################
    # NOTE: Berechnung von möglichen Moves wie folgt.
    #
    # Fall 1: (i-1,j-1) ---> (i,j)
    # ==> Stringvergleich ändert sich wie folgt:
    #     s1          s1 x
    #    ----  --->  ------
    #     s2          s2 y
    #
    # Fall 2: (i,j-1) ---> (i,j)
    # ==> Stringvergleich ändert sich wie folgt:
    #     s1          s1 GAP
    #    ----  --->  -------
    #     s2          s2  y
    #
    # Fall 3: (i-1,j) ---> (i,j)
    # ==> Stringvergleich ändert sich wie folgt:
    #     s1          s1  x
    #    ----  --->  -------
    #     s2          s2 GAP
    #
    # Diese Fälle berücksichtigen wir:
    ################################
    edges = [];
    if i > 0 and j > 0:
        edges.append((
            Directions.DIAGONAL,
            Costs[i-1, j-1] + missmatch_penalty(x, y),
        ));
    if j > 0:
        edges.append((
            Directions.HORIZONTAL,
            Costs[i, j-1] + gap_penalty(y),
        ));
    if i > 0:
        edges.append((
            Directions.VERTICAL,
            Costs[i-1, j] + gap_penalty(x),
        ));

    if len(edges) > 0:
        # Sortiere nach Priorität (festgelegt in Enum):
        edges = sorted(edges, key=lambda x: x[0].value);
        # Wähle erste Möglichkeit mit minimalen Kosten:
        index = np.argmin([ cost for _, cost in edges]);
        Moves[i, j], Costs[i, j] = edges[index];
    return;