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

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

from __future__ import annotations;

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

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

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

__all__ = [
    'FunctionDiagramWidget',
];

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

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

    def __init__(
        self,
        fnames: list[str],
        setnames: list[str],
        N: Optional[int] = None,
    ):
        self.state = None;
        self.N = N;
        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;

    def run(self):
        if self.N is None:
            control_nr = widgets.IntSlider(description=f'# Funktionen', min=1, max=self.N_max);
            display(control_nr);
        else:
            control_nr = widgets.IntSlider(value=self.N);
        self.state = widgets.interactive_output(self.handler_wrapper, controls={'N': control_nr});
        display(self.state);

    def handler_plot(self, N: int, show_labels: bool, **kwargs):
        cardinalities = [ kwargs[f'card_{k}'] for k in range(N+1)];
        injective = [ kwargs[f'injective_{k}'] for k in range(N)];
        surjective = [ kwargs[f'surjective_{k}'] for k in range(N)];
        X = [ randomset_alphabet(N=card) for card in cardinalities ];
        comp = Functions(*[
            Function(
                name     = (f'{self.fnames[k]}', f'{self.setnames[k]}', f'{self.setnames[k+1]}'),
                domain   = X[k],
                codomain = X[k+1],
                fct      = random_function(X[k], X[k+1], injective=injective[k], surjective=surjective[k], force=False)
            )
            for k in range(N)
        ]);
        fig, axs = comp.draw(show_labels=show_labels);
        # for k in range(N):
        #     [m, n] = cardinalities[k:][:2];
        #     inj = injective[k];
        #     surj = surjective[k];
        #     if m < n and surj:
        #         print(f'\x1b[1m{self.fnames[k]}\x1b[0m kann nicht surjektiv sein');
        #     if m > n and inj:
        #         print(f'\x1b[1m{self.fnames[k]}\x1b[0m kann nicht injektiv sein');
        #     if n == 1 and not surj:
        #         print(f'\x1b[1m{self.fnames[k]}\x1b[0m kann nicht nicht-surjektiv sein');
        #     if m == 1 and not inj:
        #         print(f'\x1b[1m{self.fnames[k]}\x1b[0m kann nicht nicht-injektiv sein');
        return;

    def handler_wrapper(self, N: int):
        show_labels = widgets.Checkbox(description='Labels anzeigen?', style={'description_width': 'initial'}, visible=True);
        control_nr = widgets.IntSlider(value=N);
        controls_card = [
            widgets.IntSlider(
                value       = None,
                description = f'|{self.setnames[k]}|',
                min         = 1,
                max         = 24,
            )
            for k in range(N+1)
        ];
        controls_injective = [
            widgets.Checkbox(
                value       = None,
                description = f'{self.fnames[k]} injektiv?',
                style       = {
                    'description_width': 'initial',
                },
                visible     = True,
            )
            for k in range(N)
        ];
        controls_surjective = [
            widgets.Checkbox(
                value       = None,
                description = f'{self.fnames[k]} surjectiv?',
                style       = {
                    'description_width': 'initial',
                },
                visible     = True,
            )
            for k in range(N)
        ];
        # controls_func = [None for _ in range(2*N)];
        # controls_func[::2] = controls_injective;
        # controls_func[1::2] = controls_surjective;
        # controls_func = [
        #     widgets.RadioButtons(
        #         options     = [ None, 'inj', 'surj', 'bij' ],
        #         # default:
        #         value       = 'inj',
        #         layout      = {
        #             'description_width': 'initial',
        #             'width': 'max-content',
        #         },
        #         description = f'{self.fnames[k]}',
        #         disabled    = False,
        #     )
        #     for k in range(N)
        # ];
        ui = widgets.VBox(
            [show_labels] \
            + controls_card \
            + [
                widgets.HBox([control_injective, control_surjective])
                for control_injective, control_surjective in zip(controls_injective, controls_surjective)
            ]
        );
        display(ui);
        display(widgets.interactive_output(
            f = self.handler_plot,
            controls = { 'N': control_nr, 'show_labels': show_labels } \
                | { f'card_{k}': controls_card[k] for k in range(N+1) } \
                | { f'injective_{k}': controls_injective[k] for k in range(N) } \
                | { f'surjective_{k}': controls_surjective[k] for k in range(N) }
        ));
        return;