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

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

from __future__ import annotations;

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

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

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

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

ALPHA = 'abcdefghijklmnopqrstuvwxyz';
GREEK = 'αβγδεζηθικλμνξοπρςτυφχψω';
T1 = TypeVar('T1');
T2 = TypeVar('T2');

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

def randomset_integers(N: int = -1, low: int = 1, high: int = 1) -> list[int]:
    if N == -1:
        N = random.randint(low, high);
    return list(range(1, N+1));

def randomset_alphabet(N: int = -1, low: int = 1, high: int = 1) -> list[int]:
    if N == -1:
        N = random.randint(low, high);
    return list([a for k, a in enumerate(ALPHA) if k < N]);

def randomset_greek(N: int = -1, low: int = 1, high: int = 1) -> list[int]:
    if N == -1:
        N = random.randint(low, high);
    return list([a for k, a in enumerate(GREEK) if k < 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 ];