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

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

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

from models.generated.config import *;
from src.core.utils import *;
from src.models.rucksack import *;
from src.models.stacks import *;
from src.algorithms.rucksack.display import *;

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

__all__ = [
    'rucksack_greedy_algorithm',
    'rucksack_branch_and_bound_algorithm',
];

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# METHOD greedy algorithm
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

def rucksack_greedy_algorithm(
    max_cost: float,
    costs: np.ndarray,
    values: np.ndarray,
    items: np.ndarray,
    fractional: bool,
    verbose: bool,
) -> Solution:
    '''
    Durch den Greedy-Algorithm wird der optimale Wert eines Rucksacks
    unter Rücksicht der Kapizitätsschranke eingeschätzt.

    NOTE: Wenn man `fractional = True` verwendet, liefert der Algorithmus
    eine obere Schranke des maximalen Wertes beim Originalproblem.
    '''
    # sortiere daten:
    order = get_sort_order(costs=costs, values=values);

    # verbose output hier behandeln (irrelevant für Algorithmus):
    if verbose:
        repr = display_order(order=order, costs=costs, values=values, items=items, one_based=True);
        print('');
        print('\x1b[1mRucksack Problem - Greedy\x1b[0m');
        print('');
        print(repr);
        print('');

    # führe greedy aus:
    n = len(costs);
    cost_total = 0;
    choice = [ Fraction(0) for _ in range(n) ];
    for i in order:
        # füge Item i hinzu, solange das Gesamtgewicht noch <= Schranke
        if cost_total + costs[i] <= max_cost:
            cost_total += costs[i];
            choice[i] = Fraction(1);
        # falls Bruchteile erlaubt sind, füge einen Bruchteil des i. Items hinzu und abbrechen
        elif fractional:
            choice[i] = Fraction(Fraction(max_cost - cost_total)/Fraction(costs[i]), _normalize=False);
            break;
        # ansonsten weiter machen:
        else:
            continue;

    # Aspekte der Lösung speichern:
    rucksack = [i for i, v in enumerate(choice) if v > 0]; # Indexes von Items im Rucksack
    soln = Solution(
        order = order,
        choice = choice,
        items = items[rucksack].tolist(),
        costs = costs[rucksack].tolist(),
        values = values[rucksack].tolist(),
    );

    # verbose output hier behandeln (irrelevant für Algorithmus):
    if verbose:
        repr_rucksack = display_rucksack(items=items, costs=costs, values=values, choice=choice);
        print('\x1b[1mEingeschätzte Lösung\x1b[0m');
        print('');
        print(f'Mask: [{", ".join(map(str, soln.choice))}]');
        print('Rucksack:')
        print(repr_rucksack);
        print('');

    # Lösung ausgeben
    return soln;

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# METHOD branch and bound algorithm
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

def rucksack_branch_and_bound_algorithm(
    max_cost: float,
    costs: np.ndarray,
    values: np.ndarray,
    items: np.ndarray,
    verbose: bool,
) -> Solution:
    '''
    Durch Branch & Bound wird der optimale Wert eines Rucksacks
    unter Rücksicht der Kapizitätsschranke exakt und effizienter bestimmt.
    '''

    order = get_sort_order(costs=costs, values=values);

    # verbose output hier behandeln (irrelevant für Algorithmus):
    if verbose:
        repr = display_order(order=order, costs=costs, values=values, items=items, one_based=True);
        print('');
        print('\x1b[1mRucksack Problem - Branch & Bound\x1b[0m');
        print('');
        print(repr);
        print('');

    logged_steps = [];
    step: Step;
    mask = empty_mask(n=len(costs));
    bound = np.inf;
    S = Stack();
    S.push(mask);
    while not S.empty():
        # top-Element auslesen und Bound berechnen:
        A: Mask = S.top();
        bound_subtree, choice, order_, state = estimate_lower_bound(mask=A, max_cost=max_cost, costs=costs, values=values, items=items);

        # für logging (irrelevant für Algorithmus):
        if verbose:
            step = Step(bound=bound, bound_subtree=bound_subtree, stack_str=str(S), choice=choice, order=order_, indexes=A.indexes_unset, solution=state);
            if bound_subtree < bound:
                if state is not None:
                    step.move = EnumBranchAndBoundMove.BOUND;
                    step.bound = bound_subtree;
                else:
                    step.move = EnumBranchAndBoundMove.BRANCH;
            logged_steps.append(step);

        S.pop();
        # Update nur nötig, wenn die (eingeschätzte) untere Schranke von A das bisherige Minimum verbessert:
        if bound_subtree < bound:
            # Bound aktualisieren, wenn sich A nicht weiter aufteilen od. wenn sich A wie eine einelementige Option behandeln läst:
            if state is not None:
                bound =  bound_subtree;
                mask = state;
            # Branch sonst
            else:
                B, C = A.split();
                S.push(B);
                # Nur dann C auf Stack legen, wenn mind. eine Möglichkeit in C die Kapazitätsschranke erfüllt:
                if sum(costs[C.indexes_one]) <= max_cost:
                    S.push(C);

    # Aspekte der Lösung speichern
    rucksack = mask.indexes_one; # Indexes von Items im Rucksack
    soln = Solution(
        order = order,
        choice = mask.choice,
        items = items[rucksack].tolist(),
        values = values[rucksack].tolist(),
        costs = costs[rucksack].tolist(),
    );

    # verbose output hier behandeln (irrelevant für Algorithmus):
    if verbose:
        repr = display_branch_and_bound(values=values, steps=logged_steps);
        repr_rucksack = display_rucksack(items=items, costs=costs, values=values, choice=mask.choice);
        print(repr);
        print('');
        print('\x1b[1mLösung\x1b[0m');
        print('');
        print(f'Mask: [{", ".join(map(str, soln.choice))}]');
        print('Rucksack:');
        print(repr_rucksack);
        print('');

    # Lösung ausgeben
    return soln;

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# AUXILIARY METHOD resort
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

def get_sort_order(costs: np.ndarray, values: np.ndarray) -> List[int]:
    '''
    Sortiert Daten absteigend nach values/costs.
    '''
    n = len(costs);
    indexes = list(range(n));
    margin = [ value/cost for cost, value in zip(costs, values) ];
    order = sorted(indexes, key=lambda i: -margin[i]);
    return order;

def estimate_lower_bound(
    mask: Mask,
    max_cost: float,
    costs: np.ndarray,
    values: np.ndarray,
    items: np.ndarray,
) -> Tuple[float, List[Fraction], List[int], Optional[Mask]]:
    '''
    Wenn partielle Information über den Rucksack festgelegt ist,
    kann man bei dem unbekannten Teil das Rucksack-Problem
    mit Greedy-Algorithmus »lösen«,
    um schnell eine gute Einschätzung zu bestimmen.

    NOTE: Diese Funktion wird `g(mask)` im Skript bezeichnet.
    '''
    indexes_one = mask.indexes_one;
    indexes_unset = mask.indexes_unset;
    n = len(mask);
    choice = np.zeros(shape=(n,), dtype=Fraction);
    order = np.asarray(range(n));

    # Berechnungen bei Items mit bekanntem Status in Rucksack:
    value_rucksack = sum(values[indexes_one]);
    cost_rucksack = sum(costs[indexes_one]);
    choice[indexes_one] = Fraction(1);

    # Für Rest des Rucksacks (Items mit unbekanntem Status):
    cost_rest = max_cost - cost_rucksack;
    state = None;
    # Prüfe, ob man als Lösung alles/nichts hinzufügen kann:
    if len(indexes_unset) == 0:
        state = mask;
        value_rest = 0;
    elif sum(costs[indexes_unset]) <= cost_rest:
        state = mask.pad(MaskValue.ONE);
        choice[indexes_unset] = Fraction(1);
        value_rest = sum(values[indexes_unset]);
    elif min(costs[indexes_unset]) > cost_rest:
        state = mask.pad(MaskValue.ZERO);
        choice[indexes_unset] = Fraction(0);
        value_rest = 0;
    # Sonst mit Greedy-Algorithmus lösen:
    # NOTE: Lösung ist eine Überschätzung des max-Wertes.
    else:
        soln_rest = rucksack_greedy_algorithm(
            max_cost = cost_rest, # <- Kapazität = Restgewicht
            costs = costs[indexes_unset],
            values = values[indexes_unset],
            items = items[indexes_unset],
            fractional = True,
            verbose = False,
        );
        choice[indexes_unset] = soln_rest.choice;
        value_rest = soln_rest.total_value;
        # Berechne Permutation für Teilrucksack
        permute_part(order, indexes=indexes_unset, order=soln_rest.order, in_place=True);

    # Einschätzung des max-Wertes:
    value_max_est = value_rucksack + value_rest;

    # Ausgabe mit -1 multiplizieren (weil maximiert wird):
    return -value_max_est, choice.tolist(), order.tolist(), state;