from statistics import NormalDist
from itertools import accumulate
from functools import cache
from operator import or_, and_
from math import comb

from .shared import *


class AbstractBHV:
    # recall
    # information_entropy(p) = -p*log2(p) - (1 - p)*log2(1 - p)
    # so for some active fractions, it should be a lot cheaper than for others
    # hence a specialized method for 2^-n
    @classmethod
    def rand(cls) -> Self:
        raise NotImplementedError()

    @classmethod
    def nrand(cls, n: int) -> list[Self]:
        return [cls.rand() for _ in range(n)]

    @classmethod
    def rand2(cls, power: int) -> Self:
        assert power >= 0
        r = cls.rand()
        return r if power == 0 else r & cls.rand2(power - 1)

    @classmethod
    def nrand2(cls, n: int, power: int) -> list[Self]:
        assert power >= 0
        return [cls.rand2(power) for _ in range(n)]

    @classmethod
    def random(cls, active: float) -> Self:
        raise NotImplementedError()

    @classmethod
    def nrandom(cls, n: int, active: float) -> list[Self]:
        assert 0. <= active <= 1.
        return [cls.random(active) for _ in range(n)]

    def select_rand(self, other: Self) -> Self:
        return self.rand().select(self, other)

    def select_rand2(self, other: Self, power_left: int) -> Self:
        return self.rand2(power_left).select(self, other)

    def select_random(self, other: Self, frac_left: float) -> Self:
        return self.random(frac_left).select(self, other)

    def randomize_half(self) -> Self:
        return self.rand().select_rand(self)

    def randomize_pow(self, pow_randomize: int) -> Self:
        return self.rand().select_rand2(self, pow_randomize)

    def randomize_frac(self, frac_randomize: float) -> Self:
        return self.rand().select_random(self, frac_randomize)

    def flip_half(self) -> Self:
        return self ^ self.rand()

    def flip_pow(self, pow_flip: int) -> Self:
        return self ^ self.rand2(pow_flip)

    def flip_pow_on(self, pow_flip_on: int) -> Self:
        return self | ~self.rand2(pow_flip_on)

    def flip_pow_off(self, pow_flip_off: int) -> Self:
        return self & self.rand2(pow_flip_off)

    def flip_frac(self, frac_flip: float) -> Self:
        return self ^ self.random(frac_flip)

    def flip_frac_on(self, frac_flip_on: float) -> Self:
        return self | self.random(frac_flip_on)

    def flip_frac_off(self, frac_flip_off: float) -> Self:
        return self & self.random(1. - frac_flip_off)

    def permute(self, permutation_id: int) -> Self:
        raise NotImplementedError()

    def swap_halves(self) -> Self:
        raise NotImplementedError()

    def rehash(self) -> Self:
        raise NotImplementedError()

    _FEISTAL_ROUNDS = 8

    @classmethod
    def _feistal_round(cls, block: Self, round_key: Self) -> Self:
        L = block & cls.HALF
        R = (block ^ (L ^ round_key).rehash().swap_halves()) & ~cls.HALF

        return (L | R).swap_halves()

    def feistal(self, k: Self, inv=False) -> Self:
        block = self
        rounds = range(self._FEISTAL_ROUNDS)

        for r in reversed(rounds) if inv else rounds:
            block = self._feistal_round(block, self._FEISTAL_SUBKEYS[r] & k)

        return block.swap_halves()

    @classmethod
    def majority3(cls, x, y, z) -> Self:
        return cls.majority([x, y, z])

    @classmethod
    def majority(cls, vs: list[Self]) -> Self:
        n = len(vs)

        if n == 0:
            return cls.rand()
        elif n == 1:
            return vs[0]
        elif n % 2 == 0:
            return cls.majority([cls.rand()] + vs)
        else:
            return cls._true_majority(vs)

    @classmethod
    def _true_majority(cls, vs: list[Self]) -> Self:
        assert len(vs) % 2 == 1,  "The true majority function which is only defined on uneven length inputs, see `majority`"
        half = len(vs)//2
        disjunction = list(accumulate(vs[:half-1:-1], or_))[1:]
        conjunction = list(accumulate(vs[:half-1:-1], and_))[1:]

        @cache
        def rec(ons, offs):
            if ons + 1 > half:
                return disjunction[-offs - 1]
            if offs + 1 > half:
                return conjunction[-ons - 1]
            return vs[ons + offs].select(
                rec(ons + 1, offs),
                rec(ons, offs + 1)
            )

        return rec(0, 0)

    def __eq__(self, other: Self) -> bool:
        raise NotImplementedError()

    def __xor__(self, other: Self) -> Self:
        raise NotImplementedError()

    def __and__(self, other: Self) -> Self:
        raise NotImplementedError()

    def __or__(self, other: Self) -> Self:
        raise NotImplementedError()

    def __invert__(self) -> Self:
        return self ^ self.ONE

    def select(self, when1: Self, when0: Self) -> Self:
        return when0 ^ (self & (when0 ^ when1))

    def active(self) -> int:
        raise NotImplementedError()

    def active_fraction(self) -> float:
        return self.active()/DIMENSION

    def hamming(self, other: Self) -> int:
        return (self ^ other).active()

    def bit_error_rate(self, other: Self) -> float:
        return (self ^ other).active_fraction()

    def jaccard(self, other: Self, distance=False) -> float:
        union_active = (self | other).active()
        if union_active == 0:
            raise RuntimeError("Jaccard index where both vectors are zero")
        res = (self & other).active() / union_active
        return 1. - res if distance else res

    def cosine(self, other: Self, distance=False) -> float:
        s_active = self.active()
        o_active = other.active()
        if not s_active or not o_active:
            raise RuntimeError("Cosine similarity where one of the two vectors is zero")
        res = (self & other).active() / (s_active*o_active)**.5
        return 1. - res if distance else res

    def std_apart(self, other: Self, invert=False) -> float:
        return self.frac_to_std(self.bit_error_rate(other), invert)

    @staticmethod
    def normal(mean=0., p=.5):
        return NormalDist(mean, (DIMENSION*p*(1 - p))**.5)

    @staticmethod
    def maj_ber(n):
        return comb(n - 1, (n - 1)//2)/2**n

    @staticmethod
    def frac_to_std(frac, invert=False):
        n = AbstractBHV.normal(0.0)
        estdvs = n.zscore(0.5*DIMENSION)
        stdvs = n.zscore(frac*DIMENSION)
        return estdvs - stdvs if invert else stdvs

    @staticmethod
    def std_to_frac(std, invert=False):
        n = AbstractBHV.normal(0.0)
        frac = (std*n.stdev + n.mean)/DIMENSION
        return 1. - frac if invert else frac

    def zscore(self) -> float:
        n = AbstractBHV.normal(0.5*DIMENSION)
        return n.zscore(self.active())

    def pvalue(self) -> float:
        n = AbstractBHV.normal(0.5*DIMENSION)
        s = n.cdf(self.active())
        return 2.*min(s, 1. - s)

    def related(self, other: Self, stdvs=6) -> bool:
        return abs(self.std_apart(other)) <= stdvs

    def unrelated(self, other: Self, stdvs=6) -> bool:
        return abs(self.std_apart(other, invert=True)) <= stdvs

    def bias_rel(self, other: Self, rel: Self) -> float:
        rel_l = (rel & self).active()
        rel_r = (rel & other).active()
        return rel_l/(rel_l + rel_r)

    # Alternative implementations

    @classmethod
    def _majority3_select(cls, a: Self, b: Self, c: Self) -> Self:
        # C:  1 0 0 1 0 1 1

        # |:  1 1 1 1 0 1 0
        # &:  1 0 0 0 0 0 0

        # M:  1 0 0 1 0 1 0

        return a.select(b | c, b & c)

    @classmethod
    def _majority5_select(cls, x4: Self, x3: Self, x2: Self, x1: Self, x0: Self) -> Self:
        t1 = x1 | x0
        b1 = x1 & x0
        m2 = x2.select(t1, b1)
        t2 = x2 | t1
        b2 = x2 & b1
        m3h = x3.select(t2, m2)
        m3l = x3.select(m2, b2)
        m4 = x4.select(m3h, m3l)
        return m4

    @classmethod
    def _majority7_select(cls, x6: Self, x5: Self, x4: Self, x3: Self, x2: Self, x1: Self, x0: Self) -> Self:
        t1 = x1 | x0
        b1 = x1 & x0
        m2 = x2.select(t1, b1)
        t2 = x2 | t1
        b2 = x2 & b1
        m3h = x3.select(t2, m2)
        m3l = x3.select(m2, b2)
        m4 = x4.select(m3h, m3l)
        t3 = x3 | t2
        b3 = x3 & b2
        m4h = x4.select(t3, m3h)
        m4l = x4.select(m3l, b3)
        m5h = x5.select(m4h, m4)
        m5l = x5.select(m4, m4l)
        m6 = x6.select(m5h, m5l)
        return m6

    @classmethod
    def _majority9_select(cls, x8: Self, x7: Self, x6: Self, x5: Self, x4: Self, x3: Self, x2: Self, x1: Self, x0: Self) -> Self:
        t1 = x0 | x1
        b1 = x0 & x1
        m2 = x2.select(t1, b1)
        t2 = x2 | t1
        b2 = x2 & b1
        m3h = x3.select(t2, m2)
        m3l = x3.select(m2, b2)
        t3 = x3 | t2
        b3 = x3 & b2
        m4h = x4.select(t3, m3h)
        m4 = x4.select(m3h, m3l)
        m4l = x4.select(m3l, b3)
        t4 = x4 | t3
        b4 = x4 & b3
        m5hh = x5.select(t4, m4h)
        m5h = x5.select(m4h, m4)
        m5l = x5.select(m4, m4l)
        m5ll = x5.select(m4l, b4)
        m6h = x6.select(m5hh, m5h)
        m6 = x6.select(m5h, m5l)
        m6l = x6.select(m5l, m5ll)
        m7h = x7.select(m6h, m6)
        m7l = x7.select(m6, m6l)
        m8 = x8.select(m7h, m7l)
        return m8

    @classmethod
    def _majority3(cls, a: Self, b: Self, c: Self) -> Self:
        # a:  1 0 0 1 0 1 1
        # b:  1 1 0 1 0 1 0
        # c:  1 0 1 0 0 0 0
        # M:  1 0 0 1 0 1 0

        # at least 2/3 agreeing on TRUE
        abh = a & b
        bch = b & c
        cah = c & a
        h = abh | bch | cah
        return h

    @classmethod
    def _majority5(cls, a: Self, b: Self, c: Self, d: Self, e: Self) -> Self:
        # at least 3/5 agreeing on TRUE

        # 2*10 AND
        # 9*1 OR
        # (b ∧ d ∧ e) ∨ (a ∧ d ∧ e) ∨ (b ∧ c ∧ e) ∨ (a ∧ c ∧ e) ∨ (b ∧ c ∧ d) ∨ (a ∧ c ∧ d) ∨ (c ∧ d ∧ e) ∨ (a ∧ b ∧ e) ∨ (a ∧ b ∧ d) ∨ (a ∧ b ∧ c)

        ab = a & b
        cd = c & d
        de = d & e
        ce = c & e

        bde = b & de
        ade = a & de

        bce = b & ce
        ace = a & ce

        bcd = b & cd
        acd = a & cd
        cde = e & cd

        abe = ab & e
        abd = ab & d
        abc = ab & c

        h = bde | ade | bce | ace | bcd | acd | cde | abe | abd | abc
        return h

    @classmethod
    def _majority5_via_3(cls, a: Self, b: Self, c: Self, d: Self, e: Self) -> Self:
        return cls._majority3(a, cls._majority3(b, c, d), cls._majority3(e, d, cls._majority3(c, b, a)))

    @classmethod
    def _majority7_via_3(cls, a: Self, b: Self, c: Self, d: Self, e: Self, f: Self, g: Self) -> Self:
        mdef = cls._majority3(d, e, f)
        mabc = cls._majority3(a, b, c)
        return cls._majority3(g, cls._majority3(c, mdef, cls._majority3(a, b, mdef)), cls._majority3(f, mabc, cls._majority3(d, e, mabc)))

    @classmethod
    def _majority9_via_3(cls, a: Self, b: Self, c: Self, d: Self, e: Self, f: Self, g: Self, h: Self, i: Self) -> Self:
        mabc = cls._majority3(a, b, c)
        mdef = cls._majority3(d, e, f)
        mghi = cls._majority3(g, h, i)

        l_ = cls._majority3(cls._majority3(c, mdef, b), mdef, a)
        l = cls._majority3(cls._majority3(i, l_, h), l_, g)

        r_ = cls._majority3(cls._majority3(c, mghi, b), mghi, a)
        r = cls._majority3(cls._majority3(f, r_, e), r_, d)

        m = cls._majority3(mabc, mdef, mghi)
        return cls._majority3(l, m, r)

    @classmethod
    def _majority_via_custom(cls, vs: list[Self]):
        n = len(vs)
        if n == 0:
            return cls.rand()
        elif n == 1:
            return vs[0]
        elif n == 2:
            return vs[0].select_rand(vs[1])
        elif n == 3:
            return cls._majority3(vs[0], vs[1], vs[2])
        elif n == 4:
            return cls._majority5_via_3(vs[0], vs[1], vs[2], vs[3], cls.rand())
        elif n == 5:
            return cls._majority5_via_3(vs[0], vs[1], vs[2], vs[3], vs[4])
        elif n == 6:
            return cls._majority7_via_3(vs[0], vs[1], vs[2], vs[3], vs[4], vs[5], cls.rand())
        elif n == 7:
            return cls._majority7_via_3(vs[0], vs[1], vs[2], vs[3], vs[4], vs[5], vs[6])
        elif n == 8:
            return cls._majority9_via_3(vs[0], vs[1], vs[2], vs[3], vs[4], vs[5], vs[6], vs[7], cls.rand())
        elif n == 9:
            return cls._majority9_via_3(vs[0], vs[1], vs[2], vs[3], vs[4], vs[5], vs[6], vs[7], vs[8])
        else:
            raise RuntimeError(f"No optimal implementations beyond MAJORITY-9 (tried to take MAJORITY-{n})")

    ZERO: Self
    ONE: Self
    HALF: Self


class Permutation:
    @classmethod
    def nrandom(cls, n) -> list[Self]:
        return [cls.random() for _ in range(n)]

    @classmethod
    def random(cls) -> Self:
        raise NotImplementedError()

    def __mul__(self, other: Self) -> Self:
        raise NotImplementedError()

    def __invert__(self) -> Self:
        raise NotImplementedError()

    def __pow__(self, power):
        permutation = self
        for _ in range(power - 1):
            permutation = permutation*permutation
        return permutation

    def __call__(self, hv: 'AbstractBHV') -> 'AbstractBHV':
        raise NotImplementedError()


class MemoizedPermutation(Permutation):
    _permutations: 'dict[int | tuple[int, ...], Self]'

    @classmethod
    def _get_singular(cls, permutation_id: int) -> Self:
        if permutation_id in cls._permutations:
            return cls._permutations[permutation_id]
        elif -permutation_id in cls._permutations:
            inv_permutation = cls._permutations[-permutation_id]
            permutation = ~inv_permutation
            cls._permutations[permutation_id] = permutation
            return permutation
        else:
            permutation = cls.random()
            cls._permutations[permutation_id] = permutation
            return permutation

    @classmethod
    def _get_composite(cls, permutation_id: 'tuple[int, ...]') -> Self:
        # this can be optimized a lot by looking for partial compositions
        composite_permutation = cls._get_singular(permutation_id[0])
        for component_permutation_id in permutation_id[1:]:
            component_permutation = cls.get(component_permutation_id)
            composite_permutation = composite_permutation * component_permutation

        cls._permutations[permutation_id] = composite_permutation
        return composite_permutation

    @classmethod
    def get(cls, permutation_id: 'int | tuple[int, ...]') -> Self:
        permutation = cls._get_singular(permutation_id) if isinstance(permutation_id, int) else \
            cls._get_composite(permutation_id)
        return permutation
