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

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

from __future__ import annotations;

from src.thirdparty.code import *;
from src.thirdparty.types import *;
from src.thirdparty.plots import *;
from src.thirdparty.maths import *;
from src.thirdparty.render import *;

from src.maths.diagrams import *;
from src.maths.sets import *;

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

__all__ = [
    'FunctionDiagramWidget',
];

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

class FunctionProperty(Enum):
    NOTHING = '—';
    INJECTIVE = 'injektiv';
    NOT_INJECTIVE = '¬ injektiv';
    SURJECTIVE = 'surjektiv';
    NOT_SURJECTIVE = '¬ surjektiv';
    BOTH = 'injektiv + surjektiv';
    INJECTIVE_NOT_SURJECTIVE = 'injektiv + ¬ surjektiv';
    NOT_INJECTIVE_SURJECTIVE = '¬ injektiv + surjektiv';
    NEITHER = '¬ injektiv + ¬ surjektiv';

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# Class
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

class FunctionDiagramWidget():
    state: Optional[widgets.Output];
    N: Optional[int];
    N_max: int;
    fnames: list[str];
    setnames: list[str];
    setsfrom: list[list[Any]];
    card: list[int];

    ctrl_nr: widgets.IntSlider;
    btn_refresh: widgets.ToggleButton;
    ctrls_card: list[widgets.IntSlider];
    ctrls_property: list[widgets.IntSlider];
    fcts: Functions;

    def __init__(
        self,
        fnames: list[str],
        setnames: list[str],
        setsfrom: dict[str, list[Any]] | list[list[Any]] = [],
        N: Optional[int] = None,
    ):
        self.state = None;
        self.N = N or 1;
        self.N_max = len(fnames);
        assert len(setnames) == len(fnames) + 1, f'The number of sets must be {self.N_max+1}.';
        self.fnames = fnames;
        self.setnames = setnames;
        # fix set contents:
        if isinstance(setsfrom, dict):
            default = randomset_alphabet(mode=Letters.ROMAN, shuffle=True, full=True);
            setsfrom = [ setsfrom.get(name, default) for name in setnames ];
        if len(setsfrom) == 0:
            setsfrom = [
                randomset_alphabet(mode=Letters.SYMBOLS, shuffle=True, full=True),
                randomset_alphabet(mode=Letters.ROMAN, shuffle=True, full=True),
                randomset_alphabet(mode=Letters.HEBREW, shuffle=True, full=True),
                randomset_alphabet(mode=Letters.GREEK, shuffle=True, full=True),
            ];
        n = len(setsfrom);
        self.setsfrom = [ setsfrom[k % n] for k, name in enumerate(setnames) ];
        self.card = [ min(3, len(XX)) for XX in self.setsfrom ];

    def run(self):
        '''
        Runs the widget.
        '''
        self.create_widgets();
        display(widgets.interactive_output(
            f = self.handler_main,
            controls = {'N': self.ctrl_nr},
        ));
        return;

    def create_widgets(self):
        '''
        Initialises the widgets.
        '''
        N = self.N_max;
        self.btn_refresh = widgets.ToggleButton(description='Neu laden');
        self.btn_show_labels = widgets.Checkbox(
            description = 'Labels anzeigen?',
            value = True,
            style       = {
                'description_width': 'initial',
            },
            visible     = True
        );
        self.ctrl_nr = widgets.IntSlider(value=self.N, description=f'# Funktionen', min=1, max=self.N_max);
        cards_max = [ len(XX) for XX in self.setsfrom ];
        card_max_max = max(cards_max);
        self.ctrls_card = [
            widgets.IntSlider(
                value       = card,
                description = f'|{setname}|',
                min         = 1,
                max         = card_max,
                layout = {
                    'width': 'initial',
                    # FIXME: scales wrong:
                    # 'width': f'{100*cards_max/card_max_max}%',
                },
            )
            for setname, XX, card, card_max in zip(
                self.setnames[:(N+1)],
                self.setsfrom[:(N+1)],
                self.card[:(N+1)],
                cards_max[:(N+1)],
            )
        ];
        self.ctrls_property = [
            widgets.RadioButtons(
                options     = [ e.value for e in possible_properties(self.card[k], self.card[k+1]) ],
                value       = FunctionProperty.NOTHING.value,
                layout      = {
                    'description_width': 'initial',
                    'width': 'initial',
                },
                disabled    = False,
            )
            for k in range(N)
        ];
        for k, ctrl_card in enumerate(self.ctrls_card):
            ctrl_card.observe(partial(self.handler_upd_card, k=k), names='value');
        # FIXME: this does not behave correctly:
        self.btn_show_labels.observe(partial(self.handler_plot, create=False));
        return;

    def handler_main(self, *_, **__):
        '''
        Main handler for updating widgets.
        '''
        N = self.N = self.ctrl_nr.value;
        ui = widgets.VBox([] \
            + [ self.btn_refresh ] \
            + [ self.btn_show_labels ] \
            + [ widgets.Text(f'Größen den Mengen:') ] \
            + self.ctrls_card[:(N+1)] \
            + [widgets.HBox(
                [
                    widgets.VBox(
                        [
                            widgets.Text(f'Eigenschaften von {fname}'),
                            ctrl_property
                        ],
                        layout = {
                            'description_width': 'initial',
                            'width': 'initial',
                            'display': 'flex',
                            'align_items': 'stretch',
                            'overflow': 'hidden',
                        }
                    )
                    for fname, ctrl_property in zip(self.fnames[:N], self.ctrls_property[:N])
                ],
            )],
        );
        display(ui);
        display(widgets.interactive_output(
            f = self.handler_plot,
            controls = dict(
                N           = self.ctrl_nr,
                refresh     = self.btn_refresh,
                # FIXME: shoud not have to rely on this:
                show_labels = self.btn_show_labels,
                **{
                    f'card[{k}]': ctrl_card
                    for k, ctrl_card in enumerate(self.ctrls_card)
                },
                **{
                    f'property[{k}]': ctrl_property
                    for k, ctrl_property in enumerate(self.ctrls_property)
                },
            ),
        ));
        return;

    def handler_plot(self, *_, create: bool = True, **__):
        '''
        Method to create sets, functions, and plot.
        '''
        self.card = [ ctrl_card.value for ctrl_card in self.ctrls_card ];
        if create:
            N = self.N;
            options = list(map(lambda ctrl: ctrl.value, self.ctrls_property[:N]));
            properties = list(map(option_to_property, options));
            injective = list(map(is_injective, properties));
            surjective = list(map(is_surjective, properties));
            sets = [ XX[:card] for XX, card in zip(self.setsfrom[:(N+1)], self.card[:(N+1)]) ];
            self.fcts = Functions(*[
                Function(
                    name     = (fname, nameX, nameY),
                    domain   = X,
                    codomain = Y,
                    fct      = random_function(X, Y, injective=inj, surjective=surj, force=False),
                )
                for fname, nameX, nameY, X, Y, inj, surj in zip(
                    self.fnames,
                    self.setnames,
                    self.setnames[1:],
                    sets,
                    sets[1:],
                    injective,
                    surjective,
                )
            ]);
        # FIXME:
        # upon toggling show labels button, this even triggers multiple times,
        # and only the old value registers.
        show_labels = self.btn_show_labels.value;
        fig, axs = self.fcts.draw(show_labels=show_labels);
        return;

    def handler_upd_card(self, change, k: int):
        '''
        Handler for altering function property options
        upon change of set cardinality.
        '''
        self.card[k] = change['new'];

        def upd(j: int):
            m = self.card[j];
            n = self.card[j+1];
            value = self.ctrls_property[j].value;
            options = [ e.value for e in possible_properties(m, n) ];
            if value not in options:
                value = FunctionProperty.NOTHING.value;
            self.ctrls_property[j].options = options;
            self.ctrls_property[j].value = value;
            return;

        if k > 0:
            upd(k-1);
        if k < self.N:
            upd(k);
        return;

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

def option_to_property(option: str) -> Optional[FunctionProperty]:
    try:
        return FunctionProperty(option);
    except:
        return None;

def possible_properties(m: int = 1, n: int = 1) -> list[FunctionProperty]:
    if m == n:
        if m == 1:
            return [
                e for e in FunctionProperty
                if (is_injective(e) not in [False])
                and (is_surjective(e) not in [False])
            ];
        else:
            return [
                e for e in FunctionProperty
                if not (is_surjective(e) in [True] and is_injective(e) in [False])
                and not (is_surjective(e) in [False] and is_injective(e) in [True])
            ];
    elif m == 1 and n > 1:
        return [
            e for e in FunctionProperty
            if (is_injective(e) not in [False])
            and (is_surjective(e) not in [True])
        ];
    elif n == 1 and m > 1:
        return [
            e for e in FunctionProperty
            if (is_surjective(e) not in [False])
            and (is_injective(e) not in [True])
        ];
    elif m > n:
        return [
            e for e in FunctionProperty
            if (is_injective(e) not in [True])
        ];
    elif n > m:
        return [
            e for e in FunctionProperty
            if (is_surjective(e) not in [True])
        ];
    return [ e for e in FunctionProperty ];

def is_injective(property: Optional[FunctionProperty]) -> Optional[bool]:
    match property:
        case FunctionProperty.INJECTIVE | FunctionProperty.INJECTIVE_NOT_SURJECTIVE | FunctionProperty.BOTH:
            return True;
        case FunctionProperty.NOT_INJECTIVE | FunctionProperty.NOT_INJECTIVE_SURJECTIVE | FunctionProperty.NEITHER:
            return False;
    return None;

def is_surjective(property: Optional[FunctionProperty]) -> Optional[bool]:
    match property:
        case FunctionProperty.SURJECTIVE | FunctionProperty.NOT_INJECTIVE_SURJECTIVE | FunctionProperty.BOTH:
            return True;
        case FunctionProperty.NOT_SURJECTIVE | FunctionProperty.INJECTIVE_NOT_SURJECTIVE | FunctionProperty.NEITHER:
            return False;
    return None;
