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

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

from __future__ import annotations;

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

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

__all__ = [
    'empty_mask',
    'MaskValue',
    'Mask',
];

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# ENUMS
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

class MaskValue(Enum):
    ZERO = 0;
    ONE = 1;
    UNSET = '*';

class Mask():
    index: int;
    values: List[MaskValue];

    def __init__(self, values: List[MaskValue]):
        self.values = values;
        if MaskValue.UNSET in values:
            self.index = values.index(MaskValue.UNSET);
        else:
            self.index = -1;
        return;

    def __len__(self) -> int:
        return len(self.values);

    def __str__(self) -> str:
        return ''.join([ str(m.value) for m in self.values ]);

    @property
    def choice(self) -> List[Fraction]:
        assert all(x != MaskValue.UNSET for x in self.values);
        return [ Fraction(x.value) for x in self.values ];

    @property
    def indexes_set(self) -> List[int]:
        return [i for i, value in enumerate(self.values) if value != MaskValue.UNSET];

    @property
    def indexes_one(self) -> List[int]:
        return [i for i, value in enumerate(self.values) if value == MaskValue.ONE];

    @property
    def indexes_zero(self) -> List[int]:
        return [i for i, value in enumerate(self.values) if value == MaskValue.ZERO];

    @property
    def indexes_unset(self) -> List[int]:
        return [i for i, value in enumerate(self.values) if value == MaskValue.UNSET];

    def splittable(self) -> bool:
        return self.index >= 0;

    def split(self) -> Tuple[Mask, Mask]:
        vector1 = self.values[:];
        vector1[self.index] = MaskValue.ZERO;
        vector2 = self.values[:];
        vector2[self.index] = MaskValue.ONE;
        return Mask(vector1), Mask(vector2);

    def pad(self, x: MaskValue) -> Mask:
        '''
        Pads unset values with a give by given value.
        '''
        return Mask([ x if u == MaskValue.UNSET else u for u in self.values ]);

    @property
    def support(self) -> List[int]:
        return [ i for i, v in enumerate(self.values) if v == MaskValue.ONE ];

def empty_mask(n: int):
    return Mask([MaskValue.UNSET for _ in range(n)]);