"""
fta.lattice

The Lattice class — complete divisibility structure of a positive integer.
"""

from __future__ import annotations

from typing import Dict, Tuple

from .arithmetic import factorize, to_int, _divisors, Coordinate


def _is_prime(n: int) -> bool:
    """True if n is prime."""
    if n < 2:
        return False
    if n == 2:
        return True
    if n % 2 == 0:
        return False
    i = 3
    while i * i <= n:
        if n % i == 0:
            return False
        i += 2
    return True


class Lattice:
    """
    The divisibility lattice of a positive integer n.

    Every positive divisor of n is a point in this lattice. The partial order
    is divisibility: a <= b if a divides b. Under the prime exponent coordinate
    system, this order is coordinate-wise <=.

    Parameters
    ----------
    n : int
        The integer whose divisibility structure to represent. Must be >= 1.

    Examples
    --------
    >>> L = Lattice(12)
    >>> L.members()
    (1, 2, 3, 4, 6, 12)
    >>> L.covers(3)
    (6,)
    >>> L.meet(4, 6)
    2
    >>> L.join(4, 6)
    12
    """

    __slots__ = ("_n", "_members", "_coord")

    def __init__(self, n: int) -> None:
        if not isinstance(n, int) or n < 1:
            raise ValueError(f"Lattice requires a positive integer, got {n!r}")
        self._n = n
        self._members: Tuple[int, ...] = _divisors(n)
        self._coord: Dict[int, Coordinate] = {
            m: factorize(m) for m in self._members
        }

    # ------------------------------------------------------------------
    # Core properties
    # ------------------------------------------------------------------

    @property
    def n(self) -> int:
        """The integer this lattice is built on."""
        return self._n

    @property
    def top(self) -> int:
        """The greatest element — n itself."""
        return self._n

    @property
    def bottom(self) -> int:
        """The least element — always 1."""
        return 1

    # ------------------------------------------------------------------
    # Membership
    # ------------------------------------------------------------------

    def members(self) -> Tuple[int, ...]:
        """
        All points in the lattice — the positive divisors of n, ascending.

        >>> Lattice(12).members()
        (1, 2, 3, 4, 6, 12)
        """
        return self._members

    def __contains__(self, m: int) -> bool:
        """True if m is a member of this lattice (i.e. m divides n)."""
        return m in self._coord

    def _check(self, m: int) -> None:
        if m not in self._coord:
            raise ValueError(f"{m} is not a member of Lattice({self._n})")

    # ------------------------------------------------------------------
    # Hasse diagram
    # ------------------------------------------------------------------

    def covers(self, m: int) -> Tuple[int, ...]:
        """
        Immediate parents of m in the Hasse diagram.

        d covers m iff: m divides d, d divides n, and d/m is prime
        (no intermediary between m and d in the lattice).

        >>> Lattice(12).covers(3)
        (6,)
        >>> Lattice(12).covers(1)
        (2, 3)
        >>> Lattice(12).covers(12)
        ()
        """
        self._check(m)
        return tuple(
            d for d in self._members
            if d > m and d % m == 0 and _is_prime(d // m)
        )

    def covered_by(self, m: int) -> Tuple[int, ...]:
        """
        Immediate children of m in the Hasse diagram.

        d is covered by m iff: d divides m, d divides n, and m/d is prime.

        >>> Lattice(12).covered_by(6)
        (2, 3)
        >>> Lattice(12).covered_by(1)
        ()
        """
        self._check(m)
        return tuple(
            d for d in self._members
            if d < m and m % d == 0 and _is_prime(m // d)
        )

    def hasse_edges(self) -> Tuple[Tuple[int, int], ...]:
        """
        All edges in the Hasse diagram as (child, parent) pairs, ascending.

        >>> Lattice(6).hasse_edges()
        ((1, 2), (1, 3), (2, 6), (3, 6))
        """
        edges = []
        for m in self._members:
            for parent in self.covers(m):
                edges.append((m, parent))
        return tuple(edges)

    # ------------------------------------------------------------------
    # Coordinate system
    # ------------------------------------------------------------------

    def coordinate(self, m: int) -> Coordinate:
        """
        Prime exponent coordinate of m.

        coordinate(12)  ->  {2: 2, 3: 1}
        coordinate(1)   ->  {}
        """
        self._check(m)
        return self._coord[m]

    def primes(self) -> Tuple[int, ...]:
        """
        The prime axes of this lattice — the prime factors of n, ascending.

        >>> Lattice(12).primes()
        (2, 3)
        >>> Lattice(360).primes()
        (2, 3, 5)
        """
        return tuple(sorted(factorize(self._n).keys()))

    def from_coordinate(self, coord: Coordinate) -> int:
        """
        Reconstruct a lattice member from its prime exponent coordinate.

        Raises ValueError if the result is not a member of this lattice.
        """
        m = to_int(coord)
        self._check(m)
        return m

    # ------------------------------------------------------------------
    # Lattice operations
    # ------------------------------------------------------------------

    def meet(self, a: int, b: int) -> int:
        """
        Greatest lower bound of a and b — their gcd.

        In coordinate terms: coordinate-wise min.

        >>> Lattice(12).meet(4, 6)
        2
        """
        self._check(a)
        self._check(b)
        from math import gcd
        return gcd(a, b)

    def join(self, a: int, b: int) -> int:
        """
        Least upper bound of a and b — their lcm.

        In coordinate terms: coordinate-wise max.
        Raises ValueError if lcm(a, b) is not a member of this lattice.

        >>> Lattice(12).join(4, 6)
        12
        """
        self._check(a)
        self._check(b)
        from math import lcm
        result = lcm(a, b)
        self._check(result)
        return result

    def divides(self, a: int, b: int) -> bool:
        """
        True if a divides b (a <= b in lattice order).

        >>> Lattice(12).divides(3, 6)
        True
        >>> Lattice(12).divides(4, 6)
        False
        """
        self._check(a)
        self._check(b)
        return b % a == 0

    # ------------------------------------------------------------------
    # Navigation
    # ------------------------------------------------------------------

    def members_lte(self, m: int) -> Tuple[int, ...]:
        """
        All lattice members that divide m (downward closure from m).

        >>> Lattice(12).members_lte(6)
        (1, 2, 3, 6)
        """
        self._check(m)
        return tuple(d for d in self._members if m % d == 0)

    def members_gte(self, m: int) -> Tuple[int, ...]:
        """
        All lattice members that m divides (upward closure from m).

        >>> Lattice(12).members_gte(3)
        (3, 6, 12)
        """
        self._check(m)
        return tuple(d for d in self._members if d % m == 0)

    def depth(self, m: int) -> int:
        """
        Number of edges from bottom (1) to m in the Hasse diagram.

        Equals the total number of prime factors of m (with multiplicity).

        >>> Lattice(12).depth(1)
        0
        >>> Lattice(12).depth(4)   # 4 = 2^2
        2
        >>> Lattice(12).depth(12)  # 12 = 2^2 * 3
        3
        """
        self._check(m)
        return sum(self._coord[m].values())

    def height(self, m: int) -> int:
        """
        Number of edges from m to top (n) in the Hasse diagram.

        Equals depth(n) - depth(m).

        >>> Lattice(12).height(12)
        0
        >>> Lattice(12).height(1)
        3
        """
        self._check(m)
        return self.depth(self._n) - self.depth(m)

    # ------------------------------------------------------------------
    # Dunder
    # ------------------------------------------------------------------

    def __repr__(self) -> str:
        return f"Lattice({self._n})"

    def __eq__(self, other: object) -> bool:
        return isinstance(other, Lattice) and self._n == other._n

    def __hash__(self) -> int:
        return hash(self._n)
