LeGrandChien/legrandchien.py

from enum import Enum
import numpy as np

class Operation(Enum):
    VAR = 1
    CONST = 2
    ADD = 3
    SUB = 4
    MULT = 5
    DIV = 6
    POW = 7
    SIN = 8
    COS = 9
    LN = 10
    EXP = 11

# const class
class Const:
    def __init__(self, value) -> None:
        self.value = value
        self.type = "Const"

    
    def evaluate(self, dico):
        return self.value

    def copy(self):
        return Const(self.value)

    def __add__(self, other):
        if other.type == "Const":
            return Equation(
                Operation.CONST,
                None,
                None,
                self.copy()
            ) + Equation(
                Operation.CONST,
                None,
                None,
                other
            )
        if other.type == "Var":
            return Equation(
                Operation.CONST,
                None,
                None,
                self.copy()
            ) + Equation(
                Operation.VAR,
                None,
                None,
                other
            )
        if other.type == "Equation":
            return Equation(
                Operation.CONST,
                None,
                None,
                self.copy()
            ) + other
        
    def __radd__(self, other):
        if isinstance(other, (int, float)):
            return Const(other) + self

    def __sub__(self, other):
        if other.type == "Const":
            return Equation(
                Operation.CONST,
                None,
                None,
                self.copy()
            ) - Equation(
                Operation.CONST,
                None,
                None,
                other
            )
        if other.type == "Var":
            return Equation(
                Operation.CONST,
                None,
                None,
                self.copy()
            ) - Equation(
                Operation.VAR,
                None,
                None,
                other
            )
        if other.type == "Equation":
            return Equation(
                Operation.CONST,
                None,
                None,
                self.copy()
            ) - other
        
    def __rsub__(self, other):
        if isinstance(other, (int, float)):
            return Const(other) - self
        
    def __mul__(self, other):
        if other.type == "Const":
            return Equation(
                Operation.CONST,
                None,
                None,
                self.copy()
            ) * Equation(
                Operation.CONST,
                None,
                None,
                other
            )
        if other.type == "Var":
            return Equation(
                Operation.CONST,
                None,
                None,
                self.copy()
            ) * Equation(
                Operation.VAR,
                None,
                None,
                other
            )
        if other.type == "Equation":
            return Equation(
                Operation.CONST,
                None,
                None,
                self.copy()
            ) * other
        
    def __rmul__(self, other):
        if isinstance(other, (int, float)):
            return Const(other) * self
        
    def __truediv__(self, other):
        if other.type == "Const":
            return Equation(
                Operation.CONST,
                None,
                None,
                self.copy()
            ) / Equation(
                Operation.CONST,
                None,
                None,
                other
            )
        if other.type == "Var":
            return Equation(
                Operation.CONST,
                None,
                None,
                self.copy()
            ) / Equation(
                Operation.VAR,
                None,
                None,
                other
            )
        if other.type == "Equation":
            return Equation(
                Operation.CONST,
                None,
                None,
                self.copy()
            ) / other
        
    def __truediv__(self, other):
        if isinstance(other, (int, float)):
            return Const(other) / self

# var class
class Var:
    def __init__(self, name, degree=0) -> None:
        self.name = name
        self.type = "Var"
        self.degree = degree

    def copy(self):
        return Var(self.name, self.degree)

    def diff(self, n=1):
        return Var("d_"+self.name, self.degree+n)
    
    def evaluate(self, dico):
        return dico[self.name]

    def __add__(self, other):
        if other.type == "Const":
            return Equation(
                Operation.VAR,
                None,
                None,
                self.copy()
            ) + Equation(
                Operation.CONST,
                None,
                None,
                other
            )
        if other.type == "Var":
            return Equation(
                Operation.VAR,
                None,
                None,
                self.copy()
            ) + Equation(
                Operation.VAR,
                None,
                None,
                other
            )
        if other.type == "Equation":
            return Equation(
                Operation.VAR,
                None,
                None,
                self.copy()
            ) + other
        
    def __radd__(self, other):
        if isinstance(other, (int, float)):
            return Const(other) + self
        
    def __sub__(self, other):
        if other.type == "Const":
            return Equation(
                Operation.VAR,
                None,
                None,
                self.copy()
            ) - Equation(
                Operation.CONST,
                None,
                None,
                other
            )
        if other.type == "Var":
            return Equation(
                Operation.VAR,
                None,
                None,
                self.copy()
            ) - Equation(
                Operation.VAR,
                None,
                None,
                other
            )
        if other.type == "Equation":
            return Equation(
                Operation.VAR,
                None,
                None,
                self.copy()
            ) - other
        
    def __rsub__(self, other):
        if isinstance(other, (int, float)):
            return Const(other) - self
        
    def __mul__(self, other):
        if other.type == "Const":
            return Equation(
                Operation.VAR,
                None,
                None,
                self.copy()
            ) * Equation(
                Operation.CONST,
                None,
                None,
                other
            )
        if other.type == "Var":
            return Equation(
                Operation.VAR,
                None,
                None,
                self.copy()
            ) * Equation(
                Operation.VAR,
                None,
                None,
                other
            )
        if other.type == "Equation":
            return Equation(
                Operation.VAR,
                None,
                None,
                self.copy()
            ) * other
        
    def __rmul__(self, other):
        if isinstance(other, (int, float)):
            return Const(other) * self
        
    def __truediv__(self, other):
        if other.type == "Const":
            return Equation(
                Operation.VAR,
                None,
                None,
                self.copy()
            ) / Equation(
                Operation.CONST,
                None,
                None,
                other
            )
        if other.type == "Var":
            return Equation(
                Operation.VAR,
                None,
                None,
                self.copy()
            ) / Equation(
                Operation.VAR,
                None,
                None,
                other
            )
        if other.type == "Equation":
            return Equation(
                Operation.VAR,
                None,
                None,
                self.copy()
            ) / other
        
    def __rtruediv__(self, other):
        if isinstance(other, (int, float)):
            return Const(other) / self

# equation class
class Equation:
    def __init__(self, op, left, right, value):
        self.op = op
        self.left = left
        self.right = right
        self.value = value
        self.type = "Equation"

    def copy(self):
        if self.op == Operation.CONST:
            return Const(self.value.value)
        if self.op == Operation.VAR:
            return Var(self.value.name, self.value.degree)
        return Equation(self.op, self.left.copy(), self.right.copy(), None)

    def diff(self):
        if self.op == Operation.CONST:
            return Equation(Operation.CONST, None, None, Const(0))
        if self.op == Operation.VAR:
            return Equation(Operation.VAR, None, None, Var("d_"+self.value.name, self.value.degree+1))
        if self.op == Operation.ADD:
            return self.left.diff() + self.right.diff()
        if self.op == Operation.SUB:
            return self.left.diff() - self.right.diff()
        if self.op == Operation.MULT:
            return self.left.diff()*self.right.copy() + self.left.copy()*self.right.diff()
        if self.op == Operation.DIV:
            return (self.left.diff()*self.right.copy() - self.left.copy()*self.right.diff()) \
                / (self.right.copy()*self.right.copy())
        
    def partial(self, var_name):
        if self.op == Operation.CONST:
            return Equation(Operation.CONST, None, None, Const(0))
        if self.op == Operation.VAR:
            if self.value.name == var_name:
                return Equation(Operation.CONST, None, None, Const(1))
            else:
                return Equation(Operation.CONST, None, None, Const(0))
        if self.op == Operation.ADD:
            return self.left.partial(var_name) + self.right.partial(var_name)
        if self.op == Operation.SUB:
            return self.left.partial(var_name) - self.right.partial(var_name)
        if self.op == Operation.MULT:
            return self.left.partial(var_name)*self.right.copy() + self.left.copy()*self.right.partial(var_name)
        if self.op == Operation.DIV:
            return (self.left.partial(var_name)*self.right.copy() - self.left.copy()*self.right.partial(var_name)) \
                / (self.right.copy()*self.right.copy())
        
    def evaluate(self, dico):
        if self.op == Operation.CONST:
            return self.value.value
        if self.op == Operation.VAR:
            return dico[self.value.name]
        if self.op == Operation.ADD:
            return self.left.evaluate(dico) + self.right.evaluate(dico)
        if self.op == Operation.SUB:
            return self.left.evaluate(dico) - self.right.evaluate(dico)
        if self.op == Operation.MULT:
            return self.left.evaluate(dico) * self.right.evaluate(dico)
        if self.op == Operation.DIV:
            return self.left.evaluate(dico) / self.right.evaluate(dico)
        
        
    def __add__(self, other):
        if other.type == "Const":
            return self + Equation(
                Operation.CONST,
                None,
                None,
                other
            )
        if other.type == "Var":
            return self + Equation(
                Operation.VAR,
                None,
                None,
                other
            )
        if other.type == "Equation":
            return Equation(Operation.ADD, self, other, None)
        
    def __radd__(self, other):
        if isinstance(other, (int, float)):
            return Const(other) + self
        
    def __sub__(self, other):
        if other.type == "Const":
            return self - Equation(
                Operation.CONST,
                None,
                None,
                other
            )
        if other.type == "Var":
            return self - Equation(
                Operation.VAR,
                None,
                None,
                other
            )
        if other.type == "Equation":
            return Equation(Operation.SUB, self, other, None)
        
    def __rsub__(self, other):
        if isinstance(other, (int, float)):
            return Const(other) - self
        
    def __mul__(self, other):
        if other.type == "Const":
            return self * Equation(
                Operation.CONST,
                None,
                None,
                other
            )
        if other.type == "Var":
            return self * Equation(
                Operation.VAR,
                None,
                None,
                other
            )
        if other.type == "Equation":
            return Equation(Operation.MULT, self, other, None)
        
    def __rmul__(self, other):
        if isinstance(other, (int, float)):
            return Const(other) * self
        
    def __truediv__(self, other):
        if other.type == "Const":
            return self / Equation(
                Operation.CONST,
                None,
                None,
                other
            )
        if other.type == "Var":
            return self / Equation(
                Operation.VAR,
                None,
                None,
                other
            )
        if other.type == "Equation":
            return Equation(Operation.DIV, self, other, None)
        
    def __rtruediv__(self, other):
        if isinstance(other, (int, float)):
            return Const(other) / self