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

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

from __future__ import annotations;

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

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

__all__ = [
    'Letters',
    'randomset_integers',
    'randomset_alphabet',
    'random_function',
];

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# CONSTANTS / VARIABLES
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

# local usage only
T1 = TypeVar('T1');
T2 = TypeVar('T2');

class Letters(Enum):
    # 'abcdefghijklmnopqrstuvwxyz'
    ROMAN = [chr(97 + n) for n in range(26)];
    # 'αβγδεζηθικλμνξοπρςτυφχψω'
    GREEK = [chr(945 + n) for n in range(25)]
    # 'אבגדהוזחטיךכלםמןנסעףפץצקרשת' (but <—)
    HEBREW = [chr(1488 + n) for n in range(27)];
    SYMBOLS = [
        r'$|0\rangle$',
        r'$|\uparrow\rangle$',
        r'$|\downarrow\rangle$',
        r'$\sqrt{2}$',
        r'$\pi$',
        r'$e$',
        r'$\frac{1}{137}$',
        r'$\infty$',
        # r'$-\infty$',
        r'$\clubsuit$',
        # r'$\diamondsuit$',
        r'$\heartsuit$',
        # r'$\spadesuit$',
        r'$\hbar$',
    ];

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# METHODS
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

def randomset_integers(
    N:       int  = -1,
    low:     int  = 1,
    high:    int  = 1,
    shuffle: bool = False,
) -> list[int]:
    if N == -1:
        N = random.randint(low, high);
    values = list(range(1, N+1));
    if shuffle:
        return sample(values, size=N, replace=False);
    return values;

def randomset_alphabet(
    mode: Letters,
    N:       int  = -1,
    low:     int  = 1,
    high:    int  = 1,
    shuffle: bool = False,
    full:    bool = False,
) -> list[str]:
    letters: list[str] = mode.value;
    if full:
        N = len(letters);
    elif N == -1:
        high = min(low+1, high)
        N = random.randint(low, high);
    if shuffle:
        return sample(letters, size=N, replace=False);
    return letters[:N];

def random_function(
    X: list[T1],
    Y: list[T2],
    injective: Optional[bool] = None,
    surjective: Optional[bool] = None,
    force: bool = False,
) -> list[tuple[T1, T2]]:
    m = len(X);
    n = len(Y);
    if m == 0:
        return [];
    if n == 0:
        raise Exception(f'Impossible to create a function with {m} elements in the domain and {n} in the codomain.');
    if not force:
        if injective and m > n:
            injective = None;
        if surjective and m < n:
            surjective = None;
        if not injective and m == 1:
            injective = None;
        if not surjective and n == 1:
            surjective = None;
    match (injective, surjective):
        case (True, _):
            assert m <= n, f'Impossible to create an injective function with {m} elements in the domain and {n} in the codomain.';
            Y = random.sample(Y, m);
            return [(x, y) for x, y in zip(X, Y)];
        case (_, True):
            assert m >= n, f'Impossible to create an surjective function with {m} elements in the domain and {n} in the codomain.';
            indexes = random.sample(list(range(m)), n);
            g = [ (indexes[j], Y[j]) for j in range(n) ] \
                + [
                    (i, random.choice(Y))
                    for i in range(m)
                    if not i in indexes
                ];
            g = sorted(g, key=lambda o: o[0]);
            return [ (X[i], y) for (i, y) in g ];
        case (False, _):
            assert m > 1, f'Impossible to create a non-injective function with {m} elements in the domain.';
            indexes = random.sample(list(range(m)), m);
            g = random_function(indexes, Y);
            [(i0, y0), (i1, y1)] = g[:2];
            g[0] = (i0, y1);
            g = sorted(g, key=lambda o: o[0]);
            return [ (X[i], y) for (i, y) in g ];
        case (_, False):
            assert n > 1, f'Impossible to create a non-surjective function with {n} elements in the codomain.';
            Y = random.sample(Y, n-1);
            return random_function(X, Y);
        case _:
            return [ (x, random.choice(Y)) for x in X ];
