#!/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(
    capacity: float,
    weights: 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 = resort_by_value_per_weight(weights=weights, values=values, items=items);
    uorder = iperm(order);

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

    # führe greedy aus:
    n = len(weights);
    weight_total = 0;
    vector = [ 0 for _ in range(n) ];
    for i in range(n):
        # füge Item i hinzu, solange das Gesamtgewicht noch <= Schranke
        if weight_total + weights[i] <= capacity:
            weight_total += weights[i];
            vector[i] = 1;
        # sonst abbrechen. Falls Bruchteile erlaubt, füge einen Bruchteil des i. Items hinzu
        else:
            if fractional:
                vector[i] = (capacity - weight_total)/weights[i];
            break;

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

    # verbose output hier behandeln (irrelevant für Algorithmus):
    if verbose:
        # Umsortierung rückgängig machen:
        vector = [ soln.vector[i] for i in order ];
        rucksack = [ uorder[r] for r in rucksack ];
        permute_data(weights=weights, values=values, items=items, perm=uorder);
        # Ausdrücke bestimmen:
        expr_value = display_sum(vector=vector, values=values);
        expr_weight = display_sum(vector=vector, values=weights);
        print('\x1b[1mEingeschätzte Lösung\x1b[0m');
        print('');
        if fractional:
            print(f'Mask: {soln.vector}');
            print(f'      ---> {vector} (unter urspr. Sortierung)');
        print(f'Rucksack: {", ".join(items[rucksack])}.');
        print(f'max. Value ≈ {expr_value}');
        print(f'∑ Weights  = {expr_weight}');
        print('');

    # Lösung ausgeben
    return soln;

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

def rucksack_branch_and_bound_algorithm(
    capacity: float,
    weights: 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 = resort_by_value_per_weight(weights=weights, values=values, items=items);
    uorder = iperm(order);

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

    logged_steps = [];
    vector = empty_mask(n=len(weights));
    lb_estimate = np.inf;
    S = Stack();
    S.push(vector);
    while not S.empty():
        lb, u = estimate_lower_bound(mask=S.top(), capacity=capacity, weights=weights, values=values, items=items);
        if verbose:
            logged_steps.append((lb_estimate, lb, u, str(S)));
        # Update nur nötig, wenn die (eingeschätzte) untere Schranke von A das bisherige Minimum verbessert:
        A: Mask = S.pop();
        if lb < lb_estimate:
            # Branch:
            if A.splittable():
                B, C = A.split();
                # NOTE: per Wahl erfüllt mind. eine Möglichkeit in B die Kapazitätsschranke.
                S.push(B);
                # Nur dann C auf Stack legen, wenn mind. eine Möglichkeit in C die Kapazitätsschranke erfüllt:
                if sum(weights[C.indexes_one]) <= capacity:
                    S.push(C);
            # Bound, wenn die Maske den Rucksack komplett bestimmt:
            else:
                lb_estimate = lb;
                vector = A;

    # Aspekte der Lösung speichern
    rucksack = vector.indexes_one; # Indexes von Items im Rucksack
    soln = Solution(
        vector = vector.decision,
        items = items[rucksack].tolist(),
        values = values[rucksack].tolist(),
        weights = weights[rucksack].tolist(),
    );

    # verbose output hier behandeln (irrelevant für Algorithmus):
    if verbose:
        # NOTE: Information in Tabelle gemäß permutierten Daten:
        repr = display_branch_and_bound(values=values, steps=logged_steps);
        # Umsortierung rückgängig machen:
        vector = [ soln.vector[i] for i in order ];
        rucksack = [ uorder[r] for r in rucksack ];
        permute_data(weights=weights, values=values, items=items, perm=uorder);
        # Ausdrücke bestimmen:
        expr_value = display_sum(vector=vector, values=values);
        expr_weight = display_sum(vector=vector, values=weights);
        print('\x1b[1mLösung\x1b[0m');
        print('');
        print(repr);
        print('');
        print(f'Mask: {soln.vector}');
        print(f'      ---> {vector} (unter urspr. Sortierung)');
        print(f'Rucksack: {", ".join(items[rucksack])}.');
        print(f'max. Value ≈ {expr_value}');
        print(f'∑ Weights  = {expr_weight}');
        print('');

    # Lösung ausgeben
    return soln;

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

def resort_by_value_per_weight(
    weights: np.ndarray,
    values: np.ndarray,
    items: np.ndarray,
) -> List[int]:
    '''
    Sortiert Daten absteigend nach values/weights.
    '''
    n = len(weights);
    indexes = list(range(n));
    order = sorted(indexes, key=lambda i: -values[i]/weights[i]);
    permute_data(weights=weights, values=values, items=items, perm=order);
    return order;

def permute_data(
    weights: np.ndarray,
    values: np.ndarray,
    items: np.ndarray,
    perm: List[int],
):
    weights[:] = weights[perm];
    values[:] = values[perm];
    items[:] = items[perm];
    return;

def estimate_lower_bound(
    mask: Mask,
    capacity: float,
    weights: np.ndarray,
    values: np.ndarray,
    items: np.ndarray,
) -> Tuple[float, List[float]]:
    '''
    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(vector)` im Skript bezeichnet.
    '''

    # Berechnungen bei Items mit bekanntem Status in Rucksack:
    indexes_one = mask.indexes_one;
    weight_rucksack = sum(weights[indexes_one]);
    value_rucksack = sum(values[indexes_one]);

    # Löse mit Greedy-Algorithmus auf Items mit unbekanntem Status.
    # NOTE: Lösung ist eine Überschätzung des max-Wertes.
    indexes_unset = mask.indexes_unset;
    soln_rest = rucksack_greedy_algorithm(
        capacity = capacity - weight_rucksack, # <- Kapazität = Restgewicht
        weights = weights[indexes_unset],
        values = values[indexes_unset],
        items = items[indexes_unset],
        fractional = True,
        verbose = False,
    );
    value_rest = soln_rest.total_value;

    # Lösung des [0,1]-Vektors speichern:
    n = len(mask);
    vector = np.zeros(shape=(n,), dtype=float);
    vector[indexes_one] = 1;
    vector[indexes_unset] = soln_rest.vector;

    # Einschätzung des max-Wertes (Ausgabe mit -1 multiplizieren):
    value_max_est = value_rucksack + value_rest;
    return -value_max_est, vector.tolist();