#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright © 2017 lafrite <lafrite@Poivre>
#
# Distributed under terms of the MIT license.

"""
Tokens representing polynomials functions 

"""
from ..expression import Expression
from functools import partial
from .token import Token
from ...core.MO import MO
from ...core.MO.atoms import moify

__all__ = ["Polynomial", "Quadratic", "Linear"]


class Polynomial(Token):

    """ Token representing a polynomial 

    :examples:
    >>> from ...core.MO.polynomial import MOpolynomial
    >>> P = Polynomial(MOpolynomial('x', [1, 2, 3]))
    >>> P
    <Polynomial 3x^2 + 2x + 1>
    """

    def __init__(self, a, name="", ancestor=None):
        """ Initiate Polynomial with a MO"""
        if not isinstance(a, MO):
            if isinstance(a, str):
                raise TypeError
            else:
                raise TypeError
        else:
            mo = a

        Token.__init__(self, mo, name, ancestor)
        self._mathtype = "polynome"

    @classmethod
    def from_mo(cls, mo, name="", ancestor=None):

        return cls(mo, name, ancestor)

    @classmethod
    def from_coefficients(cls, coefficients):
        """ Initiate polynomial from list of coefficients """
        pass

    @classmethod
    def random(cls):
        raise NotImplementedError

    def __setitem__(self, key, item):
        """ Use Polynomial like if they were a dictionnary to set coefficients """
        raise NotImplementedError("Can't set coefficient of a polynomial")

    def __getitem__(self, key):
        """ Use Polynomial like if they were a dictionnary to get coefficients 

        :examples:
        >>> from ...core.MO.polynomial import MOpolynomial
        >>> P = Polynomial(MOpolynomial('x', [1, 2, 3]))
        >>> P[0]
        <MOnumber 1>
        >>> P[1]
        <MOnumber 2>
        >>> P[2]
        <MOnumber 3>
        >>> P[3]
        Traceback (most recent call last):
            ...
        KeyError: 3
        """
        return self._mo.coefficients[key]

    def __call__(self, value):
        """ Call a Polynomial to evaluate itself on value 

        :examples:
        >>> from ...core.MO.polynomial import MOpolynomial
        >>> P = Polynomial(MOpolynomial('x', [1, 2, 3]))
        >>> for s in P(2).explain():
        ...     print(s)
        3 * 2^2 + 2 * 2 + 1
        3 * 4 + 4 + 1
        12 + 5
        17
        """
        tree = self._mo.tree
        variable = (set(tree.get_leafs(extract_variable)) - {None}).pop()

        dest = moify(value)
        replace_var = partial(replace, origin=variable, dest=dest)
        tree = tree.map_on_leaf(replace_var)

        return Expression(tree).simplify()
    
    def differentiate(self):
        """ Differentiate a polynome

        :example:
        >>> from ...core.MO.polynomial import MOpolynomial
        >>> P = Polynomial(MOpolynomial('x', [1, 2, 3]))
        >>> P
        <Polynomial 3x^2 + 2x + 1>
        >>> P.differentiate()
        <Linear 2 + 6x>
        >>> for s in P.differentiate().explain():
        ...     print(s)
        0 + 2 + 3 * 2x
        2 + 3 * 2 * x
        2 + 6x
        """
        return Expression(self._mo.differentiate()).simplify()
    @property
    def roots(self):
        """ Get roots of the Polynomial """
        raise NotImplementedError("Can't compute roots not specific polynomial")


class Linear(Polynomial):

    """ Token representing a linear ax + b

    :examples:
    >>> from ...core.MO.polynomial import MOpolynomial, MOMonomial
    >>> P = Linear(MOpolynomial('x', [1, 2]))
    >>> P
    <Linear 2x + 1>
    >>> P.a
    <MOnumber 2>
    >>> P.b
    <MOnumber 1>
    >>> P.differentiate()
    <Integer 2>
    >>> P.roots
    [<Fraction - 2 / 1>]

    """

    def __init__(self, mo, name="", ancestor=None):
        """ Initiate Linear with MO

        :examples:
        >>> from ...core.MO.polynomial import MOpolynomial, MOMonomial
        >>> P = Linear(MOpolynomial('x', [1, 2]))
        >>> P
        <Linear 2x + 1>
        >>> Q = Linear(MOMonomial(3, 'x', 1))
        >>> Q
        <Linear 3x>
        """
        Polynomial.__init__(self, mo, name, ancestor)
        self._mathtype = "affine"

    @classmethod
    def random(cls):
        raise NotImplementedError

    @property
    def a(self):
        return self[1]

    @property
    def b(self):
        return self[0]

    @property
    def roots(self):
        """ Get the root of the polynomial

        :examples:
        >>> from ...core.MO.polynomial import MOpolynomial, MOMonomial
        >>> P = Linear(MOpolynomial('x', [1, 2]))
        >>> P.roots
        [<Fraction - 2 / 1>]
        >>> #P = Linear(MOpolynomial('x', [1, -2]))
        >>> #P.roots
        """

        try:
            return [Expression.from_str(f"-{self.a}/{self.b}").simplify()]
        except AttributeError:
            return [Expression.from_str(f"-{self.a}/{self.b}")]


class Quadratic(Polynomial):

    """ Token representing a quadratic ax^2 + bx + c

    :examples:
    >>> from ...core.MO.polynomial import MOpolynomial
    >>> P = Quadratic(MOpolynomial('x', [1, 2, 3]))
    >>> P
    <Quadratic 3x^2 + 2x + 1>
    >>> P.a
    <MOnumber 3>
    >>> P.b
    <MOnumber 2>
    >>> P.c
    <MOnumber 1>
    >>> P.delta
    <Integer - 8>
    >>> for s in P.delta.explain():
    ...    print(s)
    2^2 - 4 * 3 * 1
    4 - 12 * 1
    4 - 12
    - 8
    >>> P.differentiate()
    <Linear 2 + 6x>
    >>> P.roots
    []

    """

    def __init__(self, mo, name="", ancestor=None):
        """ Initiate Quadratic from MO

        >>> from ...core.MO.polynomial import MOpolynomial
        >>> P = Quadratic(MOpolynomial('x', [1, 2, 3]))
        >>> P
        <Quadratic 3x^2 + 2x + 1>
        """

        Polynomial.__init__(self, mo, name, ancestor)
        self._mathtype = "polynome du 2nd degré"

    @classmethod
    def random(cls):
        raise NotImplementedError

    @property
    def a(self):
        return self[2]

    @property
    def b(self):
        return self[1]

    @property
    def c(self):
        return self[0]

    @property
    def delta(self):
        return Expression.from_str(f"{self.b}^2-4*{self.a}*{self.c}").simplify()

    @property
    def roots(self):
        if self.delta._mo < 0:
            return []
        elif self.delta._mo == 0:
            return [Expression.from_str(f"-{self.b}/(2*{self.a})").simplify()]
        else:
            raise NotImplementedError("Todo!")



def extract_variable(leaf):
    try:
        return leaf.variable
    except AttributeError:
        return None

def replace(leaf, origin, dest):
    """ Recursively replace origin to dest in leaf """
    try:
        leaf.tree
    except AttributeError:
        if leaf == origin:
            return dest
        return leaf

    replace_var = partial(replace, origin=origin, dest=dest)
    return leaf.tree.map_on_leaf(replace_var)
# -----------------------------
# Reglages pour 'vim'
# vim:set autoindent expandtab tabstop=4 shiftwidth=4:
# cursor: 16 del