from enum import Enum
import numpy as np
from tqdm import tqdm
import inspect
import dis

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 simplify(self):
        pass

    def diff(self):
        return Equation(
            Operation.CONST,
            None,
            None,
            Const(0)
        )
    
    def partial(self, _):
        return Equation(
            Operation.CONST,
            None,
            None,
            Const(0)
        )

    def __str__(self):
        return str(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
        
    def __pow__(self, other):
        if isinstance(other, (int, float)):
            return Equation(
                Operation.CONST,
                None,
                None,
                self.value ** other
            )

# var class
class Var:
    def __init__(self, name=None, degree=0) -> None:
        self.type = "Var"
        self.degree = degree
        if name == None:
            frame = inspect.currentframe().f_back
            instructions = list(dis.get_instructions(frame.f_code))
            lasti = frame.f_lasti
            for i, instr in enumerate(instructions):
                if instr.offset > lasti:
                    if instr.opname in ("STORE_NAME", "STORE_FAST"):
                        name = instr.argval
                    break
        self.name = name

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

    def diff(self, n=1):
        return Var("d_"*n+self.name, self.degree+n)
    
    def evaluate(self, dico):
        return dico[self.name]
    
    def __eq__(self, other):
        if isinstance(other, str):
            return False
        if not (self.type == "Var" and other.type == "Var"):
            return False
        return self.name == other.name and self.degree == other.degree
    
    def __hash__(self):
        return hash((self.name, self.degree))
    
    def __str__(self):
        return 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
        
    def __pow__(self, other):
        if isinstance(other, (int, float)):
            return Equation(
                Operation.POW,
                Equation(Operation.VAR, None, None, self.copy()),
                Equation(Operation.CONST, None, None, Const(other)),
                None
            )

# 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"
        self.compiled_function = None

    def copy(self):
        if self.op == Operation.CONST:
            return Equation(Operation.CONST, None, None, Const(self.value.value))
        if self.op == Operation.VAR:
            return Equation(Operation.VAR, None, None, Var(self.value.name, self.value.degree))
        if self.op in {Operation.SIN, Operation.COS}:
            return Equation(self.op, None, None, self.value.copy())
        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())
        if self.op == Operation.SIN:
            return self.value.diff() * Equation(Operation.COS, None, None, self.value.copy())
        if self.op == Operation.COS:
            return Const(-1) * self.value.diff() * Equation(Operation.SIN, None, None, self.value.copy())
        if self.op == Operation.POW:
            n = self.right.value.value
            return Equation(Operation.CONST, None, None, Const(n)) *\
                Equation(
                    Operation.POW,
                    self.left.copy(),
                    Equation(Operation.CONST, None, None, Const(n - 1)),
                    None
                ) *\
                self.left.diff()
        
    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())
        if self.op == Operation.SIN:
            return self.value.partial(var_name) * Equation(Operation.COS, None, None, self.value.copy())
        if self.op == Operation.COS:
            return Const(-1) * self.value.partial(var_name) * Equation(Operation.SIN, None, None, self.value.copy())
        if self.op == Operation.POW:
            n = self.right.value.value
            return Equation(Operation.CONST, None, None, Const(n)) *\
                Equation(
                    Operation.POW,
                    self.left.copy(),
                    Equation(Operation.CONST, None, None, Const(n - 1)),
                    None
                ) *\
                self.left.partial(var_name)

    def compile(self):
        if self.compiled_function is not None:
            return self.compiled_function
        
        code  = self.generate_code()
        self.compiled_function = eval(code)
        return self.compiled_function

    def generate_code(self):
        code, var_count = self._generate_code_recursive(0)
        return f"lambda dico: {code}"
    
    def _generate_code_recursive(self, var_count):
        if self.op == Operation.CONST:
            val = self.value.value
            return (str(val), var_count)
        
        elif self.op == Operation.VAR:
            val = self.value.name
            return (f"dico['{val}']", var_count)

        elif self.op == Operation.ADD:
            left_code, var_count = self.left._generate_code_recursive(var_count)
            right_code, var_count = self.right._generate_code_recursive(var_count)
            return (f"({left_code}) + ({right_code})", var_count)

        elif self.op == Operation.SUB:
            left_code, var_count = self.left._generate_code_recursive(var_count)
            right_code, var_count = self.right._generate_code_recursive(var_count)
            return (f"({left_code}) - ({right_code})", var_count)
        
        elif self.op == Operation.MULT:
            left_code, var_count = self.left._generate_code_recursive(var_count)
            right_code, var_count = self.right._generate_code_recursive(var_count)
            return (f"({left_code}) * ({right_code})", var_count)
        
        elif self.op == Operation.DIV:
            left_code, var_count = self.left._generate_code_recursive(var_count)
            right_code, var_count = self.right._generate_code_recursive(var_count)
            return (f"({left_code}) / ({right_code})", var_count)
        
        elif self.op == Operation.SIN:
            val_code, var_count = self.value._generate_code_recursive(var_count)
            return (f"np.sin({val_code})", var_count)
        
        elif self.op == Operation.COS:
            val_code, var_count = self.value._generate_code_recursive(var_count)
            return (f"np.cos({val_code})", var_count)
        
        elif self.op == Operation.POW:
            left_code, var_count = self.left._generate_code_recursive(var_count)
            right_code, var_count = self.right._generate_code_recursive(var_count)
            return (f"({left_code}) ** ({right_code})", var_count)
        
        return ("0", var_count)
        
    def evaluate(self, dico):
        if self.compiled_function is None:
            self.compile()
        return self.compiled_function(dico)
        
    def solve(self, init, tmax=10, dt=0.01, progress_bar=True, F_ext = Const(0)):
        dico = {}
        for v_name, v_init in init.items():
            for i in range(len(v_init)):
                dico["d_"*i + str(v_name)] = v_init[i]

        self.simplify()
        variables = self.getAllVar()

        equations = []
        # print([v.name for v in variables])
        for var in variables:
            equations.append(
                self.partial(var.diff().name).diff()\
                - self.partial(var.name)\
                + F_ext.partial(var.diff().name)
            )

        unknown = self.getUnknown(dico)
        for eq in equations:
            eq.simplify()
            unknown = unknown.union(eq.getUnknown(dico))

        t = 0
        values = {t : {var.name : dico[var.name] for var in variables}}

        for var in unknown:
            dico[var.name] = 0.

        # print([var.name for var in unknown])
        # print(equations[0])
        assert(len(unknown) == len(equations))
        n = len(unknown)

        variables = list(variables)
        unknown = list(unknown)
        equations = list(equations)

        J_template = np.empty((n, n), dtype=object)
        F_template = np.empty(n, dtype=object)
        for i in range(n):
            f = equations[i]
            f.compile()
            F_template[i] = f
            for j in range(n):
                eq = equations[i].partial(unknown[j].name)
                eq.compile()
                J_template[i][j] = eq

        if progress_bar:
            bar = tqdm([n * dt for n in range(int(tmax/dt))])
        else:
            bar = [n * dt for n in range(int(tmax/dt))]
        for t in bar:
            err = 1
            iterations = 0
            max_iteration = 50

            while err > 10**-6 and iterations < max_iteration:
                iterations += 1
                # jacobian
                J = np.zeros((n, n))
                F = np.zeros(n)
                for i in range(n):
                    F[i] = F_template[i].evaluate(dico)
                    for j in range(n):
                        # print(equations[0])
                        J[i][j] = J_template[i][j].evaluate(dico)

                # print(J, F)

                err = np.linalg.norm(F)
                dX = -np.linalg.inv(J).dot(F)

                for i in range(n):
                    dico[unknown[i].name] += dX[i]

            values[t] = {var.name : dico[var.name] for var in variables}
            t += dt

            # update all values
            for i in range(n):
                var = unknown[i]

                d = var.degree
                name = var.name
                while d > 0:
                    ivar_name = name[2:]
                    dico[ivar_name] += dt * dico[name]

                    d -= 1
                    name = ivar_name

        return values

    def getAllVar(self):
        if self.op == Operation.CONST:
            return set()
        if self.op == Operation.VAR:
            if self.value.degree == 0:
                return {self.value.copy()}
            else:
                return {Var(self.value.name[2*self.value.degree:], 0)}
        if self.op in {Operation.SIN, Operation.COS}:
            return self.value.getAllVar()
        return self.left.getAllVar().union(self.right.getAllVar())
    
    def getUnknown(self, dico):
        # print(self)
        if self.op == Operation.CONST:
            return set()
        if self.op == Operation.VAR:
            if self.value.name not in dico:
                return {self.value}
            return set()
        if self.op in {Operation.SIN, Operation.COS}:
            return self.value.getUnknown(dico)
        return self.left.getUnknown(dico).union(self.right.getUnknown(dico))
    
    def simplify(self):
        if self.left != None:
            self.left.simplify()
        if self.right != None:
            self.right.simplify()
        if self.op in {Operation.SIN, Operation.COS}:
            self.value.simplify()

        elif self.op == Operation.CONST:
            pass
        elif self.op == Operation.VAR:
            pass

        elif self.op == Operation.ADD:
            if self.left.op == Operation.CONST and self.right.op == Operation.CONST:
                self.op = Operation.CONST
                self.value = Const(self.left.value.value + self.right.value.value)
                self.left = None
                self.right = None
            elif self.left.op == Operation.CONST and self.left.value.value == 0:
                self.op = self.right.op
                self.left = self.right.left
                self.value = self.right.value
                self.right = self.right.right
            elif self.right.op == Operation.CONST and self.right.value.value == 0:
                self.op = self.left.op
                self.right = self.left.right
                self.value = self.left.value
                self.left = self.left.left

        elif self.op == Operation.SUB:
            if self.left.op == Operation.CONST and self.right.op == Operation.CONST:
                self.op = Operation.CONST
                self.value = Const(self.left.value.value - self.right.value.value)
                self.left = None
                self.right = None
            elif self.right.op == Operation.CONST and self.right.value.value == 0:
                self = self.left

        elif self.op == Operation.MULT:
            if self.left.op == Operation.CONST and self.right.op == Operation.CONST:
                self.op = Operation.CONST
                self.value = Const(self.left.value.value * self.right.value.value)
                self.left = None
                self.right = None
            elif self.left.op == Operation.CONST and self.left.value.value == 0:
                self.op = Operation.CONST
                self.left = None
                self.right = None
                self.value = Const(0)
            elif self.right.op == Operation.CONST and self.right.value.value == 0:
                self.op = Operation.CONST
                self.left = None
                self.right = None
                self.value = Const(0)
            elif self.left.op == Operation.CONST and self.left.value.value == 1:
                self.op = self.right.op
                self.left = self.right.left
                self.value = self.right.value
                self.right = self.right.right
            elif self.right.op == Operation.CONST and self.right.value.value == 1:
                self.op = self.left.op
                self.right = self.left.right
                self.value = self.left.value
                self.left = self.left.left

        elif self.op == Operation.SIN:
            if self.value.op == Operation.CONST:
                self.op = Operation.CONST
                self.left = None
                self.right = None
                self.value = Const(np.sin(self.value.value))

        elif self.op == Operation.COS:
            if self.value.op == Operation.CONST:
                self.op = Operation.CONST
                self.left = None
                self.right = None
                self.value = Const(np.cos(self.value.value))

        elif self.op == Operation.POW:
            if self.right.value.value == 0:
                self.op = Operation.CONST
                self.left = None
                self.right = None
                self.value = Const(1)
            elif self.right.value.value == 1:
                self.op = self.left.op
                self.right = self.left.right
                self.value = self.left.value
                self.left = self.left.left
            elif self.left.op == Operation.CONST:
                self.op = Operation.CONST
                self.value = Const(self.left.value.value ** self.right.value.value)
                self.left = None
                self.right = None
    
    def __str__(self):
        if self.op == Operation.CONST:
            return str(self.value.value)
        if self.op == Operation.VAR:
            return self.value.name
        if self.op == Operation.ADD:
            return f"{str(self.left)} + {str(self.right)}"
        if self.op == Operation.SUB:
            return f"{str(self.left)} - {str(self.right)}"
        if self.op == Operation.MULT:
            return f"({str(self.left)})*({str(self.right)})"
        if self.op == Operation.DIV:
            return f"({str(self.left)})/({str(self.right)})"
        if self.op == Operation.SIN:
            return f"sin({str(self.value)})"
        if self.op == Operation.COS:
            return f"cos({str(self.value)})"
        if self.op == Operation.POW:
            return f"({str(self.left)})**({str(self.right)})"
        
    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
    
    def __pow__(self, other):
        if isinstance(other, (int, float)):
            return Equation(
                Operation.POW,
                self.copy(),
                Equation(Operation.CONST, None, None, Const(other)),
                None
            )

def cos(value):
    if isinstance(value, (float, int)):
        return Equation(Operation.CONST, None, None, np.cos(value))
    if value.type == "Const":
        return Equation(Operation.CONST, None, None, np.cos(value.value))
    if value.type == "Var":
        return Equation(Operation.COS, None, None, Equation(Operation.VAR, None, None, value.copy()))
    if value.type == "Equation":
        return Equation(Operation.COS, None, None, value.copy())
    
def sin(value):
    if isinstance(value, (float, int)):
        return Equation(Operation.CONST, None, None, np.sin(value))
    if value.type == "Const":
        return Equation(Operation.CONST, None, None, np.sin(value.value))
    if value.type == "Var":
        return Equation(Operation.SIN, None, None, Equation(Operation.VAR, None, None, value.copy()))
    if value.type == "Equation":
        return Equation(Operation.SIN, None, None, value.copy())