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

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

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

from models.generated.config import *;
from models.generated.commands import *;
from src.core.log import *;
from src.core.utils import *;
from src.models.random_walk import *;
from src.algorithms.random_walk.display import *;

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

__all__ = [
    'adaptive_walk_algorithm',
    'gradient_walk_algorithm',
    'metropolis_walk_algorithm',
];

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# METHOD adaptive walk
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

def adaptive_walk_algorithm(
    landscape: Landscape,
    r: float,
    coords_init: tuple,
    optimise: EnumOptimiseMode,
    verbose: bool,
):
    '''
    Führt den Adapative-Walk Algorithmus aus, um ein lokales Minimum zu bestimmen.
    '''

    # lege Fitness- und Umgebungsfunktionen fest:
    match optimise:
        case EnumOptimiseMode.max:
            f = lambda x: -landscape.fitness(*x);
        case _:
            f = lambda x: landscape.fitness(*x);
    nbhd = lambda x: landscape.neighbourhood(*x, r=r, strict=True);
    label = lambda x: landscape.label(*x);

    # initialisiere
    x = coords_init;
    fx = f(x);
    fy = fx;
    N = nbhd(x);

    # führe walk aus:
    while True:
        # Wähle zufälligen Punkt und berechne fitness-Wert:
        y = uniform_random_choice(N);
        fy = f(y);

        # Nur dann aktualisieren, wenn sich f-Wert verbessert:
        if fy < fx:
            # Punkt + Umgebung + f-Wert aktualisieren
            x = y;
            fx = fy;
            N = nbhd(x);
        else:
            # Nichts machen!
            pass;

        # Nur dann (erfolgreich) abbrechen, wenn f-Wert lokal Min:
        if fx <= min([f(y) for y in N], default=fx):
            break;

    return x;

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# METHOD gradient walk
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

def gradient_walk_algorithm(
    landscape: Landscape,
    r: float,
    coords_init: tuple,
    optimise: EnumOptimiseMode,
    verbose: bool,
):
    '''
    Führt den Gradient-Descent (bzw. Ascent) Algorithmus aus, um ein lokales Minimum zu bestimmen.
    '''

    # lege Fitness- und Umgebungsfunktionen fest:
    match optimise:
        case EnumOptimiseMode.max:
            f = lambda x: -landscape.fitness(*x);
        case _:
            f = lambda x: landscape.fitness(*x);
    nbhd = lambda x: landscape.neighbourhood(*x, r=r, strict=True);
    label = lambda x: landscape.label(*x);

    # initialisiere
    x = coords_init;
    fx = landscape.fitness(*x);
    fy = fx;
    N = nbhd(x);
    f_values = [f(y) for y in N];
    fmin = min(f_values);
    Z = [y for y, fy in zip(N, f_values) if fy == fmin];

    # führe walk aus:
    while True:
        # Wähle zufälligen Punkt mit steilstem Abstieg und berechne fitness-Wert:
        y = uniform_random_choice(Z);
        fy = fmin;

        # Nur dann aktualisieren, wenn sich f-Wert verbessert:
        if fy < fx:
            # Punkt + Umgebung + f-Wert aktualisieren
            x = y;
            fx = fy;
            N = nbhd(y);
            f_values = [f(y) for y in N];
            fmin = min(f_values);
            Z = [y for y, fy in zip(N, f_values) if fy == fmin];
        else:
            # Nichts machen!
            pass;

        # Nur dann (erfolgreich) abbrechen, wenn f-Wert lokal Min:
        if fx <= min([f(y) for y in N], default=fx):
            break;

    return x;

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# METHOD metropolis walk
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

def metropolis_walk_algorithm(
    landscape: Landscape,
    r: float,
    coords_init: tuple,
    T: float,
    annealing: bool,
    optimise: EnumOptimiseMode,
    verbose: bool,
):
    '''
    Führt den Metropolis-Walk Algorithmus aus, um ein lokales Minimum zu bestimmen.
    '''

    # lege Fitness- und Umgebungsfunktionen fest:
    match optimise:
        case EnumOptimiseMode.max:
            f = lambda x: -landscape.fitness(*x);
        case _:
            f = lambda x: landscape.fitness(*x);
    nbhd = lambda x: landscape.neighbourhood(*x, r=r, strict=True);
    label = lambda x: landscape.label(*x);

    # initialisiere
    x = coords_init;
    fx = f(x);
    fy = fx;
    nbhd_x = nbhd(x);

    # führe walk aus:
    k = 0;
    while True:
        # Wähle zufälligen Punkt und berechne fitness-Wert:
        y = uniform_random_choice(nbhd_x);
        r = uniform(0,1);
        fy = f(y);

        # Nur dann aktualisieren, wenn sich f-Wert verbessert:
        if fy < fx or r < math.exp(-(fy-fx)/T):
            # Punkt + Umgebung + f-Wert aktualisieren
            x = y;
            fx = fy;
            nbhd_x = nbhd(x);
        else:
            # Nichts machen!
            pass;

        # »Temperatur« ggf. abkühlen:
        if annealing:
            T = cool_temperature(T, k);

        # Nur dann (erfolgreich) abbrechen, wenn f-Wert lokal Min:
        if fx <= min([f(y) for y in nbhd_x], default=fx):
            break;

        k += 1;

    return x;

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# AUXILIARY METHODS
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

def cool_temperature(T: float, k: int, const: float = 1.) -> float:
    harm = const*(k + 1);
    return T/(1 + T/harm);