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 __add__(self, other):
        if other.type == "Const":
            return Equation(
                Operation.CONST,
                None,
                None,
                self.value
            ) + Equation(
                Operation.CONST,
                None,
                None,
                other
            )
        if other.type == "Var":
            return Equation(
                Operation.CONST,
                None,
                None,
                self.value
            ) + Equation(
                Operation.VAR,
                None,
                None,
                other
            )
        if other.type == "Equation":
            return Equation(
                Operation.CONST,
                None,
                None,
                self.value
            ) + other
        
    def __sub__(self, other):
        if other.type == "Const":
            return Equation(
                Operation.CONST,
                None,
                None,
                self.value
            ) - Equation(
                Operation.CONST,
                None,
                None,
                other
            )
        if other.type == "Var":
            return Equation(
                Operation.CONST,
                None,
                None,
                self.value
            ) - Equation(
                Operation.VAR,
                None,
                None,
                other
            )
        if other.type == "Equation":
            return Equation(
                Operation.CONST,
                None,
                None,
                self.value
            ) - other
        
    def __mul__(self, other):
        if other.type == "Const":
            return Equation(
                Operation.CONST,
                None,
                None,
                self.value
            ) * Equation(
                Operation.CONST,
                None,
                None,
                other
            )
        if other.type == "Var":
            return Equation(
                Operation.CONST,
                None,
                None,
                self.value
            ) * Equation(
                Operation.VAR,
                None,
                None,
                other
            )
        if other.type == "Equation":
            return Equation(
                Operation.CONST,
                None,
                None,
                self.value
            ) * other
        
    def __truediv__(self, other):
        if other.type == "Const":
            return Equation(
                Operation.CONST,
                None,
                None,
                self.value
            ) / Equation(
                Operation.CONST,
                None,
                None,
                other
            )
        if other.type == "Var":
            return Equation(
                Operation.CONST,
                None,
                None,
                self.value
            ) / Equation(
                Operation.VAR,
                None,
                None,
                other
            )
        if other.type == "Equation":
            return Equation(
                Operation.CONST,
                None,
                None,
                self.value
            ) / other

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

    def __add__(self, other):
        if other.type == "Const":
            return Equation(
                Operation.VAR,
                None,
                None,
                self.name
            ) + Equation(
                Operation.CONST,
                None,
                None,
                other
            )
        if other.type == "Var":
            return Equation(
                Operation.VAR,
                None,
                None,
                self.name
            ) + Equation(
                Operation.VAR,
                None,
                None,
                other
            )
        if other.type == "Equation":
            return Equation(
                Operation.VAR,
                None,
                None,
                self.name
            ) + other
        
    def __sub__(self, other):
        if other.type == "Const":
            return Equation(
                Operation.VAR,
                None,
                None,
                self.name
            ) - Equation(
                Operation.CONST,
                None,
                None,
                other
            )
        if other.type == "Var":
            return Equation(
                Operation.VAR,
                None,
                None,
                self.name
            ) - Equation(
                Operation.VAR,
                None,
                None,
                other
            )
        if other.type == "Equation":
            return Equation(
                Operation.VAR,
                None,
                None,
                self.name
            ) - other
        
    def __mul__(self, other):
        if other.type == "Const":
            return Equation(
                Operation.VAR,
                None,
                None,
                self.name
            ) * Equation(
                Operation.CONST,
                None,
                None,
                other
            )
        if other.type == "Var":
            return Equation(
                Operation.VAR,
                None,
                None,
                self.name
            ) * Equation(
                Operation.VAR,
                None,
                None,
                other
            )
        if other.type == "Equation":
            return Equation(
                Operation.VAR,
                None,
                None,
                self.name
            ) * other
        
    def __truediv__(self, other):
        if other.type == "Const":
            return Equation(
                Operation.VAR,
                None,
                None,
                self.name
            ) / Equation(
                Operation.CONST,
                None,
                None,
                other
            )
        if other.type == "Var":
            return Equation(
                Operation.VAR,
                None,
                None,
                self.name
            ) / Equation(
                Operation.VAR,
                None,
                None,
                other
            )
        if other.type == "Equation":
            return Equation(
                Operation.VAR,
                None,
                None,
                self.name
            ) / other

# 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 Operation(self.op, self.left.copy(), self.right.copy(), None)

    def diff(self):
        if self.op == Operation.CONST:
            return Const(0)
        if self.op == Operation.VAR:
            return Var(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 __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 __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 __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 __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)