'''
Created on 15-Dec-2020

@author: bthakur
'''
import numpy as np   
from sympy import (Eq, symbols, solve, Rational, 
                   Add, sympify, cancel, latex, Matrix)
import random
import sympy
import traceback
import math
import os

from xv.km.managers import BaseManager
from xv.km.kelements import KProblemTemplate, KPrintElement

try:
    from IPython.display import display
    pass
except:
    display = print
    pass
class SingleVariateFirstOrder(BaseManager):
    """
    Solve equations of Single Variable First Order only
    Example:
        ke = SingleVariateFirstOrder()
        ke.getRandomProblem()
        ke.getRandomProblem(problem_type = 2)
        ke.getRandomProblem(problem_type = 3)
        ...
        ke.getEquation()
        ke.printProblem()
        ke.printAnswer()
        ke.printSolution()
        
        ke.getNotebookTemplates()
        ke.generateNotebook(
            templateName = 'SOME_TEMPLATE_NAME',
            saveToFolder = 'output'
            )        
    """
    def __new__(cls, *args, **kwargs):
        instance = super(SingleVariateFirstOrder, cls).__new__(cls)
        for base_class in SingleVariateFirstOrder.__bases__:
            base_class.__init__(instance)

        return instance
    
    def __init__(self, *args, **kwargs):
        self.x = symbols('x')
        self._steps = None
        self._matrices = None
        difficultyLevel = kwargs.get('difficultyLevel', None)
        if not difficultyLevel is None:
            kwargs['problem_type'] = difficultyLevel
            kwargs['show'] = True
            self.getRandomProblem(*args, **kwargs)
        pass
    
    def setContent(self, *args, **kwargs):
        self._content = "A first order expression of single variable"
        pass

    def help(self):
        return "TODO"
        pass
    
    def getRandomProblem(self, *args, **kwargs):
        self._initProblem(*args, **kwargs)
        show = kwargs.get('show', True)
        if show:
            return self.printProblem()
        return self._problem
        pass

    def _initProblem(self, *args, **kwargs):
        self.minV = kwargs.get('minV', -9)
        self.maxV = kwargs.get('maxV', 9)
        self._matrix = self._getRandomData()
        
        #reset matrix according to problem_type
        solution = self.solve( show = False)
        

        #get problem_type, if not set something 
        self.problem_type = kwargs.get('problem_type', 
                                       random.randint(0, len(solution) - 2))
        self.problem_type = max(0, self.problem_type)
        self.problem_type = min(len(solution) - 2, self.problem_type)
        
        prob_step_index = -1 * (self.problem_type + 1)
        self._steps = solution[prob_step_index:]
        problem_step = self._steps[0]
        self._matrix = problem_step.get('mPrev', None)
        
        statement = 'Solve the following expression:'
        expr = KPrintElement(self.getEquation(self._matrix), 'sympify/html')
        
        self._problem = []
        self._problem.append(statement) 
        self._problem.append(expr) 
        self._problem.append(f"Difficulty Level = {self.problem_type}") 

        self._answer = solve(self.getEquation())
        self._solution = [
                self._getStepsStr(),
            ]
        
        pass

    def _getRandomData(self):
        arr = np.random.randint(low = self.minV,
                                high = self.maxV, size=(2,4) 
                                ).astype('int')
                                
        arr[arr == 0] = random.randint(self.minV, self.maxV)
        
        return Matrix(arr)
    

    def _getStepsStr(self):
        sArr = []
        for i in range(len(self._steps)):
            step = self._steps[i]
            mPrev = step.get('mPrev', None)
            m = step.get('m', None)
            description = step.get('description', None)
            factor = step.get('factor', None)
            if (not mPrev is None and i == 0):
                sArr.append(
                    KPrintElement(latex(self.getEquation(mPrev)), 'text/html')
                    )
            if not description is None:
                sArr.append(description)
                
            if not factor is None:
                sArr.append(
                    KPrintElement(latex(self.getEquation(mPrev, factor = factor)), 
                                  'text/html')
                    )

            if not m is None:
                sArr.append(
                    KPrintElement(latex(self.getEquation(m)), 
                                  'sympify/html')
                    )

            sArr.append(KPrintElement(f"&nbsp;", 'text/html'))
            pass
        return sArr
        pass
    
    
    def initProblem(self, *args, **kwargs):
        raise Exception("NOT IMPLEMENTED")
        pass
    
    def _init_problem(self, minV = 1, maxV = 9):
        arr = np.random.randint(low=minV,high=maxV,size=(2,4) ).astype('int')
        arr[arr == 0] = random.randint(minV, maxV)
        return sympy.Matrix(arr)
        pass

    def _getEqSide(self, i, m):
        x = self.x
        #print(m)
        
        if m[i, 0] == 0:
            t1 = ""
        else:
            t1 = f"{m[i, 0]}"
        
        
        t2 = m[i, 1] * x + m[i, 2]
        #t2 = f"{t2}"
        
        if m[i, 3] != 1:
            t2 = f"({t2})/{m[i, 3]}"
            
        return f"{t1} + {t2}"

    def getEquation(self, m = None, factor = None):
        try:
            _ = self._matrix
        except:
            self.getRandomProblem()
            
        if m is None:
            m = self._matrix
        
        lhs = self._getEqSide(0, m)   
        #print(lhs) 
        rhs = self._getEqSide(1, m)   
        #print(rhs) 
        
        if factor:
            lhs = f"({lhs}) {factor}"
            rhs = f"({rhs}) {factor}"
            pass
        
        lhs = sympify(lhs, evaluate=False)
        rhs = sympify(rhs, evaluate=False)
          
        eq = Eq(lhs, rhs)
        return eq
        pass


    def _show(self):
        print(f"Difficulty Level = {self.problem_type}")
        for i in range(len(self._steps)):
            self._showStep(i)
            pass
        pass

    def _showStep(self, index = None, showMatrix = False):
        if index is None:
            index = -1
        lastStep = self._steps[index]
        
        description = lastStep.get('description', None)
        if description:
            print(f"\n{description}")

        factor = lastStep.get('factor', None)
        if factor:
            prevStep = self._steps[index-1]
            mPrev = prevStep.get('m')
            display(self.getEquation(mPrev, factor = factor))
        
        m = lastStep.get('m')
        if showMatrix:
            display(m)
            
        display(self.getEquation(m))
        
        pass

    def solve(self, m = None, show = False):
        if m is None:
            m = self._matrix
            self._steps = []
            self._steps.append({"m": m})
            self._matrices = []
            #self._showStep()

        #aggregate independent coeffs
        m = self._solve_0(m.copy())
        #self._showStep()

        #remove denominators
        m = self._solve_3(m.copy())
        #self._showStep()
        
        #aggregate coeffs of x
        m = self._solve_1(m.copy())
        #self._showStep()

        #simplify_free_coeff
        m = self._solve_simplify_free_coeff(m.copy())
        #self._showStep()
        
        #move x coeff to left side
        m = self._solve_swap(m.copy())
        #self._showStep()

        #move x coeff to left side
        m = self._solve_free_coeff(m.copy())
        #self._showStep()

        #move x coeff to left side
        m = self._solve_x(m.copy())
        #self._showStep()

        return self._steps
        pass

    def _solve_x(self, m):
        mPrev = m.copy()
        if (m[0, 1] == 0 ):
            return m
        
        description = f"divide both sides by {m[0, 1]}"
        factor = f"* 1/{m[0, 1]}"
        
        m[1, 0] = Rational(m[1, 0] , m[0, 1])
        m[0, 1] = 1
        
        self._steps.append({"mPrev": mPrev,
                            "m": m,
                            "factor": factor,
                            "description": description})
        
        
        return m
        pass

    
    def _solve_simplify_free_coeff(self, m):
        changeF = False
        mPrev = m.copy()
        
        if (m[0, 0] != 0 and m[0, 2] != 0):
            changeF = True
        m[0, 0] = m[0, 0] + m[0, 2]
        m[0, 2] = 0 
        
        if (m[1, 0] != 0 and m[1, 2] != 0):
            changeF = True
        m[1, 0] = m[1, 0] + m[1, 2]
        m[1, 2] = 0 
    
    
        if changeF:
            description = "Simplify"
            self._steps.append({"mPrev": mPrev,
                            "m": m,
                            "description": description})
        return m
        pass
    
    def _solve_free_coeff(self, m):
        mPrev = m.copy()
        if (m[0, 0] == 0 ):
            return m
        
        t = m[0, 0]
        
        if t > 0:
            description = f"Deducting {t} from both sides"
            factor = f" - {t}"
        else:
            description = f"Adding {0-t} from both sides"
            factor = f" + {0-t}"

        
        
        m[0, 0] = 0
        m[1, 0] = m[1, 0] - t

        self._steps.append({"mPrev": mPrev,
                            "m": m,
                            "factor": factor,
                            "description": description})
        return m
        pass


    def _solve_swap(self, m):
        mPrev = m.copy()
        if m[1, 1] == 0:
            return m

        description = f"Switching sides"
        m = sympy.Matrix([m.row(1), m.row(0)])
        
        self._steps.append({"mPrev": mPrev,
                            "m": m,
                            "description": description})
        return m
        pass

    def _solve_0(self, m):
        mPrev = m.copy()
        if not (m[0, 0] != 0 and m[1, 0] != 0):
            return m

        t = min(m[0, 0], m[1, 0])
        
        if t > 0:
            description = f"Deducting {t} from both sides"
            factor = f" - {t}"
        else:
            description = f"Adding {0-t} from both sides"
            factor = f" + {0-t}"
        

        m[0, 0] = m[0, 0] - t
        m[1, 0] = m[1, 0] - t

        self._steps.append({"mPrev": mPrev,
                            "m": m,
                            "factor": factor,
                            "description": description})
        
        return m
        pass
    
    def _solve_1(self, m):
        mPrev = m.copy()
        if not (m[0, 1] != 0 and m[1, 1] != 0):
            return m

        t = min(m[0, 1], m[1, 1])
        
        if t > 0:
            description = f"Deducting {t}{self.x} from both sides"
            factor = f" - {t}*{self.x}"
        else:
            description = f"Adding {0-t}{self.x} from both sides"
            factor = f" + {0-t}*{self.x}"
        
        m[0, 1] = m[0, 1] - t
        m[1, 1] = m[1, 1] - t

        self._steps.append({"mPrev": mPrev,
                            "m": m,
                            "factor": factor,
                            "description": description})
        return m
        pass
    
    def _solve_3(self, m):
        mPrev = m.copy()
        #if denominators on both sides are 1, don't do anything
        if (m[0, 3] == 1 and m[1, 3] == 1):
            return m

        t = m[0, 3] * m[1, 3] // math.gcd(m[0, 3], m[1, 3])
        
        leftT = t // m[0, 3]
        rightT = t // m[1, 3]
        
        
        description = f"Multiplying {t}, LCM of {m[0, 3]} and {m[1, 3]} from both sides"

        
        factor = f" * {t}"

        
        m[0, 3] = 1
        m[1, 3] = 1

        m[0, 1] = m[0, 1] * leftT
        m[1, 1] = m[1, 1] * rightT

        m[0, 2] = m[0, 2] * leftT
        m[1, 2] = m[1, 2] * rightT

        m[0, 0] = m[0, 0] * t
        m[1, 0] = m[1, 0] * t

        self._steps.append({"mPrev": mPrev,
                            "m": m,
                            "factor": factor,
                            "description": description})
        return m
        pass
    
    pass
#class SingleVariateFirstOrder:
    

