"""Sparse rational function fields."""

from __future__ import annotations

import math

from ..core import Expr, Symbol
from ..core.sympify import CantSympify, sympify
from ..domains.compositedomain import CompositeDomain
from ..domains.domainelement import DomainElement
from ..domains.field import Field
from .orderings import lex
from .polyerrors import CoercionFailedError, GeneratorsError
from .rings import PolyElement, PolynomialRing


def field(symbols, domain, order=lex):
    """Construct new rational function field returning (field, x1, ..., xn)."""
    _field = FractionField(domain, symbols, order)
    return (_field,) + _field.gens


class FractionField(Field, CompositeDomain):
    """A class for representing multivariate rational function fields."""

    is_FractionField = True

    has_assoc_Ring = True

    def __new__(cls, domain, symbols, order=lex):
        ring = PolynomialRing(domain, symbols, order)
        symbols = ring.symbols
        ngens = ring.ngens
        domain = ring.domain
        order = ring.order

        key = cls.__name__, symbols, ngens, domain, order
        obj = _field_cache.get(key)

        if obj is None:
            obj = object.__new__(cls)
            obj._hash = hash(key)
            obj.dtype = type('FracElement', (FracElement,), {'field': obj})
            obj.symbols = symbols
            obj.ngens = ngens
            obj.domain = domain
            obj.order = order

            obj.zero = obj.dtype(ring.zero)
            obj.one = obj.dtype(ring.one)

            obj.gens = obj._gens()

            obj.rep = str(domain) + '(' + ','.join(map(str, symbols)) + ')'

            for symbol, generator in zip(obj.symbols, obj.gens):
                if isinstance(symbol, Symbol):
                    name = symbol.name

                    if not hasattr(obj, name):
                        setattr(obj, name, generator)

            _field_cache[key] = obj

        return obj

    def __getnewargs_ex__(self):
        return (self.domain, self.symbols), {'order': self.order}

    @property
    def characteristic(self):
        return self.domain.characteristic

    def _gens(self):
        """Return a list of polynomial generators."""
        return tuple(self.dtype(gen) for gen in self.ring.gens)

    def __hash__(self):
        return self._hash

    def __eq__(self, other):
        return self is other

    def clone(self, symbols=None, domain=None, order=None):
        return self.__class__(domain or self.domain, symbols or self.symbols, order or self.order)

    def raw_new(self, numer, denom=None):
        return self.dtype(numer, denom)

    def domain_new(self, element):
        return self.domain.convert(element)

    def ground_new(self, element):
        try:
            return self(self.ring.ground_new(element))
        except CoercionFailedError as exc:
            domain = self.domain

            if not domain.is_Field and hasattr(domain, 'field'):
                ring = self.ring
                ground_field = domain.field
                element = ground_field.convert(element)
                numer = ring.ground_new(element.numerator)
                denom = ring.ground_new(element.denominator)
                return self.raw_new(numer, denom)
            raise NotImplementedError from exc

    def __call__(self, element):
        if isinstance(element, FracElement):
            if self == element.field:
                return element
            raise NotImplementedError('conversion')
        if isinstance(element, PolyElement):
            denom, numer = element.clear_denoms()
            numer = numer.set_ring(self.ring)
            denom = self.ring.ground_new(denom)
            return self.raw_new(numer, denom)
        if isinstance(element, tuple) and len(element) == 2:
            numer, denom = list(map(self.ring.__call__, element))
            numer, denom = numer.cancel(denom)
            return self.raw_new(numer, denom)
        if isinstance(element, str):
            raise NotImplementedError('parsing')
        if isinstance(element, Expr):
            return self.convert(element)
        return self.ground_new(element)

    def from_expr(self, expr):
        expr = sympify(expr)
        domain = self.domain
        mapping = dict(zip(self.symbols, self.gens))

        def _rebuild(expr):
            if generator := mapping.get(expr):
                return generator
            if 1/expr in mapping:
                return 1/mapping[1/expr]
            if expr.is_Add:
                return sum(map(_rebuild, expr.args))
            if expr.is_Mul:
                return math.prod(map(_rebuild, expr.args))
            if expr.is_Pow:
                c, a = expr.exp.as_coeff_Mul(rational=True)
                if c.is_Integer and c != 1:
                    return _rebuild(expr.base**a)**int(c)

            if not domain.is_Field and hasattr(domain, 'field'):
                frac = domain.field.convert(expr)
            else:
                frac = domain.convert(expr)

            return self(frac)

        try:
            return _rebuild(expr)
        except CoercionFailedError as exc:
            raise ValueError('expected an expression convertible to a '
                             f'rational function in {self}, '
                             f'got {expr}') from exc

    def to_ring(self):
        return self.domain.poly_ring(*self.symbols, order=self.order)

    def to_expr(self, element):
        ring = self.ring
        return ring.to_expr(element.numerator)/ring.to_expr(element.denominator)

    def _from_PythonIntegerRing(self, a, K0):
        return self(self.domain.convert(a, K0))
    _from_GMPYIntegerRing = _from_PythonIntegerRing
    _from_PythonRationalField = _from_PythonIntegerRing
    _from_GMPYRationalField = _from_PythonIntegerRing
    _from_RealField = _from_PythonIntegerRing
    _from_ComplexField = _from_PythonIntegerRing

    def _from_PolynomialRing(self, a, K0):
        try:
            return self(a)
        except (CoercionFailedError, GeneratorsError):
            return

    def _from_FractionField(self, a, K0):
        try:
            return a.set_field(self)
        except (CoercionFailedError, GeneratorsError):
            return

    @property
    def ring(self):
        return self.to_ring()

    def is_normal(self, a):
        return self.domain.is_normal(a.numerator.LC)


_field_cache: dict[tuple, FractionField] = {}


class FracElement(DomainElement, CantSympify):
    """Element of multivariate distributed rational function field.

    See Also
    ========

    FractionField

    """

    def __init__(self, numer, denom=None):
        """Initialize self."""
        if denom is None:
            denom = self.field.ring.one
        elif not denom:
            raise ZeroDivisionError('zero denominator')

        self._numerator = numer
        self._denominator = denom

    def __reduce__(self):
        return self.parent.__call__, ((self.numerator, self.denominator),)

    def raw_new(self, numer, denom):
        return self.__class__(numer, denom)

    def new(self, numer, denom):
        return self.raw_new(*numer.cancel(denom))

    def to_poly(self):
        if self.denominator != 1:
            raise ValueError('self.denominator should be 1')
        return self.numerator

    @property
    def numerator(self):
        return self._numerator

    @property
    def denominator(self):
        return self._denominator

    @property
    def parent(self):
        return self.field

    _hash = None

    def __hash__(self):
        _hash = self._hash
        if _hash is None:
            self._hash = _hash = hash((self.field, self.numerator, self.denominator))
        return _hash

    def copy(self):
        return self.raw_new(self.numerator.copy(), self.denominator.copy())

    def set_field(self, new_field):
        if self.field == new_field:
            return self
        new_ring = new_field.ring
        numer = self.numerator.set_ring(new_ring)
        denom = self.denominator.set_ring(new_ring)
        return new_field((numer, denom))

    def __eq__(self, other):
        if isinstance(other, self.field.dtype):
            return self.numerator == other.numerator and self.denominator == other.denominator
        return self.numerator == other and self.denominator == 1

    def __bool__(self):
        return bool(self.numerator)

    def __pos__(self):
        return self.raw_new(self.numerator, self.denominator)

    def __neg__(self):
        """Negate all coefficients in ``self``."""
        return self.raw_new(-self.numerator, self.denominator)

    def _extract_ground(self, element):
        domain = self.field.domain

        try:
            element = domain.convert(element)
        except CoercionFailedError:
            ground_field = domain.field

            try:
                element = ground_field.convert(element)
            except CoercionFailedError:
                return 0, None, None
            return -1, element.numerator, element.denominator
        return 1, element, None

    def __add__(self, other):
        """Add rational functions ``self`` and ``other``."""
        field = self.field

        if not other:
            return self
        if not self:
            return other
        if isinstance(other, field.dtype):
            if self.denominator == other.denominator:
                return self.new(self.numerator + other.numerator, self.denominator)
            return self.new(self.numerator*other.denominator + self.denominator*other.numerator,
                            self.denominator*other.denominator)
        if isinstance(other, field.ring.dtype):
            return self.new(self.numerator + self.denominator*other, self.denominator)
        if isinstance(other, FracElement):
            if isinstance(field.domain, FractionField) and field.domain.field == other.field:
                pass
            elif isinstance(other.field.domain, FractionField) and other.field.domain.field == field:
                return other.__radd__(self)
            else:
                return NotImplemented
        elif isinstance(other, PolyElement):
            if isinstance(field.domain, PolynomialRing) and field.domain.ring == other.ring:
                pass
            else:
                return other.__radd__(self)

        return self.__radd__(other)

    def __radd__(self, other):
        op, other_numer, other_denom = self._extract_ground(other)

        if op == 1:
            return self.new(self.numerator + self.denominator*other_numer, self.denominator)
        if not op:
            return NotImplemented
        return self.new(self.numerator*other_denom + self.denominator*other_numer,
                        self.denominator*other_denom)

    def __sub__(self, other):
        """Subtract rational functions ``self`` and ``other``."""
        field = self.field

        if not other:
            return self
        if not self:
            return -other
        if isinstance(other, field.dtype):
            if self.denominator == other.denominator:
                return self.new(self.numerator - other.numerator, self.denominator)
            return self.new(self.numerator*other.denominator - self.denominator*other.numerator,
                            self.denominator*other.denominator)
        if isinstance(other, field.ring.dtype):
            return self.new(self.numerator - self.denominator*other, self.denominator)
        if isinstance(other, FracElement):
            if isinstance(field.domain, FractionField) and field.domain.field == other.field:
                pass
            elif isinstance(other.field.domain, FractionField) and other.field.domain.field == field:
                return other.__rsub__(self)
            else:
                return NotImplemented
        if isinstance(other, PolyElement):
            if isinstance(field.domain, PolynomialRing) and field.domain.ring == other.ring:
                pass
            else:
                return other.__rsub__(self)

        op, other_numer, other_denom = self._extract_ground(other)

        if op == 1:
            return self.new(self.numerator - self.denominator*other_numer, self.denominator)
        if not op:
            return NotImplemented
        return self.new(self.numerator*other_denom - self.denominator*other_numer,
                        self.denominator*other_denom)

    def __rsub__(self, other):
        op, other_numer, other_denom = self._extract_ground(other)

        if op == 1:
            return self.new(-self.numerator + self.denominator*other_numer, self.denominator)
        if not op:
            return NotImplemented
        return self.new(-self.numerator*other_denom + self.denominator*other_numer,
                        self.denominator*other_denom)

    def __mul__(self, other):
        """Multiply rational functions ``self`` and ``other``."""
        field = self.field

        if not self or not other:
            return field.zero
        if isinstance(other, field.dtype):
            return self.new(self.numerator*other.numerator, self.denominator*other.denominator)
        if isinstance(other, field.ring.dtype):
            return self.new(self.numerator*other, self.denominator)
        if isinstance(other, FracElement):
            if isinstance(field.domain, FractionField) and field.domain.field == other.field:
                pass
            elif isinstance(other.field.domain, FractionField) and other.field.domain.field == field:
                return other.__rmul__(self)
            else:
                return NotImplemented
        elif isinstance(other, PolyElement):
            if isinstance(field.domain, PolynomialRing) and field.domain.ring == other.ring:
                pass
            else:
                return other.__rmul__(self)

        return self.__rmul__(other)

    def __rmul__(self, other):
        op, other_numer, other_denom = self._extract_ground(other)

        if op == 1:
            return self.new(self.numerator*other_numer, self.denominator)
        if not op:
            return NotImplemented
        return self.new(self.numerator*other_numer, self.denominator*other_denom)

    def __truediv__(self, other):
        """Computes quotient of fractions ``self`` and ``other``."""
        field = self.field

        if not other:
            raise ZeroDivisionError
        if isinstance(other, field.dtype):
            return self.new(self.numerator*other.denominator, self.denominator*other.numerator)
        if isinstance(other, field.ring.dtype):
            return self.new(self.numerator, self.denominator*other)
        if isinstance(other, FracElement):
            if isinstance(field.domain, FractionField) and field.domain.field == other.field:
                pass
            elif isinstance(other.field.domain, FractionField) and other.field.domain.field == field:
                return other.__rtruediv__(self)
            else:
                return NotImplemented
        elif isinstance(other, PolyElement):
            if isinstance(field.domain, PolynomialRing) and field.domain.ring == other.ring:
                pass
            else:
                return NotImplemented

        op, other_numer, other_denom = self._extract_ground(other)

        if op == 1:
            return self.new(self.numerator, self.denominator*other_numer)
        if not op:
            return NotImplemented
        return self.new(self.numerator*other_denom, self.denominator*other_numer)

    def __rtruediv__(self, other):
        if not self:
            raise ZeroDivisionError
        if isinstance(other, self.field.ring.dtype):
            return self.new(self.denominator*other, self.numerator)

        op, other_numer, other_denom = self._extract_ground(other)

        if op == 1:
            return self.new(self.denominator*other_numer, self.numerator)
        if not op:
            return NotImplemented
        return self.new(self.denominator*other_numer, self.numerator*other_denom)

    def __pow__(self, n):
        """Raise ``self`` to a non-negative power ``n``."""
        if n >= 0:
            return self.raw_new(self.numerator**n, self.denominator**n)
        if not self:
            raise ZeroDivisionError
        return self.raw_new(self.denominator**-n, self.numerator**-n)

    def diff(self, x):
        """Computes partial derivative in ``x``.

        Examples
        ========

        >>> _, x, y, z = field('x y z', ZZ)
        >>> ((x**2 + y)/(z + 1)).diff(x)
        2*x/(z + 1)

        """
        x = x.to_poly()
        return self.new(self.numerator.diff(x)*self.denominator -
                        self.numerator*self.denominator.diff(x), self.denominator**2)

    def __call__(self, *values):
        if 0 < len(values) <= self.field.ngens:
            return self.eval(list(zip(self.field.gens, values)))
        raise ValueError(f'expected at least 1 and at most {self.field.ngens} values, got {len(values)}')

    def eval(self, x, a=None):
        if isinstance(x, list) and a is None:
            x = [(X.to_poly(), a) for X, a in x]
            numer, denom = self.numerator.eval(x), self.denominator.eval(x)
        else:
            x = x.to_poly()
            numer, denom = self.numerator.eval(x, a), self.denominator.eval(x, a)

        if self._extract_ground(denom) == (1, 1, None):
            return numer
        if isinstance(numer, PolyElement):
            field = numer.ring.field
        else:
            field = self.field
        return field((field.ring(numer), field.ring(denom)))

    def compose(self, x, a=None):
        """Computes the functional composition."""
        field = self.field

        if isinstance(x, list) and a is None:
            x = [(X.to_poly(), a) for X, a in x]
            numer = (self.numerator.compose([(X, a.numerator) for X, a in x]) *
                     self.denominator.compose([(X, a.denominator) for X, a in x]))
            denom = (self.numerator.compose([(X, a.denominator) for X, a in x]) *
                     self.denominator.compose([(X, a.numerator) for X, a in x]))
        else:
            x = x.to_poly()
            numer = (self.numerator.compose(x, a.numerator) *
                     self.denominator.compose(x, a.denominator))
            denom = (self.numerator.compose(x, a.denominator) *
                     self.denominator.compose(x, a.numerator))

        return field((field.ring(numer), field.ring(denom)))
